Files
ai-gateway-go/internal/identity/repository.go
T
superidou 5759c1862e AI Gateway Go 0.10.0 源码快照 + 旗舰版需求规划报告
M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。
含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-12 11:45:54 +08:00

336 lines
12 KiB
Go

package identity
import (
"context"
"crypto/subtle"
"encoding/json"
"errors"
"fmt"
"strings"
"time"
platformid "aigateway.local/core/internal/platform/id"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
)
type Repository struct {
pool *pgxpool.Pool
}
func NewRepository(pool *pgxpool.Pool) *Repository {
return &Repository{pool: pool}
}
func (r *Repository) FindAdminByLogin(ctx context.Context, login string) (Account, error) {
if r.pool == nil {
return Account{}, ErrUnavailable
}
var account Account
account.Kind = KindAdmin
err := r.pool.QueryRow(ctx, `
SELECT id::text, username, display_name, role, permissions, password_hash, active,
failed_logins, locked_until, totp_enabled, encrypted_totp_secret,
totp_kek_version, totp_last_step, totp_backup_codes
FROM gateway.admin_accounts
WHERE lower(username) = lower($1)`, strings.TrimSpace(login)).Scan(
&account.ID, &account.Login, &account.DisplayName, &account.Role, &account.Permissions,
&account.PasswordHash, &account.Active, &account.FailedLogins,
&account.LockedUntil, &account.TOTPEnabled, &account.EncryptedTOTPSecret,
&account.TOTPKekVersion, &account.TOTPLastStep, &account.TOTPBackupCodes,
)
return account, mapRepositoryError(err)
}
func (r *Repository) FindAdminByID(ctx context.Context, id string) (Account, error) {
if r.pool == nil {
return Account{}, ErrUnavailable
}
var account Account
account.Kind = KindAdmin
err := r.pool.QueryRow(ctx, `
SELECT id::text, username, display_name, role, permissions, password_hash, active,
failed_logins, locked_until, totp_enabled, encrypted_totp_secret,
totp_kek_version, totp_last_step, totp_backup_codes
FROM gateway.admin_accounts
WHERE id = $1`, id).Scan(
&account.ID, &account.Login, &account.DisplayName, &account.Role, &account.Permissions,
&account.PasswordHash, &account.Active, &account.FailedLogins,
&account.LockedUntil, &account.TOTPEnabled, &account.EncryptedTOTPSecret,
&account.TOTPKekVersion, &account.TOTPLastStep, &account.TOTPBackupCodes,
)
return account, mapRepositoryError(err)
}
func (r *Repository) FindPortalByLogin(ctx context.Context, login string) (Account, error) {
if r.pool == nil {
return Account{}, ErrUnavailable
}
var account Account
account.Kind = KindPortal
err := r.pool.QueryRow(ctx, `
SELECT id::text, account, name, role, permissions, COALESCE(password_hash, ''), auth_source, active,
failed_logins, locked_until, totp_enabled, encrypted_totp_secret,
totp_kek_version, totp_last_step, totp_backup_codes, department_id::text
FROM gateway.portal_users
WHERE lower(account) = lower($1)`, strings.TrimSpace(login)).Scan(
&account.ID, &account.Login, &account.DisplayName, &account.Role, &account.Permissions, &account.PasswordHash,
&account.AuthSource, &account.Active, &account.FailedLogins,
&account.LockedUntil, &account.TOTPEnabled, &account.EncryptedTOTPSecret,
&account.TOTPKekVersion, &account.TOTPLastStep, &account.TOTPBackupCodes, &account.DepartmentID,
)
return account, mapRepositoryError(err)
}
func (r *Repository) FindPortalByID(ctx context.Context, id string) (Account, error) {
if r.pool == nil {
return Account{}, ErrUnavailable
}
var account Account
account.Kind = KindPortal
err := r.pool.QueryRow(ctx, `
SELECT id::text, account, name, role, permissions, COALESCE(password_hash, ''), auth_source, active,
failed_logins, locked_until, totp_enabled, encrypted_totp_secret,
totp_kek_version, totp_last_step, totp_backup_codes, department_id::text
FROM gateway.portal_users
WHERE id = $1`, id).Scan(
&account.ID, &account.Login, &account.DisplayName, &account.Role, &account.Permissions, &account.PasswordHash,
&account.AuthSource, &account.Active, &account.FailedLogins,
&account.LockedUntil, &account.TOTPEnabled, &account.EncryptedTOTPSecret,
&account.TOTPKekVersion, &account.TOTPLastStep, &account.TOTPBackupCodes, &account.DepartmentID,
)
return account, mapRepositoryError(err)
}
func (r *Repository) RecordFailure(ctx context.Context, account Account, maximum int, lockDuration time.Duration) (*time.Time, error) {
if r.pool == nil {
return nil, ErrUnavailable
}
table := "gateway.admin_accounts"
if account.Kind == KindPortal {
table = "gateway.portal_users"
}
query := fmt.Sprintf(`
UPDATE %s
SET locked_until = CASE
WHEN failed_logins + 1 >= $2
THEN clock_timestamp() + make_interval(secs => $3)
ELSE locked_until
END,
failed_logins = CASE WHEN failed_logins + 1 >= $2 THEN 0 ELSE failed_logins + 1 END,
updated_at = clock_timestamp()
WHERE id = $1
RETURNING locked_until`, table)
var lockedUntil *time.Time
err := r.pool.QueryRow(ctx, query, account.ID, maximum, int64(lockDuration.Seconds())).Scan(&lockedUntil)
return lockedUntil, mapRepositoryError(err)
}
func (r *Repository) CompleteLogin(ctx context.Context, account Account, upgradedHash *string) error {
if r.pool == nil {
return ErrUnavailable
}
table := "gateway.admin_accounts"
if account.Kind == KindPortal {
table = "gateway.portal_users"
}
query := fmt.Sprintf(`
UPDATE %s
SET failed_logins = 0,
locked_until = NULL,
last_login = clock_timestamp(),
password_hash = COALESCE($2::text, password_hash),
updated_at = clock_timestamp()
WHERE id = $1`, table)
result, err := r.pool.Exec(ctx, query, account.ID, upgradedHash)
if err != nil {
return fmt.Errorf("%w: %v", ErrUnavailable, err)
}
if result.RowsAffected() != 1 {
return ErrNotFound
}
return nil
}
func (r *Repository) SetPassword(ctx context.Context, account Account, passwordHash string) error {
query := fmt.Sprintf(`UPDATE %s SET password_hash=$2, failed_logins=0, locked_until=NULL, updated_at=clock_timestamp() WHERE id=$1`, identityTable(account.Kind))
return expectOne(r.pool, ctx, query, account.ID, passwordHash)
}
func (r *Repository) SetTOTPSecret(ctx context.Context, account Account, encrypted []byte, version int) error {
query := fmt.Sprintf(`
UPDATE %s
SET encrypted_totp_secret = $2, totp_kek_version = $3,
totp_enabled = false, totp_last_step = NULL,
totp_backup_codes = '[]'::jsonb, totp_confirmed_at = NULL,
updated_at = clock_timestamp()
WHERE id = $1`, identityTable(account.Kind))
return expectOne(r.pool, ctx, query, account.ID, encrypted, version)
}
func (r *Repository) EnableTOTP(ctx context.Context, account Account, step int64, records []BackupCodeRecord) error {
payload, err := json.Marshal(records)
if err != nil {
return err
}
query := fmt.Sprintf(`
UPDATE %s
SET totp_enabled = true, totp_last_step = $2,
totp_backup_codes = $3::jsonb, totp_confirmed_at = clock_timestamp(),
updated_at = clock_timestamp()
WHERE id = $1 AND encrypted_totp_secret IS NOT NULL AND NOT totp_enabled
AND (totp_last_step IS NULL OR totp_last_step < $2)`, identityTable(account.Kind))
return expectOne(r.pool, ctx, query, account.ID, step, payload)
}
func (r *Repository) ConsumeTOTPStep(ctx context.Context, account Account, step int64) (bool, error) {
if r.pool == nil {
return false, ErrUnavailable
}
query := fmt.Sprintf(`
UPDATE %s SET totp_last_step = $2, updated_at = clock_timestamp()
WHERE id = $1 AND totp_enabled
AND (totp_last_step IS NULL OR totp_last_step < $2)`, identityTable(account.Kind))
result, err := r.pool.Exec(ctx, query, account.ID, step)
if err != nil {
return false, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
return result.RowsAffected() == 1, nil
}
func (r *Repository) ConsumeBackupCode(ctx context.Context, account Account, hash string) (bool, error) {
if r.pool == nil {
return false, ErrUnavailable
}
tx, err := r.pool.BeginTx(ctx, pgx.TxOptions{})
if err != nil {
return false, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
defer func() { _ = tx.Rollback(ctx) }()
query := fmt.Sprintf(`SELECT totp_backup_codes FROM %s WHERE id = $1 AND totp_enabled FOR UPDATE`, identityTable(account.Kind))
var payload []byte
if err := tx.QueryRow(ctx, query, account.ID).Scan(&payload); err != nil {
return false, mapRepositoryError(err)
}
var records []BackupCodeRecord
if err := json.Unmarshal(payload, &records); err != nil {
return false, fmt.Errorf("invalid stored TOTP backup codes: %w", err)
}
found := -1
for index := range records {
if records[index].UsedAt == nil && subtle.ConstantTimeCompare([]byte(records[index].Hash), []byte(hash)) == 1 {
found = index
}
}
if found < 0 {
return false, nil
}
now := time.Now().UTC()
records[found].UsedAt = &now
payload, err = json.Marshal(records)
if err != nil {
return false, err
}
update := fmt.Sprintf(`UPDATE %s SET totp_backup_codes = $2::jsonb, updated_at = clock_timestamp() WHERE id = $1`, identityTable(account.Kind))
if _, err := tx.Exec(ctx, update, account.ID, payload); err != nil {
return false, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
if err := tx.Commit(ctx); err != nil {
return false, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
return true, nil
}
func (r *Repository) DisableTOTP(ctx context.Context, account Account) error {
query := fmt.Sprintf(`
UPDATE %s SET encrypted_totp_secret = NULL, totp_kek_version = NULL,
totp_enabled = false, totp_last_step = NULL,
totp_backup_codes = '[]'::jsonb, totp_confirmed_at = NULL,
updated_at = clock_timestamp()
WHERE id = $1 AND totp_enabled`, identityTable(account.Kind))
return expectOne(r.pool, ctx, query, account.ID)
}
func (r *Repository) ReplaceBackupCodes(ctx context.Context, account Account, records []BackupCodeRecord) error {
payload, err := json.Marshal(records)
if err != nil {
return err
}
query := fmt.Sprintf(`UPDATE %s SET totp_backup_codes = $2::jsonb, updated_at = clock_timestamp() WHERE id = $1 AND totp_enabled`, identityTable(account.Kind))
return expectOne(r.pool, ctx, query, account.ID, payload)
}
func identityTable(kind Kind) string {
if kind == KindPortal {
return "gateway.portal_users"
}
return "gateway.admin_accounts"
}
func expectOne(pool *pgxpool.Pool, ctx context.Context, query string, arguments ...any) error {
if pool == nil {
return ErrUnavailable
}
result, err := pool.Exec(ctx, query, arguments...)
if err != nil {
return fmt.Errorf("%w: %v", ErrUnavailable, err)
}
if result.RowsAffected() != 1 {
return ErrNotFound
}
return nil
}
func (r *Repository) CreateAdmin(ctx context.Context, login, displayName, role, passwordHash string) (string, error) {
if r.pool == nil {
return "", ErrUnavailable
}
id, err := platformid.NewUUID()
if err != nil {
return "", err
}
result, err := r.pool.Exec(ctx, `
INSERT INTO gateway.admin_accounts
(id, username, display_name, role, password_hash)
VALUES ($1, $2, $3, $4, $5)
ON CONFLICT DO NOTHING`, id, strings.ToLower(strings.TrimSpace(login)), displayName, role, passwordHash)
if err != nil {
return "", err
}
if result.RowsAffected() != 1 {
return "", errors.New("administrator already exists")
}
return id, nil
}
func (r *Repository) CreatePortalUser(ctx context.Context, login, displayName, passwordHash string) (string, error) {
if r.pool == nil {
return "", ErrUnavailable
}
id, err := platformid.NewUUID()
if err != nil {
return "", err
}
result, err := r.pool.Exec(ctx, `
INSERT INTO gateway.portal_users (id, account, name, password_hash, auth_source)
VALUES ($1, $2, $3, $4, 'local')
ON CONFLICT DO NOTHING`, id, strings.ToLower(strings.TrimSpace(login)), displayName, passwordHash)
if err != nil {
return "", err
}
if result.RowsAffected() != 1 {
return "", errors.New("portal user already exists")
}
return id, nil
}
func mapRepositoryError(err error) error {
if errors.Is(err, pgx.ErrNoRows) {
return ErrNotFound
}
if err != nil {
return fmt.Errorf("%w: %v", ErrUnavailable, err)
}
return nil
}