ResolveSpec/pkg/security/totp_provider_database.go

package security

import (
	"context"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"sync"
)

// DatabaseTwoFactorProvider implements TwoFactorAuthProvider using PostgreSQL stored procedures
// Procedure names are configurable via SQLNames (see DefaultSQLNames for defaults)
// See totp_database_schema.sql for procedure definitions
type DatabaseTwoFactorProvider struct {
	db         *sql.DB
	dbMu       sync.RWMutex
	dbFactory  func() (*sql.DB, error)
	totpGen    *TOTPGenerator
	sqlNames   *SQLNames
	tableNames *TableNames
	queryMode  QueryMode
	capability *dbCapability
}

// NewDatabaseTwoFactorProvider creates a new database-backed 2FA provider
func NewDatabaseTwoFactorProvider(db *sql.DB, config *TwoFactorConfig, names ...*SQLNames) *DatabaseTwoFactorProvider {
	if config == nil {
		config = DefaultTwoFactorConfig()
	}
	return &DatabaseTwoFactorProvider{
		db:         db,
		totpGen:    NewTOTPGenerator(config),
		sqlNames:   resolveSQLNames(names...),
		tableNames: DefaultTableNames(),
		capability: newDBCapability(),
	}
}

// WithDBFactory configures a factory used to reopen the database connection if it is closed.
func (p *DatabaseTwoFactorProvider) WithDBFactory(factory func() (*sql.DB, error)) *DatabaseTwoFactorProvider {
	p.dbFactory = factory
	return p
}

// WithTableNames configures Direct-mode table names. If names is nil, defaults are used.
func (p *DatabaseTwoFactorProvider) WithTableNames(names *TableNames) *DatabaseTwoFactorProvider {
	p.tableNames = resolveTableNames(names)
	return p
}

// WithQueryMode selects stored-procedure vs Direct-mode SQL (default ModeAuto).
func (p *DatabaseTwoFactorProvider) WithQueryMode(mode QueryMode) *DatabaseTwoFactorProvider {
	p.queryMode = mode
	return p
}

func (p *DatabaseTwoFactorProvider) getDB() *sql.DB {
	p.dbMu.RLock()
	defer p.dbMu.RUnlock()
	return p.db
}

func (p *DatabaseTwoFactorProvider) reconnectDB() error {
	if p.dbFactory == nil {
		return fmt.Errorf("no db factory configured for reconnect")
	}
	newDB, err := p.dbFactory()
	if err != nil {
		return err
	}
	p.dbMu.Lock()
	p.db = newDB
	p.dbMu.Unlock()
	if p.capability != nil {
		p.capability.reset()
	}
	return nil
}

func (p *DatabaseTwoFactorProvider) runDBOpWithReconnect(run func(*sql.DB) error) error {
	db := p.getDB()
	if db == nil {
		return fmt.Errorf("database connection is nil")
	}
	err := run(db)
	if isDBClosed(err) {
		if reconnErr := p.reconnectDB(); reconnErr == nil {
			err = run(p.getDB())
		}
	}
	return err
}

// Generate2FASecret creates a new secret for a user
func (p *DatabaseTwoFactorProvider) Generate2FASecret(userID int, issuer, accountName string) (*TwoFactorSecret, error) {
	secret, err := p.totpGen.GenerateSecret()
	if err != nil {
		return nil, fmt.Errorf("failed to generate secret: %w", err)
	}

	qrURL := p.totpGen.GenerateQRCodeURL(secret, issuer, accountName)

	backupCodes, err := GenerateBackupCodes(10)
	if err != nil {
		return nil, fmt.Errorf("failed to generate backup codes: %w", err)
	}

	return &TwoFactorSecret{
		Secret:      secret,
		QRCodeURL:   qrURL,
		BackupCodes: backupCodes,
		Issuer:      issuer,
		AccountName: accountName,
	}, nil
}

// Validate2FACode verifies a TOTP code
func (p *DatabaseTwoFactorProvider) Validate2FACode(secret string, code string) (bool, error) {
	return p.totpGen.ValidateCode(secret, code)
}

// Enable2FA activates 2FA for a user
func (p *DatabaseTwoFactorProvider) Enable2FA(userID int, secret string, backupCodes []string) error {
	// Hash backup codes for secure storage
	hashedCodes := make([]string, len(backupCodes))
	for i, code := range backupCodes {
		hash := sha256.Sum256([]byte(code))
		hashedCodes[i] = hex.EncodeToString(hash[:])
	}

	// Convert to JSON array
	codesJSON, err := json.Marshal(hashedCodes)
	if err != nil {
		return fmt.Errorf("failed to marshal backup codes: %w", err)
	}

	ctx := context.Background()
	if !p.capability.ShouldUseProcedure(ctx, p.queryMode, p.getDB(), p.sqlNames.TOTPEnable) {
		return p.enable2FADirect(ctx, userID, secret, hashedCodes)
	}

	// Call stored procedure
	var success bool
	var errorMsg sql.NullString

	query := fmt.Sprintf(`SELECT p_success, p_error FROM %s($1, $2, $3::jsonb)`, p.sqlNames.TOTPEnable)
	err = p.getDB().QueryRow(query, userID, secret, string(codesJSON)).Scan(&success, &errorMsg)
	if err != nil {
		return fmt.Errorf("enable 2FA query failed: %w", err)
	}

	if !success {
		if errorMsg.Valid {
			return fmt.Errorf("%s", errorMsg.String)
		}
		return fmt.Errorf("failed to enable 2FA")
	}

	return nil
}

