llmgw/internal/store/users.go

package store

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/jackc/pgx/v5/pgconn"
)

// UserState is where an account is in its life.
type UserState string

// The account states. A new account is pending until an operator approves it,
// which is what keeps a public signup form from being a way in.
const (
	UserPending   UserState = "pending"
	UserActive    UserState = "active"
	UserSuspended UserState = "suspended"
)

// ErrEmailTaken says an address already has an account.
//
// It is returned to callers that need to know — the CLI — and deliberately not
// surfaced to a signup form, where it would confirm which addresses have
// accounts here.
var ErrEmailTaken = errors.New("that address already has an account")

// User is a person.
type User struct {
	ID           string
	Email        string
	Name         string
	PasswordHash string
	State        UserState
	// ProjectID is assigned at approval and decides which limits the person's
	// keys inherit. Empty while pending.
	ProjectID   string
	CreatedAt   time.Time
	ApprovedAt  time.Time
	ApprovedBy  string
	LastLoginAt time.Time
	// IsAdmin lets this account administer: every page under /admin, and
	// every change there.
	IsAdmin bool
	// Role is a lesser part in running the gateway: RoleViewer looks at
	// every administration page, RoleOperator also runs the cluster from them
	// (models, the mirror, nodes), RoleBilling sees usage and quotas. Empty
	// for everybody else.
	Role string
}

// Staff roles.
const (
	RoleViewer   = "viewer"
	RoleOperator = "operator"
	RoleBilling  = "billing"
	RoleAdmin    = "admin"
)

// StaffRoles are the roles SetStaffRole takes; "" takes them all away.
var StaffRoles = []string{RoleViewer, RoleOperator, RoleBilling, RoleAdmin}

// Staff reports whether a person has any part in running the gateway.
func (u User) Staff() bool { return u.IsAdmin || u.Role != "" }

// StaffRole is the role a person holds, admin included.
func (u User) StaffRole() string {
	if u.IsAdmin {
		return RoleAdmin
	}
	return u.Role
}

// CanSignIn reports whether this account may open a session.
func (u User) CanSignIn() bool { return u.State == UserActive }

// NormalizeEmail is applied at every boundary, so two accounts cannot differ
// only in case — which would be two accounts to the database and one to every
// human who looks at them.
func NormalizeEmail(email string) string {
	return strings.ToLower(strings.TrimSpace(email))
}

const userColumns = `id, email, name, password_hash, state, COALESCE(project_id, ''),
	created_at, approved_at, approved_by, last_login_at, is_admin, staff_role`

// CreateUser records a new, unapproved account.
func (s *DB) CreateUser(ctx context.Context, u User) error {
	u.Email = NormalizeEmail(u.Email)
	_, err := s.db.ExecContext(ctx,
		`INSERT INTO users (id, email, name, password_hash, state, project_id, created_at)
		 VALUES (?, ?, ?, ?, ?, ?, ?)`,
		u.ID, u.Email, u.Name, u.PasswordHash, string(u.State),
		nullString(u.ProjectID), millis(u.CreatedAt))
	if isUniqueViolation(err) {
		return ErrEmailTaken
	}
	return err
}

// isUniqueViolation reports whether err is a row refused for a value another
// row already has. Each database says it its own way: SQLite in words
// ("UNIQUE constraint failed"), PostgreSQL by code, 23505, in a message with
// "unique" in lower case — which a check for the upper-case word missed, so a
// second sign-up with an address on PostgreSQL was "something went wrong".
func isUniqueViolation(err error) bool {
	if err == nil {
		return false
	}
	var pe *pgconn.PgError
	if errors.As(err, &pe) {
		return pe.Code == "23505"
	}
	m := err.Error()
	return strings.Contains(m, "UNIQUE constraint failed") || strings.Contains(m, "SQLSTATE 23505")
}

// UserByEmail finds an account by its address.
func (s *DB) UserByEmail(ctx context.Context, email string) (User, error) {
	return s.scanUser(s.db.QueryRowContext(ctx,
		`SELECT `+userColumns+` FROM users WHERE email = ?`, NormalizeEmail(email)))
}

