feat(auth): persist sessions and accounts via PostgresAuthStore

Add PgStore with hashed token sessions, IAM bootstrap, and store factory wiring; remove in-memory token map and allow_dev_switch from login payload.
This commit is contained in:
weichen
2026-09-12 01:39:31 +00:00
parent bc37debc31
commit f97ed2dd98
7 changed files with 657 additions and 22 deletions
+114
View File
@@ -0,0 +1,114 @@
package auth
import (
"context"
"fmt"
"github.com/jackc/pgx/v5/pgxpool"
"chat-agent-system/internal/config"
)
// BootstrapAdminFromConfig 确保配置的管理员账户在 PG 中存在且具备 role-admin。
func BootstrapAdminFromConfig(ctx context.Context, store AuthStore, cfg *config.Config) error {
if store == nil || cfg == nil {
return nil
}
adminUser := cfg.AdminUsername
adminPass := cfg.AdminPassword
if adminUser == "" {
adminUser = "admin"
}
if adminPass == "" {
return nil
}
existing, err := store.GetUserByUsername(ctx, adminUser)
if err == nil && existing != nil {
if existing.RoleID != "role-admin" {
_, _ = store.UpdateUserRole(ctx, existing.ID, "role-admin", "系统管理员")
}
if existing.PasswordHash == "" {
hash, hashErr := HashPassword(adminPass)
if hashErr != nil {
return hashErr
}
_, err = store.UpdateUserPassword(ctx, existing.ID, hash, false)
return err
}
return nil
}
hash, err := HashPassword(adminPass)
if err != nil {
return fmt.Errorf("bootstrap admin hash: %w", err)
}
_, err = store.CreateUser(ctx, &UserAccount{
Username: adminUser,
Name: "系统管理员",
PasswordHash: hash,
Status: StatusActive,
MustChangePassword: false,
RoleID: "role-admin",
Position: "系统管理员",
Department: "系统运维部",
EmployeeID: "SYS-000",
})
return err
}
// BootstrapDevSeedUsers 在 dev/test 环境补齐演示账号(与 MemoryStore 种子对齐)。
func BootstrapDevSeedUsers(ctx context.Context, store AuthStore, cfg *config.Config) error {
if store == nil || cfg == nil {
return nil
}
if cfg.Env != "dev" && cfg.Env != "test" && cfg.Env != "" {
return nil
}
seeds := []struct {
username string
name string
roleID string
position string
dept string
empID string
}{
{"zhanggong", "张工", "role-engineer", "方案工程师", "骨科耗材研发组", "RD-001"},
{"lizhuguan", "李主管", "role-tech-lead", "技术主管", "骨科耗材研发组", "RD-002"},
{"wangzong", "王总监", "role-director", "研发总监", "研发中心", "RD-003"},
{"guest", "访客", "role-guest", "待指派", "外协中心", "EXT-999"},
}
passHash, err := HashPassword("123456")
if err != nil {
return err
}
for _, su := range seeds {
if _, err := store.GetUserByUsername(ctx, su.username); err == nil {
continue
}
_, err := store.CreateUser(ctx, &UserAccount{
Username: su.username,
Name: su.name,
PasswordHash: passHash,
Status: StatusActive,
RoleID: su.roleID,
Position: su.position,
Department: su.dept,
EmployeeID: su.empID,
})
if err != nil {
return fmt.Errorf("seed user %s: %w", su.username, err)
}
}
return nil
}
// BootstrapIAM 在 PG IAM 就绪后执行管理员与 dev 种子初始化。
func BootstrapIAM(ctx context.Context, pool *pgxpool.Pool, store AuthStore, cfg *config.Config) error {
if !IAMTablesReady(ctx, pool) {
return nil
}
if err := BootstrapAdminFromConfig(ctx, store, cfg); err != nil {
return err
}
return BootstrapDevSeedUsers(ctx, store, cfg)
}
-1
View File
@@ -39,5 +39,4 @@ type SessionInfo struct {
User *UserAccount `json:"user"`
MustChangePassword bool `json:"must_change_password"`
Env string `json:"env"`
AllowDevSwitch bool `json:"allow_dev_switch"`
}
+391
View File
@@ -0,0 +1,391 @@
package auth
import (
"context"
"encoding/json"
"errors"
"fmt"
"strings"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
"chat-agent-system/internal/config"
)
// PgStore PostgreSQL 实现的账户、凭证与会话存储。
type PgStore struct {
pool *pgxpool.Pool
cfg *config.Config
}
// NewPgStore 构造 PG AuthStore;pool 必填且 IAM 表已迁移。
func NewPgStore(pool *pgxpool.Pool, cfg *config.Config) (*PgStore, error) {
if pool == nil {
return nil, errors.New("pg auth store: pool required")
}
return &PgStore{pool: pool, cfg: cfg}, nil
}
// IAMTablesReady 检查 Phase 1 IAM 表是否可用。
func IAMTablesReady(ctx context.Context, pool *pgxpool.Pool) bool {
if pool == nil {
return false
}
var ok bool
err := pool.QueryRow(ctx, `
SELECT EXISTS (
SELECT 1 FROM information_schema.tables
WHERE table_schema = 'public' AND table_name = 'user_credentials'
)
`).Scan(&ok)
return err == nil && ok
}
func (s *PgStore) scanUser(row pgx.Row) (*UserAccount, error) {
var u UserAccount
var passwordHash string
var createdAt, updatedAt time.Time
err := row.Scan(
&u.ID,
&u.Username,
&u.Name,
&passwordHash,
&u.Status,
&u.MustChangePassword,
&u.RoleID,
&u.Position,
&u.Department,
&u.EmployeeID,
&u.Email,
&createdAt,
&updatedAt,
)
if err != nil {
return nil, err
}
u.PasswordHash = passwordHash
u.CreatedAt = createdAt.Format("2006-01-02 15:04:05")
u.UpdatedAt = updatedAt.Format("2006-01-02 15:04:05")
return &u, nil
}
const userSelectSQL = `
SELECT u.id,
COALESCE(NULLIF(u.username, ''), u.account_id) AS username,
u.name,
COALESCE(uc.password_hash, '') AS password_hash,
u.status,
u.must_change_password,
u.role_id,
u.position,
u.department,
u.employee_id,
u.email,
u.created_at,
u.updated_at
FROM users u
LEFT JOIN user_credentials uc ON uc.user_id = u.id
`
func (s *PgStore) GetUserByUsername(ctx context.Context, username string) (*UserAccount, error) {
username = strings.TrimSpace(username)
row := s.pool.QueryRow(ctx, userSelectSQL+`
WHERE lower(COALESCE(NULLIF(u.username, ''), u.account_id)) = lower($1)
LIMIT 1
`, username)
u, err := s.scanUser(row)
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return nil, fmt.Errorf("user not found: %s", username)
}
return nil, err
}
return u, nil
}
func (s *PgStore) GetUserByID(ctx context.Context, id int64) (*UserAccount, error) {
row := s.pool.QueryRow(ctx, userSelectSQL+` WHERE u.id = $1`, id)
u, err := s.scanUser(row)
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return nil, fmt.Errorf("user not found: %d", id)
}
return nil, err
}
return u, nil
}
func (s *PgStore) CreateUser(ctx context.Context, user *UserAccount) (*UserAccount, error) {
if user == nil || strings.TrimSpace(user.Username) == "" {
return nil, errors.New("username required")
}
tx, err := s.pool.Begin(ctx)
if err != nil {
return nil, err
}
defer tx.Rollback(ctx)
roleLegacy := "engineer"
switch user.RoleID {
case "role-director":
roleLegacy = "director"
case "role-tech-lead":
roleLegacy = "regulatory"
case "role-guest":
roleLegacy = "sales"
}
var id int64
err = tx.QueryRow(ctx, `
INSERT INTO users (
account_id, name, role, username, status, must_change_password,
role_id, position, department, employee_id, email
) VALUES ($1, $2, $3, $1, $4, $5, $6, $7, $8, $9, $10)
RETURNING id
`, user.Username, user.Name, roleLegacy, user.Status, user.MustChangePassword,
user.RoleID, user.Position, user.Department, user.EmployeeID, user.Email).Scan(&id)
if err != nil {
return nil, fmt.Errorf("insert user: %w", err)
}
if user.PasswordHash != "" {
_, err = tx.Exec(ctx, `
INSERT INTO user_credentials (user_id, password_hash)
VALUES ($1, $2)
`, id, user.PasswordHash)
if err != nil {
return nil, fmt.Errorf("insert credentials: %w", err)
}
}
if user.RoleID != "" {
_, _ = tx.Exec(ctx, `
INSERT INTO user_role_bindings (user_id, role_id, scope_type, scope_id, assigned_by)
VALUES ($1, $2, 'system', 0, $1)
ON CONFLICT DO NOTHING
`, id, user.RoleID)
}
if err := tx.Commit(ctx); err != nil {
return nil, err
}
return s.GetUserByID(ctx, id)
}
func (s *PgStore) UpdateUser(ctx context.Context, user *UserAccount) (*UserAccount, error) {
_, err := s.pool.Exec(ctx, `
UPDATE users SET
name = $2,
department = $3,
employee_id = $4,
email = $5,
updated_at = now()
WHERE id = $1
`, user.ID, user.Name, user.Department, user.EmployeeID, user.Email)
if err != nil {
return nil, err
}
return s.GetUserByID(ctx, user.ID)
}
func (s *PgStore) UpdateUserRole(ctx context.Context, id int64, roleID, position string) (*UserAccount, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
return nil, err
}
defer tx.Rollback(ctx)
_, err = tx.Exec(ctx, `
UPDATE users SET role_id = $2, position = $3, updated_at = now() WHERE id = $1
`, id, roleID, position)
if err != nil {
return nil, err
}
_, err = tx.Exec(ctx, `
DELETE FROM user_role_bindings
WHERE user_id = $1 AND scope_type = 'system' AND scope_id = 0 AND role_id <> $2
`, id, roleID)
if err != nil {
return nil, err
}
_, err = tx.Exec(ctx, `
INSERT INTO user_role_bindings (user_id, role_id, scope_type, scope_id, assigned_by)
VALUES ($1, $2, 'system', 0, $1)
ON CONFLICT DO NOTHING
`, id, roleID)
if err != nil {
return nil, err
}
if err := tx.Commit(ctx); err != nil {
return nil, err
}
return s.GetUserByID(ctx, id)
}
func (s *PgStore) UpdateUserStatus(ctx context.Context, id int64, status AccountStatus) (*UserAccount, error) {
_, err := s.pool.Exec(ctx, `
UPDATE users SET status = $2, updated_at = now() WHERE id = $1
`, id, status)
if err != nil {
return nil, err
}
if status == StatusDisabled {
_ = s.RevokeAllForUser(ctx, id)
}
return s.GetUserByID(ctx, id)
}
func (s *PgStore) UpdateUserPassword(ctx context.Context, id int64, passwordHash string, mustChangePassword bool) (*UserAccount, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
return nil, err
}
defer tx.Rollback(ctx)
_, err = tx.Exec(ctx, `
INSERT INTO user_credentials (user_id, password_hash, password_changed_at, failed_count, locked_until)
VALUES ($1, $2, now(), 0, NULL)
ON CONFLICT (user_id) DO UPDATE SET
password_hash = EXCLUDED.password_hash,
password_changed_at = now(),
failed_count = 0,
locked_until = NULL
`, id, passwordHash)
if err != nil {
return nil, err
}
_, err = tx.Exec(ctx, `
UPDATE users SET must_change_password = $2, updated_at = now() WHERE id = $1
`, id, mustChangePassword)
if err != nil {
return nil, err
}
if err := tx.Commit(ctx); err != nil {
return nil, err
}
return s.GetUserByID(ctx, id)
}
func (s *PgStore) ListUsers(ctx context.Context) ([]*UserAccount, error) {
rows, err := s.pool.Query(ctx, userSelectSQL+` ORDER BY u.id ASC`)
if err != nil {
return nil, err
}
defer rows.Close()
list := make([]*UserAccount, 0)
for rows.Next() {
u, err := s.scanUser(rows)
if err != nil {
return nil, err
}
u.PasswordHash = ""
list = append(list, u)
}
return list, rows.Err()
}
func (s *PgStore) GetConfig(ctx context.Context) (SystemAuthConfig, error) {
var raw []byte
err := s.pool.QueryRow(ctx, `
SELECT config_value FROM iam_system_config WHERE config_key = 'auth'
`).Scan(&raw)
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return SystemAuthConfig{AllowRegistration: true, RequireApproval: false}, nil
}
return SystemAuthConfig{}, err
}
var cfg SystemAuthConfig
if err := json.Unmarshal(raw, &cfg); err != nil {
return SystemAuthConfig{}, err
}
return cfg, nil
}
func (s *PgStore) UpdateConfig(ctx context.Context, cfg SystemAuthConfig) error {
raw, err := json.Marshal(cfg)
if err != nil {
return err
}
_, err = s.pool.Exec(ctx, `
INSERT INTO iam_system_config (config_key, config_value, updated_at)
VALUES ('auth', $1, now())
ON CONFLICT (config_key) DO UPDATE SET config_value = EXCLUDED.config_value, updated_at = now()
`, raw)
return err
}
func (s *PgStore) CreateSession(ctx context.Context, userID int64, tokenHash string, meta SessionMeta) error {
now := time.Now()
_, err := s.pool.Exec(ctx, `
INSERT INTO auth_sessions (
token_hash, user_id, created_at, last_seen_at, idle_expires_at, absolute_expires_at, ip_hash, user_agent
) VALUES ($1, $2, $3, $3, $4, $5, $6, $7)
`, tokenHash, userID, now, now.Add(defaultIdleSessionTTL), now.Add(defaultAbsoluteSessionTTL), meta.IPHash, meta.UserAgent)
return err
}
func (s *PgStore) ResolveUserID(ctx context.Context, tokenHash string) (int64, error) {
var userID int64
var idleExp, absExp time.Time
var revokedAt *time.Time
err := s.pool.QueryRow(ctx, `
SELECT user_id, idle_expires_at, absolute_expires_at, revoked_at
FROM auth_sessions WHERE token_hash = $1
`, tokenHash).Scan(&userID, &idleExp, &absExp, &revokedAt)
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return 0, errSessionNotFound
}
return 0, err
}
if revokedAt != nil {
return 0, errSessionNotFound
}
now := time.Now()
if now.After(idleExp) || now.After(absExp) {
return 0, errSessionNotFound
}
var status AccountStatus
if err := s.pool.QueryRow(ctx, `SELECT status FROM users WHERE id = $1`, userID).Scan(&status); err != nil {
return 0, errSessionNotFound
}
if status == StatusDisabled || status == StatusPendingApproval {
return 0, errSessionNotFound
}
_, _ = s.pool.Exec(ctx, `
UPDATE auth_sessions SET last_seen_at = now(), idle_expires_at = $2 WHERE token_hash = $1
`, tokenHash, now.Add(defaultIdleSessionTTL))
return userID, nil
}
func (s *PgStore) TouchSession(ctx context.Context, tokenHash string) error {
_, err := s.pool.Exec(ctx, `
UPDATE auth_sessions SET last_seen_at = now(), idle_expires_at = $2
WHERE token_hash = $1 AND revoked_at IS NULL
`, tokenHash, time.Now().Add(defaultIdleSessionTTL))
return err
}
func (s *PgStore) RevokeByTokenHash(ctx context.Context, tokenHash string) error {
_, err := s.pool.Exec(ctx, `
UPDATE auth_sessions SET revoked_at = now() WHERE token_hash = $1 AND revoked_at IS NULL
`, tokenHash)
return err
}
func (s *PgStore) RevokeAllForUser(ctx context.Context, userID int64) error {
_, err := s.pool.Exec(ctx, `
UPDATE auth_sessions SET revoked_at = now() WHERE user_id = $1 AND revoked_at IS NULL
`, userID)
return err
}
+16 -21
View File
@@ -7,8 +7,6 @@ import (
"errors"
"fmt"
"strings"
"sync"
"chat-agent-system/internal/config"
)
@@ -43,20 +41,22 @@ type CreateUserRequest struct {
// Service 认证与账号治理业务逻辑
type Service struct {
store AuthStore
sessions SessionStore
cfg *config.Config
tokensMu sync.RWMutex
tokens map[string]int64 // token -> userID
}
// NewService 实例化认证服务
func NewService(store AuthStore, cfg *config.Config) *Service {
// NewService 实例化认证服务;sessionStore 为空时使用进程内 MemorySessionStore。
func NewService(store AuthStore, sessionStore SessionStore, cfg *config.Config) *Service {
if store == nil {
store = NewMemoryStore(cfg)
}
if sessionStore == nil {
sessionStore = NewMemorySessionStore()
}
return &Service{
store: store,
cfg: cfg,
tokens: make(map[string]int64),
store: store,
sessions: sessionStore,
cfg: cfg,
}
}
@@ -91,9 +91,10 @@ func (s *Service) Login(ctx context.Context, username, password string) (*Sessio
}
token := generateToken()
s.tokensMu.Lock()
s.tokens[token] = user.ID
s.tokensMu.Unlock()
tokenHash := HashToken(token)
if err := s.sessions.CreateSession(ctx, user.ID, tokenHash, SessionMeta{}); err != nil {
return nil, fmt.Errorf("创建会话失败: %w", err)
}
// 敏感信息脱敏
user.PasswordHash = ""
@@ -108,7 +109,6 @@ func (s *Service) Login(ctx context.Context, username, password string) (*Sessio
User: user,
MustChangePassword: user.MustChangePassword,
Env: env,
AllowDevSwitch: env == "dev",
}, nil
}
@@ -191,11 +191,8 @@ func (s *Service) ChangePassword(ctx context.Context, userID int64, oldPass, new
// GetUserByToken 根据会话 Token 查找登录用户
func (s *Service) GetUserByToken(ctx context.Context, token string) (*UserAccount, error) {
s.tokensMu.RLock()
userID, ok := s.tokens[token]
s.tokensMu.RUnlock()
if !ok {
userID, err := s.sessions.ResolveUserID(ctx, HashToken(token))
if err != nil {
return nil, errors.New("会话已过期或未登录")
}
@@ -220,9 +217,7 @@ func (s *Service) GetUserByID(ctx context.Context, userID int64) (*UserAccount,
// Logout 注销会话
func (s *Service) Logout(token string) {
s.tokensMu.Lock()
defer s.tokensMu.Unlock()
delete(s.tokens, token)
_ = s.sessions.RevokeByTokenHash(context.Background(), HashToken(token))
}
// ListUsers 列出所有用户
+77
View File
@@ -0,0 +1,77 @@
package auth
import (
"context"
"crypto/sha256"
"encoding/hex"
"errors"
"time"
)
var errSessionNotFound = errors.New("session not found or expired")
const (
defaultIdleSessionTTL = 24 * time.Hour
defaultAbsoluteSessionTTL = 7 * 24 * time.Hour
)
// SessionMeta 创建会话时的客户端元数据(可选)。
type SessionMeta struct {
IPHash string
UserAgent string
}
// SessionStore 持久化或内存会话存储。
type SessionStore interface {
CreateSession(ctx context.Context, userID int64, tokenHash string, meta SessionMeta) error
ResolveUserID(ctx context.Context, tokenHash string) (int64, error)
TouchSession(ctx context.Context, tokenHash string) error
RevokeByTokenHash(ctx context.Context, tokenHash string) error
RevokeAllForUser(ctx context.Context, userID int64) error
}
// HashToken 将会话 token 转为存储用 SHA-256 十六进制摘要。
func HashToken(token string) string {
sum := sha256.Sum256([]byte(token))
return hex.EncodeToString(sum[:])
}
// MemorySessionStore 进程内 token 映射(仅 dev/stub fallback)。
type MemorySessionStore struct {
tokens map[string]int64
}
func NewMemorySessionStore() *MemorySessionStore {
return &MemorySessionStore{tokens: make(map[string]int64)}
}
func (s *MemorySessionStore) CreateSession(_ context.Context, userID int64, tokenHash string, _ SessionMeta) error {
s.tokens[tokenHash] = userID
return nil
}
func (s *MemorySessionStore) ResolveUserID(_ context.Context, tokenHash string) (int64, error) {
userID, ok := s.tokens[tokenHash]
if !ok {
return 0, errSessionNotFound
}
return userID, nil
}
func (s *MemorySessionStore) TouchSession(_ context.Context, _ string) error {
return nil
}
func (s *MemorySessionStore) RevokeByTokenHash(_ context.Context, tokenHash string) error {
delete(s.tokens, tokenHash)
return nil
}
func (s *MemorySessionStore) RevokeAllForUser(_ context.Context, userID int64) error {
for tok, uid := range s.tokens {
if uid == userID {
delete(s.tokens, tok)
}
}
return nil
}
+29
View File
@@ -0,0 +1,29 @@
package auth_test
import (
"context"
"testing"
"chat-agent-system/internal/auth"
)
func TestHashTokenAndMemorySessionStore(t *testing.T) {
store := auth.NewMemorySessionStore()
ctx := context.Background()
token := "sample-token-abc"
hash := auth.HashToken(token)
if err := store.CreateSession(ctx, 42, hash, auth.SessionMeta{}); err != nil {
t.Fatalf("CreateSession: %v", err)
}
userID, err := store.ResolveUserID(ctx, hash)
if err != nil || userID != 42 {
t.Fatalf("ResolveUserID got (%d, %v)", userID, err)
}
if err := store.RevokeByTokenHash(ctx, hash); err != nil {
t.Fatalf("Revoke: %v", err)
}
if _, err := store.ResolveUserID(ctx, hash); err == nil {
t.Fatal("expected revoked session to fail")
}
}
+30
View File
@@ -0,0 +1,30 @@
package auth
import (
"context"
"github.com/jackc/pgx/v5/pgxpool"
"chat-agent-system/internal/config"
)
// Stores 认证账户存储与会话存储组合。
type Stores struct {
Auth AuthStore
Sessions SessionStore
}
// NewStores 根据数据库与 IAM 迁移状态选择 PG 或内存实现。
func NewStores(ctx context.Context, pool *pgxpool.Pool, cfg *config.Config) (*Stores, error) {
if pool != nil && IAMTablesReady(ctx, pool) {
pg, err := NewPgStore(pool, cfg)
if err != nil {
return nil, err
}
if err := BootstrapIAM(ctx, pool, pg, cfg); err != nil {
return nil, err
}
return &Stores{Auth: pg, Sessions: pg}, nil
}
mem := NewMemoryStore(cfg)
return &Stores{Auth: mem, Sessions: NewMemorySessionStore()}, nil
}