knowledge-governance/internal/platform/postgres/auth.go

package postgres

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"golang.org/x/crypto/bcrypt"
)

// Roles for the governance control plane. V1 keeps exactly three; the release
// double-gate additionally requires admin so production authority is a role,
// not a typed name in a request body.
const (
	RoleViewer   = "viewer"
	RoleEngineer = "engineer"
	RoleAdmin    = "admin"
)

func ValidRole(role string) bool {
	return role == RoleViewer || role == RoleEngineer || role == RoleAdmin
}

var (
	ErrAuthUserNotFound      = errors.New("auth user not found")
	ErrAuthInvalidCredential = errors.New("invalid username or password")
	ErrAuthUserDisabled      = errors.New("account is disabled")
	ErrAuthLocked            = errors.New("account temporarily locked")
	ErrAuthSessionInvalid    = errors.New("auth session invalid")
	ErrAuthDuplicateUsername = errors.New("username already exists")
	ErrAuthInvalidUser       = errors.New("invalid user definition")
	ErrAuthLastAdmin         = errors.New("cannot remove or demote the last active admin")
)

const authMaxFailedAttempts = 8
const authLockDuration = 15 * time.Minute

type AuthUser struct {
	ID                 string
	Username           string
	DisplayName        string
	Role               string
	MustChangePassword bool
	State              string
	CreatedAt          time.Time
	LastLoginAt        *time.Time
}

// AuthSession is the server-side session behind the HttpOnly cookie. The
// caller only ever sees the opaque token; the store keeps its SHA-256.
type AuthSession struct {
	ID        string
	User      AuthUser
	ExpiresAt time.Time
}

func isUniqueViolation(err error) bool {
	var pgErr *pgconn.PgError
	return errors.As(err, &pgErr) && pgErr.Code == "23505"
}

func hashSessionToken(token string) string {
	sum := sha256.Sum256([]byte(token))
	return hex.EncodeToString(sum[:])
}

func validUsername(value string) bool {
	if len(value) < 2 || len(value) > 64 {
		return false
	}
	for index, char := range value {
		ok := (char >= 'a' && char <= 'z') || (char >= '0' && char <= '9') ||
			(index > 0 && (char == '.' || char == '_' || char == '-'))
		if !ok {
			return false
		}
	}
	return (value[0] >= 'a' && value[0] <= 'z') || (value[0] >= '0' && value[0] <= '9')
}

func validPasswordPolicy(password string) bool {
	return len(password) >= 8 && len(password) <= 128
}

// CreateUser stores only a bcrypt hash; the plaintext password never leaves
// the caller. createdBy may be empty for the bootstrap admin.
func (s *Store) CreateUser(ctx context.Context, username, displayName, role, password, createdBy string, now time.Time) (AuthUser, error) {
	if !validUsername(username) || !ValidRole(role) || !validPasswordPolicy(password) || now.IsZero() ||
		displayName != strings.TrimSpace(displayName) || len(displayName) > 128 {
		return AuthUser{}, ErrAuthInvalidUser
	}
	hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		return AuthUser{}, errors.New("hash password")
	}
	id, err := newOpaqueID("usr-")
	if err != nil {
		return AuthUser{}, err
	}
	_, err = s.db.Exec(ctx, `
INSERT INTO users(id,username,display_name,role,password_hash,must_change_password,state,created_by,created_at,updated_at)
VALUES($1,$2,$3,$4,$5,$6,'active',$7,$8,$8)`,
		id, username, displayName, role, string(hash), true, createdBy, now.UTC())
	if err != nil {
		if isUniqueViolation(err) {
			return AuthUser{}, ErrAuthDuplicateUsername
		}
		return AuthUser{}, fmt.Errorf("create auth user: %w", err)
	}
	return AuthUser{ID: id, Username: username, DisplayName: displayName, Role: role, State: "active", CreatedAt: now.UTC()}, nil
}

func (s *Store) GetAuthUserByName(ctx context.Context, username string) (AuthUser, string, error) {
	var user AuthUser
	var hash string
	err := s.db.QueryRow(ctx, `
SELECT id,username,display_name,role,password_hash,must_change_password,state,created_at,updated_at,last_login_at
FROM users WHERE username=$1`, username).Scan(
		&user.ID, &user.Username, &user.DisplayName, &user.Role, &hash,
		&user.MustChangePassword, &user.State, &user.CreatedAt, &user.CreatedAt, &user.LastLoginAt)
	if errors.Is(err, pgx.ErrNoRows) {
		return AuthUser{}, "", ErrAuthUserNotFound
	}
	if err != nil {
		return AuthUser{}, "", fmt.Errorf("load auth user: %w", err)
	}
	return user, hash, nil
}

