hydra-webauthn/internal/storage/sqlite/user_repository.go

683 lines
15 KiB
Go

package sqlite
import (
"context"
"database/sql"
"encoding/json"
"fmt"
"strings"
"time"
"forge.cadoles.com/wpetit/hydra-webauthn/internal/storage"
"github.com/go-webauthn/webauthn/webauthn"
"github.com/pkg/errors"
"gitlab.com/wpetit/goweb/logger"
_ "embed"
_ "modernc.org/sqlite"
)
type UserRepository struct {
db *sql.DB
sqliteBusyRetryMaxAttempts int
}
// Create implements storage.UserRepository.
func (r *UserRepository) CreateUser(ctx context.Context, username string, attributes map[string]any) (*storage.User, error) {
user := storage.NewUser(username, attributes)
err := r.withTxRetry(ctx, func(tx *sql.Tx) error {
query := `SELECT COUNT(id) FROM users WHERE username = $1`
args := []any{
user.Username,
}
row := tx.QueryRowContext(ctx, query, args...)
var count int64
if err := row.Scan(&count); err != nil {
return errors.WithStack(err)
}
if err := row.Err(); err != nil {
return errors.WithStack(err)
}
if count > 0 {
return errors.WithStack(storage.ErrAlreadyExist)
}
query = `
INSERT INTO users (id, username, attributes, created_at, updated_at)
VALUES ($1, $2, $3, $4, $4)
`
rawAttributes, err := json.Marshal(user.Attributes)
if err != nil {
return errors.WithStack(err)
}
args = []any{
user.ID,
user.Username,
rawAttributes,
user.CreatedAt,
user.UpdatedAt,
}
if _, err := tx.ExecContext(ctx, query, args...); err != nil {
return errors.WithStack(err)
}
return nil
})
if err != nil {
return nil, errors.WithStack(err)
}
return user, nil
}
// DeleteUserByID implements storage.UserRepository.
func (*UserRepository) DeleteUserByID(ctx context.Context, username string) error {
panic("unimplemented")
}
// FindUserByUsername implements storage.UserRepository.
func (r *UserRepository) FindUserByUsername(ctx context.Context, username string) (*storage.User, error) {
user, err := r.findUserBy(ctx, "username", username)
if err != nil {
return nil, errors.WithStack(err)
}
return user, nil
}
// FindUserByID implements storage.UserRepository.
func (r *UserRepository) FindUserByID(ctx context.Context, userID string) (*storage.User, error) {
user, err := r.findUserBy(ctx, "id", userID)
if err != nil {
return nil, errors.WithStack(err)
}
return user, nil
}
func (r *UserRepository) findUserBy(ctx context.Context, column string, value any) (*storage.User, error) {
user := &storage.User{}
err := r.withTxRetry(ctx, func(tx *sql.Tx) error {
query := fmt.Sprintf(`SELECT id, username, attributes, created_at, updated_at FROM users WHERE %s = $1`, column)
args := []any{
value,
}
var rawAttributes []byte
row := tx.QueryRowContext(ctx, query, args...)
if err := row.Scan(&user.ID, &user.Username, &rawAttributes, &user.CreatedAt, &user.UpdatedAt); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return errors.WithStack(storage.ErrNotFound)
}
return errors.WithStack(err)
}
if err := row.Err(); err != nil {
return errors.WithStack(err)
}
if err := json.Unmarshal(rawAttributes, &user.Attributes); err != nil {
return errors.WithStack(err)
}
if user.Attributes == nil {
user.Attributes = make(map[string]any)
}
query = `SELECT credential FROM user_credentials WHERE user_id = $1`
args = []any{
user.ID,
}
rows, err := tx.QueryContext(ctx, query, args...)
if err != nil {
return errors.WithStack(err)
}
defer func() {
if err := rows.Close(); err != nil {
logger.Error(ctx, "could not close rows", logger.E(errors.WithStack(err)))
}
}()
user.Credentials = make([]webauthn.Credential, 0)
for rows.Next() {
var (
rawCredential []byte
credential webauthn.Credential
)
if err := rows.Scan(&rawCredential); err != nil {
return errors.WithStack(err)
}
if err := json.Unmarshal(rawCredential, &credential); err != nil {
return errors.WithStack(err)
}
user.Credentials = append(user.Credentials, credential)
}
if err := rows.Err(); err != nil {
return errors.WithStack(err)
}
return nil
})
if err != nil {
return nil, errors.WithStack(err)
}
return user, nil
}
// List implements storage.UserRepository.
func (r *UserRepository) ListUsers(ctx context.Context) ([]storage.UserHeader, error) {
users := make([]storage.UserHeader, 0)
err := r.withTxRetry(ctx, func(tx *sql.Tx) error {
query := `SELECT id, username, created_at, updated_at FROM users`
rows, err := tx.QueryContext(ctx, query)
if err != nil {
return errors.WithStack(err)
}
defer func() {
if err := rows.Close(); err != nil {
logger.Error(ctx, "could not close rows", logger.E(errors.WithStack(err)))
}
}()
for rows.Next() {
user := storage.UserHeader{}
if err := rows.Scan(&user.ID, &user.Username, &user.CreatedAt, &user.UpdatedAt); err != nil {
return errors.WithStack(err)
}
users = append(users, user)
}
if err := rows.Err(); err != nil {
return errors.WithStack(err)
}
return nil
})
if err != nil {
return nil, errors.WithStack(err)
}
return users, nil
}
// UpdateUsername implements storage.UserRepository.
func (r *UserRepository) UpdateUserUsername(ctx context.Context, userID string, username string) (*storage.User, error) {
var user *storage.User
err := r.withTxRetry(ctx, func(tx *sql.Tx) error {
query := `SELECT COUNT(id) FROM users WHERE id = $1`
args := []any{
userID,
}
row := tx.QueryRowContext(ctx, query, args...)
var count int64
if err := row.Scan(&count); err != nil {
return errors.WithStack(err)
}
if err := row.Err(); err != nil {
return errors.WithStack(err)
}
if count == 0 {
return errors.WithStack(storage.ErrNotFound)
}
query = `
UPDATE users SET username = $1, updated_at = $2 WHERE id = $3
RETURNING id, username, attributes, created_at, updated_at
`
args = []any{
username,
time.Now(),
userID,
}
var rawAttributes []byte
user = &storage.User{}
row = tx.QueryRowContext(ctx, query, args...)
if err := row.Scan(&user.ID, &user.Username, &rawAttributes, &user.CreatedAt, &user.UpdatedAt); err != nil {
return errors.WithStack(err)
}
if err := row.Err(); err != nil {
return errors.WithStack(err)
}
if err := json.Unmarshal(rawAttributes, &user.Attributes); err != nil {
return errors.WithStack(err)
}
if user.Attributes == nil {
user.Attributes = make(map[string]any)
}
return nil
})
if err != nil {
return nil, errors.WithStack(err)
}
return user, nil
}
// Update implements storage.UserRepository.
func (r *UserRepository) UpdateUserAttributes(ctx context.Context, userID string, attributes map[string]any) (*storage.User, error) {
var user *storage.User
err := r.withTxRetry(ctx, func(tx *sql.Tx) error {
query := `SELECT COUNT(id) FROM users WHERE id = $1`
args := []any{
userID,
}
row := tx.QueryRowContext(ctx, query, args...)
var count int64
if err := row.Scan(&count); err != nil {
return errors.WithStack(err)
}
if err := row.Err(); err != nil {
return errors.WithStack(err)
}
if count == 0 {
return errors.WithStack(storage.ErrNotFound)
}
query = `
UPDATE users SET attributes = $1, updated_at = $2 WHERE id = $3
RETURNING id, username, attributes, created_at, updated_at
`
rawAttributes, err := json.Marshal(attributes)
if err != nil {
return errors.WithStack(err)
}
args = []any{
rawAttributes,
time.Now(),
userID,
}
user = &storage.User{}
row = tx.QueryRowContext(ctx, query, args...)
if err := row.Scan(&user.ID, &user.Username, &rawAttributes, &user.CreatedAt, &user.UpdatedAt); err != nil {
return errors.WithStack(err)
}
if err := row.Err(); err != nil {
return errors.WithStack(err)
}
if err := json.Unmarshal(rawAttributes, &user.Attributes); err != nil {
return errors.WithStack(err)
}
if user.Attributes == nil {
user.Attributes = make(map[string]any)
}
return nil
})
if err != nil {
return nil, errors.WithStack(err)
}
return user, nil
}
// AddCredential implements storage.UserRepository.
func (r *UserRepository) AddUserCredential(ctx context.Context, userID string, credential *webauthn.Credential) (string, error) {
credentialID := storage.NewID()
err := r.withTxRetry(ctx, func(tx *sql.Tx) error {
exists, err := r.userExists(ctx, tx, userID)
if err != nil {
return errors.WithStack(err)
}
if !exists {
return errors.WithStack(storage.ErrNotFound)
}
query := `
INSERT INTO user_credentials (id, user_id, credential, created_at)
VALUES ($1, $2, $3, $4)
`
rawCredential, err := json.Marshal(credential)
if err != nil {
return errors.WithStack(err)
}
args := []any{
credentialID,
userID,
rawCredential,
time.Now(),
}
if _, err := tx.ExecContext(ctx, query, args...); err != nil {
return errors.WithStack(err)
}
return nil
})
if err != nil {
return "", errors.WithStack(err)
}
return credentialID, nil
}
// RemoveCredential implements storage.UserRepository.
func (r *UserRepository) RemoveUserCredential(ctx context.Context, userID string, credentialID string) error {
err := r.withTxRetry(ctx, func(tx *sql.Tx) error {
exists, err := r.userExists(ctx, tx, userID)
if err != nil {
return errors.WithStack(err)
}
if !exists {
return errors.WithStack(storage.ErrNotFound)
}
query := `
DELETE FROM user_credentials WHERE id = $1 AND user_id = $2
`
args := []any{
credentialID,
userID,
}
if _, err := tx.ExecContext(ctx, query, args...); err != nil {
return errors.WithStack(err)
}
return nil
})
if err != nil {
return errors.WithStack(err)
}
return nil
}
// GenerateRegistrationLink implements storage.UserRepository.
func (r *UserRepository) GenerateRegistrationLink(ctx context.Context, userID string) (*storage.RegistrationLink, error) {
registrationLink := storage.NewRegistrationLink(userID)
err := r.withTxRetry(ctx, func(tx *sql.Tx) error {
exists, err := r.userExists(ctx, tx, userID)
if err != nil {
return errors.WithStack(err)
}
if !exists {
return errors.WithStack(storage.ErrNotFound)
}
query := `DELETE FROM registration_links WHERE user_id = $1`
args := []any{
registrationLink.UserID,
}
if _, err := tx.ExecContext(ctx, query, args...); err != nil {
return errors.WithStack(err)
}
query = `
INSERT INTO registration_links (token, user_id, created_at)
VALUES ($1, $2, $3)
`
args = []any{
registrationLink.Token,
registrationLink.UserID,
registrationLink.CreatedAt,
}
if _, err := tx.ExecContext(ctx, query, args...); err != nil {
return errors.WithStack(err)
}
return nil
})
if err != nil {
return nil, errors.WithStack(err)
}
return registrationLink, nil
}
// ClearRegistrationLink implements storage.UserRepository.
func (*UserRepository) ClearRegistrationLink(ctx context.Context, userID string) error {
panic("unimplemented")
}
// GetRegistrationLink implements storage.UserRepository.
func (r *UserRepository) GetRegistrationLink(ctx context.Context, userID string) (*storage.RegistrationLink, error) {
registrationLink := &storage.RegistrationLink{}
err := r.withTxRetry(ctx, func(tx *sql.Tx) error {
query := `SELECT token, user_id, created_at FROM registration_links WHERE user_id = $1`
args := []any{
userID,
}
row := tx.QueryRowContext(ctx, query, args...)
if err := row.Scan(&registrationLink.Token, &registrationLink.UserID, &registrationLink.CreatedAt); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return errors.WithStack(storage.ErrNotFound)
}
return errors.WithStack(err)
}
if err := row.Err(); err != nil {
return errors.WithStack(err)
}
return nil
})
if err != nil {
return nil, errors.WithStack(err)
}
return registrationLink, nil
}
// GetRegistrationLinkByToken implements storage.UserRepository.
func (r *UserRepository) GetRegistrationLinkByToken(ctx context.Context, token string) (*storage.RegistrationLink, error) {
registrationLink := &storage.RegistrationLink{}
err := r.withTxRetry(ctx, func(tx *sql.Tx) error {
query := `SELECT token, user_id, created_at FROM registration_links WHERE token = $1`
args := []any{
token,
}
row := tx.QueryRowContext(ctx, query, args...)
if err := row.Scan(&registrationLink.Token, &registrationLink.UserID, &registrationLink.CreatedAt); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return errors.WithStack(storage.ErrNotFound)
}
return errors.WithStack(err)
}
if err := row.Err(); err != nil {
return errors.WithStack(err)
}
return nil
})
if err != nil {
return nil, errors.WithStack(err)
}
return registrationLink, nil
}
func (r *UserRepository) userExists(ctx context.Context, tx *sql.Tx, userID string) (bool, error) {
query := `SELECT COUNT(id) FROM users WHERE id = $1`
args := []any{
userID,
}
row := tx.QueryRowContext(ctx, query, args...)
var count int64
if err := row.Scan(&count); err != nil {
return false, errors.WithStack(err)
}
if err := row.Err(); err != nil {
return false, errors.WithStack(err)
}
return count >= 1, nil
}
func (r *UserRepository) withTxRetry(ctx context.Context, fn func(*sql.Tx) error) error {
attempts := 0
max := r.sqliteBusyRetryMaxAttempts
ctx = logger.With(ctx, logger.F("max", max))
var err error
for {
ctx = logger.With(ctx)
if attempts >= max {
logger.Debug(ctx, "transaction retrying failed", logger.F("attempts", attempts))
return errors.Wrapf(err, "transaction failed after %d attempts", max)
}
err = r.withTx(ctx, fn)
if err != nil {
if !strings.Contains(err.Error(), "(5) (SQLITE_BUSY)") {
return errors.WithStack(err)
}
err = errors.WithStack(err)
logger.Warn(ctx, "database is busy", logger.E(err))
wait := time.Duration(8<<(attempts+1)) * time.Millisecond
logger.Debug(
ctx, "database is busy, waiting before retrying transaction",
logger.F("wait", wait.String()),
logger.F("attempts", attempts),
)
timer := time.NewTimer(wait)
select {
case <-timer.C:
attempts++
continue
case <-ctx.Done():
if err := ctx.Err(); err != nil {
return errors.WithStack(err)
}
return nil
}
}
return nil
}
}
func (r *UserRepository) withTx(ctx context.Context, fn func(*sql.Tx) error) error {
tx, err := r.db.BeginTx(ctx, nil)
if err != nil {
return errors.WithStack(err)
}
defer func() {
if err := tx.Rollback(); err != nil {
if errors.Is(err, sql.ErrTxDone) {
return
}
err = errors.WithStack(err)
logger.Error(ctx, "could not rollback transaction", logger.CapturedE(err))
}
}()
if err := fn(tx); err != nil {
return errors.WithStack(err)
}
if err := tx.Commit(); err != nil {
return errors.WithStack(err)
}
return nil
}
func NewUserRepository(dsn string) (*UserRepository, error) {
db, err := sql.Open("sqlite", dsn)
if err != nil {
return nil, errors.WithStack(err)
}
if err := applyUserRepositoryMigration(db); err != nil {
return nil, errors.Wrap(err, "could not migrate schema")
}
return &UserRepository{db, 5}, nil
}
var _ storage.UserRepository = &UserRepository{}
//go:embed user_repository.sql
var userRepositoryMigrationScript string
func applyUserRepositoryMigration(db *sql.DB) error {
if err := db.Ping(); err != nil {
return errors.WithStack(err)
}
if _, err := db.Exec(userRepositoryMigrationScript); err != nil {
return errors.WithStack(err)
}
return nil
}