// Disable2FA deactivates 2FA for a user
func (p *DatabaseTwoFactorProvider) Disable2FA(userID int) error {
	ctx := context.Background()
	if !p.capability.ShouldUseProcedure(ctx, p.queryMode, p.getDB(), p.sqlNames.TOTPDisable) {
		return p.disable2FADirect(ctx, userID)
	}

	var success bool
	var errorMsg sql.NullString

	query := fmt.Sprintf(`SELECT p_success, p_error FROM %s($1)`, p.sqlNames.TOTPDisable)
	err := p.getDB().QueryRow(query, userID).Scan(&success, &errorMsg)
	if err != nil {
		return fmt.Errorf("disable 2FA query failed: %w", err)
	}

	if !success {
		if errorMsg.Valid {
			return fmt.Errorf("%s", errorMsg.String)
		}
		return fmt.Errorf("failed to disable 2FA")
	}

	return nil
}

// Get2FAStatus checks if user has 2FA enabled
func (p *DatabaseTwoFactorProvider) Get2FAStatus(userID int) (bool, error) {
	ctx := context.Background()
	if !p.capability.ShouldUseProcedure(ctx, p.queryMode, p.getDB(), p.sqlNames.TOTPGetStatus) {
		return p.get2FAStatusDirect(ctx, userID)
	}

	var success bool
	var errorMsg sql.NullString
	var enabled bool

	query := fmt.Sprintf(`SELECT p_success, p_error, p_enabled FROM %s($1)`, p.sqlNames.TOTPGetStatus)
	err := p.getDB().QueryRow(query, userID).Scan(&success, &errorMsg, &enabled)
	if err != nil {
		return false, fmt.Errorf("get 2FA status query failed: %w", err)
	}

	if !success {
		if errorMsg.Valid {
			return false, fmt.Errorf("%s", errorMsg.String)
		}
		return false, fmt.Errorf("failed to get 2FA status")
	}

	return enabled, nil
}

// Get2FASecret retrieves the user's 2FA secret
func (p *DatabaseTwoFactorProvider) Get2FASecret(userID int) (string, error) {
	ctx := context.Background()
	if !p.capability.ShouldUseProcedure(ctx, p.queryMode, p.getDB(), p.sqlNames.TOTPGetSecret) {
		return p.get2FASecretDirect(ctx, userID)
	}

	var success bool
	var errorMsg sql.NullString
	var secret sql.NullString

	query := fmt.Sprintf(`SELECT p_success, p_error, p_secret FROM %s($1)`, p.sqlNames.TOTPGetSecret)
	err := p.getDB().QueryRow(query, userID).Scan(&success, &errorMsg, &secret)
	if err != nil {
		return "", fmt.Errorf("get 2FA secret query failed: %w", err)
	}

	if !success {
		if errorMsg.Valid {
			return "", fmt.Errorf("%s", errorMsg.String)
		}
		return "", fmt.Errorf("failed to get 2FA secret")
	}

	if !secret.Valid {
		return "", fmt.Errorf("2FA secret not found")
	}

	return secret.String, nil
}

// GenerateBackupCodes creates backup codes for 2FA
func (p *DatabaseTwoFactorProvider) GenerateBackupCodes(userID int, count int) ([]string, error) {
	codes, err := GenerateBackupCodes(count)
	if err != nil {
		return nil, fmt.Errorf("failed to generate backup codes: %w", err)
	}

	// Hash backup codes for storage
	hashedCodes := make([]string, len(codes))
	for i, code := range codes {
		hash := sha256.Sum256([]byte(code))
		hashedCodes[i] = hex.EncodeToString(hash[:])
	}

	ctx := context.Background()
	if !p.capability.ShouldUseProcedure(ctx, p.queryMode, p.getDB(), p.sqlNames.TOTPRegenerateBackup) {
		if err := p.regenerateBackupCodesDirect(ctx, userID, hashedCodes); err != nil {
			return nil, err
		}
		return codes, nil
	}

	// Convert to JSON array
	codesJSON, err := json.Marshal(hashedCodes)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal backup codes: %w", err)
	}

	// Call stored procedure
	var success bool
	var errorMsg sql.NullString

	query := fmt.Sprintf(`SELECT p_success, p_error FROM %s($1, $2::jsonb)`, p.sqlNames.TOTPRegenerateBackup)
	err = p.getDB().QueryRow(query, userID, string(codesJSON)).Scan(&success, &errorMsg)
	if err != nil {
		return nil, fmt.Errorf("regenerate backup codes query failed: %w", err)
	}

	if !success {
		if errorMsg.Valid {
			return nil, fmt.Errorf("%s", errorMsg.String)
		}
		return nil, fmt.Errorf("failed to regenerate backup codes")
	}

	// Return unhashed codes to user (only time they see them)
	return codes, nil
}

// ValidateBackupCode checks and consumes a backup code
func (p *DatabaseTwoFactorProvider) ValidateBackupCode(userID int, code string) (bool, error) {
	// Hash the code
	hash := sha256.Sum256([]byte(code))
	codeHash := hex.EncodeToString(hash[:])

	ctx := context.Background()
	if !p.capability.ShouldUseProcedure(ctx, p.queryMode, p.getDB(), p.sqlNames.TOTPValidateBackupCode) {
		return p.validateBackupCodeDirect(ctx, userID, codeHash)
	}

	var success bool
	var errorMsg sql.NullString
	var valid bool

	query := fmt.Sprintf(`SELECT p_success, p_error, p_valid FROM %s($1, $2)`, p.sqlNames.TOTPValidateBackupCode)
	err := p.getDB().QueryRow(query, userID, codeHash).Scan(&success, &errorMsg, &valid)
	if err != nil {
		return false, fmt.Errorf("validate backup code query failed: %w", err)
	}

	if !success {
		if errorMsg.Valid {
			return false, fmt.Errorf("%s", errorMsg.String)
		}
		return false, nil
	}

	return valid, nil
}