func (s *Store) GetAuthUserByID(ctx context.Context, id string) (AuthUser, error) {
	var user AuthUser
	err := s.db.QueryRow(ctx, `
SELECT id,username,display_name,role,must_change_password,state,created_at,last_login_at
FROM users WHERE id=$1`, id).Scan(
		&user.ID, &user.Username, &user.DisplayName, &user.Role,
		&user.MustChangePassword, &user.State, &user.CreatedAt, &user.LastLoginAt)
	if errors.Is(err, pgx.ErrNoRows) {
		return AuthUser{}, ErrAuthUserNotFound
	}
	if err != nil {
		return AuthUser{}, fmt.Errorf("load auth user by id: %w", err)
	}
	return user, nil
}

func (s *Store) ListAuthUsers(ctx context.Context) ([]AuthUser, error) {
	rows, err := s.db.Query(ctx, `
SELECT id,username,display_name,role,must_change_password,state,created_at,last_login_at
FROM users ORDER BY created_at, username`)
	if err != nil {
		return nil, fmt.Errorf("list auth users: %w", err)
	}
	defer rows.Close()
	users := []AuthUser{}
	for rows.Next() {
		var user AuthUser
		if err := rows.Scan(&user.ID, &user.Username, &user.DisplayName, &user.Role,
			&user.MustChangePassword, &user.State, &user.CreatedAt, &user.LastLoginAt); err != nil {
			return nil, fmt.Errorf("scan auth user: %w", err)
		}
		users = append(users, user)
	}
	return users, rows.Err()
}

// CountAuthUsers backs the day-one bootstrap guard: the initial admin exists
// only while the account table is empty.
func (s *Store) CountAuthUsers(ctx context.Context) (int, error) {
	var count int
	if err := s.db.QueryRow(ctx, `SELECT COUNT(*) FROM users`).Scan(&count); err != nil {
		return 0, fmt.Errorf("count auth users: %w", err)
	}
	return count, nil
}

// CreateInitialAdmin installs the fixed day-one administrator with the
// operator-agreed bootstrap password and must_change_password=TRUE enforced
// at the first login. It refuses to run once any account exists so this code
// path can never mint a second superuser outside the UI.
func (s *Store) CreateInitialAdmin(ctx context.Context, username, bootstrapPassword string, now time.Time) (AuthUser, error) {
	if username == "" || bootstrapPassword == "" || now.IsZero() {
		return AuthUser{}, ErrAuthInvalidUser
	}
	count, err := s.CountAuthUsers(ctx)
	if err != nil {
		return AuthUser{}, err
	}
	if count > 0 {
		return AuthUser{}, ErrAuthDuplicateUsername
	}
	hash, err := bcrypt.GenerateFromPassword([]byte(bootstrapPassword), bcrypt.DefaultCost)
	if err != nil {
		return AuthUser{}, errors.New("hash password")
	}
	id, err := newOpaqueID("usr-")
	if err != nil {
		return AuthUser{}, err
	}
	_, err = s.db.Exec(ctx, `
INSERT INTO users(id,username,display_name,role,password_hash,must_change_password,state,created_by,created_at,updated_at)
VALUES($1,$2,'系统管理员','admin',$3,TRUE,'active','bootstrap',$4,$4)`, id, username, string(hash), now.UTC())
	if err != nil {
		if isUniqueViolation(err) {
			return AuthUser{}, ErrAuthDuplicateUsername
		}
		return AuthUser{}, fmt.Errorf("create initial admin: %w", err)
	}
	return AuthUser{ID: id, Username: username, DisplayName: "系统管理员", Role: RoleAdmin, MustChangePassword: true, State: "active", CreatedAt: now.UTC()}, nil
}

// TouchSession rotates the expiry of a currently valid session; the auth
// middleware uses this for sliding renewal while keeping the same cookie.
func (s *Store) TouchSession(ctx context.Context, sessionID string, now, expiresAt time.Time) error {
	_, err := s.db.Exec(ctx, `UPDATE user_sessions SET last_seen_at=$2, expires_at=$3 WHERE id=$1 AND revoked_at IS NULL`, sessionID, now.UTC(), expiresAt.UTC())
	if err != nil {
		return fmt.Errorf("touch auth session: %w", err)
	}
	return nil
}