// UserByID finds an account.
func (s *DB) UserByID(ctx context.Context, id string) (User, error) {
	return s.scanUser(s.db.QueryRowContext(ctx,
		`SELECT `+userColumns+` FROM users WHERE id = ?`, id))
}

// Users lists accounts, optionally filtered by state.
func (s *DB) Users(ctx context.Context, state UserState) ([]User, error) {
	query := `SELECT ` + userColumns + ` FROM users`
	var args []any
	if state != "" {
		query += ` WHERE state = ?`
		args = append(args, string(state))
	}
	query += ` ORDER BY created_at`
	return queryAll(ctx, s.db, query, args, func(rows *sql.Rows) (User, error) {
		return s.scanUser(rows)
	})
}

// ApproveUser activates an account, and puts it in a project when one is given.
//
// A project is not required: a person need not belong to anything to use the
// gateway on their own, bounded by their personal limits. Memberships can be
// added at approval or at any time after.
func (s *DB) ApproveUser(ctx context.Context, id, projectID, by string, at time.Time) error {
	return s.inTx(ctx, func(c conn) error {
		res, err := c.ExecContext(ctx,
			`UPDATE users SET state = ?, approved_at = ?, approved_by = ? WHERE id = ? AND state = ?`,
			string(UserActive), millis(at), by, id, string(UserPending))
		if err != nil {
			return err
		}
		if n, _ := res.RowsAffected(); n == 0 {
			return ErrNotFound
		}
		if projectID == "" {
			return nil
		}
		_, err = c.ExecContext(ctx, `
			INSERT INTO memberships (user_id, scope_type, scope_id, role, added_at, added_by)
			VALUES (?, 'project', ?, 'member', ?, ?)
			ON CONFLICT (user_id, scope_type, scope_id) DO NOTHING`,
			id, projectID, millis(at), nullString(by))
		return err
	})
}

// SetUserState suspends or reinstates an account.
//
// Only an account that was once approved can be reinstated. Otherwise a
// signup suspended while pending and then reinstated would become active
// without anybody approving it.
func (s *DB) SetUserState(ctx context.Context, id string, state UserState) error {
	if state == UserActive {
		var approved sql.NullInt64
		if err := s.db.QueryRowContext(ctx, `SELECT approved_at FROM users WHERE id = ?`, id).
			Scan(&approved); err == nil && !approved.Valid {
			return errors.New("this account was never approved; approve it instead")
		}
	}
	res, err := s.db.ExecContext(ctx, `UPDATE users SET state = ? WHERE id = ?`, string(state), id)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNotFound
	}
	return nil
}

// SetAdmin grants or withdraws access to the cluster dashboard.
//
// Withdrawing does not close their sessions: the page is read-only, so the
// worst a stale session buys is a few more seconds of looking. Suspending the
// account is what closes sessions, and that is the lever for someone who should
// not be here at all.
func (s *DB) SetAdmin(ctx context.Context, id string, admin bool) error {
	// An integer, not a bool: the column is an integer on both dialects,
	// and PostgreSQL refuses a bool for it where SQLite took one — so making
	// an administrator worked in every test and failed on a fresh install.
	flag := 0
	if admin {
		flag = 1
	}
	res, err := s.db.ExecContext(ctx, `UPDATE users SET is_admin = ? WHERE id = ?`, flag, id)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNotFound
	}
	return nil
}

// SetStaffRole gives a person a part in running the gateway — viewer,
// operator, billing or admin — or, with "", none.
func (s *DB) SetStaffRole(ctx context.Context, id, role string) error {
	admin, rest := 0, role
	switch role {
	case RoleAdmin:
		admin, rest = 1, ""
	case "", RoleViewer, RoleOperator, RoleBilling:
	default:
		return fmt.Errorf("role %q is not one of viewer, operator, billing, admin", role)
	}
	res, err := s.db.ExecContext(ctx, `UPDATE users SET is_admin = ?, staff_role = ? WHERE id = ?`, admin, rest, id)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNotFound
	}
	return nil
}

