yovision/Bell/server/app/auth/store.go

package auth

import (
	"context"
	"fmt"
	"time"

	"git.ilapage.cn/ila/yovision/Bell/server/app/rbac"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

type Store struct{ DB *pgxpool.Pool }

func (s Store) UserByUsername(ctx context.Context, username string) (User, error) {
	var user User
	err := s.DB.QueryRow(ctx, `SELECT id::text,username,display_name,password_hash,enabled FROM bell_users WHERE username=$1`, username).Scan(&user.ID, &user.Username, &user.DisplayName, &user.PasswordHash, &user.Enabled)
	if err != nil {
		return User{}, err
	}
	if err = s.loadAccess(ctx, &user); err != nil {
		return User{}, err
	}
	return user, nil
}

func (s Store) UserBySessionDigest(ctx context.Context, digest []byte) (User, error) {
	var user User
	err := s.DB.QueryRow(ctx, `SELECT u.id::text,u.username,u.display_name,u.password_hash,u.enabled FROM bell_sessions s JOIN bell_users u ON u.id=s.user_id WHERE s.token_digest=$1 AND s.revoked_at IS NULL AND s.expires_at>now()`, digest).Scan(&user.ID, &user.Username, &user.DisplayName, &user.PasswordHash, &user.Enabled)
	if err != nil {
		return User{}, err
	}
	if err = s.loadAccess(ctx, &user); err != nil {
		return User{}, err
	}
	return user, nil
}

func (s Store) loadAccess(ctx context.Context, user *User) error {
	rows, err := s.DB.Query(ctx, `SELECT DISTINCT r.code,p.permission_code FROM bell_user_roles ur JOIN bell_roles r ON r.id=ur.role_id LEFT JOIN bell_role_permissions p ON p.role_id=r.id WHERE ur.user_id=$1 ORDER BY r.code,p.permission_code`, user.ID)
	if err != nil {
		return err
	}
	defer rows.Close()
	user.Permissions = rbac.Set{}
	roleSeen := map[string]bool{}
	for rows.Next() {
		var role string
		var permission *string
		if err := rows.Scan(&role, &permission); err != nil {
			return err
		}
		if !roleSeen[role] {
			user.Roles = append(user.Roles, role)
			roleSeen[role] = true
		}
		if permission != nil {
			user.Permissions[*permission] = struct{}{}
		}
	}
	return rows.Err()
}

func (s Store) CreateSession(ctx context.Context, userID string, digest []byte, expires time.Time) error {
	_, err := s.DB.Exec(ctx, `INSERT INTO bell_sessions(user_id,token_digest,expires_at) VALUES($1,$2,$3)`, userID, digest, expires)
	return err
}

func (s Store) RevokeSession(ctx context.Context, digest []byte) error {
	_, err := s.DB.Exec(ctx, `UPDATE bell_sessions SET revoked_at=COALESCE(revoked_at,now()) WHERE token_digest=$1`, digest)
	return err
}

func (s Store) BootstrapAdministrator(ctx context.Context, username, displayName, passwordHash string) (string, error) {
	tx, err := s.DB.Begin(ctx)
	if err != nil {
		return "", err
	}
	defer func() { _ = tx.Rollback(ctx) }()
	var id string
	err = tx.QueryRow(ctx, `INSERT INTO bell_users(username,display_name,password_hash) VALUES($1,$2,$3) ON CONFLICT(username) DO NOTHING RETURNING id::text`, username, displayName, passwordHash).Scan(&id)
	if err != nil {
		if err == pgx.ErrNoRows {
			return "", fmt.Errorf("administrator %q already exists", username)
		}
		return "", err
	}
	if _, err = tx.Exec(ctx, `INSERT INTO bell_user_roles(user_id,role_id) SELECT $1,id FROM bell_roles WHERE code='administrator'`, id); err != nil {
		return "", err
	}
	if err = tx.Commit(ctx); err != nil {
		return "", err
	}
	return id, nil
}

func (s Store) ListUsers(ctx context.Context) ([]User, error) {
	rows, err := s.DB.Query(ctx, `SELECT id::text,username,display_name,enabled FROM bell_users ORDER BY username LIMIT 200`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	users := []User{}
	for rows.Next() {
		var u User
		if err := rows.Scan(&u.ID, &u.Username, &u.DisplayName, &u.Enabled); err != nil {
			return nil, err
		}
		users = append(users, u)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	for i := range users {
		if err := s.loadAccess(ctx, &users[i]); err != nil {
			return nil, err
		}
	}
	return users, nil
}

func (s Store) CreateUser(ctx context.Context, username, displayName, passwordHash, role string) (string, error) {
	tx, err := s.DB.Begin(ctx)
	if err != nil {
		return "", err
	}
	defer func() { _ = tx.Rollback(ctx) }()
	var id string
	if err = tx.QueryRow(ctx, `INSERT INTO bell_users(username,display_name,password_hash) VALUES($1,$2,$3) RETURNING id::text`, username, displayName, passwordHash).Scan(&id); err != nil {
		return "", err
	}
	result, err := tx.Exec(ctx, `INSERT INTO bell_user_roles(user_id,role_id) SELECT $1,id FROM bell_roles WHERE code=$2`, id, role)
	if err != nil {
		return "", err
	}
	if result.RowsAffected() != 1 {
		return "", fmt.Errorf("unknown role %q", role)
	}
	if err = tx.Commit(ctx); err != nil {
		return "", err
	}
	return id, nil
}

func (s Store) ReplaceRoles(ctx context.Context, userID string, roles []string) error {
	tx, err := s.DB.Begin(ctx)
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback(ctx) }()
	if _, err = tx.Exec(ctx, `DELETE FROM bell_user_roles WHERE user_id=$1`, userID); err != nil {
		return err
	}
	for _, role := range roles {
		result, err := tx.Exec(ctx, `INSERT INTO bell_user_roles(user_id,role_id) SELECT $1,id FROM bell_roles WHERE code=$2`, userID, role)
		if err != nil {
			return err
		}
		if result.RowsAffected() != 1 {
			return fmt.Errorf("unknown role %q", role)
		}
	}
	return tx.Commit(ctx)
}