func (s *Store) DeleteSession(ctx context.Context, sessionID string, now time.Time) error {
	return s.Logout(ctx, sessionID, now)
}

type LoginOutcome struct {
	SessionToken string
	SessionID    string
	User         AuthUser
}

// Login verifies credentials, enforces the failed-attempt lockout, and mints
// a server-side session. Every failure path returns the same error so the UI
// cannot enumerate usernames.
func (s *Store) Login(ctx context.Context, username, password, userAgent string, now time.Time, ttl time.Duration) (LoginOutcome, error) {
	if ttl <= 0 {
		ttl = 12 * time.Hour
	}
	user, hash, err := s.GetAuthUserByName(ctx, username)
	if err != nil {
		if errors.Is(err, ErrAuthUserNotFound) {
			// Constant-ish work even for unknown users (bcrypt compare against
			// a dummy hash) to slow enumeration.
			_ = bcrypt.CompareHashAndPassword([]byte("$2a$10$CwTycUXWue0Thq9StjUM0uJ8fH9l0d0d0d0d0d0d0d0d0d0d0d0dO"), []byte(password))
			return LoginOutcome{}, ErrAuthInvalidCredential
		}
		return LoginOutcome{}, err
	}
	if user.State != "active" {
		return LoginOutcome{}, ErrAuthUserDisabled
	}
	if lockedUntil, locked, lockErr := s.userLockedUntil(ctx, user.ID); lockErr != nil {
		return LoginOutcome{}, lockErr
	} else if locked && lockedUntil.After(now) {
		return LoginOutcome{}, ErrAuthLocked
	}
	if bcrypt.CompareHashAndPassword([]byte(hash), []byte(password)) != nil {
		if err := s.recordFailedLogin(ctx, user.ID, now.UTC()); err != nil {
			return LoginOutcome{}, err
		}
		return LoginOutcome{}, ErrAuthInvalidCredential
	}
	token, err := newSessionToken()
	if err != nil {
		return LoginOutcome{}, err
	}
	sessionID, err := newOpaqueID("ses-")
	if err != nil {
		return LoginOutcome{}, err
	}
	var mustChange bool
	err = s.db.QueryRow(ctx, `
UPDATE users SET failed_attempts=0, locked_until=NULL, last_login_at=$2, updated_at=$2
WHERE id=$1 AND state='active'
RETURNING must_change_password`, user.ID, now.UTC()).Scan(&mustChange)
	if err != nil {
		return LoginOutcome{}, fmt.Errorf("complete auth login: %w", err)
	}
	var insertedID string
	err = s.db.QueryRow(ctx, `
INSERT INTO user_sessions(id,token_hash,user_id,created_at,last_seen_at,expires_at,user_agent)
VALUES($1,$2,$3,$4,$4,$5,$6) RETURNING id`,
		sessionID, hashSessionToken(token), user.ID, now.UTC(), now.Add(ttl), truncateUserAgent(userAgent)).Scan(&insertedID)
	if err != nil {
		return LoginOutcome{}, fmt.Errorf("create auth session: %w", err)
	}
	user.MustChangePassword = mustChange
	return LoginOutcome{SessionToken: token, SessionID: insertedID, User: user}, nil
}

func (s *Store) userLockedUntil(ctx context.Context, userID string) (time.Time, bool, error) {
	var lockedUntil *time.Time
	err := s.db.QueryRow(ctx, `SELECT locked_until FROM users WHERE id=$1`, userID).Scan(&lockedUntil)
	if errors.Is(err, pgx.ErrNoRows) {
		return time.Time{}, false, ErrAuthUserNotFound
	}
	if err != nil {
		return time.Time{}, false, fmt.Errorf("load auth lock: %w", err)
	}
	return derefTime(lockedUntil), lockedUntil != nil, nil
}

