mirror of
https://github.com/bitechdev/ResolveSpec.git
synced 2026-02-01 07:24:25 +00:00
feat(security): ✨ Add two-factor authentication support
* Implement TwoFactorAuthenticator for 2FA login. * Create DatabaseTwoFactorProvider for PostgreSQL integration. * Add MemoryTwoFactorProvider for in-memory testing. * Develop TOTPGenerator for generating and validating codes. * Include tests for all new functionalities. * Ensure backup codes are securely hashed and validated.
This commit is contained in:
@@ -6,6 +6,7 @@ Type-safe, composable security system for ResolveSpec with support for authentic
|
|||||||
|
|
||||||
- ✅ **Interface-Based** - Type-safe providers instead of callbacks
|
- ✅ **Interface-Based** - Type-safe providers instead of callbacks
|
||||||
- ✅ **Login/Logout Support** - Built-in authentication lifecycle
|
- ✅ **Login/Logout Support** - Built-in authentication lifecycle
|
||||||
|
- ✅ **Two-Factor Authentication (2FA)** - Optional TOTP support for enhanced security
|
||||||
- ✅ **Composable** - Mix and match different providers
|
- ✅ **Composable** - Mix and match different providers
|
||||||
- ✅ **No Global State** - Each handler has its own security configuration
|
- ✅ **No Global State** - Each handler has its own security configuration
|
||||||
- ✅ **Testable** - Easy to mock and test
|
- ✅ **Testable** - Easy to mock and test
|
||||||
@@ -212,6 +213,23 @@ auth := security.NewJWTAuthenticator("secret-key", db)
|
|||||||
// Note: Requires JWT library installation for token signing/verification
|
// Note: Requires JWT library installation for token signing/verification
|
||||||
```
|
```
|
||||||
|
|
||||||
|
**TwoFactorAuthenticator** - Wraps any authenticator with TOTP 2FA:
|
||||||
|
```go
|
||||||
|
baseAuth := security.NewDatabaseAuthenticator(db)
|
||||||
|
|
||||||
|
// Use in-memory provider (for testing)
|
||||||
|
tfaProvider := security.NewMemoryTwoFactorProvider(nil)
|
||||||
|
|
||||||
|
// Or use database provider (for production)
|
||||||
|
tfaProvider := security.NewDatabaseTwoFactorProvider(db, nil)
|
||||||
|
// Requires: users table with totp fields, user_totp_backup_codes table
|
||||||
|
// Requires: resolvespec_totp_* stored procedures (see totp_database_schema.sql)
|
||||||
|
|
||||||
|
auth := security.NewTwoFactorAuthenticator(baseAuth, tfaProvider, nil)
|
||||||
|
// Supports: TOTP codes, backup codes, QR code generation
|
||||||
|
// Compatible with Google Authenticator, Microsoft Authenticator, Authy, etc.
|
||||||
|
```
|
||||||
|
|
||||||
### Column Security Providers
|
### Column Security Providers
|
||||||
|
|
||||||
**DatabaseColumnSecurityProvider** - Loads rules from database:
|
**DatabaseColumnSecurityProvider** - Loads rules from database:
|
||||||
@@ -334,7 +352,182 @@ func handleRefresh(securityList *security.SecurityList) http.HandlerFunc {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
http.Error(w, err.Error(), http.StatusUnauthorized)
|
http.Error(w, err.Error(), http.StatusUnauthorized)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
## Two-Factor Authentication (2FA)
|
||||||
|
|
||||||
|
### Overview
|
||||||
|
|
||||||
|
- **Optional per-user** - Enable/disable 2FA individually
|
||||||
|
- **TOTP standard** - Compatible with Google Authenticator, Microsoft Authenticator, Authy, 1Password, etc.
|
||||||
|
- **Configurable** - SHA1/SHA256/SHA512, 6/8 digits, custom time periods
|
||||||
|
- **Backup codes** - One-time recovery codes with secure hashing
|
||||||
|
- **Clock skew** - Handles time differences between client/server
|
||||||
|
|
||||||
|
### Setup
|
||||||
|
|
||||||
|
```go
|
||||||
|
// 1. Wrap existing authenticator with 2FA support
|
||||||
|
baseAuth := security.NewDatabaseAuthenticator(db)
|
||||||
|
tfaProvider := security.NewMemoryTwoFactorProvider(nil) // Use custom DB implementation in production
|
||||||
|
tfaAuth := security.NewTwoFactorAuthenticator(baseAuth, tfaProvider, nil)
|
||||||
|
|
||||||
|
// 2. Use as normal authenticator
|
||||||
|
provider := security.NewCompositeSecurityProvider(tfaAuth, colSec, rowSec)
|
||||||
|
securityList := security.NewSecurityList(provider)
|
||||||
|
```
|
||||||
|
|
||||||
|
### Enable 2FA for User
|
||||||
|
|
||||||
|
```go
|
||||||
|
// 1. Initiate 2FA setup
|
||||||
|
secret, err := tfaAuth.Setup2FA(userID, "MyApp", "user@example.com")
|
||||||
|
// Returns: secret.Secret, secret.QRCodeURL, secret.BackupCodes
|
||||||
|
|
||||||
|
// 2. User scans QR code with authenticator app
|
||||||
|
// Display secret.QRCodeURL as QR code image
|
||||||
|
|
||||||
|
// 3. User enters verification code from app
|
||||||
|
code := "123456" // From authenticator app
|
||||||
|
err = tfaAuth.Enable2FA(userID, secret.Secret, code)
|
||||||
|
// 2FA is now enabled for this user
|
||||||
|
|
||||||
|
// 4. Store backup codes securely and show to user once
|
||||||
|
// Display: secret.BackupCodes (10 codes)
|
||||||
|
```
|
||||||
|
|
||||||
|
### Login Flow with 2FA
|
||||||
|
|
||||||
|
```go
|
||||||
|
// 1. User provides credentials
|
||||||
|
req := security.LoginRequest{
|
||||||
|
Username: "user@example.com",
|
||||||
|
Password: "password",
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := tfaAuth.Login(ctx, req)
|
||||||
|
|
||||||
|
// 2. Check if 2FA required
|
||||||
|
if resp.Requires2FA {
|
||||||
|
// Prompt user for 2FA code
|
||||||
|
code := getUserInput() // From authenticator app or backup code
|
||||||
|
|
||||||
|
// 3. Login again with 2FA code
|
||||||
|
req.TwoFactorCode = code
|
||||||
|
resp, err = tfaAuth.Login(ctx, req)
|
||||||
|
|
||||||
|
// 4. Success - token is returned
|
||||||
|
token := resp.Token
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
### Manage 2FA
|
||||||
|
|
||||||
|
```go
|
||||||
|
// Disable 2FA
|
||||||
|
err := tfaAuth.Disable2FA(userID)
|
||||||
|
|
||||||
|
// Regenerate backup codes
|
||||||
|
newCodes, err := tfaAuth.RegenerateBackupCodes(userID, 10)
|
||||||
|
|
||||||
|
// Check status
|
||||||
|
has2FA, err := tfaProvider.Get2FAStatus(userID)
|
||||||
|
```
|
||||||
|
|
||||||
|
### Custom 2FA Storage
|
||||||
|
|
||||||
|
**Option 1: Use DatabaseTwoFactorProvider (Recommended)**
|
||||||
|
|
||||||
|
```go
|
||||||
|
// Uses PostgreSQL stored procedures for all operations
|
||||||
|
db := setupDatabase()
|
||||||
|
|
||||||
|
// Run migrations from totp_database_schema.sql
|
||||||
|
// - Add totp_secret, totp_enabled, totp_enabled_at to users table
|
||||||
|
// - Create user_totp_backup_codes table
|
||||||
|
// - Create resolvespec_totp_* stored procedures
|
||||||
|
|
||||||
|
tfaProvider := security.NewDatabaseTwoFactorProvider(db, nil)
|
||||||
|
tfaAuth := security.NewTwoFactorAuthenticator(baseAuth, tfaProvider, nil)
|
||||||
|
```
|
||||||
|
|
||||||
|
**Option 2: Implement Custom Provider**
|
||||||
|
|
||||||
|
Implement `TwoFactorAuthProvider` for custom storage:
|
||||||
|
|
||||||
|
```go
|
||||||
|
type DBTwoFactorProvider struct {
|
||||||
|
db *gorm.DB
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *DBTwoFactorProvider) Enable2FA(userID int, secret string, backupCodes []string) error {
|
||||||
|
// Store secret and hashed backup codes in database
|
||||||
|
return p.db.Exec("UPDATE users SET totp_secret = ?, backup_codes = ? WHERE id = ?",
|
||||||
|
secret, hashCodes(backupCodes), userID).Error
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *DBTwoFactorProvider) Get2FASecret(userID int) (string, error) {
|
||||||
|
var secret string
|
||||||
|
err := p.db.Raw("SELECT totp_secret FROM users WHERE id = ?", userID).Scan(&secret).Error
|
||||||
|
return secret, err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Implement remaining methods: Generate2FASecret, Validate2FACode, Disable2FA,
|
||||||
|
// Get2FAStatus, GenerateBackupCodes, ValidateBackupCode
|
||||||
|
```
|
||||||
|
|
||||||
|
### Configuration
|
||||||
|
|
||||||
|
```go
|
||||||
|
config := &security.TwoFactorConfig{
|
||||||
|
Algorithm: "SHA256", // SHA1, SHA256, SHA512
|
||||||
|
Digits: 8, // 6 or 8
|
||||||
|
Period: 30, // Seconds per code
|
||||||
|
SkewWindow: 2, // Accept codes ±2 periods
|
||||||
|
}
|
||||||
|
|
||||||
|
totp := security.NewTOTPGenerator(config)
|
||||||
|
tfaAuth := security.NewTwoFactorAuthenticator(baseAuth, tfaProvider, config)
|
||||||
|
```
|
||||||
|
|
||||||
|
### API Response Structure
|
||||||
|
|
||||||
|
```go
|
||||||
|
// LoginResponse with 2FA
|
||||||
|
type LoginResponse struct {
|
||||||
|
Token string `json:"token"`
|
||||||
|
Requires2FA bool `json:"requires_2fa"`
|
||||||
|
TwoFactorSetupData *TwoFactorSecret `json:"two_factor_setup,omitempty"`
|
||||||
|
User *UserContext `json:"user"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// TwoFactorSecret for setup
|
||||||
|
type TwoFactorSecret struct {
|
||||||
|
Secret string `json:"secret"` // Base32 encoded
|
||||||
|
QRCodeURL string `json:"qr_code_url"` // otpauth://totp/...
|
||||||
|
BackupCodes []string `json:"backup_codes"` // 10 recovery codes
|
||||||
|
}
|
||||||
|
|
||||||
|
// UserContext includes 2FA status
|
||||||
|
type UserContext struct {
|
||||||
|
UserID int `json:"user_id"`
|
||||||
|
TwoFactorEnabled bool `json:"two_factor_enabled"`
|
||||||
|
// ... other fields
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
### Security Best Practices
|
||||||
|
|
||||||
|
- **Store secrets encrypted** - Never store TOTP secrets in plain text
|
||||||
|
- **Hash backup codes** - Use SHA-256 before storing
|
||||||
|
- **Rate limit** - Limit 2FA verification attempts
|
||||||
|
- **Require password** - Always verify password before disabling 2FA
|
||||||
|
- **Show backup codes once** - Display only during setup/regeneration
|
||||||
|
- **Log 2FA events** - Track enable/disable/failed attempts
|
||||||
|
- **Mark codes as used** - Backup codes are single-use only
|
||||||
|
|
||||||
|
|
||||||
json.NewEncoder(w).Encode(resp)
|
json.NewEncoder(w).Encode(resp)
|
||||||
} else {
|
} else {
|
||||||
http.Error(w, "Refresh not supported", http.StatusNotImplemented)
|
http.Error(w, "Refresh not supported", http.StatusNotImplemented)
|
||||||
|
|||||||
@@ -15,7 +15,11 @@ CREATE TABLE IF NOT EXISTS users (
|
|||||||
last_login_at TIMESTAMP,
|
last_login_at TIMESTAMP,
|
||||||
-- OAuth2 fields
|
-- OAuth2 fields
|
||||||
remote_id VARCHAR(255), -- Provider's user ID (e.g., Google sub, GitHub id)
|
remote_id VARCHAR(255), -- Provider's user ID (e.g., Google sub, GitHub id)
|
||||||
auth_provider VARCHAR(50) -- 'local', 'google', 'github', 'microsoft', 'facebook', etc.
|
auth_provider VARCHAR(50), -- 'local', 'google', 'github', 'microsoft', 'facebook', etc.
|
||||||
|
-- Two-Factor Authentication fields
|
||||||
|
totp_secret VARCHAR(255), -- Base32 encoded TOTP secret (encrypted recommended)
|
||||||
|
totp_enabled BOOLEAN DEFAULT false,
|
||||||
|
totp_enabled_at TIMESTAMP
|
||||||
);
|
);
|
||||||
|
|
||||||
-- User sessions table for DatabaseAuthenticator and OAuth2Authenticator
|
-- User sessions table for DatabaseAuthenticator and OAuth2Authenticator
|
||||||
@@ -52,6 +56,19 @@ CREATE TABLE IF NOT EXISTS token_blacklist (
|
|||||||
CREATE INDEX IF NOT EXISTS idx_token ON token_blacklist(token);
|
CREATE INDEX IF NOT EXISTS idx_token ON token_blacklist(token);
|
||||||
CREATE INDEX IF NOT EXISTS idx_blacklist_expires_at ON token_blacklist(expires_at);
|
CREATE INDEX IF NOT EXISTS idx_blacklist_expires_at ON token_blacklist(expires_at);
|
||||||
|
|
||||||
|
-- Two-Factor Authentication backup codes table
|
||||||
|
CREATE TABLE IF NOT EXISTS user_totp_backup_codes (
|
||||||
|
id SERIAL PRIMARY KEY,
|
||||||
|
user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||||
|
code_hash VARCHAR(64) NOT NULL, -- SHA-256 hash of backup code
|
||||||
|
used BOOLEAN DEFAULT false,
|
||||||
|
used_at TIMESTAMP,
|
||||||
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||||
|
);
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_totp_user_id ON user_totp_backup_codes(user_id);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_totp_code_hash ON user_totp_backup_codes(code_hash);
|
||||||
|
|
||||||
-- Example: Seed admin user (password should be hashed with bcrypt)
|
-- Example: Seed admin user (password should be hashed with bcrypt)
|
||||||
-- INSERT INTO users (username, email, password, user_level, roles, is_active)
|
-- INSERT INTO users (username, email, password, user_level, roles, is_active)
|
||||||
-- VALUES ('admin', 'admin@example.com', '$2a$10$...', 10, 'admin,user', true);
|
-- VALUES ('admin', 'admin@example.com', '$2a$10$...', 10, 'admin,user', true);
|
||||||
@@ -849,3 +866,212 @@ $$ LANGUAGE plpgsql;
|
|||||||
|
|
||||||
-- Test get user
|
-- Test get user
|
||||||
-- SELECT * FROM resolvespec_oauth_getuser(1);
|
-- SELECT * FROM resolvespec_oauth_getuser(1);
|
||||||
|
|
||||||
|
|
||||||
|
-- ============================================
|
||||||
|
-- Stored Procedures for Two-Factor Authentication
|
||||||
|
-- ============================================
|
||||||
|
|
||||||
|
-- 1. resolvespec_totp_enable - Enable 2FA for a user
|
||||||
|
-- Input: p_user_id (integer), p_secret (text), p_backup_codes (jsonb array)
|
||||||
|
-- Output: p_success (bool), p_error (text)
|
||||||
|
CREATE OR REPLACE FUNCTION resolvespec_totp_enable(
|
||||||
|
p_user_id INTEGER,
|
||||||
|
p_secret TEXT,
|
||||||
|
p_backup_codes jsonb
|
||||||
|
)
|
||||||
|
RETURNS TABLE(p_success boolean, p_error text) AS $$
|
||||||
|
DECLARE
|
||||||
|
v_code TEXT;
|
||||||
|
v_code_hash TEXT;
|
||||||
|
BEGIN
|
||||||
|
-- Update user record with TOTP secret
|
||||||
|
UPDATE users
|
||||||
|
SET totp_secret = p_secret,
|
||||||
|
totp_enabled = true,
|
||||||
|
totp_enabled_at = CURRENT_TIMESTAMP
|
||||||
|
WHERE id = p_user_id;
|
||||||
|
|
||||||
|
IF NOT FOUND THEN
|
||||||
|
RETURN QUERY SELECT false, 'User not found'::text;
|
||||||
|
RETURN;
|
||||||
|
END IF;
|
||||||
|
|
||||||
|
-- Delete old backup codes
|
||||||
|
DELETE FROM user_totp_backup_codes WHERE user_id = p_user_id;
|
||||||
|
|
||||||
|
-- Insert new backup codes
|
||||||
|
FOR i IN 0..jsonb_array_length(p_backup_codes)-1 LOOP
|
||||||
|
v_code_hash := p_backup_codes->>i;
|
||||||
|
INSERT INTO user_totp_backup_codes (user_id, code_hash)
|
||||||
|
VALUES (p_user_id, v_code_hash);
|
||||||
|
END LOOP;
|
||||||
|
|
||||||
|
RETURN QUERY SELECT true, NULL::text;
|
||||||
|
END;
|
||||||
|
$$ LANGUAGE plpgsql;
|
||||||
|
|
||||||
|
-- 2. resolvespec_totp_disable - Disable 2FA for a user
|
||||||
|
-- Input: p_user_id (integer)
|
||||||
|
-- Output: p_success (bool), p_error (text)
|
||||||
|
CREATE OR REPLACE FUNCTION resolvespec_totp_disable(p_user_id INTEGER)
|
||||||
|
RETURNS TABLE(p_success boolean, p_error text) AS $$
|
||||||
|
BEGIN
|
||||||
|
-- Clear TOTP secret and disable 2FA
|
||||||
|
UPDATE users
|
||||||
|
SET totp_secret = NULL,
|
||||||
|
totp_enabled = false
|
||||||
|
WHERE id = p_user_id;
|
||||||
|
|
||||||
|
IF NOT FOUND THEN
|
||||||
|
RETURN QUERY SELECT false, 'User not found'::text;
|
||||||
|
RETURN;
|
||||||
|
END IF;
|
||||||
|
|
||||||
|
-- Delete all backup codes
|
||||||
|
DELETE FROM user_totp_backup_codes WHERE user_id = p_user_id;
|
||||||
|
|
||||||
|
RETURN QUERY SELECT true, NULL::text;
|
||||||
|
END;
|
||||||
|
$$ LANGUAGE plpgsql;
|
||||||
|
|
||||||
|
-- 3. resolvespec_totp_get_status - Check if user has 2FA enabled
|
||||||
|
-- Input: p_user_id (integer)
|
||||||
|
-- Output: p_success (bool), p_error (text), p_enabled (bool)
|
||||||
|
CREATE OR REPLACE FUNCTION resolvespec_totp_get_status(p_user_id INTEGER)
|
||||||
|
RETURNS TABLE(p_success boolean, p_error text, p_enabled boolean) AS $$
|
||||||
|
DECLARE
|
||||||
|
v_enabled BOOLEAN;
|
||||||
|
BEGIN
|
||||||
|
SELECT totp_enabled
|
||||||
|
INTO v_enabled
|
||||||
|
FROM users
|
||||||
|
WHERE id = p_user_id;
|
||||||
|
|
||||||
|
IF NOT FOUND THEN
|
||||||
|
RETURN QUERY SELECT false, 'User not found'::text, false;
|
||||||
|
RETURN;
|
||||||
|
END IF;
|
||||||
|
|
||||||
|
RETURN QUERY SELECT true, NULL::text, COALESCE(v_enabled, false);
|
||||||
|
END;
|
||||||
|
$$ LANGUAGE plpgsql;
|
||||||
|
|
||||||
|
-- 4. resolvespec_totp_get_secret - Get user's TOTP secret
|
||||||
|
-- Input: p_user_id (integer)
|
||||||
|
-- Output: p_success (bool), p_error (text), p_secret (text)
|
||||||
|
CREATE OR REPLACE FUNCTION resolvespec_totp_get_secret(p_user_id INTEGER)
|
||||||
|
RETURNS TABLE(p_success boolean, p_error text, p_secret text) AS $$
|
||||||
|
DECLARE
|
||||||
|
v_secret TEXT;
|
||||||
|
v_enabled BOOLEAN;
|
||||||
|
BEGIN
|
||||||
|
SELECT totp_secret, totp_enabled
|
||||||
|
INTO v_secret, v_enabled
|
||||||
|
FROM users
|
||||||
|
WHERE id = p_user_id;
|
||||||
|
|
||||||
|
IF NOT FOUND THEN
|
||||||
|
RETURN QUERY SELECT false, 'User not found'::text, NULL::text;
|
||||||
|
RETURN;
|
||||||
|
END IF;
|
||||||
|
|
||||||
|
IF NOT COALESCE(v_enabled, false) THEN
|
||||||
|
RETURN QUERY SELECT false, 'TOTP not enabled for user'::text, NULL::text;
|
||||||
|
RETURN;
|
||||||
|
END IF;
|
||||||
|
|
||||||
|
RETURN QUERY SELECT true, NULL::text, v_secret;
|
||||||
|
END;
|
||||||
|
$$ LANGUAGE plpgsql;
|
||||||
|
|
||||||
|
-- 5. resolvespec_totp_regenerate_backup_codes - Generate new backup codes
|
||||||
|
-- Input: p_user_id (integer), p_backup_codes (jsonb array of hashed codes)
|
||||||
|
-- Output: p_success (bool), p_error (text)
|
||||||
|
CREATE OR REPLACE FUNCTION resolvespec_totp_regenerate_backup_codes(
|
||||||
|
p_user_id INTEGER,
|
||||||
|
p_backup_codes jsonb
|
||||||
|
)
|
||||||
|
RETURNS TABLE(p_success boolean, p_error text) AS $$
|
||||||
|
DECLARE
|
||||||
|
v_code_hash TEXT;
|
||||||
|
BEGIN
|
||||||
|
-- Verify user exists and has 2FA enabled
|
||||||
|
IF NOT EXISTS (SELECT 1 FROM users WHERE id = p_user_id AND totp_enabled = true) THEN
|
||||||
|
RETURN QUERY SELECT false, 'User not found or TOTP not enabled'::text;
|
||||||
|
RETURN;
|
||||||
|
END IF;
|
||||||
|
|
||||||
|
-- Delete old backup codes
|
||||||
|
DELETE FROM user_totp_backup_codes WHERE user_id = p_user_id;
|
||||||
|
|
||||||
|
-- Insert new backup codes
|
||||||
|
FOR i IN 0..jsonb_array_length(p_backup_codes)-1 LOOP
|
||||||
|
v_code_hash := p_backup_codes->>i;
|
||||||
|
INSERT INTO user_totp_backup_codes (user_id, code_hash)
|
||||||
|
VALUES (p_user_id, v_code_hash);
|
||||||
|
END LOOP;
|
||||||
|
|
||||||
|
RETURN QUERY SELECT true, NULL::text;
|
||||||
|
END;
|
||||||
|
$$ LANGUAGE plpgsql;
|
||||||
|
|
||||||
|
-- 6. resolvespec_totp_validate_backup_code - Validate and mark backup code as used
|
||||||
|
-- Input: p_user_id (integer), p_code_hash (text)
|
||||||
|
-- Output: p_success (bool), p_error (text), p_valid (bool)
|
||||||
|
CREATE OR REPLACE FUNCTION resolvespec_totp_validate_backup_code(
|
||||||
|
p_user_id INTEGER,
|
||||||
|
p_code_hash TEXT
|
||||||
|
)
|
||||||
|
RETURNS TABLE(p_success boolean, p_error text, p_valid boolean) AS $$
|
||||||
|
DECLARE
|
||||||
|
v_code_id INTEGER;
|
||||||
|
v_used BOOLEAN;
|
||||||
|
BEGIN
|
||||||
|
-- Find the backup code
|
||||||
|
SELECT id, used
|
||||||
|
INTO v_code_id, v_used
|
||||||
|
FROM user_totp_backup_codes
|
||||||
|
WHERE user_id = p_user_id AND code_hash = p_code_hash;
|
||||||
|
|
||||||
|
IF NOT FOUND THEN
|
||||||
|
RETURN QUERY SELECT true, NULL::text, false;
|
||||||
|
RETURN;
|
||||||
|
END IF;
|
||||||
|
|
||||||
|
-- Check if already used
|
||||||
|
IF v_used THEN
|
||||||
|
RETURN QUERY SELECT false, 'Backup code already used'::text, false;
|
||||||
|
RETURN;
|
||||||
|
END IF;
|
||||||
|
|
||||||
|
-- Mark as used
|
||||||
|
UPDATE user_totp_backup_codes
|
||||||
|
SET used = true, used_at = CURRENT_TIMESTAMP
|
||||||
|
WHERE id = v_code_id;
|
||||||
|
|
||||||
|
RETURN QUERY SELECT true, NULL::text, true;
|
||||||
|
END;
|
||||||
|
$$ LANGUAGE plpgsql;
|
||||||
|
|
||||||
|
-- ============================================
|
||||||
|
-- Example: Test TOTP stored procedures
|
||||||
|
-- ============================================
|
||||||
|
|
||||||
|
-- Enable 2FA
|
||||||
|
-- SELECT * FROM resolvespec_totp_enable(1, 'JBSWY3DPEHPK3PXP', '["abc123", "def456"]'::jsonb);
|
||||||
|
|
||||||
|
-- Disable 2FA
|
||||||
|
-- SELECT * FROM resolvespec_totp_disable(1);
|
||||||
|
|
||||||
|
-- Get 2FA status
|
||||||
|
-- SELECT * FROM resolvespec_totp_get_status(1);
|
||||||
|
|
||||||
|
-- Get TOTP secret
|
||||||
|
-- SELECT * FROM resolvespec_totp_get_secret(1);
|
||||||
|
|
||||||
|
-- Regenerate backup codes
|
||||||
|
-- SELECT * FROM resolvespec_totp_regenerate_backup_codes(1, '["new123", "new456"]'::jsonb);
|
||||||
|
|
||||||
|
-- Validate backup code
|
||||||
|
-- SELECT * FROM resolvespec_totp_validate_backup_code(1, 'abc123');
|
||||||
|
|||||||
@@ -17,12 +17,14 @@ type UserContext struct {
|
|||||||
Email string `json:"email"`
|
Email string `json:"email"`
|
||||||
Claims map[string]any `json:"claims"`
|
Claims map[string]any `json:"claims"`
|
||||||
Meta map[string]any `json:"meta"` // Additional metadata that can hold any JSON-serializable values
|
Meta map[string]any `json:"meta"` // Additional metadata that can hold any JSON-serializable values
|
||||||
|
TwoFactorEnabled bool `json:"two_factor_enabled"` // Indicates if 2FA is enabled for this user
|
||||||
}
|
}
|
||||||
|
|
||||||
// LoginRequest contains credentials for login
|
// LoginRequest contains credentials for login
|
||||||
type LoginRequest struct {
|
type LoginRequest struct {
|
||||||
Username string `json:"username"`
|
Username string `json:"username"`
|
||||||
Password string `json:"password"`
|
Password string `json:"password"`
|
||||||
|
TwoFactorCode string `json:"two_factor_code,omitempty"` // TOTP or backup code
|
||||||
Claims map[string]any `json:"claims"` // Additional login data
|
Claims map[string]any `json:"claims"` // Additional login data
|
||||||
Meta map[string]any `json:"meta"` // Additional metadata to be set on user context
|
Meta map[string]any `json:"meta"` // Additional metadata to be set on user context
|
||||||
}
|
}
|
||||||
@@ -44,6 +46,8 @@ type LoginResponse struct {
|
|||||||
RefreshToken string `json:"refresh_token"`
|
RefreshToken string `json:"refresh_token"`
|
||||||
User *UserContext `json:"user"`
|
User *UserContext `json:"user"`
|
||||||
ExpiresIn int64 `json:"expires_in"` // Token expiration in seconds
|
ExpiresIn int64 `json:"expires_in"` // Token expiration in seconds
|
||||||
|
Requires2FA bool `json:"requires_2fa"` // True if 2FA code is required
|
||||||
|
TwoFactorSetupData *TwoFactorSecret `json:"two_factor_setup,omitempty"` // Present when setting up 2FA
|
||||||
Meta map[string]any `json:"meta"` // Additional metadata to be set on user context
|
Meta map[string]any `json:"meta"` // Additional metadata to be set on user context
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
188
pkg/security/totp.go
Normal file
188
pkg/security/totp.go
Normal file
@@ -0,0 +1,188 @@
|
|||||||
|
package security
|
||||||
|
|
||||||
|
import (
|
||||||
|
"crypto/hmac"
|
||||||
|
"crypto/rand"
|
||||||
|
"crypto/sha1"
|
||||||
|
"crypto/sha256"
|
||||||
|
"crypto/sha512"
|
||||||
|
"encoding/base32"
|
||||||
|
"encoding/binary"
|
||||||
|
"fmt"
|
||||||
|
"hash"
|
||||||
|
"math"
|
||||||
|
"net/url"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// TwoFactorAuthProvider defines interface for 2FA operations
|
||||||
|
type TwoFactorAuthProvider interface {
|
||||||
|
// Generate2FASecret creates a new secret for a user
|
||||||
|
Generate2FASecret(userID int, issuer, accountName string) (*TwoFactorSecret, error)
|
||||||
|
|
||||||
|
// Validate2FACode verifies a TOTP code
|
||||||
|
Validate2FACode(secret string, code string) (bool, error)
|
||||||
|
|
||||||
|
// Enable2FA activates 2FA for a user (store secret in your database)
|
||||||
|
Enable2FA(userID int, secret string, backupCodes []string) error
|
||||||
|
|
||||||
|
// Disable2FA deactivates 2FA for a user
|
||||||
|
Disable2FA(userID int) error
|
||||||
|
|
||||||
|
// Get2FAStatus checks if user has 2FA enabled
|
||||||
|
Get2FAStatus(userID int) (bool, error)
|
||||||
|
|
||||||
|
// Get2FASecret retrieves the user's 2FA secret
|
||||||
|
Get2FASecret(userID int) (string, error)
|
||||||
|
|
||||||
|
// GenerateBackupCodes creates backup codes for 2FA
|
||||||
|
GenerateBackupCodes(userID int, count int) ([]string, error)
|
||||||
|
|
||||||
|
// ValidateBackupCode checks and consumes a backup code
|
||||||
|
ValidateBackupCode(userID int, code string) (bool, error)
|
||||||
|
}
|
||||||
|
|
||||||
|
// TwoFactorSecret contains 2FA setup information
|
||||||
|
type TwoFactorSecret struct {
|
||||||
|
Secret string `json:"secret"` // Base32 encoded secret
|
||||||
|
QRCodeURL string `json:"qr_code_url"` // URL for QR code generation
|
||||||
|
BackupCodes []string `json:"backup_codes"` // One-time backup codes
|
||||||
|
Issuer string `json:"issuer"` // Application name
|
||||||
|
AccountName string `json:"account_name"` // User identifier (email/username)
|
||||||
|
}
|
||||||
|
|
||||||
|
// TwoFactorConfig holds TOTP configuration
|
||||||
|
type TwoFactorConfig struct {
|
||||||
|
Algorithm string // SHA1, SHA256, SHA512
|
||||||
|
Digits int // Number of digits in code (6 or 8)
|
||||||
|
Period int // Time step in seconds (default 30)
|
||||||
|
SkewWindow int // Number of time steps to check before/after (default 1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// DefaultTwoFactorConfig returns standard TOTP configuration
|
||||||
|
func DefaultTwoFactorConfig() *TwoFactorConfig {
|
||||||
|
return &TwoFactorConfig{
|
||||||
|
Algorithm: "SHA1",
|
||||||
|
Digits: 6,
|
||||||
|
Period: 30,
|
||||||
|
SkewWindow: 1,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TOTPGenerator handles TOTP code generation and validation
|
||||||
|
type TOTPGenerator struct {
|
||||||
|
config *TwoFactorConfig
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewTOTPGenerator creates a new TOTP generator with config
|
||||||
|
func NewTOTPGenerator(config *TwoFactorConfig) *TOTPGenerator {
|
||||||
|
if config == nil {
|
||||||
|
config = DefaultTwoFactorConfig()
|
||||||
|
}
|
||||||
|
return &TOTPGenerator{
|
||||||
|
config: config,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// GenerateSecret creates a random base32-encoded secret
|
||||||
|
func (t *TOTPGenerator) GenerateSecret() (string, error) {
|
||||||
|
secret := make([]byte, 20)
|
||||||
|
_, err := rand.Read(secret)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("failed to generate random secret: %w", err)
|
||||||
|
}
|
||||||
|
return base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(secret), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GenerateQRCodeURL creates a URL for QR code generation
|
||||||
|
func (t *TOTPGenerator) GenerateQRCodeURL(secret, issuer, accountName string) string {
|
||||||
|
params := url.Values{}
|
||||||
|
params.Set("secret", secret)
|
||||||
|
params.Set("issuer", issuer)
|
||||||
|
params.Set("algorithm", t.config.Algorithm)
|
||||||
|
params.Set("digits", fmt.Sprintf("%d", t.config.Digits))
|
||||||
|
params.Set("period", fmt.Sprintf("%d", t.config.Period))
|
||||||
|
|
||||||
|
label := url.PathEscape(fmt.Sprintf("%s:%s", issuer, accountName))
|
||||||
|
return fmt.Sprintf("otpauth://totp/%s?%s", label, params.Encode())
|
||||||
|
}
|
||||||
|
|
||||||
|
// GenerateCode creates a TOTP code for a given time
|
||||||
|
func (t *TOTPGenerator) GenerateCode(secret string, timestamp time.Time) (string, error) {
|
||||||
|
// Decode secret
|
||||||
|
key, err := base32.StdEncoding.WithPadding(base32.NoPadding).DecodeString(strings.ToUpper(secret))
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("invalid secret: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Calculate counter (time steps since Unix epoch)
|
||||||
|
counter := uint64(timestamp.Unix()) / uint64(t.config.Period)
|
||||||
|
|
||||||
|
// Generate HMAC
|
||||||
|
h := t.getHashFunc()
|
||||||
|
mac := hmac.New(h, key)
|
||||||
|
|
||||||
|
// Convert counter to 8-byte array
|
||||||
|
buf := make([]byte, 8)
|
||||||
|
binary.BigEndian.PutUint64(buf, counter)
|
||||||
|
mac.Write(buf)
|
||||||
|
|
||||||
|
sum := mac.Sum(nil)
|
||||||
|
|
||||||
|
// Dynamic truncation
|
||||||
|
offset := sum[len(sum)-1] & 0x0f
|
||||||
|
truncated := binary.BigEndian.Uint32(sum[offset:]) & 0x7fffffff
|
||||||
|
|
||||||
|
// Generate code with specified digits
|
||||||
|
code := truncated % uint32(math.Pow10(t.config.Digits))
|
||||||
|
|
||||||
|
format := fmt.Sprintf("%%0%dd", t.config.Digits)
|
||||||
|
return fmt.Sprintf(format, code), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ValidateCode checks if a code is valid for the secret
|
||||||
|
func (t *TOTPGenerator) ValidateCode(secret, code string) (bool, error) {
|
||||||
|
now := time.Now()
|
||||||
|
|
||||||
|
// Check current time and skew window
|
||||||
|
for i := -t.config.SkewWindow; i <= t.config.SkewWindow; i++ {
|
||||||
|
timestamp := now.Add(time.Duration(i*t.config.Period) * time.Second)
|
||||||
|
expected, err := t.GenerateCode(secret, timestamp)
|
||||||
|
if err != nil {
|
||||||
|
return false, err
|
||||||
|
}
|
||||||
|
|
||||||
|
if code == expected {
|
||||||
|
return true, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return false, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// getHashFunc returns the hash function based on algorithm
|
||||||
|
func (t *TOTPGenerator) getHashFunc() func() hash.Hash {
|
||||||
|
switch strings.ToUpper(t.config.Algorithm) {
|
||||||
|
case "SHA256":
|
||||||
|
return sha256.New
|
||||||
|
case "SHA512":
|
||||||
|
return sha512.New
|
||||||
|
default:
|
||||||
|
return sha1.New
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// GenerateBackupCodes creates random backup codes
|
||||||
|
func GenerateBackupCodes(count int) ([]string, error) {
|
||||||
|
codes := make([]string, count)
|
||||||
|
for i := 0; i < count; i++ {
|
||||||
|
code := make([]byte, 4)
|
||||||
|
_, err := rand.Read(code)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to generate backup code: %w", err)
|
||||||
|
}
|
||||||
|
codes[i] = fmt.Sprintf("%08X", binary.BigEndian.Uint32(code))
|
||||||
|
}
|
||||||
|
return codes, nil
|
||||||
|
}
|
||||||
399
pkg/security/totp_integration_test.go
Normal file
399
pkg/security/totp_integration_test.go
Normal file
@@ -0,0 +1,399 @@
|
|||||||
|
package security_test
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"net/http"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/bitechdev/ResolveSpec/pkg/security"
|
||||||
|
)
|
||||||
|
|
||||||
|
var ErrInvalidCredentials = errors.New("invalid credentials")
|
||||||
|
|
||||||
|
// MockAuthenticator is a simple authenticator for testing 2FA
|
||||||
|
type MockAuthenticator struct {
|
||||||
|
users map[string]*security.UserContext
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewMockAuthenticator() *MockAuthenticator {
|
||||||
|
return &MockAuthenticator{
|
||||||
|
users: map[string]*security.UserContext{
|
||||||
|
"testuser": {
|
||||||
|
UserID: 1,
|
||||||
|
UserName: "testuser",
|
||||||
|
Email: "test@example.com",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *MockAuthenticator) Login(ctx context.Context, req security.LoginRequest) (*security.LoginResponse, error) {
|
||||||
|
user, exists := m.users[req.Username]
|
||||||
|
if !exists || req.Password != "password" {
|
||||||
|
return nil, ErrInvalidCredentials
|
||||||
|
}
|
||||||
|
|
||||||
|
return &security.LoginResponse{
|
||||||
|
Token: "mock-token",
|
||||||
|
RefreshToken: "mock-refresh-token",
|
||||||
|
User: user,
|
||||||
|
ExpiresIn: 3600,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *MockAuthenticator) Logout(ctx context.Context, req security.LogoutRequest) error {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *MockAuthenticator) Authenticate(r *http.Request) (*security.UserContext, error) {
|
||||||
|
return m.users["testuser"], nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTwoFactorAuthenticator_Setup(t *testing.T) {
|
||||||
|
baseAuth := NewMockAuthenticator()
|
||||||
|
provider := security.NewMemoryTwoFactorProvider(nil)
|
||||||
|
tfaAuth := security.NewTwoFactorAuthenticator(baseAuth, provider, nil)
|
||||||
|
|
||||||
|
// Setup 2FA
|
||||||
|
secret, err := tfaAuth.Setup2FA(1, "TestApp", "test@example.com")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Setup2FA() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if secret.Secret == "" {
|
||||||
|
t.Error("Setup2FA() returned empty secret")
|
||||||
|
}
|
||||||
|
|
||||||
|
if secret.QRCodeURL == "" {
|
||||||
|
t.Error("Setup2FA() returned empty QR code URL")
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(secret.BackupCodes) == 0 {
|
||||||
|
t.Error("Setup2FA() returned no backup codes")
|
||||||
|
}
|
||||||
|
|
||||||
|
if secret.Issuer != "TestApp" {
|
||||||
|
t.Errorf("Setup2FA() Issuer = %s, want TestApp", secret.Issuer)
|
||||||
|
}
|
||||||
|
|
||||||
|
if secret.AccountName != "test@example.com" {
|
||||||
|
t.Errorf("Setup2FA() AccountName = %s, want test@example.com", secret.AccountName)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTwoFactorAuthenticator_Enable2FA(t *testing.T) {
|
||||||
|
baseAuth := NewMockAuthenticator()
|
||||||
|
provider := security.NewMemoryTwoFactorProvider(nil)
|
||||||
|
tfaAuth := security.NewTwoFactorAuthenticator(baseAuth, provider, nil)
|
||||||
|
|
||||||
|
// Setup 2FA
|
||||||
|
secret, err := tfaAuth.Setup2FA(1, "TestApp", "test@example.com")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Setup2FA() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Generate valid code
|
||||||
|
totp := security.NewTOTPGenerator(nil)
|
||||||
|
code, err := totp.GenerateCode(secret.Secret, time.Now())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GenerateCode() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Enable 2FA with valid code
|
||||||
|
err = tfaAuth.Enable2FA(1, secret.Secret, code)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Enable2FA() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify 2FA is enabled
|
||||||
|
status, err := provider.Get2FAStatus(1)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Get2FAStatus() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if !status {
|
||||||
|
t.Error("Enable2FA() did not enable 2FA")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTwoFactorAuthenticator_Enable2FA_InvalidCode(t *testing.T) {
|
||||||
|
baseAuth := NewMockAuthenticator()
|
||||||
|
provider := security.NewMemoryTwoFactorProvider(nil)
|
||||||
|
tfaAuth := security.NewTwoFactorAuthenticator(baseAuth, provider, nil)
|
||||||
|
|
||||||
|
// Setup 2FA
|
||||||
|
secret, err := tfaAuth.Setup2FA(1, "TestApp", "test@example.com")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Setup2FA() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try to enable with invalid code
|
||||||
|
err = tfaAuth.Enable2FA(1, secret.Secret, "000000")
|
||||||
|
if err == nil {
|
||||||
|
t.Error("Enable2FA() should fail with invalid code")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify 2FA is not enabled
|
||||||
|
status, _ := provider.Get2FAStatus(1)
|
||||||
|
if status {
|
||||||
|
t.Error("Enable2FA() should not enable 2FA with invalid code")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTwoFactorAuthenticator_Login_Without2FA(t *testing.T) {
|
||||||
|
baseAuth := NewMockAuthenticator()
|
||||||
|
provider := security.NewMemoryTwoFactorProvider(nil)
|
||||||
|
tfaAuth := security.NewTwoFactorAuthenticator(baseAuth, provider, nil)
|
||||||
|
|
||||||
|
req := security.LoginRequest{
|
||||||
|
Username: "testuser",
|
||||||
|
Password: "password",
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := tfaAuth.Login(context.Background(), req)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Login() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if resp.Requires2FA {
|
||||||
|
t.Error("Login() should not require 2FA when not enabled")
|
||||||
|
}
|
||||||
|
|
||||||
|
if resp.Token == "" {
|
||||||
|
t.Error("Login() should return token when 2FA not required")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTwoFactorAuthenticator_Login_With2FA_NoCode(t *testing.T) {
|
||||||
|
baseAuth := NewMockAuthenticator()
|
||||||
|
provider := security.NewMemoryTwoFactorProvider(nil)
|
||||||
|
tfaAuth := security.NewTwoFactorAuthenticator(baseAuth, provider, nil)
|
||||||
|
|
||||||
|
// Setup and enable 2FA
|
||||||
|
secret, _ := tfaAuth.Setup2FA(1, "TestApp", "test@example.com")
|
||||||
|
totp := security.NewTOTPGenerator(nil)
|
||||||
|
code, _ := totp.GenerateCode(secret.Secret, time.Now())
|
||||||
|
tfaAuth.Enable2FA(1, secret.Secret, code)
|
||||||
|
|
||||||
|
// Try to login without 2FA code
|
||||||
|
req := security.LoginRequest{
|
||||||
|
Username: "testuser",
|
||||||
|
Password: "password",
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := tfaAuth.Login(context.Background(), req)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Login() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if !resp.Requires2FA {
|
||||||
|
t.Error("Login() should require 2FA when enabled")
|
||||||
|
}
|
||||||
|
|
||||||
|
if resp.Token != "" {
|
||||||
|
t.Error("Login() should not return token when 2FA required but not provided")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTwoFactorAuthenticator_Login_With2FA_ValidCode(t *testing.T) {
|
||||||
|
baseAuth := NewMockAuthenticator()
|
||||||
|
provider := security.NewMemoryTwoFactorProvider(nil)
|
||||||
|
tfaAuth := security.NewTwoFactorAuthenticator(baseAuth, provider, nil)
|
||||||
|
|
||||||
|
// Setup and enable 2FA
|
||||||
|
secret, _ := tfaAuth.Setup2FA(1, "TestApp", "test@example.com")
|
||||||
|
totp := security.NewTOTPGenerator(nil)
|
||||||
|
code, _ := totp.GenerateCode(secret.Secret, time.Now())
|
||||||
|
tfaAuth.Enable2FA(1, secret.Secret, code)
|
||||||
|
|
||||||
|
// Generate new valid code for login
|
||||||
|
newCode, _ := totp.GenerateCode(secret.Secret, time.Now())
|
||||||
|
|
||||||
|
// Login with 2FA code
|
||||||
|
req := security.LoginRequest{
|
||||||
|
Username: "testuser",
|
||||||
|
Password: "password",
|
||||||
|
TwoFactorCode: newCode,
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := tfaAuth.Login(context.Background(), req)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Login() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if resp.Requires2FA {
|
||||||
|
t.Error("Login() should not require 2FA when valid code provided")
|
||||||
|
}
|
||||||
|
|
||||||
|
if resp.Token == "" {
|
||||||
|
t.Error("Login() should return token when 2FA validated")
|
||||||
|
}
|
||||||
|
|
||||||
|
if !resp.User.TwoFactorEnabled {
|
||||||
|
t.Error("Login() should set TwoFactorEnabled on user")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTwoFactorAuthenticator_Login_With2FA_InvalidCode(t *testing.T) {
|
||||||
|
baseAuth := NewMockAuthenticator()
|
||||||
|
provider := security.NewMemoryTwoFactorProvider(nil)
|
||||||
|
tfaAuth := security.NewTwoFactorAuthenticator(baseAuth, provider, nil)
|
||||||
|
|
||||||
|
// Setup and enable 2FA
|
||||||
|
secret, _ := tfaAuth.Setup2FA(1, "TestApp", "test@example.com")
|
||||||
|
totp := security.NewTOTPGenerator(nil)
|
||||||
|
code, _ := totp.GenerateCode(secret.Secret, time.Now())
|
||||||
|
tfaAuth.Enable2FA(1, secret.Secret, code)
|
||||||
|
|
||||||
|
// Try to login with invalid code
|
||||||
|
req := security.LoginRequest{
|
||||||
|
Username: "testuser",
|
||||||
|
Password: "password",
|
||||||
|
TwoFactorCode: "000000",
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err := tfaAuth.Login(context.Background(), req)
|
||||||
|
if err == nil {
|
||||||
|
t.Error("Login() should fail with invalid 2FA code")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTwoFactorAuthenticator_Login_WithBackupCode(t *testing.T) {
|
||||||
|
baseAuth := NewMockAuthenticator()
|
||||||
|
provider := security.NewMemoryTwoFactorProvider(nil)
|
||||||
|
tfaAuth := security.NewTwoFactorAuthenticator(baseAuth, provider, nil)
|
||||||
|
|
||||||
|
// Setup and enable 2FA
|
||||||
|
secret, _ := tfaAuth.Setup2FA(1, "TestApp", "test@example.com")
|
||||||
|
totp := security.NewTOTPGenerator(nil)
|
||||||
|
code, _ := totp.GenerateCode(secret.Secret, time.Now())
|
||||||
|
tfaAuth.Enable2FA(1, secret.Secret, code)
|
||||||
|
|
||||||
|
// Get backup codes
|
||||||
|
backupCodes, _ := tfaAuth.RegenerateBackupCodes(1, 10)
|
||||||
|
|
||||||
|
// Login with backup code
|
||||||
|
req := security.LoginRequest{
|
||||||
|
Username: "testuser",
|
||||||
|
Password: "password",
|
||||||
|
TwoFactorCode: backupCodes[0],
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := tfaAuth.Login(context.Background(), req)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Login() with backup code error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if resp.Token == "" {
|
||||||
|
t.Error("Login() should return token when backup code validated")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try to use same backup code again
|
||||||
|
req2 := security.LoginRequest{
|
||||||
|
Username: "testuser",
|
||||||
|
Password: "password",
|
||||||
|
TwoFactorCode: backupCodes[0],
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err = tfaAuth.Login(context.Background(), req2)
|
||||||
|
if err == nil {
|
||||||
|
t.Error("Login() should fail when reusing backup code")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTwoFactorAuthenticator_Disable2FA(t *testing.T) {
|
||||||
|
baseAuth := NewMockAuthenticator()
|
||||||
|
provider := security.NewMemoryTwoFactorProvider(nil)
|
||||||
|
tfaAuth := security.NewTwoFactorAuthenticator(baseAuth, provider, nil)
|
||||||
|
|
||||||
|
// Setup and enable 2FA
|
||||||
|
secret, _ := tfaAuth.Setup2FA(1, "TestApp", "test@example.com")
|
||||||
|
totp := security.NewTOTPGenerator(nil)
|
||||||
|
code, _ := totp.GenerateCode(secret.Secret, time.Now())
|
||||||
|
tfaAuth.Enable2FA(1, secret.Secret, code)
|
||||||
|
|
||||||
|
// Disable 2FA
|
||||||
|
err := tfaAuth.Disable2FA(1)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Disable2FA() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify 2FA is disabled
|
||||||
|
status, _ := provider.Get2FAStatus(1)
|
||||||
|
if status {
|
||||||
|
t.Error("Disable2FA() did not disable 2FA")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Login should not require 2FA
|
||||||
|
req := security.LoginRequest{
|
||||||
|
Username: "testuser",
|
||||||
|
Password: "password",
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := tfaAuth.Login(context.Background(), req)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Login() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if resp.Requires2FA {
|
||||||
|
t.Error("Login() should not require 2FA after disabling")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTwoFactorAuthenticator_RegenerateBackupCodes(t *testing.T) {
|
||||||
|
baseAuth := NewMockAuthenticator()
|
||||||
|
provider := security.NewMemoryTwoFactorProvider(nil)
|
||||||
|
tfaAuth := security.NewTwoFactorAuthenticator(baseAuth, provider, nil)
|
||||||
|
|
||||||
|
// Setup and enable 2FA
|
||||||
|
secret, _ := tfaAuth.Setup2FA(1, "TestApp", "test@example.com")
|
||||||
|
totp := security.NewTOTPGenerator(nil)
|
||||||
|
code, _ := totp.GenerateCode(secret.Secret, time.Now())
|
||||||
|
tfaAuth.Enable2FA(1, secret.Secret, code)
|
||||||
|
|
||||||
|
// Get initial backup codes
|
||||||
|
codes1, err := tfaAuth.RegenerateBackupCodes(1, 10)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("RegenerateBackupCodes() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(codes1) != 10 {
|
||||||
|
t.Errorf("RegenerateBackupCodes() returned %d codes, want 10", len(codes1))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Regenerate backup codes
|
||||||
|
codes2, err := tfaAuth.RegenerateBackupCodes(1, 10)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("RegenerateBackupCodes() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Old codes should not work
|
||||||
|
req := security.LoginRequest{
|
||||||
|
Username: "testuser",
|
||||||
|
Password: "password",
|
||||||
|
TwoFactorCode: codes1[0],
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err = tfaAuth.Login(context.Background(), req)
|
||||||
|
if err == nil {
|
||||||
|
t.Error("Login() should fail with old backup code after regeneration")
|
||||||
|
}
|
||||||
|
|
||||||
|
// New codes should work
|
||||||
|
req2 := security.LoginRequest{
|
||||||
|
Username: "testuser",
|
||||||
|
Password: "password",
|
||||||
|
TwoFactorCode: codes2[0],
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := tfaAuth.Login(context.Background(), req2)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Login() with new backup code error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if resp.Token == "" {
|
||||||
|
t.Error("Login() should return token with new backup code")
|
||||||
|
}
|
||||||
|
}
|
||||||
134
pkg/security/totp_middleware.go
Normal file
134
pkg/security/totp_middleware.go
Normal file
@@ -0,0 +1,134 @@
|
|||||||
|
package security
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
)
|
||||||
|
|
||||||
|
// TwoFactorAuthenticator wraps an Authenticator and adds 2FA support
|
||||||
|
type TwoFactorAuthenticator struct {
|
||||||
|
baseAuth Authenticator
|
||||||
|
totp *TOTPGenerator
|
||||||
|
provider TwoFactorAuthProvider
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewTwoFactorAuthenticator creates a new 2FA-enabled authenticator
|
||||||
|
func NewTwoFactorAuthenticator(baseAuth Authenticator, provider TwoFactorAuthProvider, config *TwoFactorConfig) *TwoFactorAuthenticator {
|
||||||
|
if config == nil {
|
||||||
|
config = DefaultTwoFactorConfig()
|
||||||
|
}
|
||||||
|
return &TwoFactorAuthenticator{
|
||||||
|
baseAuth: baseAuth,
|
||||||
|
totp: NewTOTPGenerator(config),
|
||||||
|
provider: provider,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Login authenticates with 2FA support
|
||||||
|
func (t *TwoFactorAuthenticator) Login(ctx context.Context, req LoginRequest) (*LoginResponse, error) {
|
||||||
|
// First, perform standard authentication
|
||||||
|
resp, err := t.baseAuth.Login(ctx, req)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if user has 2FA enabled
|
||||||
|
if resp.User == nil {
|
||||||
|
return resp, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
has2FA, err := t.provider.Get2FAStatus(resp.User.UserID)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to check 2FA status: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if !has2FA {
|
||||||
|
// User doesn't have 2FA enabled, return normal response
|
||||||
|
return resp, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// User has 2FA enabled
|
||||||
|
if req.TwoFactorCode == "" {
|
||||||
|
// No 2FA code provided, require it
|
||||||
|
resp.Requires2FA = true
|
||||||
|
resp.Token = "" // Don't return token until 2FA is verified
|
||||||
|
resp.RefreshToken = ""
|
||||||
|
return resp, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Validate 2FA code
|
||||||
|
secret, err := t.provider.Get2FASecret(resp.User.UserID)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to get 2FA secret: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try TOTP code first
|
||||||
|
valid, err := t.totp.ValidateCode(secret, req.TwoFactorCode)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to validate 2FA code: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if !valid {
|
||||||
|
// Try backup code
|
||||||
|
valid, err = t.provider.ValidateBackupCode(resp.User.UserID, req.TwoFactorCode)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to validate backup code: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if !valid {
|
||||||
|
return nil, fmt.Errorf("invalid 2FA code")
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2FA verified, return full response with token
|
||||||
|
resp.User.TwoFactorEnabled = true
|
||||||
|
return resp, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Logout delegates to base authenticator
|
||||||
|
func (t *TwoFactorAuthenticator) Logout(ctx context.Context, req LogoutRequest) error {
|
||||||
|
return t.baseAuth.Logout(ctx, req)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Authenticate delegates to base authenticator
|
||||||
|
func (t *TwoFactorAuthenticator) Authenticate(r *http.Request) (*UserContext, error) {
|
||||||
|
return t.baseAuth.Authenticate(r)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Setup2FA initiates 2FA setup for a user
|
||||||
|
func (t *TwoFactorAuthenticator) Setup2FA(userID int, issuer, accountName string) (*TwoFactorSecret, error) {
|
||||||
|
return t.provider.Generate2FASecret(userID, issuer, accountName)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Enable2FA completes 2FA setup after user confirms with a valid code
|
||||||
|
func (t *TwoFactorAuthenticator) Enable2FA(userID int, secret, verificationCode string) error {
|
||||||
|
// Verify the code before enabling
|
||||||
|
valid, err := t.totp.ValidateCode(secret, verificationCode)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to validate code: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if !valid {
|
||||||
|
return fmt.Errorf("invalid verification code")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Generate backup codes
|
||||||
|
backupCodes, err := t.provider.GenerateBackupCodes(userID, 10)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to generate backup codes: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Enable 2FA
|
||||||
|
return t.provider.Enable2FA(userID, secret, backupCodes)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Disable2FA removes 2FA from a user account
|
||||||
|
func (t *TwoFactorAuthenticator) Disable2FA(userID int) error {
|
||||||
|
return t.provider.Disable2FA(userID)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RegenerateBackupCodes creates new backup codes for a user
|
||||||
|
func (t *TwoFactorAuthenticator) RegenerateBackupCodes(userID int, count int) ([]string, error) {
|
||||||
|
return t.provider.GenerateBackupCodes(userID, count)
|
||||||
|
}
|
||||||
229
pkg/security/totp_provider_database.go
Normal file
229
pkg/security/totp_provider_database.go
Normal file
@@ -0,0 +1,229 @@
|
|||||||
|
package security
|
||||||
|
|
||||||
|
import (
|
||||||
|
"crypto/sha256"
|
||||||
|
"database/sql"
|
||||||
|
"encoding/hex"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
)
|
||||||
|
|
||||||
|
// DatabaseTwoFactorProvider implements TwoFactorAuthProvider using PostgreSQL stored procedures
|
||||||
|
// Requires stored procedures: resolvespec_totp_enable, resolvespec_totp_disable,
|
||||||
|
// resolvespec_totp_get_status, resolvespec_totp_get_secret,
|
||||||
|
// resolvespec_totp_regenerate_backup_codes, resolvespec_totp_validate_backup_code
|
||||||
|
// See totp_database_schema.sql for procedure definitions
|
||||||
|
type DatabaseTwoFactorProvider struct {
|
||||||
|
db *sql.DB
|
||||||
|
totpGen *TOTPGenerator
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewDatabaseTwoFactorProvider creates a new database-backed 2FA provider
|
||||||
|
func NewDatabaseTwoFactorProvider(db *sql.DB, config *TwoFactorConfig) *DatabaseTwoFactorProvider {
|
||||||
|
if config == nil {
|
||||||
|
config = DefaultTwoFactorConfig()
|
||||||
|
}
|
||||||
|
return &DatabaseTwoFactorProvider{
|
||||||
|
db: db,
|
||||||
|
totpGen: NewTOTPGenerator(config),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Call stored procedure
|
||||||
|
var success bool
|
||||||
|
var errorMsg sql.NullString
|
||||||
|
|
||||||
|
query := `SELECT p_success, p_error FROM resolvespec_totp_enable($1, $2, $3::jsonb)`
|
||||||
|
err = p.db.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 {
|
||||||
|
var success bool
|
||||||
|
var errorMsg sql.NullString
|
||||||
|
|
||||||
|
query := `SELECT p_success, p_error FROM resolvespec_totp_disable($1)`
|
||||||
|
err := p.db.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) {
|
||||||
|
var success bool
|
||||||
|
var errorMsg sql.NullString
|
||||||
|
var enabled bool
|
||||||
|
|
||||||
|
query := `SELECT p_success, p_error, p_enabled FROM resolvespec_totp_get_status($1)`
|
||||||
|
err := p.db.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) {
|
||||||
|
var success bool
|
||||||
|
var errorMsg sql.NullString
|
||||||
|
var secret sql.NullString
|
||||||
|
|
||||||
|
query := `SELECT p_success, p_error, p_secret FROM resolvespec_totp_get_secret($1)`
|
||||||
|
err := p.db.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[:])
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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 := `SELECT p_success, p_error FROM resolvespec_totp_regenerate_backup_codes($1, $2::jsonb)`
|
||||||
|
err = p.db.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[:])
|
||||||
|
|
||||||
|
var success bool
|
||||||
|
var errorMsg sql.NullString
|
||||||
|
var valid bool
|
||||||
|
|
||||||
|
query := `SELECT p_success, p_error, p_valid FROM resolvespec_totp_validate_backup_code($1, $2)`
|
||||||
|
err := p.db.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
|
||||||
|
}
|
||||||
218
pkg/security/totp_provider_database_test.go
Normal file
218
pkg/security/totp_provider_database_test.go
Normal file
@@ -0,0 +1,218 @@
|
|||||||
|
package security_test
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/bitechdev/ResolveSpec/pkg/security"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Note: These tests require a PostgreSQL database with the schema from totp_database_schema.sql
|
||||||
|
// Set TEST_DATABASE_URL environment variable or skip tests
|
||||||
|
|
||||||
|
func setupTestDB(t *testing.T) *sql.DB {
|
||||||
|
// Skip if no test database configured
|
||||||
|
t.Skip("Database tests require TEST_DATABASE_URL environment variable")
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDatabaseTwoFactorProvider_Enable2FA(t *testing.T) {
|
||||||
|
db := setupTestDB(t)
|
||||||
|
if db == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
provider := security.NewDatabaseTwoFactorProvider(db, nil)
|
||||||
|
|
||||||
|
// Generate secret and backup codes
|
||||||
|
secret, err := provider.Generate2FASecret(1, "TestApp", "test@example.com")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Generate2FASecret() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Enable 2FA
|
||||||
|
err = provider.Enable2FA(1, secret.Secret, secret.BackupCodes)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Enable2FA() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify enabled
|
||||||
|
enabled, err := provider.Get2FAStatus(1)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Get2FAStatus() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if !enabled {
|
||||||
|
t.Error("Get2FAStatus() = false, want true")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDatabaseTwoFactorProvider_Disable2FA(t *testing.T) {
|
||||||
|
db := setupTestDB(t)
|
||||||
|
if db == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
provider := security.NewDatabaseTwoFactorProvider(db, nil)
|
||||||
|
|
||||||
|
// Enable first
|
||||||
|
secret, _ := provider.Generate2FASecret(1, "TestApp", "test@example.com")
|
||||||
|
provider.Enable2FA(1, secret.Secret, secret.BackupCodes)
|
||||||
|
|
||||||
|
// Disable
|
||||||
|
err := provider.Disable2FA(1)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Disable2FA() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify disabled
|
||||||
|
enabled, err := provider.Get2FAStatus(1)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Get2FAStatus() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if enabled {
|
||||||
|
t.Error("Get2FAStatus() = true, want false")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDatabaseTwoFactorProvider_GetSecret(t *testing.T) {
|
||||||
|
db := setupTestDB(t)
|
||||||
|
if db == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
provider := security.NewDatabaseTwoFactorProvider(db, nil)
|
||||||
|
|
||||||
|
// Enable 2FA
|
||||||
|
secret, _ := provider.Generate2FASecret(1, "TestApp", "test@example.com")
|
||||||
|
provider.Enable2FA(1, secret.Secret, secret.BackupCodes)
|
||||||
|
|
||||||
|
// Retrieve secret
|
||||||
|
retrieved, err := provider.Get2FASecret(1)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Get2FASecret() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if retrieved != secret.Secret {
|
||||||
|
t.Errorf("Get2FASecret() = %v, want %v", retrieved, secret.Secret)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDatabaseTwoFactorProvider_ValidateBackupCode(t *testing.T) {
|
||||||
|
db := setupTestDB(t)
|
||||||
|
if db == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
provider := security.NewDatabaseTwoFactorProvider(db, nil)
|
||||||
|
|
||||||
|
// Enable 2FA
|
||||||
|
secret, _ := provider.Generate2FASecret(1, "TestApp", "test@example.com")
|
||||||
|
provider.Enable2FA(1, secret.Secret, secret.BackupCodes)
|
||||||
|
|
||||||
|
// Validate backup code
|
||||||
|
valid, err := provider.ValidateBackupCode(1, secret.BackupCodes[0])
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("ValidateBackupCode() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if !valid {
|
||||||
|
t.Error("ValidateBackupCode() = false, want true")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try to use same code again
|
||||||
|
valid, err = provider.ValidateBackupCode(1, secret.BackupCodes[0])
|
||||||
|
if err == nil {
|
||||||
|
t.Error("ValidateBackupCode() should error on reuse")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try invalid code
|
||||||
|
valid, err = provider.ValidateBackupCode(1, "INVALID")
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("ValidateBackupCode() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if valid {
|
||||||
|
t.Error("ValidateBackupCode() = true for invalid code")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDatabaseTwoFactorProvider_RegenerateBackupCodes(t *testing.T) {
|
||||||
|
db := setupTestDB(t)
|
||||||
|
if db == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
provider := security.NewDatabaseTwoFactorProvider(db, nil)
|
||||||
|
|
||||||
|
// Enable 2FA
|
||||||
|
secret, _ := provider.Generate2FASecret(1, "TestApp", "test@example.com")
|
||||||
|
provider.Enable2FA(1, secret.Secret, secret.BackupCodes)
|
||||||
|
|
||||||
|
// Regenerate codes
|
||||||
|
newCodes, err := provider.GenerateBackupCodes(1, 10)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("GenerateBackupCodes() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(newCodes) != 10 {
|
||||||
|
t.Errorf("GenerateBackupCodes() returned %d codes, want 10", len(newCodes))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Old codes should not work
|
||||||
|
valid, _ := provider.ValidateBackupCode(1, secret.BackupCodes[0])
|
||||||
|
if valid {
|
||||||
|
t.Error("Old backup code should not work after regeneration")
|
||||||
|
}
|
||||||
|
|
||||||
|
// New codes should work
|
||||||
|
valid, err = provider.ValidateBackupCode(1, newCodes[0])
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("ValidateBackupCode() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if !valid {
|
||||||
|
t.Error("ValidateBackupCode() = false for new code")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDatabaseTwoFactorProvider_Generate2FASecret(t *testing.T) {
|
||||||
|
db := setupTestDB(t)
|
||||||
|
if db == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
provider := security.NewDatabaseTwoFactorProvider(db, nil)
|
||||||
|
|
||||||
|
secret, err := provider.Generate2FASecret(1, "TestApp", "test@example.com")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Generate2FASecret() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if secret.Secret == "" {
|
||||||
|
t.Error("Generate2FASecret() returned empty secret")
|
||||||
|
}
|
||||||
|
|
||||||
|
if secret.QRCodeURL == "" {
|
||||||
|
t.Error("Generate2FASecret() returned empty QR code URL")
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(secret.BackupCodes) != 10 {
|
||||||
|
t.Errorf("Generate2FASecret() returned %d backup codes, want 10", len(secret.BackupCodes))
|
||||||
|
}
|
||||||
|
|
||||||
|
if secret.Issuer != "TestApp" {
|
||||||
|
t.Errorf("Generate2FASecret() Issuer = %v, want TestApp", secret.Issuer)
|
||||||
|
}
|
||||||
|
|
||||||
|
if secret.AccountName != "test@example.com" {
|
||||||
|
t.Errorf("Generate2FASecret() AccountName = %v, want test@example.com", secret.AccountName)
|
||||||
|
}
|
||||||
|
}
|
||||||
156
pkg/security/totp_provider_memory.go
Normal file
156
pkg/security/totp_provider_memory.go
Normal file
@@ -0,0 +1,156 @@
|
|||||||
|
package security
|
||||||
|
|
||||||
|
import (
|
||||||
|
"crypto/sha256"
|
||||||
|
"encoding/hex"
|
||||||
|
"fmt"
|
||||||
|
"sync"
|
||||||
|
)
|
||||||
|
|
||||||
|
// MemoryTwoFactorProvider is an in-memory implementation of TwoFactorAuthProvider for testing/examples
|
||||||
|
type MemoryTwoFactorProvider struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
secrets map[int]string // userID -> secret
|
||||||
|
backupCodes map[int]map[string]bool // userID -> backup codes (code -> used)
|
||||||
|
totpGen *TOTPGenerator
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewMemoryTwoFactorProvider creates a new in-memory 2FA provider
|
||||||
|
func NewMemoryTwoFactorProvider(config *TwoFactorConfig) *MemoryTwoFactorProvider {
|
||||||
|
if config == nil {
|
||||||
|
config = DefaultTwoFactorConfig()
|
||||||
|
}
|
||||||
|
return &MemoryTwoFactorProvider{
|
||||||
|
secrets: make(map[int]string),
|
||||||
|
backupCodes: make(map[int]map[string]bool),
|
||||||
|
totpGen: NewTOTPGenerator(config),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Generate2FASecret creates a new secret for a user
|
||||||
|
func (m *MemoryTwoFactorProvider) Generate2FASecret(userID int, issuer, accountName string) (*TwoFactorSecret, error) {
|
||||||
|
secret, err := m.totpGen.GenerateSecret()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
qrURL := m.totpGen.GenerateQRCodeURL(secret, issuer, accountName)
|
||||||
|
|
||||||
|
backupCodes, err := GenerateBackupCodes(10)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return &TwoFactorSecret{
|
||||||
|
Secret: secret,
|
||||||
|
QRCodeURL: qrURL,
|
||||||
|
BackupCodes: backupCodes,
|
||||||
|
Issuer: issuer,
|
||||||
|
AccountName: accountName,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Validate2FACode verifies a TOTP code
|
||||||
|
func (m *MemoryTwoFactorProvider) Validate2FACode(secret string, code string) (bool, error) {
|
||||||
|
return m.totpGen.ValidateCode(secret, code)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Enable2FA activates 2FA for a user
|
||||||
|
func (m *MemoryTwoFactorProvider) Enable2FA(userID int, secret string, backupCodes []string) error {
|
||||||
|
m.mu.Lock()
|
||||||
|
defer m.mu.Unlock()
|
||||||
|
|
||||||
|
m.secrets[userID] = secret
|
||||||
|
|
||||||
|
// Store backup codes
|
||||||
|
if m.backupCodes[userID] == nil {
|
||||||
|
m.backupCodes[userID] = make(map[string]bool)
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, code := range backupCodes {
|
||||||
|
// Hash backup codes for security
|
||||||
|
hash := sha256.Sum256([]byte(code))
|
||||||
|
m.backupCodes[userID][hex.EncodeToString(hash[:])] = false
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Disable2FA deactivates 2FA for a user
|
||||||
|
func (m *MemoryTwoFactorProvider) Disable2FA(userID int) error {
|
||||||
|
m.mu.Lock()
|
||||||
|
defer m.mu.Unlock()
|
||||||
|
|
||||||
|
delete(m.secrets, userID)
|
||||||
|
delete(m.backupCodes, userID)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get2FAStatus checks if user has 2FA enabled
|
||||||
|
func (m *MemoryTwoFactorProvider) Get2FAStatus(userID int) (bool, error) {
|
||||||
|
m.mu.RLock()
|
||||||
|
defer m.mu.RUnlock()
|
||||||
|
|
||||||
|
_, exists := m.secrets[userID]
|
||||||
|
return exists, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get2FASecret retrieves the user's 2FA secret
|
||||||
|
func (m *MemoryTwoFactorProvider) Get2FASecret(userID int) (string, error) {
|
||||||
|
m.mu.RLock()
|
||||||
|
defer m.mu.RUnlock()
|
||||||
|
|
||||||
|
secret, exists := m.secrets[userID]
|
||||||
|
if !exists {
|
||||||
|
return "", fmt.Errorf("user does not have 2FA enabled")
|
||||||
|
}
|
||||||
|
return secret, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GenerateBackupCodes creates backup codes for 2FA
|
||||||
|
func (m *MemoryTwoFactorProvider) GenerateBackupCodes(userID int, count int) ([]string, error) {
|
||||||
|
codes, err := GenerateBackupCodes(count)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
m.mu.Lock()
|
||||||
|
defer m.mu.Unlock()
|
||||||
|
|
||||||
|
// Clear old backup codes and store new ones
|
||||||
|
m.backupCodes[userID] = make(map[string]bool)
|
||||||
|
for _, code := range codes {
|
||||||
|
hash := sha256.Sum256([]byte(code))
|
||||||
|
m.backupCodes[userID][hex.EncodeToString(hash[:])] = false
|
||||||
|
}
|
||||||
|
|
||||||
|
return codes, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ValidateBackupCode checks and consumes a backup code
|
||||||
|
func (m *MemoryTwoFactorProvider) ValidateBackupCode(userID int, code string) (bool, error) {
|
||||||
|
m.mu.Lock()
|
||||||
|
defer m.mu.Unlock()
|
||||||
|
|
||||||
|
userCodes, exists := m.backupCodes[userID]
|
||||||
|
if !exists {
|
||||||
|
return false, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Hash the provided code
|
||||||
|
hash := sha256.Sum256([]byte(code))
|
||||||
|
hashStr := hex.EncodeToString(hash[:])
|
||||||
|
|
||||||
|
used, exists := userCodes[hashStr]
|
||||||
|
if !exists {
|
||||||
|
return false, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
if used {
|
||||||
|
return false, fmt.Errorf("backup code already used")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Mark as used
|
||||||
|
userCodes[hashStr] = true
|
||||||
|
return true, nil
|
||||||
|
}
|
||||||
292
pkg/security/totp_test.go
Normal file
292
pkg/security/totp_test.go
Normal file
@@ -0,0 +1,292 @@
|
|||||||
|
package security
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestTOTPGenerator_GenerateSecret(t *testing.T) {
|
||||||
|
totp := NewTOTPGenerator(nil)
|
||||||
|
|
||||||
|
secret, err := totp.GenerateSecret()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GenerateSecret() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if secret == "" {
|
||||||
|
t.Error("GenerateSecret() returned empty secret")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Secret should be base32 encoded
|
||||||
|
if len(secret) < 16 {
|
||||||
|
t.Error("GenerateSecret() returned secret that is too short")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTOTPGenerator_GenerateQRCodeURL(t *testing.T) {
|
||||||
|
totp := NewTOTPGenerator(nil)
|
||||||
|
|
||||||
|
secret := "JBSWY3DPEHPK3PXP"
|
||||||
|
issuer := "TestApp"
|
||||||
|
accountName := "user@example.com"
|
||||||
|
|
||||||
|
url := totp.GenerateQRCodeURL(secret, issuer, accountName)
|
||||||
|
|
||||||
|
if !strings.HasPrefix(url, "otpauth://totp/") {
|
||||||
|
t.Errorf("GenerateQRCodeURL() = %v, want otpauth://totp/ prefix", url)
|
||||||
|
}
|
||||||
|
|
||||||
|
if !strings.Contains(url, "secret="+secret) {
|
||||||
|
t.Errorf("GenerateQRCodeURL() missing secret parameter")
|
||||||
|
}
|
||||||
|
|
||||||
|
if !strings.Contains(url, "issuer="+issuer) {
|
||||||
|
t.Errorf("GenerateQRCodeURL() missing issuer parameter")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTOTPGenerator_GenerateCode(t *testing.T) {
|
||||||
|
config := &TwoFactorConfig{
|
||||||
|
Algorithm: "SHA1",
|
||||||
|
Digits: 6,
|
||||||
|
Period: 30,
|
||||||
|
SkewWindow: 1,
|
||||||
|
}
|
||||||
|
totp := NewTOTPGenerator(config)
|
||||||
|
|
||||||
|
secret := "JBSWY3DPEHPK3PXP"
|
||||||
|
|
||||||
|
// Test with known time
|
||||||
|
timestamp := time.Unix(1234567890, 0)
|
||||||
|
code, err := totp.GenerateCode(secret, timestamp)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GenerateCode() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(code) != 6 {
|
||||||
|
t.Errorf("GenerateCode() returned code with length %d, want 6", len(code))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Code should be numeric
|
||||||
|
for _, c := range code {
|
||||||
|
if c < '0' || c > '9' {
|
||||||
|
t.Errorf("GenerateCode() returned non-numeric code: %s", code)
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTOTPGenerator_ValidateCode(t *testing.T) {
|
||||||
|
config := &TwoFactorConfig{
|
||||||
|
Algorithm: "SHA1",
|
||||||
|
Digits: 6,
|
||||||
|
Period: 30,
|
||||||
|
SkewWindow: 1,
|
||||||
|
}
|
||||||
|
totp := NewTOTPGenerator(config)
|
||||||
|
|
||||||
|
secret := "JBSWY3DPEHPK3PXP"
|
||||||
|
|
||||||
|
// Generate a code for current time
|
||||||
|
now := time.Now()
|
||||||
|
code, err := totp.GenerateCode(secret, now)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GenerateCode() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Validate the code
|
||||||
|
valid, err := totp.ValidateCode(secret, code)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ValidateCode() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if !valid {
|
||||||
|
t.Error("ValidateCode() = false, want true for current code")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Test with invalid code
|
||||||
|
valid, err = totp.ValidateCode(secret, "000000")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ValidateCode() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// This might occasionally pass if 000000 is the correct code, but very unlikely
|
||||||
|
if valid && code != "000000" {
|
||||||
|
t.Error("ValidateCode() = true for invalid code")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTOTPGenerator_ValidateCode_WithSkew(t *testing.T) {
|
||||||
|
config := &TwoFactorConfig{
|
||||||
|
Algorithm: "SHA1",
|
||||||
|
Digits: 6,
|
||||||
|
Period: 30,
|
||||||
|
SkewWindow: 2, // Allow 2 periods before/after
|
||||||
|
}
|
||||||
|
totp := NewTOTPGenerator(config)
|
||||||
|
|
||||||
|
secret := "JBSWY3DPEHPK3PXP"
|
||||||
|
|
||||||
|
// Generate code for 1 period ago
|
||||||
|
past := time.Now().Add(-30 * time.Second)
|
||||||
|
code, err := totp.GenerateCode(secret, past)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GenerateCode() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Should still validate with skew window
|
||||||
|
valid, err := totp.ValidateCode(secret, code)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ValidateCode() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if !valid {
|
||||||
|
t.Error("ValidateCode() = false, want true for code within skew window")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTOTPGenerator_DifferentAlgorithms(t *testing.T) {
|
||||||
|
algorithms := []string{"SHA1", "SHA256", "SHA512"}
|
||||||
|
secret := "JBSWY3DPEHPK3PXP"
|
||||||
|
|
||||||
|
for _, algo := range algorithms {
|
||||||
|
t.Run(algo, func(t *testing.T) {
|
||||||
|
config := &TwoFactorConfig{
|
||||||
|
Algorithm: algo,
|
||||||
|
Digits: 6,
|
||||||
|
Period: 30,
|
||||||
|
SkewWindow: 1,
|
||||||
|
}
|
||||||
|
totp := NewTOTPGenerator(config)
|
||||||
|
|
||||||
|
code, err := totp.GenerateCode(secret, time.Now())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GenerateCode() with %s error = %v", algo, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
valid, err := totp.ValidateCode(secret, code)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ValidateCode() with %s error = %v", algo, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if !valid {
|
||||||
|
t.Errorf("ValidateCode() with %s = false, want true", algo)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTOTPGenerator_8Digits(t *testing.T) {
|
||||||
|
config := &TwoFactorConfig{
|
||||||
|
Algorithm: "SHA1",
|
||||||
|
Digits: 8,
|
||||||
|
Period: 30,
|
||||||
|
SkewWindow: 1,
|
||||||
|
}
|
||||||
|
totp := NewTOTPGenerator(config)
|
||||||
|
|
||||||
|
secret := "JBSWY3DPEHPK3PXP"
|
||||||
|
|
||||||
|
code, err := totp.GenerateCode(secret, time.Now())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GenerateCode() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(code) != 8 {
|
||||||
|
t.Errorf("GenerateCode() returned code with length %d, want 8", len(code))
|
||||||
|
}
|
||||||
|
|
||||||
|
valid, err := totp.ValidateCode(secret, code)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ValidateCode() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if !valid {
|
||||||
|
t.Error("ValidateCode() = false, want true for 8-digit code")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGenerateBackupCodes(t *testing.T) {
|
||||||
|
count := 10
|
||||||
|
codes, err := GenerateBackupCodes(count)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GenerateBackupCodes() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(codes) != count {
|
||||||
|
t.Errorf("GenerateBackupCodes() returned %d codes, want %d", len(codes), count)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check uniqueness
|
||||||
|
seen := make(map[string]bool)
|
||||||
|
for _, code := range codes {
|
||||||
|
if seen[code] {
|
||||||
|
t.Errorf("GenerateBackupCodes() generated duplicate code: %s", code)
|
||||||
|
}
|
||||||
|
seen[code] = true
|
||||||
|
|
||||||
|
// Check format (8 hex characters)
|
||||||
|
if len(code) != 8 {
|
||||||
|
t.Errorf("GenerateBackupCodes() code length = %d, want 8", len(code))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDefaultTwoFactorConfig(t *testing.T) {
|
||||||
|
config := DefaultTwoFactorConfig()
|
||||||
|
|
||||||
|
if config.Algorithm != "SHA1" {
|
||||||
|
t.Errorf("DefaultTwoFactorConfig() Algorithm = %s, want SHA1", config.Algorithm)
|
||||||
|
}
|
||||||
|
|
||||||
|
if config.Digits != 6 {
|
||||||
|
t.Errorf("DefaultTwoFactorConfig() Digits = %d, want 6", config.Digits)
|
||||||
|
}
|
||||||
|
|
||||||
|
if config.Period != 30 {
|
||||||
|
t.Errorf("DefaultTwoFactorConfig() Period = %d, want 30", config.Period)
|
||||||
|
}
|
||||||
|
|
||||||
|
if config.SkewWindow != 1 {
|
||||||
|
t.Errorf("DefaultTwoFactorConfig() SkewWindow = %d, want 1", config.SkewWindow)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTOTPGenerator_InvalidSecret(t *testing.T) {
|
||||||
|
totp := NewTOTPGenerator(nil)
|
||||||
|
|
||||||
|
// Test with invalid base32 secret
|
||||||
|
_, err := totp.GenerateCode("INVALID!!!", time.Now())
|
||||||
|
if err == nil {
|
||||||
|
t.Error("GenerateCode() with invalid secret should return error")
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err = totp.ValidateCode("INVALID!!!", "123456")
|
||||||
|
if err == nil {
|
||||||
|
t.Error("ValidateCode() with invalid secret should return error")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Benchmark tests
|
||||||
|
func BenchmarkTOTPGenerator_GenerateCode(b *testing.B) {
|
||||||
|
totp := NewTOTPGenerator(nil)
|
||||||
|
secret := "JBSWY3DPEHPK3PXP"
|
||||||
|
now := time.Now()
|
||||||
|
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
_, _ = totp.GenerateCode(secret, now)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkTOTPGenerator_ValidateCode(b *testing.B) {
|
||||||
|
totp := NewTOTPGenerator(nil)
|
||||||
|
secret := "JBSWY3DPEHPK3PXP"
|
||||||
|
code, _ := totp.GenerateCode(secret, time.Now())
|
||||||
|
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
_, _ = totp.ValidateCode(secret, code)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user