archivdms/internal/userstore/userstore.go

// Package userstore is a PostgreSQL-backed user account store, ported from
// archivmail's internal/userstore pattern. LDAP and TOTP are intentionally
// left out of this initial scaffold (can be re-added later following the
// same pattern archivmail uses) — archivdms starts with local password auth
// only.
package userstore

import (
	"context"
	"errors"
	"fmt"
	"time"

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

const (
	RoleUser        = "user"
	RoleDomainAdmin = "domain_admin"
	RoleSuperAdmin  = "superadmin"

	// AuthSourceLocal / AuthSourceLDAP are the values of users.auth_source.
	AuthSourceLocal = "local"
	AuthSourceLDAP  = "ldap"

	bcryptCost = 12

	// dummyBcryptHash burns bcrypt time on "user not found" so that the login
	// endpoint does not leak identifier existence via timing (ported from
	// archivmail's VerifyLogin hardening).
	dummyBcryptHash = "$2a$12$C6UzMDM.H6dfI/f/IKcEeO4TW/OZ/6PdTdSU0/eV1JCJXo.0DGvTa"
)

// User represents a user account.
type User struct {
	ID        int64     `json:"id"`
	Username  string    `json:"username"`
	Email     string    `json:"email"`
	Role      string    `json:"role"`
	Active    bool      `json:"active"`
	CreatedAt time.Time `json:"created_at"`
	TenantID  *int64    `json:"tenant_id,omitempty"`
}

// CreateUserRequest holds parameters for creating a new user.
type CreateUserRequest struct {
	Username string
	Email    string
	Password string
	Role     string
	TenantID *int64
}

// UpdateUserRequest holds optional fields for updating a user.
type UpdateUserRequest struct {
	Email    *string
	Role     *string
	Active   *bool
	Password *string
}

// Store is a PostgreSQL-backed user store.
type Store struct {
	pool *pgxpool.Pool
}

// New connects to PostgreSQL using the given DSN and initialises the schema.
func New(dsn string) (*Store, error) {
	ctx := context.Background()
	pool, err := pgxpool.New(ctx, dsn)
	if err != nil {
		return nil, fmt.Errorf("userstore: connect: %w", err)
	}
	s := &Store{pool: pool}
	if err := s.initSchema(ctx); err != nil {
		pool.Close()
		return nil, fmt.Errorf("userstore: init schema: %w", err)
	}
	return s, nil
}

func (s *Store) initSchema(ctx context.Context) error {
	_, err := s.pool.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS users (
			id            BIGSERIAL PRIMARY KEY,
			username      VARCHAR(100) UNIQUE NOT NULL,
			email         VARCHAR(255) UNIQUE NOT NULL,
			password_hash VARCHAR(255) NOT NULL DEFAULT '',
			role          VARCHAR(20)  NOT NULL DEFAULT 'user',
			active        BOOLEAN      NOT NULL DEFAULT true,
			created_at    TIMESTAMPTZ  NOT NULL DEFAULT NOW(),
			last_login_at TIMESTAMPTZ,
			tenant_id     BIGINT
		);
		CREATE TABLE IF NOT EXISTS token_blacklist (
			jti        VARCHAR(255) PRIMARY KEY,
			expires_at TIMESTAMPTZ  NOT NULL
		);
		CREATE INDEX IF NOT EXISTS idx_users_tenant ON users (tenant_id);
	`)
	return err
}

// Close closes the underlying connection pool.
func (s *Store) Close() error {
	s.pool.Close()
	return nil
}

// Create inserts a new local user with a bcrypt-hashed password.
func (s *Store) Create(req CreateUserRequest) (*User, error) {
	hash, err := bcrypt.GenerateFromPassword([]byte(req.Password), bcryptCost)
	if err != nil {
		return nil, fmt.Errorf("userstore: bcrypt: %w", err)
	}
	ctx := context.Background()
	var id int64
	err = s.pool.QueryRow(ctx,
		`INSERT INTO users (username, email, password_hash, role, active, created_at, tenant_id)
		 VALUES ($1, $2, $3, $4, true, NOW(), $5) RETURNING id`,
		req.Username, req.Email, string(hash), req.Role, req.TenantID,
	).Scan(&id)
	if err != nil {
		return nil, fmt.Errorf("userstore: create: %w", err)
	}
	return s.GetByID(id)
}

// GetByID retrieves a user by their numeric ID.
func (s *Store) GetByID(id int64) (*User, error) {
	ctx := context.Background()
	row := s.pool.QueryRow(ctx,
		`SELECT id, username, email, role, active, created_at, tenant_id FROM users WHERE id = $1`, id)
	return scanUser(row)
}

// GetByUsername retrieves a user by their username.
func (s *Store) GetByUsername(username string) (*User, error) {
	ctx := context.Background()
	row := s.pool.QueryRow(ctx,
		`SELECT id, username, email, role, active, created_at, tenant_id FROM users WHERE username = $1`, username)
	return scanUser(row)
}

// GetByEmail retrieves a user by email address.
func (s *Store) GetByEmail(ctx context.Context, email string) (*User, error) {
	row := s.pool.QueryRow(ctx,
		`SELECT id, username, email, role, active, created_at, tenant_id FROM users WHERE email = $1`, email)
	return scanUser(row)
}

// VerifyLogin checks credentials by email (falling back to username for
// non-tenant users) and returns the user on success.
func (s *Store) VerifyLogin(ctx context.Context, identifier, password string) (*User, error) {
	row := s.pool.QueryRow(ctx,
		`SELECT id, username, email, role, active, created_at, tenant_id, password_hash FROM users WHERE email = $1`,
		identifier,
	)
	u, hash, err := scanUserWithHash(row)
	if err != nil {
		if !errors.Is(err, pgx.ErrNoRows) {
			return nil, fmt.Errorf("userstore: verify login (email): %w", err)
		}
		row = s.pool.QueryRow(ctx,
			`SELECT id, username, email, role, active, created_at, tenant_id, password_hash
			 FROM users WHERE username = $1 AND tenant_id IS NULL`,
			identifier,
		)
		u, hash, err = scanUserWithHash(row)
		if errors.Is(err, pgx.ErrNoRows) {
			_ = bcrypt.CompareHashAndPassword([]byte(dummyBcryptHash), []byte(password))
			return nil, errors.New("userstore: user not found")
		}
		if err != nil {
			return nil, fmt.Errorf("userstore: verify login (username): %w", err)
		}
	}
	if !u.Active {
		return nil, errors.New("userstore: account disabled")
	}
	if err := bcrypt.CompareHashAndPassword([]byte(hash), []byte(password)); err != nil {
		return nil, errors.New("userstore: wrong password")
	}
	return u, nil
}

// ErrUserNotFound is returned by FindForLogin when no matching account exists.
var ErrUserNotFound = errors.New("userstore: user not found")

// LoginRecord is the minimal set of fields the login flow needs: the user, the
// bcrypt hash (empty for LDAP accounts) and the auth_source discriminator.
type LoginRecord struct {
	User       *User
	Hash       string
	AuthSource string
}

// FindForLogin resolves a login identifier (email, or username for a
// tenant-less account) to a LoginRecord, returning ErrUserNotFound when no
// account matches. It performs the SAME lookup precedence as VerifyLogin but,
// crucially, does not itself verify the password — password/LDAP verification
// is orchestrated by the auth manager based on AuthSource.
func (s *Store) FindForLogin(ctx context.Context, identifier string) (*LoginRecord, error) {
	row := s.pool.QueryRow(ctx,
		`SELECT id, username, email, role, active, created_at, tenant_id, password_hash, auth_source
		 FROM users WHERE email = $1`, identifier)
	rec, err := scanLoginRecord(row)
	if err == nil {
		return rec, nil
	}
	if !errors.Is(err, pgx.ErrNoRows) {
		return nil, fmt.Errorf("userstore: find for login (email): %w", err)
	}
	row = s.pool.QueryRow(ctx,
		`SELECT id, username, email, role, active, created_at, tenant_id, password_hash, auth_source
		 FROM users WHERE username = $1 AND tenant_id IS NULL`, identifier)
	rec, err = scanLoginRecord(row)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil, ErrUserNotFound
	}
	if err != nil {
		return nil, fmt.Errorf("userstore: find for login (username): %w", err)
	}
	return rec, nil
}

// BurnPasswordTiming performs a throwaway bcrypt comparison to keep the "user
// not found" / "LDAP disabled" paths timing-indistinguishable from a real
// local password check (see dummyBcryptHash).
func (s *Store) BurnPasswordTiming(password string) {
	_ = bcrypt.CompareHashAndPassword([]byte(dummyBcryptHash), []byte(password))
}

// CompareLocalPassword verifies a plaintext password against a bcrypt hash.
func CompareLocalPassword(hash, password string) error {
	return bcrypt.CompareHashAndPassword([]byte(hash), []byte(password))
}

// LDAPUserRequest holds attributes for JIT-provisioning an LDAP user.
type LDAPUserRequest struct {
	Username string
	Email    string
	Role     string
	LdapUID  string
	TenantID *int64
}

// CreateLDAPUser inserts a new LDAP-backed user (auth_source='ldap',
// empty password_hash). Used by just-in-time provisioning on first LDAP login.
func (s *Store) CreateLDAPUser(ctx context.Context, req LDAPUserRequest) (*User, error) {
	var id int64
	err := s.pool.QueryRow(ctx,
		`INSERT INTO users (username, email, password_hash, role, active, created_at, tenant_id,
			auth_source, ldap_uid, ldap_synced_at)
		 VALUES ($1, $2, '', $3, true, NOW(), $4, 'ldap', $5, NOW()) RETURNING id`,
		req.Username, req.Email, req.Role, req.TenantID, req.LdapUID,
	).Scan(&id)
	if err != nil {
		return nil, fmt.Errorf("userstore: create ldap user: %w", err)
	}
	return s.GetByID(id)
}

// SyncLDAPUser re-synchronises the mutable attributes of an existing LDAP user
// on each login (email, display role — downgrades included). auth_source and
// ldap_uid are never changed here.
func (s *Store) SyncLDAPUser(ctx context.Context, id int64, email, role string) error {
	_, err := s.pool.Exec(ctx,
		`UPDATE users SET email = $1, role = $2, ldap_synced_at = NOW()
		 WHERE id = $3 AND auth_source = 'ldap'`,
		email, role, id)
	if err != nil {
		return fmt.Errorf("userstore: sync ldap user: %w", err)
	}
	return nil
}

// Update applies a partial update to a user record.
func (s *Store) Update(id int64, req UpdateUserRequest) (*User, error) {
	ctx := context.Background()
	if req.Email != nil {
		if _, err := s.pool.Exec(ctx, `UPDATE users SET email = $1 WHERE id = $2`, *req.Email, id); err != nil {
			return nil, fmt.Errorf("userstore: update email: %w", err)
		}
	}
	if req.Role != nil {
		if _, err := s.pool.Exec(ctx, `UPDATE users SET role = $1 WHERE id = $2`, *req.Role, id); err != nil {
			return nil, fmt.Errorf("userstore: update role: %w", err)
		}
	}
	if req.Active != nil {
		if _, err := s.pool.Exec(ctx, `UPDATE users SET active = $1 WHERE id = $2`, *req.Active, id); err != nil {
			return nil, fmt.Errorf("userstore: update active: %w", err)
		}
	}
	if req.Password != nil {
		hash, err := bcrypt.GenerateFromPassword([]byte(*req.Password), bcryptCost)
		if err != nil {
			return nil, fmt.Errorf("userstore: bcrypt: %w", err)
		}
		if _, err := s.pool.Exec(ctx, `UPDATE users SET password_hash = $1 WHERE id = $2`, string(hash), id); err != nil {
			return nil, fmt.Errorf("userstore: update password: %w", err)
		}
	}
	return s.GetByID(id)
}

// Delete removes a user by ID.
func (s *Store) Delete(id int64) error {
	ctx := context.Background()
	tag, err := s.pool.Exec(ctx, `DELETE FROM users WHERE id = $1`, id)
	if err != nil {
		return fmt.Errorf("userstore: delete: %w", err)
	}
	if tag.RowsAffected() == 0 {
		return fmt.Errorf("userstore: user %d not found", id)
	}
	return nil
}

// List returns all users, optionally filtered by role. Pass role="" to list all.
func (s *Store) List(role string) ([]*User, error) {
	ctx := context.Background()
	var rows pgx.Rows
	var err error
	if role == "" {
		rows, err = s.pool.Query(ctx, `SELECT id, username, email, role, active, created_at, tenant_id FROM users ORDER BY id`)
	} else {
		rows, err = s.pool.Query(ctx, `SELECT id, username, email, role, active, created_at, tenant_id FROM users WHERE role = $1 ORDER BY id`, role)
	}
	if err != nil {
		return nil, fmt.Errorf("userstore: list: %w", err)
	}
	defer rows.Close()
	users := make([]*User, 0)
	for rows.Next() {
		u, err := scanUserRow(rows)
		if err != nil {
			return nil, err
		}
		users = append(users, u)
	}
	return users, rows.Err()
}

// ListByTenant returns all users belonging to a specific tenant.
func (s *Store) ListByTenant(ctx context.Context, tenantID int64) ([]*User, error) {
	rows, err := s.pool.Query(ctx,
		`SELECT id, username, email, role, active, created_at, tenant_id FROM users WHERE tenant_id = $1 ORDER BY id`, tenantID)
	if err != nil {
		return nil, fmt.Errorf("userstore: list by tenant: %w", err)
	}
	defer rows.Close()
	users := make([]*User, 0)
	for rows.Next() {
		u, err := scanUserRow(rows)
		if err != nil {
			return nil, err
		}
		users = append(users, u)
	}
	return users, rows.Err()
}

// BlacklistToken adds a JWT ID to the token blacklist.
func (s *Store) BlacklistToken(jti string, expires time.Time) error {
	ctx := context.Background()
	_, err := s.pool.Exec(ctx,
		`INSERT INTO token_blacklist (jti, expires_at) VALUES ($1, $2)
		 ON CONFLICT (jti) DO UPDATE SET expires_at = EXCLUDED.expires_at`,
		jti, expires.UTC(),
	)
	return err
}

// IsBlacklisted returns true if the given JTI is in the blacklist.
func (s *Store) IsBlacklisted(jti string) (bool, error) {
	ctx := context.Background()
	var count int
	err := s.pool.QueryRow(ctx, `SELECT COUNT(*) FROM token_blacklist WHERE jti = $1`, jti).Scan(&count)
	return count > 0, err
}

// UpdateLastLogin sets last_login_at to now for the given user.
func (s *Store) UpdateLastLogin(id int64) error {
	ctx := context.Background()
	_, err := s.pool.Exec(ctx, `UPDATE users SET last_login_at = NOW() WHERE id = $1`, id)
	return err
}

// AdminCount returns the number of active privileged users.
func (s *Store) AdminCount() (int, error) {
	ctx := context.Background()
	var count int
	err := s.pool.QueryRow(ctx,
		`SELECT COUNT(*) FROM users WHERE role IN ('domain_admin','superadmin') AND active = true`).Scan(&count)
	return count, err
}

// --- helpers ---

func scanUser(row pgx.Row) (*User, error) {
	var u User
	err := row.Scan(&u.ID, &u.Username, &u.Email, &u.Role, &u.Active, &u.CreatedAt, &u.TenantID)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil, fmt.Errorf("userstore: not found")
	}
	if err != nil {
		return nil, fmt.Errorf("userstore: scan: %w", err)
	}
	return &u, nil
}

func scanUserWithHash(row pgx.Row) (*User, string, error) {
	var u User
	var hash string
	err := row.Scan(&u.ID, &u.Username, &u.Email, &u.Role, &u.Active, &u.CreatedAt, &u.TenantID, &hash)
	if err != nil {
		return nil, "", err
	}
	return &u, hash, nil
}

func scanLoginRecord(row pgx.Row) (*LoginRecord, error) {
	var u User
	var hash, authSource string
	err := row.Scan(&u.ID, &u.Username, &u.Email, &u.Role, &u.Active, &u.CreatedAt, &u.TenantID, &hash, &authSource)
	if err != nil {
		return nil, err
	}
	return &LoginRecord{User: &u, Hash: hash, AuthSource: authSource}, nil
}

func scanUserRow(rows pgx.Rows) (*User, error) {
	var u User
	if err := rows.Scan(&u.ID, &u.Username, &u.Email, &u.Role, &u.Active, &u.CreatedAt, &u.TenantID); err != nil {
		return nil, fmt.Errorf("userstore: scan row: %w", err)
	}
	return &u, nil
}