func (s *Store) recordFailedLogin(ctx context.Context, userID string, now time.Time) error {
	// Explicit parameter casts: PostgreSQL cannot infer the type of a bare
	// parameter inside `$4::interval` / the integer comparison, and the
	// extended-protocol Prepare would fail with "could not determine data type".
	_, err := s.db.Exec(ctx, `
UPDATE users SET
    failed_attempts = failed_attempts + 1,
    locked_until = CASE WHEN failed_attempts + 1 >= $3::int THEN $2::timestamptz + $4::interval ELSE locked_until END,
    updated_at = $2::timestamptz
WHERE id = $1::varchar`, userID, now, authMaxFailedAttempts, fmt.Sprintf("%d minutes", int(authLockDuration.Minutes())))
	if err != nil {
		return fmt.Errorf("record failed login: %w", err)
	}
	return nil
}

// LookupSession resolves a session token, sliding-expiring it on use, and
// rejects sessions whose user was disabled or whose chain was revoked.
func (s *Store) LookupSession(ctx context.Context, token string, now time.Time) (AuthSession, error) {
	if len(token) < 16 {
		return AuthSession{}, ErrAuthSessionInvalid
	}
	var session AuthSession
	var revoked bool
	err := s.db.QueryRow(ctx, `
SELECT ses.id, ses.expires_at, ses.revoked_at IS NOT NULL,
       usr.id, usr.username, usr.display_name, usr.role, usr.must_change_password, usr.state, usr.created_at, usr.last_login_at
FROM user_sessions ses
JOIN users usr ON usr.id = ses.user_id
WHERE ses.token_hash = $1`, hashSessionToken(token)).Scan(
		&session.ID, &session.ExpiresAt, &revoked,
		&session.User.ID, &session.User.Username, &session.User.DisplayName, &session.User.Role,
		&session.User.MustChangePassword, &session.User.State, &session.User.CreatedAt, &session.User.LastLoginAt)
	if errors.Is(err, pgx.ErrNoRows) {
		return AuthSession{}, ErrAuthSessionInvalid
	}
	if err != nil {
		return AuthSession{}, fmt.Errorf("lookup auth session: %w", err)
	}
	if revoked || session.User.State != "active" || !session.ExpiresAt.After(now) {
		return AuthSession{}, ErrAuthSessionInvalid
	}
	if _, err := s.db.Exec(ctx, `UPDATE user_sessions SET last_seen_at=$2 WHERE id=$1`, session.ID, now.UTC()); err != nil {
		return AuthSession{}, fmt.Errorf("touch auth session: %w", err)
	}
	return session, nil
}

func (s *Store) Logout(ctx context.Context, sessionID string, now time.Time) error {
	_, err := s.db.Exec(ctx, `UPDATE user_sessions SET revoked_at=$2 WHERE id=$1 AND revoked_at IS NULL`, sessionID, now.UTC())
	if err != nil {
		return fmt.Errorf("revoke auth session: %w", err)
	}
	return nil
}

// RevokeAllUserSessions is used on password change, role change and disable
// so a stolen or stale cookie dies immediately, not at natural expiry.
func (s *Store) RevokeAllUserSessions(ctx context.Context, userID string, now time.Time) error {
	_, err := s.db.Exec(ctx, `UPDATE user_sessions SET revoked_at=$2 WHERE user_id=$1 AND revoked_at IS NULL`, userID, now.UTC())
	if err != nil {
		return fmt.Errorf("revoke user sessions: %w", err)
	}
	return nil
}

// ChangePassword verifies the old password, installs the new hash and clears
// the must-change flag; every other session dies. Returns the new session
// token id only after success so callers can rotate the cookie.
func (s *Store) ChangePassword(ctx context.Context, userID, oldPassword, newPassword string, now time.Time) error {
	if !validPasswordPolicy(newPassword) {
		return ErrAuthInvalidUser
	}
	var hash string
	var mustChange bool
	err := s.db.QueryRow(ctx, `SELECT password_hash, must_change_password FROM users WHERE id=$1 AND state='active'`, userID).Scan(&hash, &mustChange)
	if errors.Is(err, pgx.ErrNoRows) {
		return ErrAuthUserNotFound
	}
	if err != nil {
		return fmt.Errorf("load auth user for password change: %w", err)
	}
	if bcrypt.CompareHashAndPassword([]byte(hash), []byte(oldPassword)) != nil {
		return ErrAuthInvalidCredential
	}
	if oldPassword == newPassword {
		return ErrAuthInvalidUser
	}
	fresh, err := bcrypt.GenerateFromPassword([]byte(newPassword), bcrypt.DefaultCost)
	if err != nil {
		return errors.New("hash password")
	}
	if _, err := s.db.Exec(ctx, `
UPDATE users SET password_hash=$2, must_change_password=FALSE, updated_at=$3 WHERE id=$1`, userID, string(fresh), now.UTC()); err != nil {
		return fmt.Errorf("update auth password: %w", err)
	}
	return s.RevokeAllUserSessions(ctx, userID, now)
}