// TouchLogin records a successful sign-in.
func (s *DB) TouchLogin(ctx context.Context, id string, at time.Time) error {
	_, err := s.db.ExecContext(ctx, `UPDATE users SET last_login_at = ? WHERE id = ?`,
		millis(at), id)
	return err
}

// KeysForUser lists the keys a person made for themselves.
func (s *DB) KeysForUser(ctx context.Context, userID string) ([]Key, error) {
	return queryAll(ctx, s.db,
		`SELECT `+keyColumns+` FROM api_keys WHERE user_id = ? ORDER BY created_at DESC`,
		[]any{userID}, func(rows *sql.Rows) (Key, error) { return scanKey(rows) })
}

// CreateUserKey issues a key owned by a person.
//
// With a project, the key is used there and the person must be able to use
// that project; without one it is personal, bounded by the person's own limits.
func (s *DB) CreateUserKey(ctx context.Context, userID string, k Key) error {
	if k.ProjectID != "" {
		ok, err := s.CanUseProject(ctx, userID, k.ProjectID)
		if err != nil {
			return err
		}
		if !ok {
			return ErrNotMember
		}
	}
	k.UserID = userID
	return s.CreateKey(ctx, k)
}

// RevokeUserKey revokes a key only if it belongs to this person.
//
// The ownership check is in the statement rather than in a prior read: two
// queries would leave a window where a key changes hands between them, and more
// importantly it would be one forgotten check away from letting anyone revoke
// anyone's key by id.
func (s *DB) RevokeUserKey(ctx context.Context, userID, keyID string, at time.Time) error {
	res, err := s.db.ExecContext(ctx,
		`UPDATE api_keys SET revoked_at = ?
		 WHERE id = ? AND user_id = ? AND revoked_at IS NULL`,
		millis(at), keyID, userID)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNotFound
	}
	return nil
}

func (s *DB) scanUser(sc scanner) (User, error) {
	var u User
	var state string
	var created int64
	var approved, lastLogin sql.NullInt64
	var approvedBy sql.NullString

	if err := sc.Scan(&u.ID, &u.Email, &u.Name, &u.PasswordHash, &state, &u.ProjectID,
		&created, &approved, &approvedBy, &lastLogin, &u.IsAdmin, &u.Role); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return User{}, ErrNotFound
		}
		return User{}, err
	}
	u.State = UserState(state)
	u.CreatedAt = fromMillis(created)
	u.ApprovedAt = fromNullMillis(approved)
	u.ApprovedBy = approvedBy.String
	u.LastLoginAt = fromNullMillis(lastLogin)
	return u, nil
}

// SetPassword replaces an account's password hash.
//
// It does not close sessions; that is the caller's decision, because the two
// reasons to change a password want different things. Someone rotating theirs
// wants their other devices signed out. An operator resetting a forgotten one
// wants the same, but also has nobody to keep signed in.
func (s *DB) SetPassword(ctx context.Context, userID, hash string) error {
	res, err := s.db.ExecContext(ctx, `UPDATE users SET password_hash = ? WHERE id = ?`,
		hash, userID)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNotFound
	}
	return nil
}

// AllProjects lists every project with the path that identifies it.
//
// For the approval form, which has to offer somewhere to put a new account's
// keys. Walking the tree in the page would mean three round trips per render.
func (s *DB) AllProjects(ctx context.Context) ([]ProjectPath, error) {
	return queryAll(ctx, s.db, `
		SELECT p.id, o.name || ' / ' || t.name || ' / ' || p.name
		FROM projects p
		JOIN teams t ON t.id = p.team_id
		JOIN organizations o ON o.id = t.org_id
		ORDER BY 2`, nil, func(rows *sql.Rows) (ProjectPath, error) {
		var p ProjectPath
		return p, rows.Scan(&p.ID, &p.Path)
	})
}

// ProjectPath is a project and where it sits in the tenancy tree.
type ProjectPath struct {
	ID, Path string
}