superidou
c22669c31d
- 新增迁移 000031-000034(登录日志/角色/模型配额/记忆) - 新增包: license/memory/modelquota/assistant,扫描引擎 - 全部功能后端+前端+端到端验证通过(25 包单测)
423 lines
15 KiB
Go
423 lines
15 KiB
Go
package identity
|
|
|
|
import (
|
|
"context"
|
|
"crypto/subtle"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"strconv"
|
|
"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
|
|
}
|
|
|
|
// LoginLog 是一条登录尝试记录。
|
|
type LoginLog struct {
|
|
ID string `json:"id"`
|
|
Kind string `json:"kind"`
|
|
Login string `json:"login"`
|
|
Success bool `json:"success"`
|
|
IP *string `json:"ip,omitempty"`
|
|
UserAgent string `json:"user_agent"`
|
|
Reason string `json:"reason"`
|
|
CreatedAt time.Time `json:"created_at"`
|
|
}
|
|
|
|
// RecordLoginLog 记录一次登录尝试(成功或失败)。
|
|
func (r *Repository) RecordLoginLog(ctx context.Context, kind Kind, login string, success bool, ip, userAgent, reason string) error {
|
|
if r.pool == nil {
|
|
return ErrUnavailable
|
|
}
|
|
id, err := platformid.NewUUID()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
ipValue := strings.TrimSpace(ip)
|
|
_, err = r.pool.Exec(ctx, `INSERT INTO gateway.login_logs(id,kind,login,success,ip,user_agent,reason) VALUES($1,$2,$3,$4,nullif($5,'')::inet,$6,$7)`,
|
|
id, string(kind), login, success, ipValue, truncateText(userAgent, 256), truncateText(reason, 64))
|
|
if err != nil {
|
|
return fmt.Errorf("%w: %v", ErrUnavailable, err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ListLoginLogs 查询登录记录(按时间倒序)。
|
|
func (r *Repository) ListLoginLogs(ctx context.Context, kind Kind, login string, limit int) ([]LoginLog, error) {
|
|
if r.pool == nil {
|
|
return nil, ErrUnavailable
|
|
}
|
|
if limit < 1 || limit > 500 {
|
|
limit = 50
|
|
}
|
|
query := `SELECT id::text,kind,login,success,ip::text,user_agent,reason,created_at FROM gateway.login_logs WHERE kind=$1`
|
|
args := []any{string(kind)}
|
|
if login != "" {
|
|
args = append(args, login)
|
|
query += ` AND login=$` + strconv.Itoa(len(args))
|
|
}
|
|
query += ` ORDER BY created_at DESC LIMIT $` + strconv.Itoa(len(args)+1)
|
|
args = append(args, limit)
|
|
rows, err := r.pool.Query(ctx, query, args...)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("%w: %v", ErrUnavailable, err)
|
|
}
|
|
defer rows.Close()
|
|
items := []LoginLog{}
|
|
for rows.Next() {
|
|
var item LoginLog
|
|
var ip *string
|
|
if err := rows.Scan(&item.ID, &item.Kind, &item.Login, &item.Success, &ip, &item.UserAgent, &item.Reason, &item.CreatedAt); err != nil {
|
|
return nil, err
|
|
}
|
|
if ip != nil && *ip != "" {
|
|
item.IP = ip
|
|
}
|
|
items = append(items, item)
|
|
}
|
|
return items, rows.Err()
|
|
}
|
|
|
|
func truncateText(value string, limit int) string {
|
|
value = strings.TrimSpace(value)
|
|
if len(value) <= limit {
|
|
return value
|
|
}
|
|
return value[:limit]
|
|
}
|
|
|
|
// CountIdentities 统计管理员与门户账号总数(License 账号数管控用)。
|
|
func (r *Repository) CountIdentities(ctx context.Context, total *int) error {
|
|
if r.pool == nil {
|
|
return ErrUnavailable
|
|
}
|
|
err := r.pool.QueryRow(ctx, `SELECT (SELECT count(*) FROM gateway.admin_accounts) + (SELECT count(*) FROM gateway.portal_users)`).Scan(total)
|
|
if err != nil {
|
|
return fmt.Errorf("%w: %v", ErrUnavailable, err)
|
|
}
|
|
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
|
|
}
|