// AdminResetPassword sets a fresh temporary password, forces a change on next
// login and revokes all sessions for the target user. It refuses to lock the
// platform out of its last admin.
func (s *Store) AdminResetPassword(ctx context.Context, actorID, userID, temporaryPassword string, now time.Time) error {
	if userID == actorID {
		return ErrAuthInvalidUser
	}
	if !validPasswordPolicy(temporaryPassword) {
		return ErrAuthInvalidUser
	}
	target, err := s.GetAuthUserByID(ctx, userID)
	if err != nil {
		return err
	}
	if target.State != "active" {
		return ErrAuthUserDisabled
	}
	fresh, err := bcrypt.GenerateFromPassword([]byte(temporaryPassword), bcrypt.DefaultCost)
	if err != nil {
		return errors.New("hash password")
	}
	if _, err := s.db.Exec(ctx, `
UPDATE users SET password_hash=$2, must_change_password=TRUE, failed_attempts=0, locked_until=NULL, updated_at=$3
WHERE id=$1 AND state='active'`, userID, string(fresh), now.UTC()); err != nil {
		return fmt.Errorf("reset auth password: %w", err)
	}
	return s.RevokeAllUserSessions(ctx, userID, now)
}

func (s *Store) SetUserRole(ctx context.Context, actorID, userID, role string, now time.Time) error {
	if !ValidRole(role) {
		return ErrAuthInvalidUser
	}
	target, err := s.GetAuthUserByID(ctx, userID)
	if err != nil {
		return err
	}
	if target.Role == RoleAdmin && role != RoleAdmin && target.State == "active" {
		last, err := s.isLastActiveAdmin(ctx, target.ID)
		if err != nil {
			return err
		}
		if last {
			return ErrAuthLastAdmin
		}
	}
	if _, err := s.db.Exec(ctx, `UPDATE users SET role=$2, updated_at=$3 WHERE id=$1`, userID, role, now.UTC()); err != nil {
		return fmt.Errorf("set auth role: %w", err)
	}
	return s.RevokeAllUserSessions(ctx, userID, now)
}

// SetUserState activates or disables an account. Disabling revokes sessions
// and preserves every business row and audit reference.
func (s *Store) SetUserState(ctx context.Context, actorID, userID, state string, now time.Time) error {
	if state != "active" && state != "disabled" {
		return ErrAuthInvalidUser
	}
	target, err := s.GetAuthUserByID(ctx, userID)
	if err != nil {
		return err
	}
	if state == "disabled" && target.Role == RoleAdmin && target.State == "active" {
		last, err := s.isLastActiveAdmin(ctx, target.ID)
		if err != nil {
			return err
		}
		if last {
			return ErrAuthLastAdmin
		}
	}
	if _, err := s.db.Exec(ctx, `UPDATE users SET state=$2, updated_at=$3 WHERE id=$1`, userID, state, now.UTC()); err != nil {
		return fmt.Errorf("set auth state: %w", err)
	}
	if state == "disabled" {
		return s.RevokeAllUserSessions(ctx, userID, now)
	}
	return nil
}

func (s *Store) isLastActiveAdmin(ctx context.Context, userID string) (bool, error) {
	var count int
	if err := s.db.QueryRow(ctx, `SELECT COUNT(*) FROM users WHERE role='admin' AND state='active' AND id<>$1`, userID).Scan(&count); err != nil {
		return false, fmt.Errorf("count active admins: %w", err)
	}
	return count == 0, nil
}

func newSessionToken() (string, error) {
	raw := make([]byte, 32)
	if _, err := rand.Read(raw); err != nil {
		return "", errors.New("generate session token")
	}
	return base64.RawURLEncoding.EncodeToString(raw), nil
}

func truncateUserAgent(value string) string {
	value = strings.Map(func(r rune) rune {
		if r == '\n' || r == '\r' || r == '\t' {
			return ' '
		}
		if r < 0x20 {
			return -1
		}
		return r
	}, value)
	if len(value) > 256 {
		return value[:256]
	}
	return value
}

func derefTime(value *time.Time) time.Time {
	if value == nil {
		return time.Time{}
	}
	return *value
}