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) }