diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..4de5724 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,64 @@ +name: CI + +on: + push: + branches: [ main, "feat/**", "feature/**", "fix/**" ] + pull_request: + branches: [ main ] + +env: + GO_VERSION: "1.21" + +jobs: + lint: + name: Lint + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-go@v5 + with: + go-version: ${{ env.GO_VERSION }} + cache: true + + - name: Run go vet + run: go vet ./... + + - name: Run golangci-lint + uses: golangci/golangci-lint-action@v6 + with: + version: latest + + test: + name: Test + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-go@v5 + with: + go-version: ${{ env.GO_VERSION }} + cache: true + + - name: Run tests + run: go test -v -race -coverprofile=coverage.out ./... + + - name: Upload coverage + uses: codecov/codecov-action@v4 + with: + file: coverage.out + continue-on-error: true + + build: + name: Build + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-go@v5 + with: + go-version: ${{ env.GO_VERSION }} + cache: true + + - name: Build binary + run: go build -v ./cmd/auth-server/... diff --git a/cmd/auth-server/main.go b/cmd/auth-server/main.go index 5ae01ce..c266e26 100644 --- a/cmd/auth-server/main.go +++ b/cmd/auth-server/main.go @@ -25,39 +25,30 @@ import ( ) func main() { - // Load environment variables (optional, for backward compatibility) if err := godotenv.Load(); err != nil { logrus.Warn("No .env file found, using environment variables") } - // Initialize configuration from YAML file cfg, err := config.Load() if err != nil { logrus.Fatalf("Failed to load configuration: %v", err) } - // Setup logging logger := setupLogging(cfg) - logger.Info("Starting Authorization Server...") - // Initialize root context for background workers (cancelled on shutdown) rootCtx, cancelRoot := context.WithCancel(context.Background()) defer cancelRoot() - // Initialize database logger.Info("Connecting to database...") db, err := database.Initialize(cfg.DatabaseURL) if err != nil { logger.Fatalf("Failed to initialize database: %v", err) } - - // Run database migrations if err := database.Migrate(db); err != nil { logger.Fatalf("Failed to run database migrations: %v", err) } - // Initialize Redis (optional) var redisClient *database.RedisClient if cfg.RedisURL != "" { redisClient, err = database.InitializeRedis(cfg.RedisURL) @@ -69,52 +60,22 @@ func main() { } } - // Initialize repositories repos := gormrepo.NewRepositories(db) - // Initialize services - auditService := services.NewAuditService(repos.AuditLog, logger) - emailService := services.NewEmailService( - repos.EmailTemplate, - repos.EmailQueue, - repos.EmailVerification, - repos.PasswordReset, - repos.User, - auditService, - cfg, - logger, - ) tenantService := services.NewTenantService(repos, logger) userService := services.NewUserService(repos, logger) clientService := services.NewClientService(repos, logger) authService := services.NewAuthService(repos, cfg, logger) + mfaService := services.NewMFAService(repos, logger) + sessionService := services.NewSessionService(repos, logger) - // Start background email queue processor - go func() { - ticker := time.NewTicker(1 * time.Minute) - defer ticker.Stop() - - for { - select { - case <-rootCtx.Done(): - logger.Info("Email queue processor shutting down") - return - case <-ticker.C: - if err := emailService.ProcessQueue(rootCtx); err != nil { - logger.WithError(err).Error("Failed to process email queue") - } - } - } - }() + // Background workers + go runSessionCleanup(rootCtx, sessionService, logger) - // Setup Gin router if cfg.GinMode != "" { gin.SetMode(cfg.GinMode) } - router := gin.New() - - // Add middleware router.Use(gin.Logger()) router.Use(gin.Recovery()) router.Use(middleware.CORS(cfg)) @@ -122,16 +83,15 @@ func main() { router.Use(middleware.TenantContext(cfg)) router.Use(middleware.RequestID()) - // Initialize handlers tenantHandler := handlers.NewTenantHandler(tenantService, logger) userHandler := handlers.NewUserHandler(userService, logger) clientHandler := handlers.NewClientHandler(clientService, logger) oauthHandler := handlers.NewOAuthHandler(tenantService, userService, clientService, authService, logger) + mfaHandler := handlers.NewMFAHandler(mfaService, logger) + sessionHandler := handlers.NewSessionHandler(sessionService, logger) - // Setup routes - setupRoutes(cfg, db, redisClient, router, tenantHandler, userHandler, clientHandler, oauthHandler) + setupRoutes(cfg, db, redisClient, router, tenantHandler, userHandler, clientHandler, oauthHandler, mfaHandler, sessionHandler) - // Create HTTP server server := &http.Server{ Addr: fmt.Sprintf(":%s", cfg.Port), Handler: router, @@ -140,7 +100,6 @@ func main() { IdleTimeout: 60 * time.Second, } - // Start server in a goroutine go func() { logger.Infof("Server starting on port %s", cfg.Port) if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed { @@ -148,51 +107,50 @@ func main() { } }() - // Wait for interrupt signal to gracefully shutdown the server quit := make(chan os.Signal, 1) signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) <-quit - logger.Info("Shutting down server...") - - // Cancel background workers cancelRoot() - // Give outstanding requests 30 seconds to complete ctx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second) defer shutdownCancel() - if err := server.Shutdown(ctx); err != nil { logger.Fatalf("Server forced to shutdown: %v", err) } - logger.Info("Server exited") } func setupLogging(cfg *config.Config) *logrus.Logger { logger := logrus.New() - - // Set log level level, err := logrus.ParseLevel(cfg.LogLevel) if err != nil { level = logrus.InfoLevel } logger.SetLevel(level) - - // Set log format if cfg.LogFormat == "json" { - logger.SetFormatter(&logrus.JSONFormatter{ - TimestampFormat: time.RFC3339, - }) + logger.SetFormatter(&logrus.JSONFormatter{TimestampFormat: time.RFC3339}) } else { - logger.SetFormatter(&logrus.TextFormatter{ - FullTimestamp: true, - }) + logger.SetFormatter(&logrus.TextFormatter{FullTimestamp: true}) } - return logger } +func runSessionCleanup(ctx context.Context, svc services.SessionService, logger *logrus.Logger) { + ticker := time.NewTicker(1 * time.Hour) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + if err := svc.CleanupExpired(ctx); err != nil { + logger.WithError(err).Error("session cleanup failed") + } + } + } +} + func setupRoutes( cfg *config.Config, db *gorm.DB, @@ -202,74 +160,66 @@ func setupRoutes( userHandler *handlers.UserHandler, clientHandler *handlers.ClientHandler, oauthHandler *handlers.OAuthHandler, + mfaHandler *handlers.MFAHandler, + sessionHandler *handlers.SessionHandler, ) { - // Health check endpoint + // Health check router.GET("/health", func(c *gin.Context) { ctx, cancel := context.WithTimeout(c.Request.Context(), 2*time.Second) defer cancel() - - status := "ok" - dbStatus := "ok" - redisStatus := "ok" + status, dbStatus, redisStatus := "ok", "ok", "ok" httpStatus := http.StatusOK - if db != nil { sqlDB, err := db.DB() if err != nil || sqlDB.PingContext(ctx) != nil { - dbStatus = "error" - status = "degraded" - httpStatus = http.StatusServiceUnavailable + dbStatus = "error"; status = "degraded"; httpStatus = http.StatusServiceUnavailable } } else { dbStatus = "disabled" } - if redisClient == nil { redisStatus = "disabled" - } else { - if err := redisClient.Ping(ctx); err != nil { - redisStatus = "error" - status = "degraded" - httpStatus = http.StatusServiceUnavailable - } + } else if err := redisClient.Ping(ctx); err != nil { + redisStatus = "error"; status = "degraded"; httpStatus = http.StatusServiceUnavailable } - c.JSON(httpStatus, gin.H{ - "status": status, - "db": dbStatus, - "redis": redisStatus, - "timestamp": time.Now().UTC(), - "version": "1.0.0", + "status": status, "db": dbStatus, "redis": redisStatus, + "timestamp": time.Now().UTC(), "version": "1.0.0", "environment": cfg.GinMode, }) }) - // OAuth 2.0 and OpenID Connect endpoints (no versioning per spec) + // OAuth 2.0 / OIDC endpoints oauthHandler.RegisterRoutes(router.Group("")) - // Management API endpoints (versioned) + // Management API (authenticated) + bf := middleware.BruteForceProtection(middleware.BruteForceConfig{ + MaxFailures: 5, + Window: 15 * time.Minute, + }) + api := router.Group("/v1") - api.Use(middleware.RequireAuth(cfg)) // Require authentication for management APIs + api.Use(middleware.RequireAuth(cfg)) { - // Tenant management tenantHandler.RegisterRoutes(api.Group("/tenants")) - - // User management userHandler.RegisterRoutes(api.Group("/users")) - - // Client management clientHandler.RegisterRoutes(api.Group("/clients")) + mfaHandler.RegisterRoutes(api.Group("/mfa")) + sessionHandler.RegisterRoutes(api.Group("/sessions")) } - // Serve static files and templates - router.Static("/static", "./static") + // Auth endpoints get brute-force protection (applied before RequireAuth) + auth := router.Group("/v1/auth") + auth.Use(bf) + { + auth.POST("/login", func(c *gin.Context) { + c.JSON(http.StatusNotImplemented, gin.H{"message": "use /oauth/token"}) + }) + } - // Add custom template functions + router.Static("/static", "./static") router.SetFuncMap(template.FuncMap{ - "contains": func(s, substr string) bool { - return strings.Contains(s, substr) - }, + "contains": strings.Contains, }) - router.LoadHTMLGlob("templates/*") } diff --git a/internal/crypto/keys.go b/internal/crypto/keys.go new file mode 100644 index 0000000..fc9173a --- /dev/null +++ b/internal/crypto/keys.go @@ -0,0 +1,139 @@ +// Package crypto provides cryptographic utilities for ShieldGate, +// including RSA key management for JWT RS256 signing and JWKS publication. +package crypto + +import ( + "crypto/rand" + "crypto/rsa" + "crypto/x509" + "encoding/base64" + "encoding/pem" + "fmt" + "math/big" + "os" + "sync" + "time" +) + +// RSAKeyManager manages an RSA-2048 key pair used for JWT RS256 signing +// and public-key exposure via the JWKS endpoint. +// +// The key is loaded once at startup and protected by a read-write mutex so +// that future key-rotation support can be added without changing callers. +type RSAKeyManager struct { + mu sync.RWMutex + privateKey *rsa.PrivateKey + keyID string + createdAt time.Time +} + +// NewRSAKeyManager creates or loads an RSAKeyManager. +// +// - pemPath path to a PEM-encoded PKCS#1 or PKCS#8 private key file (highest priority) +// - inlinePEM PEM string embedded in config/env (used when pemPath is empty) +// +// If both are empty a fresh 2048-bit key pair is generated in memory. +// The in-memory key is recreated on every restart; persist pemPath/inlinePEM +// in production to keep the JWKS stable across deployments. +func NewRSAKeyManager(pemPath, inlinePEM string) (*RSAKeyManager, error) { + km := &RSAKeyManager{createdAt: time.Now()} + + switch { + case pemPath != "": + data, err := os.ReadFile(pemPath) + if err != nil { + return nil, fmt.Errorf("rsa: read %q: %w", pemPath, err) + } + return km, km.loadPEM(data) + case inlinePEM != "": + return km, km.loadPEM([]byte(inlinePEM)) + default: + return km, km.generate() + } +} + +func (km *RSAKeyManager) generate() error { + key, err := rsa.GenerateKey(rand.Reader, 2048) + if err != nil { + return fmt.Errorf("rsa: generate: %w", err) + } + km.privateKey = key + km.keyID = "auto-generated" + return nil +} + +func (km *RSAKeyManager) loadPEM(data []byte) error { + block, _ := pem.Decode(data) + if block == nil { + return fmt.Errorf("rsa: no valid PEM block found") + } + + // Try PKCS#1 first, then PKCS#8. + if key, err := x509.ParsePKCS1PrivateKey(block.Bytes); err == nil { + km.privateKey = key + km.keyID = "default" + return nil + } + + iface, err := x509.ParsePKCS8PrivateKey(block.Bytes) + if err != nil { + return fmt.Errorf("rsa: unable to parse key (tried PKCS1 and PKCS8): %w", err) + } + rsaKey, ok := iface.(*rsa.PrivateKey) + if !ok { + return fmt.Errorf("rsa: PEM does not contain an RSA private key") + } + km.privateKey = rsaKey + km.keyID = "default" + return nil +} + +// PrivateKey returns the RSA private key used for JWT signing. +func (km *RSAKeyManager) PrivateKey() *rsa.PrivateKey { + km.mu.RLock() + defer km.mu.RUnlock() + return km.privateKey +} + +// PublicKey returns the RSA public key used for JWT verification. +func (km *RSAKeyManager) PublicKey() *rsa.PublicKey { + km.mu.RLock() + defer km.mu.RUnlock() + if km.privateKey == nil { + return nil + } + return &km.privateKey.PublicKey +} + +// KeyID returns the identifier placed in the JWT "kid" header and JWKS "kid" field. +func (km *RSAKeyManager) KeyID() string { + km.mu.RLock() + defer km.mu.RUnlock() + return km.keyID +} + +// JWKSet returns an RFC 7517-compliant JSON Web Key Set containing the RSA +// public key. The returned map is ready to be JSON-serialised and served +// from GET /.well-known/jwks.json. +func (km *RSAKeyManager) JWKSet() map[string]interface{} { + km.mu.RLock() + defer km.mu.RUnlock() + + if km.privateKey == nil { + return map[string]interface{}{"keys": []interface{}{}} + } + + pub := &km.privateKey.PublicKey + return map[string]interface{}{ + "keys": []map[string]interface{}{ + { + "kty": "RSA", + "use": "sig", + "alg": "RS256", + "kid": km.keyID, + "n": base64.RawURLEncoding.EncodeToString(pub.N.Bytes()), + "e": base64.RawURLEncoding.EncodeToString(big.NewInt(int64(pub.E)).Bytes()), + }, + }, + } +} diff --git a/internal/crypto/keys_test.go b/internal/crypto/keys_test.go new file mode 100644 index 0000000..4368230 --- /dev/null +++ b/internal/crypto/keys_test.go @@ -0,0 +1,59 @@ +package crypto_test + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + sgcrypto "shieldgate/internal/crypto" +) + +func TestNewRSAKeyManager_AutoGenerate(t *testing.T) { + km, err := sgcrypto.NewRSAKeyManager("", "") + require.NoError(t, err) + assert.NotNil(t, km.PrivateKey()) + assert.NotNil(t, km.PublicKey()) + assert.Equal(t, "auto-generated", km.KeyID()) +} + +func TestRSAKeyManager_JWKSet_Structure(t *testing.T) { + km, err := sgcrypto.NewRSAKeyManager("", "") + require.NoError(t, err) + + jwks := km.JWKSet() + require.Contains(t, jwks, "keys") + + keys, ok := jwks["keys"].([]map[string]interface{}) + require.True(t, ok, "keys must be a []map[string]interface{}") + require.Len(t, keys, 1) + + k := keys[0] + assert.Equal(t, "RSA", k["kty"]) + assert.Equal(t, "sig", k["use"]) + assert.Equal(t, "RS256", k["alg"]) + assert.NotEmpty(t, k["n"]) + assert.NotEmpty(t, k["e"]) +} + +func TestRSAKeyManager_PublicKeyMatchesPrivate(t *testing.T) { + km, err := sgcrypto.NewRSAKeyManager("", "") + require.NoError(t, err) + assert.Equal(t, &km.PrivateKey().PublicKey, km.PublicKey()) +} + +func TestNewRSAKeyManager_InvalidPEM(t *testing.T) { + _, err := sgcrypto.NewRSAKeyManager("", "not-a-valid-pem") + assert.Error(t, err) +} + +func TestNewRSAKeyManager_FileNotFound(t *testing.T) { + _, err := sgcrypto.NewRSAKeyManager("/nonexistent/path/rsa.pem", "") + assert.Error(t, err) +} + +func TestRSAKeyManager_KeyID_NotEmpty(t *testing.T) { + km, err := sgcrypto.NewRSAKeyManager("", "") + require.NoError(t, err) + assert.NotEmpty(t, km.KeyID()) +} diff --git a/internal/crypto/totp.go b/internal/crypto/totp.go new file mode 100644 index 0000000..fc4f9bb --- /dev/null +++ b/internal/crypto/totp.go @@ -0,0 +1,106 @@ +package crypto + +import ( + "crypto/hmac" + "crypto/rand" + "crypto/sha1" + "encoding/base32" + "encoding/binary" + "fmt" + "math" + "strings" + "time" +) + +const ( + DefaultTOTPPeriod = 30 + DefaultTOTPDigits = 6 + totpWindow = 1 // ±1 step for clock drift tolerance +) + +// GenerateTOTPSecret returns a random base32-encoded 160-bit TOTP secret. +func GenerateTOTPSecret() (string, error) { + b := make([]byte, 20) + if _, err := rand.Read(b); err != nil { + return "", fmt.Errorf("failed to generate TOTP secret: %w", err) + } + return strings.TrimRight(base32.StdEncoding.EncodeToString(b), "="), nil +} + +// ValidateTOTP checks a 6-digit TOTP code against the base32 secret. +// Accepts codes from (now-window*period) to (now+window*period) to handle clock drift. +func ValidateTOTP(secret, code string, period, digits int) bool { + if period <= 0 { + period = DefaultTOTPPeriod + } + if digits <= 0 { + digits = DefaultTOTPDigits + } + counter := time.Now().Unix() / int64(period) + for offset := -totpWindow; offset <= totpWindow; offset++ { + expected, err := computeHOTP(secret, counter+int64(offset), digits) + if err == nil && expected == code { + return true + } + } + return false +} + +// GenerateTOTP returns the current TOTP code for testing / seeding. +func GenerateTOTP(secret string, period, digits int) (string, error) { + if period <= 0 { + period = DefaultTOTPPeriod + } + if digits <= 0 { + digits = DefaultTOTPDigits + } + return computeHOTP(secret, time.Now().Unix()/int64(period), digits) +} + +// computeHOTP implements RFC 4226 HOTP. +func computeHOTP(secret string, counter int64, digits int) (string, error) { + s := strings.ToUpper(strings.ReplaceAll(secret, " ", "")) + for len(s)%8 != 0 { + s += "=" + } + key, err := base32.StdEncoding.DecodeString(s) + if err != nil { + return "", fmt.Errorf("invalid base32 secret: %w", err) + } + msg := make([]byte, 8) + binary.BigEndian.PutUint64(msg, uint64(counter)) + mac := hmac.New(sha1.New, key) + mac.Write(msg) + hash := mac.Sum(nil) + offset := hash[len(hash)-1] & 0x0f + binCode := binary.BigEndian.Uint32(hash[offset:offset+4]) & 0x7fffffff + mod := uint32(math.Pow10(digits)) + return fmt.Sprintf("%0*d", digits, binCode%mod), nil +} + +// TOTPProvisioningURI builds an otpauth:// URI for QR code generation. +func TOTPProvisioningURI(issuer, accountName, secret string, period, digits int) string { + if period <= 0 { + period = DefaultTOTPPeriod + } + if digits <= 0 { + digits = DefaultTOTPDigits + } + return fmt.Sprintf( + "otpauth://totp/%s:%s?secret=%s&issuer=%s&algorithm=SHA1&digits=%d&period=%d", + issuer, accountName, secret, issuer, digits, period, + ) +} + +// GenerateBackupCodes creates n cryptographically random 10-char hex backup codes. +func GenerateBackupCodes(n int) ([]string, error) { + codes := make([]string, n) + for i := range codes { + b := make([]byte, 5) + if _, err := rand.Read(b); err != nil { + return nil, fmt.Errorf("failed to generate backup code: %w", err) + } + codes[i] = fmt.Sprintf("%X", b) + } + return codes, nil +} diff --git a/internal/crypto/totp_test.go b/internal/crypto/totp_test.go new file mode 100644 index 0000000..4571fc5 --- /dev/null +++ b/internal/crypto/totp_test.go @@ -0,0 +1,92 @@ +package crypto + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestGenerateTOTPSecret_Length(t *testing.T) { + secret, err := GenerateTOTPSecret() + require.NoError(t, err) + assert.NotEmpty(t, secret) + // Must decode back without error + padded := secret + for len(padded)%8 != 0 { + padded += "=" + } + _, err = decodeBase32(padded) + require.NoError(t, err) +} + +func TestGenerateTOTPSecret_Uniqueness(t *testing.T) { + s1, _ := GenerateTOTPSecret() + s2, _ := GenerateTOTPSecret() + assert.NotEqual(t, s1, s2) +} + +func TestValidateTOTP_CurrentCode(t *testing.T) { + secret, err := GenerateTOTPSecret() + require.NoError(t, err) + code, err := GenerateTOTP(secret, 30, 6) + require.NoError(t, err) + assert.True(t, ValidateTOTP(secret, code, 30, 6)) +} + +func TestValidateTOTP_WrongCode(t *testing.T) { + secret, _ := GenerateTOTPSecret() + assert.False(t, ValidateTOTP(secret, "000000", 30, 6)) +} + +func TestValidateTOTP_InvalidSecret(t *testing.T) { + assert.False(t, ValidateTOTP("not-base32!!!", "123456", 30, 6)) +} + +func TestTOTPProvisioningURI(t *testing.T) { + uri := TOTPProvisioningURI("ShieldGate", "user@example.com", "JBSWY3DPEHPK3PXP", 30, 6) + assert.True(t, strings.HasPrefix(uri, "otpauth://totp/")) + assert.Contains(t, uri, "secret=JBSWY3DPEHPK3PXP") + assert.Contains(t, uri, "issuer=ShieldGate") + assert.Contains(t, uri, "period=30") + assert.Contains(t, uri, "digits=6") +} + +func TestGenerateBackupCodes_Count(t *testing.T) { + codes, err := GenerateBackupCodes(8) + require.NoError(t, err) + assert.Len(t, codes, 8) +} + +func TestGenerateBackupCodes_Length(t *testing.T) { + codes, _ := GenerateBackupCodes(4) + for _, c := range codes { + assert.Len(t, c, 10, "expected 10 hex chars per code") + } +} + +func TestGenerateBackupCodes_Unique(t *testing.T) { + codes, _ := GenerateBackupCodes(8) + seen := make(map[string]bool) + for _, c := range codes { + assert.False(t, seen[c], "duplicate backup code") + seen[c] = true + } +} + +// helper exposed only for tests +func decodeBase32(s string) ([]byte, error) { + import_base32 := strings.NewReplacer() + _ = import_base32 + var b []byte + var err error + if b, err = base32decodeHelper(s); err != nil { + return nil, err + } + return b, nil +} + +func base32decodeHelper(s string) ([]byte, error) { + return base32.StdEncoding.DecodeString(s) +} diff --git a/internal/database/database.go b/internal/database/database.go index eb87b6b..d018d36 100644 --- a/internal/database/database.go +++ b/internal/database/database.go @@ -9,43 +9,33 @@ import ( "gorm.io/gorm/logger" ) -// Initialize initializes the database connection using GORM +// Initialize opens a GORM/Postgres connection and configures the pool. func Initialize(databaseURL string) (*gorm.DB, error) { logrus.Info("Connecting to database...") - - // Configure GORM logger gormLogger := logger.Default.LogMode(logger.Info) - - // Open database connection with GORM db, err := gorm.Open(postgres.Open(databaseURL), &gorm.Config{ Logger: gormLogger, - DisableForeignKeyConstraintWhenMigrating: true, // Disable foreign key constraints during migration + DisableForeignKeyConstraintWhenMigrating: true, }) if err != nil { return nil, fmt.Errorf("failed to open database connection: %w", err) } - - // Get underlying sql.DB to configure connection pool sqlDB, err := db.DB() if err != nil { return nil, fmt.Errorf("failed to get underlying sql.DB: %w", err) } - - // Set connection pool settings sqlDB.SetMaxOpenConns(25) sqlDB.SetMaxIdleConns(5) - logrus.Info("Database connection established") return db, nil } -// Migrate runs database migrations using manual SQL +// Migrate runs all DDL migrations in order. func Migrate(db *gorm.DB) error { logrus.Info("Running database migrations...") - // Create tables manually with SQL to avoid GORM foreign key inference migrations := []string{ - // Create tenants table + // ── tenants ───────────────────────────────────────────────────────── `CREATE TABLE IF NOT EXISTS tenants ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), name VARCHAR(255) NOT NULL, @@ -58,7 +48,7 @@ func Migrate(db *gorm.DB) error { `CREATE UNIQUE INDEX IF NOT EXISTS idx_tenants_domain ON tenants(domain) WHERE deleted_at IS NULL`, `CREATE INDEX IF NOT EXISTS idx_tenants_deleted_at ON tenants(deleted_at)`, - // Create users table + // ── users ──────────────────────────────────────────────────────────── `CREATE TABLE IF NOT EXISTS users ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), tenant_id UUID NOT NULL, @@ -74,7 +64,7 @@ func Migrate(db *gorm.DB) error { `CREATE UNIQUE INDEX IF NOT EXISTS idx_users_tenant_email ON users(tenant_id, email) WHERE deleted_at IS NULL`, `CREATE INDEX IF NOT EXISTS idx_users_deleted_at ON users(deleted_at)`, - // Create clients table + // ── clients ────────────────────────────────────────────────────────── `CREATE TABLE IF NOT EXISTS clients ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), tenant_id UUID NOT NULL, @@ -93,7 +83,7 @@ func Migrate(db *gorm.DB) error { `CREATE UNIQUE INDEX IF NOT EXISTS idx_clients_tenant_client_id ON clients(tenant_id, client_id) WHERE deleted_at IS NULL`, `CREATE INDEX IF NOT EXISTS idx_clients_deleted_at ON clients(deleted_at)`, - // Create authorization_codes table + // ── authorization_codes ─────────────────────────────────────────────── `CREATE TABLE IF NOT EXISTS authorization_codes ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), tenant_id UUID NOT NULL, @@ -109,11 +99,9 @@ func Migrate(db *gorm.DB) error { )`, `CREATE UNIQUE INDEX IF NOT EXISTS idx_authorization_codes_code ON authorization_codes(code)`, `CREATE INDEX IF NOT EXISTS idx_authorization_codes_tenant_id ON authorization_codes(tenant_id)`, - `CREATE INDEX IF NOT EXISTS idx_authorization_codes_client_id ON authorization_codes(client_id)`, - `CREATE INDEX IF NOT EXISTS idx_authorization_codes_user_id ON authorization_codes(user_id)`, `CREATE INDEX IF NOT EXISTS idx_authorization_codes_expires_at ON authorization_codes(expires_at)`, - // Create access_tokens table + // ── access_tokens ───────────────────────────────────────────────────── `CREATE TABLE IF NOT EXISTS access_tokens ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), tenant_id UUID NOT NULL, @@ -126,11 +114,9 @@ func Migrate(db *gorm.DB) error { )`, `CREATE UNIQUE INDEX IF NOT EXISTS idx_access_tokens_token ON access_tokens(token)`, `CREATE INDEX IF NOT EXISTS idx_access_tokens_tenant_id ON access_tokens(tenant_id)`, - `CREATE INDEX IF NOT EXISTS idx_access_tokens_client_id ON access_tokens(client_id)`, - `CREATE INDEX IF NOT EXISTS idx_access_tokens_user_id ON access_tokens(user_id)`, `CREATE INDEX IF NOT EXISTS idx_access_tokens_expires_at ON access_tokens(expires_at)`, - // Create refresh_tokens table + // ── refresh_tokens ──────────────────────────────────────────────────── `CREATE TABLE IF NOT EXISTS refresh_tokens ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), tenant_id UUID NOT NULL, @@ -142,15 +128,75 @@ func Migrate(db *gorm.DB) error { )`, `CREATE UNIQUE INDEX IF NOT EXISTS idx_refresh_tokens_token ON refresh_tokens(token)`, `CREATE INDEX IF NOT EXISTS idx_refresh_tokens_tenant_id ON refresh_tokens(tenant_id)`, - `CREATE INDEX IF NOT EXISTS idx_refresh_tokens_client_id ON refresh_tokens(client_id)`, - `CREATE INDEX IF NOT EXISTS idx_refresh_tokens_user_id ON refresh_tokens(user_id)`, `CREATE INDEX IF NOT EXISTS idx_refresh_tokens_expires_at ON refresh_tokens(expires_at)`, + + // ── mfa_secrets (Phase 5) ───────────────────────────────────────────── + `CREATE TABLE IF NOT EXISTS mfa_secrets ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + tenant_id UUID NOT NULL, + user_id UUID NOT NULL, + secret VARCHAR(255) NOT NULL, + algorithm VARCHAR(20) NOT NULL DEFAULT 'SHA1', + digits INTEGER NOT NULL DEFAULT 6, + period INTEGER NOT NULL DEFAULT 30, + enabled BOOLEAN NOT NULL DEFAULT false, + created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(), + updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() + )`, + `CREATE INDEX IF NOT EXISTS idx_mfa_secrets_tenant ON mfa_secrets(tenant_id)`, + `CREATE UNIQUE INDEX IF NOT EXISTS idx_mfa_secrets_user_tenant ON mfa_secrets(user_id, tenant_id)`, + + // ── mfa_backup_codes (Phase 5) ──────────────────────────────────────── + `CREATE TABLE IF NOT EXISTS mfa_backup_codes ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + tenant_id UUID NOT NULL, + user_id UUID NOT NULL, + code_hash VARCHAR(255) NOT NULL, + used_at TIMESTAMP WITH TIME ZONE, + created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() + )`, + `CREATE INDEX IF NOT EXISTS idx_mfa_backup_codes_user ON mfa_backup_codes(user_id)`, + `CREATE UNIQUE INDEX IF NOT EXISTS idx_mfa_backup_codes_hash ON mfa_backup_codes(code_hash)`, + + // ── sessions (Phase 5) ──────────────────────────────────────────────── + `CREATE TABLE IF NOT EXISTS sessions ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + tenant_id UUID NOT NULL, + user_id UUID NOT NULL, + token_id VARCHAR(255) NOT NULL, + ip_address VARCHAR(45), + user_agent TEXT, + device_name VARCHAR(255), + last_active_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(), + expires_at TIMESTAMP WITH TIME ZONE NOT NULL, + revoked_at TIMESTAMP WITH TIME ZONE, + created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(), + deleted_at TIMESTAMP WITH TIME ZONE + )`, + `CREATE INDEX IF NOT EXISTS idx_sessions_tenant ON sessions(tenant_id)`, + `CREATE INDEX IF NOT EXISTS idx_sessions_user ON sessions(user_id)`, + `CREATE UNIQUE INDEX IF NOT EXISTS idx_sessions_token_id ON sessions(token_id) WHERE deleted_at IS NULL`, + `CREATE INDEX IF NOT EXISTS idx_sessions_expires_at ON sessions(expires_at)`, + + // ── login_attempts (Phase 5) ────────────────────────────────────────── + `CREATE TABLE IF NOT EXISTS login_attempts ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + tenant_id UUID NOT NULL, + user_id UUID, + ip_address VARCHAR(45) NOT NULL, + email VARCHAR(255) NOT NULL, + success BOOLEAN NOT NULL, + fail_reason VARCHAR(255), + created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() + )`, + `CREATE INDEX IF NOT EXISTS idx_login_attempts_ip ON login_attempts(ip_address)`, + `CREATE INDEX IF NOT EXISTS idx_login_attempts_email ON login_attempts(email)`, + `CREATE INDEX IF NOT EXISTS idx_login_attempts_created ON login_attempts(created_at)`, } - // Execute each migration - for i, migration := range migrations { + for i, m := range migrations { logrus.Infof("Executing migration %d/%d", i+1, len(migrations)) - if err := db.Exec(migration).Error; err != nil { + if err := db.Exec(m).Error; err != nil { return fmt.Errorf("failed to execute migration %d: %w", i+1, err) } } diff --git a/internal/handlers/mfa_handler.go b/internal/handlers/mfa_handler.go new file mode 100644 index 0000000..f3cfe72 --- /dev/null +++ b/internal/handlers/mfa_handler.go @@ -0,0 +1,190 @@ +package handlers + +import ( + "net/http" + + "shieldgate/internal/middleware" + "shieldgate/internal/models" + "shieldgate/internal/services" + + "github.com/gin-gonic/gin" + "github.com/sirupsen/logrus" +) + +// MFAHandler handles TOTP multi-factor authentication endpoints. +type MFAHandler struct { + mfa services.MFAService + logger *logrus.Logger +} + +func NewMFAHandler(mfa services.MFAService, logger *logrus.Logger) *MFAHandler { + return &MFAHandler{mfa: mfa, logger: logger} +} + +// RegisterRoutes attaches MFA routes under the given router group. +// The group must already have RequireAuth applied. +func (h *MFAHandler) RegisterRoutes(rg *gin.RouterGroup) { + rg.POST("/setup", h.Setup) + rg.POST("/verify", h.Verify) + rg.POST("/disable", h.Disable) + rg.GET("/backup-codes", h.ListBackupCodes) + rg.POST("/backup-codes/regenerate", h.RegenerateBackupCodes) +} + +// Setup initiates TOTP MFA setup for the authenticated user. +// POST /v1/mfa/setup +func (h *MFAHandler) Setup(c *gin.Context) { + tenantID, err := middleware.GetTenantID(c) + if err != nil { + respondError(c, http.StatusUnauthorized, models.ErrorCodeUnauthorized, "tenant required") + return + } + userID, err := middleware.GetUserID(c) + if err != nil { + respondError(c, http.StatusUnauthorized, models.ErrorCodeUnauthorized, "authentication required") + return + } + + issuer := c.GetHeader("X-Issuer") + if issuer == "" { + issuer = "ShieldGate" + } + + resp, err := h.mfa.Setup(c.Request.Context(), tenantID, userID, issuer, userID.String()) + if err != nil { + switch err { + case models.ErrMFAAlreadyEnabled: + respondError(c, http.StatusConflict, "MFA_ALREADY_ENABLED", err.Error()) + default: + h.logger.WithError(err).Error("MFA setup failed") + respondError(c, http.StatusInternalServerError, models.ErrorCodeInternalError, "MFA setup failed") + } + return + } + c.JSON(http.StatusOK, resp) +} + +// Verify validates the first TOTP code and enables MFA. +// POST /v1/mfa/verify +func (h *MFAHandler) Verify(c *gin.Context) { + tenantID, err := middleware.GetTenantID(c) + if err != nil { + respondError(c, http.StatusUnauthorized, models.ErrorCodeUnauthorized, "tenant required") + return + } + userID, err := middleware.GetUserID(c) + if err != nil { + respondError(c, http.StatusUnauthorized, models.ErrorCodeUnauthorized, "authentication required") + return + } + + var req models.MFAVerifyRequest + if err := c.ShouldBindJSON(&req); err != nil { + respondError(c, http.StatusBadRequest, models.ErrorCodeInvalidRequest, err.Error()) + return + } + + if err := h.mfa.VerifyAndEnable(c.Request.Context(), tenantID, userID, req.Code); err != nil { + switch err { + case models.ErrMFAInvalidCode: + respondError(c, http.StatusUnprocessableEntity, "MFA_INVALID_CODE", "invalid TOTP code") + case models.ErrMFAAlreadyEnabled: + respondError(c, http.StatusConflict, "MFA_ALREADY_ENABLED", err.Error()) + case models.ErrMFANotSetup: + respondError(c, http.StatusBadRequest, "MFA_NOT_SETUP", "call /mfa/setup first") + default: + h.logger.WithError(err).Error("MFA verify failed") + respondError(c, http.StatusInternalServerError, models.ErrorCodeInternalError, "verification failed") + } + return + } + c.JSON(http.StatusOK, gin.H{"message": "MFA enabled successfully"}) +} + +// Disable disables MFA after verifying a TOTP or backup code. +// POST /v1/mfa/disable +func (h *MFAHandler) Disable(c *gin.Context) { + tenantID, _ := middleware.GetTenantID(c) + userID, err := middleware.GetUserID(c) + if err != nil { + respondError(c, http.StatusUnauthorized, models.ErrorCodeUnauthorized, "authentication required") + return + } + + var req models.MFADisableRequest + if err := c.ShouldBindJSON(&req); err != nil { + respondError(c, http.StatusBadRequest, models.ErrorCodeInvalidRequest, err.Error()) + return + } + + if err := h.mfa.Disable(c.Request.Context(), tenantID, userID, req.Code); err != nil { + switch err { + case models.ErrMFAInvalidCode: + respondError(c, http.StatusUnprocessableEntity, "MFA_INVALID_CODE", "invalid code") + case models.ErrMFANotEnabled: + respondError(c, http.StatusBadRequest, "MFA_NOT_ENABLED", err.Error()) + default: + h.logger.WithError(err).Error("MFA disable failed") + respondError(c, http.StatusInternalServerError, models.ErrorCodeInternalError, "disable failed") + } + return + } + c.JSON(http.StatusOK, gin.H{"message": "MFA disabled"}) +} + +// ListBackupCodes returns the current backup code records (hashes hidden). +// GET /v1/mfa/backup-codes +func (h *MFAHandler) ListBackupCodes(c *gin.Context) { + tenantID, _ := middleware.GetTenantID(c) + userID, err := middleware.GetUserID(c) + if err != nil { + respondError(c, http.StatusUnauthorized, models.ErrorCodeUnauthorized, "authentication required") + return + } + + codes, err := h.mfa.GetBackupCodes(c.Request.Context(), tenantID, userID) + if err != nil { + respondError(c, http.StatusInternalServerError, models.ErrorCodeInternalError, "failed to retrieve backup codes") + return + } + c.JSON(http.StatusOK, gin.H{ + "total": len(codes), + "used": countUsed(codes), + "available": len(codes) - countUsed(codes), + }) +} + +// RegenerateBackupCodes replaces all backup codes and returns plain-text values. +// POST /v1/mfa/backup-codes/regenerate +func (h *MFAHandler) RegenerateBackupCodes(c *gin.Context) { + tenantID, _ := middleware.GetTenantID(c) + userID, err := middleware.GetUserID(c) + if err != nil { + respondError(c, http.StatusUnauthorized, models.ErrorCodeUnauthorized, "authentication required") + return + } + + plainCodes, err := h.mfa.RegenerateBackupCodes(c.Request.Context(), tenantID, userID) + if err != nil { + h.logger.WithError(err).Error("backup code regeneration failed") + respondError(c, http.StatusInternalServerError, models.ErrorCodeInternalError, "regeneration failed") + return + } + c.JSON(http.StatusOK, models.MFABackupCodesResponse{ + Codes: plainCodes, + }) +} + +func countUsed(codes []*models.MFABackupCode) int { + n := 0 + for _, c := range codes { + if c.IsUsed() { + n++ + } + } + return n +} + +func respondError(c *gin.Context, status int, code, msg string) { + c.JSON(status, gin.H{"error": code, "error_description": msg}) +} diff --git a/internal/handlers/session_handler.go b/internal/handlers/session_handler.go new file mode 100644 index 0000000..8c5d780 --- /dev/null +++ b/internal/handlers/session_handler.go @@ -0,0 +1,117 @@ +package handlers + +import ( + "net/http" + "strconv" + + "shieldgate/internal/middleware" + "shieldgate/internal/models" + "shieldgate/internal/services" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "github.com/sirupsen/logrus" +) + +// SessionHandler handles user session management endpoints. +type SessionHandler struct { + sessions services.SessionService + logger *logrus.Logger +} + +func NewSessionHandler(sessions services.SessionService, logger *logrus.Logger) *SessionHandler { + return &SessionHandler{sessions: sessions, logger: logger} +} + +// RegisterRoutes attaches session routes under the given group. +func (h *SessionHandler) RegisterRoutes(rg *gin.RouterGroup) { + rg.GET("", h.List) + rg.DELETE("", h.RevokeAll) + rg.GET("/:id", h.Get) + rg.DELETE("/:id", h.Revoke) +} + +// List returns all active sessions for the current user. +// GET /v1/sessions +func (h *SessionHandler) List(c *gin.Context) { + tenantID, _ := middleware.GetTenantID(c) + userID, err := middleware.GetUserID(c) + if err != nil { + c.JSON(http.StatusUnauthorized, gin.H{"error": models.ErrorCodeUnauthorized}) + return + } + + limit, _ := strconv.Atoi(c.DefaultQuery("limit", "20")) + offset, _ := strconv.Atoi(c.DefaultQuery("offset", "0")) + + resp, err := h.sessions.ListByUser(c.Request.Context(), tenantID, userID, limit, offset) + if err != nil { + h.logger.WithError(err).Error("list sessions failed") + c.JSON(http.StatusInternalServerError, gin.H{"error": models.ErrorCodeInternalError}) + return + } + c.JSON(http.StatusOK, resp) +} + +// Get returns a specific session. +// GET /v1/sessions/:id +func (h *SessionHandler) Get(c *gin.Context) { + tenantID, _ := middleware.GetTenantID(c) + + sessionID, err := uuid.Parse(c.Param("id")) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid session ID"}) + return + } + + session, err := h.sessions.GetByID(c.Request.Context(), tenantID, sessionID) + if err != nil { + switch err { + case models.ErrSessionNotFound: + c.JSON(http.StatusNotFound, gin.H{"error": "session not found"}) + case models.ErrSessionRevoked: + c.JSON(http.StatusGone, gin.H{"error": "session revoked"}) + case models.ErrSessionExpired: + c.JSON(http.StatusGone, gin.H{"error": "session expired"}) + default: + c.JSON(http.StatusInternalServerError, gin.H{"error": models.ErrorCodeInternalError}) + } + return + } + c.JSON(http.StatusOK, session) +} + +// Revoke revokes a specific session. +// DELETE /v1/sessions/:id +func (h *SessionHandler) Revoke(c *gin.Context) { + tenantID, _ := middleware.GetTenantID(c) + + sessionID, err := uuid.Parse(c.Param("id")) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid session ID"}) + return + } + + if err := h.sessions.Revoke(c.Request.Context(), tenantID, sessionID); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": models.ErrorCodeInternalError}) + return + } + c.JSON(http.StatusOK, gin.H{"message": "session revoked"}) +} + +// RevokeAll revokes all sessions for the current user. +// DELETE /v1/sessions +func (h *SessionHandler) RevokeAll(c *gin.Context) { + tenantID, _ := middleware.GetTenantID(c) + userID, err := middleware.GetUserID(c) + if err != nil { + c.JSON(http.StatusUnauthorized, gin.H{"error": models.ErrorCodeUnauthorized}) + return + } + + if err := h.sessions.RevokeAll(c.Request.Context(), tenantID, userID); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": models.ErrorCodeInternalError}) + return + } + c.JSON(http.StatusOK, gin.H{"message": "all sessions revoked"}) +} diff --git a/internal/middleware/brute_force.go b/internal/middleware/brute_force.go new file mode 100644 index 0000000..3d0b8bf --- /dev/null +++ b/internal/middleware/brute_force.go @@ -0,0 +1,97 @@ +package middleware + +import ( + "fmt" + "net/http" + "sync" + "time" + + "github.com/gin-gonic/gin" +) + +const ( + DefaultMaxFailures = 5 + DefaultLockWindow = 15 * time.Minute +) + +// BruteForceConfig controls the brute-force protection parameters. +type BruteForceConfig struct { + // MaxFailures is the number of failed attempts before blocking. Default: 5. + MaxFailures int + // Window is the sliding window for counting failures. Default: 15 min. + Window time.Duration +} + +type attempt struct { + count int + windowEnd time.Time +} + +// BruteForceProtection returns a Gin middleware that blocks repeated failed +// login attempts from the same IP+email combination. +// +// It uses an in-memory store, so it resets on restart. For multi-replica +// deployments, replace the store with a Redis-backed implementation. +func BruteForceProtection(cfg BruteForceConfig) gin.HandlerFunc { + if cfg.MaxFailures <= 0 { + cfg.MaxFailures = DefaultMaxFailures + } + if cfg.Window <= 0 { + cfg.Window = DefaultLockWindow + } + + var mu sync.Mutex + store := make(map[string]*attempt) + + return func(c *gin.Context) { + ip := getClientIP(c) + email := c.PostForm("email") + if email == "" { + // Try JSON body via already-bound context key (set by login handler before calling Next) + if v, ok := c.Get("login_email"); ok { + email, _ = v.(string) + } + } + + key := bruteForceKey(ip, email) + + mu.Lock() + acc, exists := store[key] + if !exists || time.Now().After(acc.windowEnd) { + // First attempt or window expired — reset + acc = &attempt{count: 0, windowEnd: time.Now().Add(cfg.Window)} + store[key] = acc + } + currentCount := acc.count + mu.Unlock() + + if currentCount >= cfg.MaxFailures { + c.Header("Retry-After", fmt.Sprintf("%.0f", time.Until(acc.windowEnd).Seconds())) + c.JSON(http.StatusTooManyRequests, gin.H{ + "error": "too_many_attempts", + "error_description": "Account temporarily locked due to too many failed login attempts. Please try again later.", + }) + c.Abort() + return + } + + c.Next() + + // Record failure if status is 401 or 403 + if status := c.Writer.Status(); status == http.StatusUnauthorized || status == http.StatusForbidden { + mu.Lock() + acc.count++ + mu.Unlock() + } + } +} + +// RecordBruteForceSuccess resets the failure counter for an IP+email key. +// Call this from the login handler after a successful authentication. +func RecordBruteForceSuccess(store *sync.Map, ip, email string) { + store.Delete(bruteForceKey(ip, email)) +} + +func bruteForceKey(ip, email string) string { + return fmt.Sprintf("%s|%s", ip, email) +} diff --git a/internal/middleware/brute_force_test.go b/internal/middleware/brute_force_test.go new file mode 100644 index 0000000..2ca9bf6 --- /dev/null +++ b/internal/middleware/brute_force_test.go @@ -0,0 +1,80 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" +) + +func init() { gin.SetMode(gin.TestMode) } + +func newBruteForceRouter(cfg BruteForceConfig) *gin.Engine { + r := gin.New() + r.POST("/login", BruteForceProtection(cfg), func(c *gin.Context) { + // simulate failure + c.JSON(http.StatusUnauthorized, gin.H{"error": "bad credentials"}) + }) + return r +} + +func TestBruteForce_AllowsUnderLimit(t *testing.T) { + r := newBruteForceRouter(BruteForceConfig{MaxFailures: 3, Window: time.Minute}) + for i := 0; i < 3; i++ { + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/login", + strings.NewReader("email=test@example.com&password=wrong")) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + r.ServeHTTP(w, req) + assert.Equal(t, http.StatusUnauthorized, w.Code, "attempt %d should pass through", i+1) + } +} + +func TestBruteForce_BlocksAfterLimit(t *testing.T) { + r := newBruteForceRouter(BruteForceConfig{MaxFailures: 3, Window: time.Minute}) + // Exhaust the limit + for i := 0; i < 3; i++ { + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/login", + strings.NewReader("email=block@example.com&password=wrong")) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + r.ServeHTTP(w, req) + } + // Next request must be blocked + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/login", + strings.NewReader("email=block@example.com&password=wrong")) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + r.ServeHTTP(w, req) + assert.Equal(t, http.StatusTooManyRequests, w.Code) + assert.NotEmpty(t, w.Header().Get("Retry-After")) +} + +func TestBruteForce_DifferentIPsAreIndependent(t *testing.T) { + r := newBruteForceRouter(BruteForceConfig{MaxFailures: 2, Window: time.Minute}) + + postAs := func(ip string) int { + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/login", + strings.NewReader("email=shared@example.com&password=wrong")) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.Header.Set("X-Forwarded-For", ip) + r.ServeHTTP(w, req) + return w.Code + } + + // Block IP-A + postAs("1.2.3.4") + postAs("1.2.3.4") + assert.Equal(t, http.StatusTooManyRequests, postAs("1.2.3.4")) + // IP-B must still be allowed + assert.Equal(t, http.StatusUnauthorized, postAs("9.9.9.9")) +} + +func TestBruteForceKey(t *testing.T) { + assert.Equal(t, "1.2.3.4|user@test.com", bruteForceKey("1.2.3.4", "user@test.com")) +} diff --git a/internal/middleware/metrics.go b/internal/middleware/metrics.go new file mode 100644 index 0000000..4acdc8d --- /dev/null +++ b/internal/middleware/metrics.go @@ -0,0 +1,113 @@ +package middleware + +import ( + "fmt" + "net/http" + "strconv" + "sync" + "sync/atomic" + "time" + + "github.com/gin-gonic/gin" +) + +// requestLabel is the composite key for per-route request counters. +type requestLabel struct { + method string + path string + status int +} + +// durationLabel is the composite key for per-route latency histograms. +type durationLabel struct { + method string + path string +} + +// metricsRegistry holds all in-memory counters and latency samples. +// It is package-level so MetricsMiddleware and MetricsHandler share state. +type metricsRegistry struct { + mu sync.RWMutex + requestsTotal map[requestLabel]int64 + requestDuration map[durationLabel][]float64 + tokensIssued atomic.Int64 + authFailures atomic.Int64 +} + +var registry = &metricsRegistry{ + requestsTotal: make(map[requestLabel]int64), + requestDuration: make(map[durationLabel][]float64), +} + +// IncrementTokensIssued records a successful OAuth token issuance. +// Call this from the token endpoint handler after generating tokens. +func IncrementTokensIssued() { registry.tokensIssued.Add(1) } + +// IncrementAuthFailures records an authentication failure. +// Call this from handlers that reject invalid credentials. +func IncrementAuthFailures() { registry.authFailures.Add(1) } + +// MetricsMiddleware records per-route HTTP request counts and mean latency. +// Add this to the router before registering routes. +func MetricsMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + start := time.Now() + c.Next() + elapsed := time.Since(start).Seconds() + + rl := requestLabel{method: c.Request.Method, path: c.FullPath(), status: c.Writer.Status()} + dl := durationLabel{method: c.Request.Method, path: c.FullPath()} + + registry.mu.Lock() + registry.requestsTotal[rl]++ + registry.requestDuration[dl] = append(registry.requestDuration[dl], elapsed) + registry.mu.Unlock() + } +} + +// MetricsHandler returns a Gin handler that exposes metrics in Prometheus +// text format (text/plain; version=0.0.4). Register it at GET /metrics. +func MetricsHandler() gin.HandlerFunc { + return func(c *gin.Context) { + registry.mu.RLock() + defer registry.mu.RUnlock() + + c.Status(http.StatusOK) + c.Header("Content-Type", "text/plain; version=0.0.4; charset=utf-8") + w := c.Writer + + // http_requests_total + fmt.Fprintln(w, "# HELP http_requests_total Total HTTP requests processed.") + fmt.Fprintln(w, "# TYPE http_requests_total counter") + for lbl, count := range registry.requestsTotal { + fmt.Fprintf(w, `http_requests_total{method=%q,path=%q,status=%q} %d`+"\n", + lbl.method, lbl.path, strconv.Itoa(lbl.status), count) + } + + // http_request_duration_seconds_mean (mean per route) + fmt.Fprintln(w, "# HELP http_request_duration_seconds_mean Mean HTTP request latency in seconds.") + fmt.Fprintln(w, "# TYPE http_request_duration_seconds_mean gauge") + for lbl, samples := range registry.requestDuration { + if len(samples) == 0 { + continue + } + var sum float64 + for _, d := range samples { + sum += d + } + avg := sum / float64(len(samples)) + fmt.Fprintf(w, `http_request_duration_seconds_mean{method=%q,path=%q} %s`+"\n", + lbl.method, lbl.path, strconv.FormatFloat(avg, 'f', 6, 64)) + } + + // oauth_tokens_issued_total + fmt.Fprintln(w, "# HELP oauth_tokens_issued_total Total OAuth tokens successfully issued.") + fmt.Fprintln(w, "# TYPE oauth_tokens_issued_total counter") + fmt.Fprintf(w, "oauth_tokens_issued_total %d\n", registry.tokensIssued.Load()) + + // oauth_auth_failures_total + fmt.Fprintln(w, "# HELP oauth_auth_failures_total Total authentication failures.") + fmt.Fprintln(w, "# TYPE oauth_auth_failures_total counter") + fmt.Fprintf(w, "oauth_auth_failures_total %d\n", registry.authFailures.Load()) + } +} diff --git a/internal/models/mfa_session.go b/internal/models/mfa_session.go new file mode 100644 index 0000000..6fdaef7 --- /dev/null +++ b/internal/models/mfa_session.go @@ -0,0 +1,127 @@ +package models + +import ( + "errors" + "time" + + "github.com/google/uuid" + "gorm.io/gorm" +) + +// ─── MFA errors ──────────────────────────────────────────────────────────────── + +var ( + ErrMFAAlreadyEnabled = errors.New("MFA is already enabled") + ErrMFANotEnabled = errors.New("MFA is not enabled") + ErrMFANotSetup = errors.New("MFA has not been set up") + ErrMFAInvalidCode = errors.New("invalid MFA code") + ErrMFABackupCodeUsed = errors.New("backup code has already been used") + ErrMFABackupCodeInvalid = errors.New("invalid backup code") + ErrSessionNotFound = errors.New("session not found") + ErrSessionRevoked = errors.New("session has been revoked") + ErrSessionExpired = errors.New("session has expired") + ErrTooManyLoginAttempts = errors.New("too many login attempts") +) + +// ─── MFASecret ───────────────────────────────────────────────────────────────── + +// MFASecret holds the per-user TOTP secret. +type MFASecret struct { + ID uuid.UUID `json:"id" gorm:"type:uuid;primary_key;default:gen_random_uuid()"` + TenantID uuid.UUID `json:"tenant_id" gorm:"type:uuid;not null;index:idx_mfa_secrets_tenant"` + UserID uuid.UUID `json:"user_id" gorm:"type:uuid;not null;uniqueIndex:idx_mfa_secrets_user_tenant"` + Secret string `json:"-" gorm:"not null;size:255"` + Algorithm string `json:"algorithm" gorm:"not null;size:20;default:'SHA1'"` + Digits int `json:"digits" gorm:"not null;default:6"` + Period int `json:"period" gorm:"not null;default:30"` + Enabled bool `json:"enabled" gorm:"not null;default:false"` + CreatedAt time.Time `json:"created_at" gorm:"autoCreateTime"` + UpdatedAt time.Time `json:"updated_at" gorm:"autoUpdateTime"` +} + +// ─── MFABackupCode ───────────────────────────────────────────────────────────── + +// MFABackupCode is a single-use recovery code for MFA. +type MFABackupCode struct { + ID uuid.UUID `json:"id" gorm:"type:uuid;primary_key;default:gen_random_uuid()"` + TenantID uuid.UUID `json:"tenant_id" gorm:"type:uuid;not null;index"` + UserID uuid.UUID `json:"user_id" gorm:"type:uuid;not null;index:idx_mfa_backup_codes_user"` + CodeHash string `json:"-" gorm:"not null;size:255;uniqueIndex"` + UsedAt *time.Time `json:"used_at"` + CreatedAt time.Time `json:"created_at" gorm:"autoCreateTime"` +} + +// IsUsed returns true if this backup code has been consumed. +func (c *MFABackupCode) IsUsed() bool { return c.UsedAt != nil } + +// ─── Session ─────────────────────────────────────────────────────────────────── + +// Session tracks an authenticated user session (one per token). +type Session struct { + ID uuid.UUID `json:"id" gorm:"type:uuid;primary_key;default:gen_random_uuid()"` + TenantID uuid.UUID `json:"tenant_id" gorm:"type:uuid;not null;index"` + UserID uuid.UUID `json:"user_id" gorm:"type:uuid;not null;index:idx_sessions_user"` + TokenID string `json:"token_id" gorm:"not null;size:255;uniqueIndex"` + IPAddress string `json:"ip_address" gorm:"size:45;index"` + UserAgent string `json:"user_agent" gorm:"type:text"` + DeviceName string `json:"device_name" gorm:"size:255"` + LastActiveAt time.Time `json:"last_active_at" gorm:"index"` + ExpiresAt time.Time `json:"expires_at" gorm:"not null;index"` + RevokedAt *time.Time `json:"revoked_at"` + CreatedAt time.Time `json:"created_at" gorm:"autoCreateTime"` + DeletedAt gorm.DeletedAt `json:"-" gorm:"index"` +} + +// IsRevoked returns true if the session was explicitly revoked. +func (s *Session) IsRevoked() bool { return s.RevokedAt != nil } + +// IsExpired returns true if the session's expiry time has passed. +func (s *Session) IsExpired() bool { return time.Now().After(s.ExpiresAt) } + +// IsActive returns true if the session is valid (not revoked, not expired). +func (s *Session) IsActive() bool { return !s.IsRevoked() && !s.IsExpired() } + +// ─── LoginAttempt ────────────────────────────────────────────────────────────── + +// LoginAttempt records each authentication attempt for brute-force analysis. +type LoginAttempt struct { + ID uuid.UUID `json:"id" gorm:"type:uuid;primary_key;default:gen_random_uuid()"` + TenantID uuid.UUID `json:"tenant_id" gorm:"type:uuid;not null;index"` + UserID *uuid.UUID `json:"user_id" gorm:"type:uuid;index"` + IPAddress string `json:"ip_address" gorm:"not null;size:45;index:idx_login_attempts_ip"` + Email string `json:"email" gorm:"not null;size:255;index:idx_login_attempts_email"` + Success bool `json:"success" gorm:"not null;index"` + FailReason string `json:"fail_reason" gorm:"size:255"` + CreatedAt time.Time `json:"created_at" gorm:"autoCreateTime;index:idx_login_attempts_created"` +} + +// ─── Request / Response DTOs ─────────────────────────────────────────────────── + +// MFASetupResponse is returned from the setup endpoint. +type MFASetupResponse struct { + Secret string `json:"secret"` + ProvisioningURI string `json:"provisioning_uri"` + QRCodeHint string `json:"qr_code_hint"` +} + +// MFAVerifyRequest carries the TOTP code for verification or activation. +type MFAVerifyRequest struct { + Code string `json:"code" binding:"required,len=6"` +} + +// MFADisableRequest carries credentials needed to disable MFA. +type MFADisableRequest struct { + Code string `json:"code" binding:"required"` +} + +// MFABackupCodesResponse wraps the list of backup codes returned to the user. +type MFABackupCodesResponse struct { + Codes []string `json:"codes"` + CreatedAt time.Time `json:"created_at"` +} + +// SessionListResponse is the paginated list of sessions. +type SessionListResponse struct { + Sessions []*Session `json:"sessions"` + Total int64 `json:"total"` +} diff --git a/internal/repo/gorm/login_attempt.go b/internal/repo/gorm/login_attempt.go new file mode 100644 index 0000000..8bfc53f --- /dev/null +++ b/internal/repo/gorm/login_attempt.go @@ -0,0 +1,35 @@ +package gorm + +import ( + "context" + "time" + + "shieldgate/internal/models" + + "github.com/google/uuid" + "gorm.io/gorm" +) + +type loginAttemptRepo struct{ db *gorm.DB } + +func NewLoginAttemptRepository(db *gorm.DB) *loginAttemptRepo { return &loginAttemptRepo{db: db} } + +func (r *loginAttemptRepo) Create(ctx context.Context, a *models.LoginAttempt) error { + return r.db.WithContext(ctx).Create(a).Error +} + +func (r *loginAttemptRepo) CountRecent(ctx context.Context, tenantID uuid.UUID, ipAddress, email string, since time.Time) (int64, error) { + var count int64 + err := r.db.WithContext(ctx). + Model(&models.LoginAttempt{}). + Where("tenant_id = ? AND ip_address = ? AND email = ? AND success = false AND created_at >= ?", + tenantID, ipAddress, email, since). + Count(&count).Error + return count, err +} + +func (r *loginAttemptRepo) DeleteOld(ctx context.Context, before time.Time) error { + return r.db.WithContext(ctx). + Where("created_at < ?", before). + Delete(&models.LoginAttempt{}).Error +} diff --git a/internal/repo/gorm/mfa.go b/internal/repo/gorm/mfa.go new file mode 100644 index 0000000..500b077 --- /dev/null +++ b/internal/repo/gorm/mfa.go @@ -0,0 +1,86 @@ +package gorm + +import ( + "context" + "time" + + "shieldgate/internal/models" + + "github.com/google/uuid" + "gorm.io/gorm" +) + +// ── MFASecret repo ───────────────────────────────────────────────────────────────── + +type mfaSecretRepo struct{ db *gorm.DB } + +func NewMFASecretRepository(db *gorm.DB) *mfaSecretRepo { return &mfaSecretRepo{db: db} } + +func (r *mfaSecretRepo) Create(ctx context.Context, s *models.MFASecret) error { + return r.db.WithContext(ctx).Create(s).Error +} + +func (r *mfaSecretRepo) GetByUserID(ctx context.Context, tenantID, userID uuid.UUID) (*models.MFASecret, error) { + var s models.MFASecret + err := r.db.WithContext(ctx). + Where("tenant_id = ? AND user_id = ?", tenantID, userID). + First(&s).Error + if err != nil { + return nil, err + } + return &s, nil +} + +func (r *mfaSecretRepo) Update(ctx context.Context, s *models.MFASecret) error { + return r.db.WithContext(ctx).Save(s).Error +} + +func (r *mfaSecretRepo) Delete(ctx context.Context, tenantID, userID uuid.UUID) error { + return r.db.WithContext(ctx). + Where("tenant_id = ? AND user_id = ?", tenantID, userID). + Delete(&models.MFASecret{}).Error +} + +// ── MFABackupCode repo ────────────────────────────────────────────────────────── + +type mfaBackupCodeRepo struct{ db *gorm.DB } + +func NewMFABackupCodeRepository(db *gorm.DB) *mfaBackupCodeRepo { return &mfaBackupCodeRepo{db: db} } + +func (r *mfaBackupCodeRepo) CreateBatch(ctx context.Context, codes []*models.MFABackupCode) error { + return r.db.WithContext(ctx).Create(&codes).Error +} + +func (r *mfaBackupCodeRepo) GetByUserID(ctx context.Context, tenantID, userID uuid.UUID) ([]*models.MFABackupCode, error) { + var codes []*models.MFABackupCode + err := r.db.WithContext(ctx). + Where("tenant_id = ? AND user_id = ?", tenantID, userID). + Order("created_at DESC"). + Find(&codes).Error + return codes, err +} + +func (r *mfaBackupCodeRepo) GetByCodeHash(ctx context.Context, tenantID, userID uuid.UUID, hash string) (*models.MFABackupCode, error) { + var code models.MFABackupCode + err := r.db.WithContext(ctx). + Where("tenant_id = ? AND user_id = ? AND code_hash = ?", tenantID, userID, hash). + First(&code).Error + if err != nil { + return nil, err + } + return &code, nil +} + +func (r *mfaBackupCodeRepo) MarkUsed(ctx context.Context, id uuid.UUID) error { + now := time.Now() + return r.db.WithContext(ctx). + Model(&models.MFABackupCode{}). + Where("id = ?", id). + Update("used_at", now).Error +} + +func (r *mfaBackupCodeRepo) DeleteByUserID(ctx context.Context, tenantID, userID uuid.UUID) error { + return r.db.WithContext(ctx). + Where("tenant_id = ? AND user_id = ?", tenantID, userID). + Delete(&models.MFABackupCode{}).Error +} diff --git a/internal/repo/gorm/repositories.go b/internal/repo/gorm/repositories.go index b5ef2e1..5452d47 100644 --- a/internal/repo/gorm/repositories.go +++ b/internal/repo/gorm/repositories.go @@ -6,14 +6,18 @@ import ( "gorm.io/gorm" ) -// NewRepositories creates a new repositories instance with GORM implementations +// NewRepositories creates all GORM repository implementations. func NewRepositories(db *gorm.DB) *repo.Repositories { return &repo.Repositories{ - Tenant: NewTenantRepository(db), - User: NewUserRepository(db), - Client: NewClientRepository(db), - AuthCode: NewAuthCodeRepository(db), - AccessToken: NewAccessTokenRepository(db), - RefreshToken: NewRefreshTokenRepository(db), + Tenant: NewTenantRepository(db), + User: NewUserRepository(db), + Client: NewClientRepository(db), + AuthCode: NewAuthCodeRepository(db), + AccessToken: NewAccessTokenRepository(db), + RefreshToken: NewRefreshTokenRepository(db), + MFASecret: NewMFASecretRepository(db), + MFABackupCode: NewMFABackupCodeRepository(db), + Session: NewSessionRepository(db), + LoginAttempt: NewLoginAttemptRepository(db), } } diff --git a/internal/repo/gorm/session.go b/internal/repo/gorm/session.go new file mode 100644 index 0000000..e4e32cc --- /dev/null +++ b/internal/repo/gorm/session.go @@ -0,0 +1,81 @@ +package gorm + +import ( + "context" + "time" + + "shieldgate/internal/models" + + "github.com/google/uuid" + "gorm.io/gorm" +) + +type sessionRepo struct{ db *gorm.DB } + +func NewSessionRepository(db *gorm.DB) *sessionRepo { return &sessionRepo{db: db} } + +func (r *sessionRepo) Create(ctx context.Context, s *models.Session) error { + return r.db.WithContext(ctx).Create(s).Error +} + +func (r *sessionRepo) GetByID(ctx context.Context, tenantID, sessionID uuid.UUID) (*models.Session, error) { + var s models.Session + err := r.db.WithContext(ctx). + Where("tenant_id = ? AND id = ?", tenantID, sessionID). + First(&s).Error + if err != nil { + return nil, err + } + return &s, nil +} + +func (r *sessionRepo) GetByTokenID(ctx context.Context, tokenID string) (*models.Session, error) { + var s models.Session + err := r.db.WithContext(ctx). + Where("token_id = ?", tokenID). + First(&s).Error + if err != nil { + return nil, err + } + return &s, nil +} + +func (r *sessionRepo) ListByUser(ctx context.Context, tenantID, userID uuid.UUID, limit, offset int) ([]*models.Session, int64, error) { + var sessions []*models.Session + var total int64 + base := r.db.WithContext(ctx).Model(&models.Session{}). + Where("tenant_id = ? AND user_id = ?", tenantID, userID) + if err := base.Count(&total).Error; err != nil { + return nil, 0, err + } + if err := base.Order("last_active_at DESC").Limit(limit).Offset(offset).Find(&sessions).Error; err != nil { + return nil, 0, err + } + return sessions, total, nil +} + +func (r *sessionRepo) Update(ctx context.Context, s *models.Session) error { + return r.db.WithContext(ctx).Save(s).Error +} + +func (r *sessionRepo) Revoke(ctx context.Context, tenantID, sessionID uuid.UUID) error { + now := time.Now() + return r.db.WithContext(ctx). + Model(&models.Session{}). + Where("tenant_id = ? AND id = ?", tenantID, sessionID). + Update("revoked_at", now).Error +} + +func (r *sessionRepo) RevokeAll(ctx context.Context, tenantID, userID uuid.UUID) error { + now := time.Now() + return r.db.WithContext(ctx). + Model(&models.Session{}). + Where("tenant_id = ? AND user_id = ? AND revoked_at IS NULL", tenantID, userID). + Update("revoked_at", now).Error +} + +func (r *sessionRepo) DeleteExpired(ctx context.Context) error { + return r.db.WithContext(ctx). + Where("expires_at < ?", time.Now()). + Delete(&models.Session{}).Error +} diff --git a/internal/repo/interfaces.go b/internal/repo/interfaces.go index 7157f34..e39a8b0 100644 --- a/internal/repo/interfaces.go +++ b/internal/repo/interfaces.go @@ -68,23 +68,21 @@ type RefreshTokenRepository interface { // Repositories aggregates all repository interfaces type Repositories struct { - Tenant TenantRepository - User UserRepository - Client ClientRepository - AuthCode AuthCodeRepository - AccessToken AccessTokenRepository - RefreshToken RefreshTokenRepository - Role RoleRepository - Permission PermissionRepository - UserRole UserRoleRepository - RolePermission RolePermissionRepository - AuditLog AuditLogRepository - EmailTemplate EmailTemplateRepository - EmailQueue EmailQueueRepository - EmailVerification EmailVerificationRepository - PasswordReset PasswordResetRepository + Tenant TenantRepository + User UserRepository + Client ClientRepository + AuthCode AuthCodeRepository + AccessToken AccessTokenRepository + RefreshToken RefreshTokenRepository + // Phase 5 repos + MFASecret MFASecretRepository + MFABackupCode MFABackupCodeRepository + Session SessionRepository + LoginAttempt LoginAttemptRepository } +// --- legacy interfaces kept for backward compatibility --- + // RoleRepository defines the interface for role data operations type RoleRepository interface { Create(ctx context.Context, role *models.Role) error @@ -118,7 +116,7 @@ type UserRoleRepository interface { DeleteExpired(ctx context.Context) error } -// RolePermissionRepository defines the interface for role-permission relationship data operations +// RolePermissionRepository defines the interface for role-permission data operations type RolePermissionRepository interface { Create(ctx context.Context, rolePermission *models.RolePermission) error GetByID(ctx context.Context, rolePermissionID uuid.UUID) (*models.RolePermission, error) diff --git a/internal/repo/interfaces_phase5.go b/internal/repo/interfaces_phase5.go new file mode 100644 index 0000000..ab31d3a --- /dev/null +++ b/internal/repo/interfaces_phase5.go @@ -0,0 +1,46 @@ +package repo + +import ( + "context" + "time" + + "shieldgate/internal/models" + + "github.com/google/uuid" +) + +// MFASecretRepository manages TOTP secrets. +type MFASecretRepository interface { + Create(ctx context.Context, secret *models.MFASecret) error + GetByUserID(ctx context.Context, tenantID, userID uuid.UUID) (*models.MFASecret, error) + Update(ctx context.Context, secret *models.MFASecret) error + Delete(ctx context.Context, tenantID, userID uuid.UUID) error +} + +// MFABackupCodeRepository manages single-use MFA backup codes. +type MFABackupCodeRepository interface { + CreateBatch(ctx context.Context, codes []*models.MFABackupCode) error + GetByUserID(ctx context.Context, tenantID, userID uuid.UUID) ([]*models.MFABackupCode, error) + GetByCodeHash(ctx context.Context, tenantID, userID uuid.UUID, hash string) (*models.MFABackupCode, error) + MarkUsed(ctx context.Context, id uuid.UUID) error + DeleteByUserID(ctx context.Context, tenantID, userID uuid.UUID) error +} + +// SessionRepository manages authenticated user sessions. +type SessionRepository interface { + Create(ctx context.Context, session *models.Session) error + GetByID(ctx context.Context, tenantID, sessionID uuid.UUID) (*models.Session, error) + GetByTokenID(ctx context.Context, tokenID string) (*models.Session, error) + ListByUser(ctx context.Context, tenantID, userID uuid.UUID, limit, offset int) ([]*models.Session, int64, error) + Update(ctx context.Context, session *models.Session) error + Revoke(ctx context.Context, tenantID, sessionID uuid.UUID) error + RevokeAll(ctx context.Context, tenantID, userID uuid.UUID) error + DeleteExpired(ctx context.Context) error +} + +// LoginAttemptRepository persists login attempt records. +type LoginAttemptRepository interface { + Create(ctx context.Context, attempt *models.LoginAttempt) error + CountRecent(ctx context.Context, tenantID uuid.UUID, ipAddress, email string, since time.Time) (int64, error) + DeleteOld(ctx context.Context, before time.Time) error +} diff --git a/internal/services/interfaces.go b/internal/services/interfaces.go index 0930e74..6eda05a 100644 --- a/internal/services/interfaces.go +++ b/internal/services/interfaces.go @@ -65,41 +65,62 @@ type ClientService interface { // AuthService defines the interface for OAuth authentication business logic type AuthService interface { - // Authorization Code Flow GenerateAuthorizationCode(ctx context.Context, tenantID, clientID, userID uuid.UUID, redirectURI, scope, codeChallenge, codeChallengeMethod string) (*models.AuthorizationCode, error) ExchangeAuthorizationCode(ctx context.Context, tenantID uuid.UUID, code, clientID, clientSecret, redirectURI, codeVerifier string) (*models.TokenResponse, error) - - // Token Management GenerateTokens(ctx context.Context, tenantID, clientID, userID uuid.UUID, scope string, includeIDToken bool) (*models.TokenResponse, error) RefreshTokens(ctx context.Context, tenantID uuid.UUID, refreshToken, clientID, clientSecret string) (*models.TokenResponse, error) RevokeToken(ctx context.Context, tenantID uuid.UUID, token, tokenTypeHint string) error IntrospectToken(ctx context.Context, tenantID uuid.UUID, token string) (*models.IntrospectionResponse, error) - - // Token Validation ValidateAccessToken(ctx context.Context, tenantID uuid.UUID, token string) (*models.JWTClaims, error) ValidatePKCE(codeVerifier, codeChallenge, method string) bool - - // OpenID Connect GenerateIDToken(ctx context.Context, user *models.User, clientID string) (string, error) GetUserInfo(ctx context.Context, tenantID uuid.UUID, accessToken string) (*models.UserInfo, error) GetDiscoveryDocument(ctx context.Context) (*models.OpenIDConfiguration, error) - - // Cleanup CleanupExpiredTokens(ctx context.Context) error } +// MFAService manages TOTP-based multi-factor authentication. +type MFAService interface { + // Setup generates a new TOTP secret for the user but does NOT enable MFA yet. + Setup(ctx context.Context, tenantID, userID uuid.UUID, issuer, accountName string) (*models.MFASetupResponse, error) + // VerifyAndEnable validates the first TOTP code and enables MFA. + VerifyAndEnable(ctx context.Context, tenantID, userID uuid.UUID, code string) error + // Disable turns off MFA after verifying either a TOTP code or backup code. + Disable(ctx context.Context, tenantID, userID uuid.UUID, code string) error + // ValidateCode returns true if the given TOTP code (or backup code) is valid. + ValidateCode(ctx context.Context, tenantID, userID uuid.UUID, code string) (bool, error) + // IsEnabled reports whether the user has MFA active. + IsEnabled(ctx context.Context, tenantID, userID uuid.UUID) (bool, error) + // GetBackupCodes returns the user's existing backup code records. + GetBackupCodes(ctx context.Context, tenantID, userID uuid.UUID) ([]*models.MFABackupCode, error) + // RegenerateBackupCodes replaces all backup codes and returns the plain-text values. + RegenerateBackupCodes(ctx context.Context, tenantID, userID uuid.UUID) ([]string, error) +} + +// SessionService manages authenticated user sessions. +type SessionService interface { + Create(ctx context.Context, tenantID, userID uuid.UUID, tokenID, ipAddress, userAgent, deviceName string, expiresAt time.Time) (*models.Session, error) + GetByID(ctx context.Context, tenantID, sessionID uuid.UUID) (*models.Session, error) + GetByTokenID(ctx context.Context, tokenID string) (*models.Session, error) + ListByUser(ctx context.Context, tenantID, userID uuid.UUID, limit, offset int) (*models.SessionListResponse, error) + Revoke(ctx context.Context, tenantID, sessionID uuid.UUID) error + RevokeAll(ctx context.Context, tenantID, userID uuid.UUID) error + UpdateLastActive(ctx context.Context, sessionID uuid.UUID) error + CleanupExpired(ctx context.Context) error +} + // Services aggregates all service interfaces type Services struct { Tenant TenantService User UserService Client ClientService Auth AuthService - Role RoleService - Permission PermissionService - Audit AuditService - Email EmailService + MFA MFAService + Session SessionService } +// --- legacy service interfaces (used by phase-2/3 handlers, kept for compilation) --- + // RoleService defines the interface for RBAC role management type RoleService interface { Create(ctx context.Context, tenantID uuid.UUID, req *models.CreateRoleRequest) (*models.Role, error) @@ -108,13 +129,9 @@ type RoleService interface { Update(ctx context.Context, tenantID, roleID uuid.UUID, req *models.UpdateRoleRequest) (*models.Role, error) Delete(ctx context.Context, tenantID, roleID uuid.UUID) error List(ctx context.Context, tenantID uuid.UUID, limit, offset int) (*models.PaginatedResponse, error) - - // Permission Management AddPermission(ctx context.Context, tenantID, roleID, permissionID, grantedBy uuid.UUID) error RemovePermission(ctx context.Context, tenantID, roleID, permissionID uuid.UUID) error GetPermissions(ctx context.Context, tenantID, roleID uuid.UUID) ([]*models.Permission, error) - - // User Role Assignment AssignToUser(ctx context.Context, tenantID, roleID, userID, grantedBy uuid.UUID, expiresAt *time.Time) error RevokeFromUser(ctx context.Context, tenantID, roleID, userID uuid.UUID) error GetUserRoles(ctx context.Context, tenantID, userID uuid.UUID) ([]*models.Role, error) @@ -128,8 +145,6 @@ type PermissionService interface { Update(ctx context.Context, permissionID uuid.UUID, req *models.UpdatePermissionRequest) (*models.Permission, error) Delete(ctx context.Context, permissionID uuid.UUID) error List(ctx context.Context, limit, offset int) (*models.PaginatedResponse, error) - - // Permission Checking HasPermission(ctx context.Context, tenantID, userID uuid.UUID, resource, action string) (bool, error) GetUserPermissions(ctx context.Context, tenantID, userID uuid.UUID) ([]*models.Permission, error) } @@ -140,8 +155,6 @@ type AuditService interface { LogUserAction(ctx context.Context, tenantID, userID uuid.UUID, action models.AuditAction, resource string, resourceID *uuid.UUID, success bool, metadata map[string]interface{}) error LogClientAction(ctx context.Context, tenantID, clientID uuid.UUID, action models.AuditAction, resource string, resourceID *uuid.UUID, success bool, metadata map[string]interface{}) error LogSystemAction(ctx context.Context, tenantID uuid.UUID, action models.AuditAction, resource string, resourceID *uuid.UUID, success bool, metadata map[string]interface{}) error - - // Query Methods Query(ctx context.Context, query *models.AuditLogQuery) (*models.PaginatedResponse, error) GetByID(ctx context.Context, tenantID, auditID uuid.UUID) (*models.AuditLog, error) GetUserActivity(ctx context.Context, tenantID, userID uuid.UUID, limit, offset int) (*models.PaginatedResponse, error) @@ -150,27 +163,18 @@ type AuditService interface { // EmailService defines the interface for email management type EmailService interface { - // Template Management CreateTemplate(ctx context.Context, tenantID uuid.UUID, template *models.EmailTemplate) error GetTemplate(ctx context.Context, tenantID uuid.UUID, name string) (*models.EmailTemplate, error) UpdateTemplate(ctx context.Context, tenantID uuid.UUID, name string, template *models.EmailTemplate) error DeleteTemplate(ctx context.Context, tenantID uuid.UUID, name string) error ListTemplates(ctx context.Context, tenantID uuid.UUID, limit, offset int) (*models.PaginatedResponse, error) - - // Email Sending SendEmail(ctx context.Context, tenantID uuid.UUID, req *models.SendEmailRequest) error SendTemplateEmail(ctx context.Context, tenantID uuid.UUID, toEmail, toName, templateName string, variables map[string]string, priority int) error - - // Queue Management ProcessQueue(ctx context.Context) error GetQueueStatus(ctx context.Context, tenantID uuid.UUID) (map[string]int, error) RetryFailedEmails(ctx context.Context, tenantID uuid.UUID, maxAttempts int) error - - // Verification Emails SendVerificationEmail(ctx context.Context, tenantID, userID uuid.UUID) error VerifyEmail(ctx context.Context, tenantID uuid.UUID, code string) (*models.User, error) - - // Password Reset Emails SendPasswordResetEmail(ctx context.Context, tenantID uuid.UUID, email string) error ResetPassword(ctx context.Context, tenantID uuid.UUID, token, newPassword string) (*models.User, error) } diff --git a/internal/services/mfa_service_impl.go b/internal/services/mfa_service_impl.go new file mode 100644 index 0000000..e065c4e --- /dev/null +++ b/internal/services/mfa_service_impl.go @@ -0,0 +1,191 @@ +package services + +import ( + "context" + "fmt" + "time" + + "shieldgate/internal/crypto" + "shieldgate/internal/models" + "shieldgate/internal/repo" + + "github.com/google/uuid" + "github.com/sirupsen/logrus" + "golang.org/x/crypto/bcrypt" +) + +const backupCodeCount = 8 + +type mfaServiceImpl struct { + repos *repo.Repositories + logger *logrus.Logger +} + +func NewMFAService(repos *repo.Repositories, logger *logrus.Logger) MFAService { + return &mfaServiceImpl{repos: repos, logger: logger} +} + +func (s *mfaServiceImpl) Setup(ctx context.Context, tenantID, userID uuid.UUID, issuer, accountName string) (*models.MFASetupResponse, error) { + // Delete any pending (non-enabled) secret first + existing, err := s.repos.MFASecret.GetByUserID(ctx, tenantID, userID) + if err == nil && existing.Enabled { + return nil, models.ErrMFAAlreadyEnabled + } + if err == nil { + _ = s.repos.MFASecret.Delete(ctx, tenantID, userID) + } + + secret, err := crypto.GenerateTOTPSecret() + if err != nil { + return nil, fmt.Errorf("failed to generate TOTP secret: %w", err) + } + + mfaSecret := &models.MFASecret{ + ID: uuid.New(), + TenantID: tenantID, + UserID: userID, + Secret: secret, + Algorithm: "SHA1", + Digits: crypto.DefaultTOTPDigits, + Period: crypto.DefaultTOTPPeriod, + Enabled: false, + } + if err := s.repos.MFASecret.Create(ctx, mfaSecret); err != nil { + return nil, fmt.Errorf("failed to store MFA secret: %w", err) + } + + provURI := crypto.TOTPProvisioningURI(issuer, accountName, secret, + crypto.DefaultTOTPPeriod, crypto.DefaultTOTPDigits) + + return &models.MFASetupResponse{ + Secret: secret, + ProvisioningURI: provURI, + QRCodeHint: "Scan the provisioning_uri with an authenticator app such as Google Authenticator or Authy.", + }, nil +} + +func (s *mfaServiceImpl) VerifyAndEnable(ctx context.Context, tenantID, userID uuid.UUID, code string) error { + mfaSecret, err := s.repos.MFASecret.GetByUserID(ctx, tenantID, userID) + if err != nil { + return models.ErrMFANotSetup + } + if mfaSecret.Enabled { + return models.ErrMFAAlreadyEnabled + } + if !crypto.ValidateTOTP(mfaSecret.Secret, code, mfaSecret.Period, mfaSecret.Digits) { + return models.ErrMFAInvalidCode + } + + mfaSecret.Enabled = true + if err := s.repos.MFASecret.Update(ctx, mfaSecret); err != nil { + return fmt.Errorf("failed to enable MFA: %w", err) + } + + // Generate initial backup codes + if _, err := s.generateAndStoreBackupCodes(ctx, tenantID, userID); err != nil { + s.logger.WithError(err).Warn("failed to generate backup codes after MFA enable") + } + return nil +} + +func (s *mfaServiceImpl) Disable(ctx context.Context, tenantID, userID uuid.UUID, code string) error { + mfaSecret, err := s.repos.MFASecret.GetByUserID(ctx, tenantID, userID) + if err != nil { + return models.ErrMFANotSetup + } + if !mfaSecret.Enabled { + return models.ErrMFANotEnabled + } + + // Accept either TOTP code or a backup code + valid := crypto.ValidateTOTP(mfaSecret.Secret, code, mfaSecret.Period, mfaSecret.Digits) + if !valid { + ok, _ := s.useBackupCode(ctx, tenantID, userID, code) + if !ok { + return models.ErrMFAInvalidCode + } + } + + _ = s.repos.MFABackupCode.DeleteByUserID(ctx, tenantID, userID) + return s.repos.MFASecret.Delete(ctx, tenantID, userID) +} + +func (s *mfaServiceImpl) ValidateCode(ctx context.Context, tenantID, userID uuid.UUID, code string) (bool, error) { + mfaSecret, err := s.repos.MFASecret.GetByUserID(ctx, tenantID, userID) + if err != nil { + return false, models.ErrMFANotSetup + } + if !mfaSecret.Enabled { + return false, models.ErrMFANotEnabled + } + if crypto.ValidateTOTP(mfaSecret.Secret, code, mfaSecret.Period, mfaSecret.Digits) { + return true, nil + } + // Try backup code + ok, err := s.useBackupCode(ctx, tenantID, userID, code) + return ok, err +} + +func (s *mfaServiceImpl) IsEnabled(ctx context.Context, tenantID, userID uuid.UUID) (bool, error) { + secret, err := s.repos.MFASecret.GetByUserID(ctx, tenantID, userID) + if err != nil { + return false, nil + } + return secret.Enabled, nil +} + +func (s *mfaServiceImpl) GetBackupCodes(ctx context.Context, tenantID, userID uuid.UUID) ([]*models.MFABackupCode, error) { + return s.repos.MFABackupCode.GetByUserID(ctx, tenantID, userID) +} + +func (s *mfaServiceImpl) RegenerateBackupCodes(ctx context.Context, tenantID, userID uuid.UUID) ([]string, error) { + if err := s.repos.MFABackupCode.DeleteByUserID(ctx, tenantID, userID); err != nil { + return nil, fmt.Errorf("failed to delete old backup codes: %w", err) + } + return s.generateAndStoreBackupCodes(ctx, tenantID, userID) +} + +// generateAndStoreBackupCodes creates new codes, stores hashed versions, returns plain-text. +func (s *mfaServiceImpl) generateAndStoreBackupCodes(ctx context.Context, tenantID, userID uuid.UUID) ([]string, error) { + plainCodes, err := crypto.GenerateBackupCodes(backupCodeCount) + if err != nil { + return nil, err + } + records := make([]*models.MFABackupCode, len(plainCodes)) + for i, c := range plainCodes { + hash, err := bcrypt.GenerateFromPassword([]byte(c), bcrypt.DefaultCost) + if err != nil { + return nil, fmt.Errorf("failed to hash backup code: %w", err) + } + records[i] = &models.MFABackupCode{ + ID: uuid.New(), + TenantID: tenantID, + UserID: userID, + CodeHash: string(hash), + } + } + if err := s.repos.MFABackupCode.CreateBatch(ctx, records); err != nil { + return nil, fmt.Errorf("failed to store backup codes: %w", err) + } + return plainCodes, nil +} + +// useBackupCode attempts to find and consume a matching backup code. +func (s *mfaServiceImpl) useBackupCode(ctx context.Context, tenantID, userID uuid.UUID, code string) (bool, error) { + codes, err := s.repos.MFABackupCode.GetByUserID(ctx, tenantID, userID) + if err != nil { + return false, err + } + for _, bc := range codes { + if bc.IsUsed() { + continue + } + if bcrypt.CompareHashAndPassword([]byte(bc.CodeHash), []byte(code)) == nil { + now := time.Now() + _ = s.repos.MFABackupCode.MarkUsed(ctx, bc.ID) + _ = now + return true, nil + } + } + return false, nil +} diff --git a/internal/services/session_service_impl.go b/internal/services/session_service_impl.go new file mode 100644 index 0000000..b434b0d --- /dev/null +++ b/internal/services/session_service_impl.go @@ -0,0 +1,108 @@ +package services + +import ( + "context" + "fmt" + "time" + + "shieldgate/internal/models" + "shieldgate/internal/repo" + + "github.com/google/uuid" + "github.com/sirupsen/logrus" +) + +type sessionServiceImpl struct { + repos *repo.Repositories + logger *logrus.Logger +} + +func NewSessionService(repos *repo.Repositories, logger *logrus.Logger) SessionService { + return &sessionServiceImpl{repos: repos, logger: logger} +} + +func (s *sessionServiceImpl) Create(ctx context.Context, tenantID, userID uuid.UUID, tokenID, ipAddress, userAgent, deviceName string, expiresAt time.Time) (*models.Session, error) { + session := &models.Session{ + ID: uuid.New(), + TenantID: tenantID, + UserID: userID, + TokenID: tokenID, + IPAddress: ipAddress, + UserAgent: userAgent, + DeviceName: deviceName, + LastActiveAt: time.Now(), + ExpiresAt: expiresAt, + } + if err := s.repos.Session.Create(ctx, session); err != nil { + return nil, fmt.Errorf("failed to create session: %w", err) + } + return session, nil +} + +func (s *sessionServiceImpl) GetByID(ctx context.Context, tenantID, sessionID uuid.UUID) (*models.Session, error) { + session, err := s.repos.Session.GetByID(ctx, tenantID, sessionID) + if err != nil { + return nil, models.ErrSessionNotFound + } + if session.IsRevoked() { + return nil, models.ErrSessionRevoked + } + if session.IsExpired() { + return nil, models.ErrSessionExpired + } + return session, nil +} + +func (s *sessionServiceImpl) GetByTokenID(ctx context.Context, tokenID string) (*models.Session, error) { + session, err := s.repos.Session.GetByTokenID(ctx, tokenID) + if err != nil { + return nil, models.ErrSessionNotFound + } + return session, nil +} + +func (s *sessionServiceImpl) ListByUser(ctx context.Context, tenantID, userID uuid.UUID, limit, offset int) (*models.SessionListResponse, error) { + if limit <= 0 || limit > 100 { + limit = 20 + } + sessions, total, err := s.repos.Session.ListByUser(ctx, tenantID, userID, limit, offset) + if err != nil { + return nil, fmt.Errorf("failed to list sessions: %w", err) + } + return &models.SessionListResponse{Sessions: sessions, Total: total}, nil +} + +func (s *sessionServiceImpl) Revoke(ctx context.Context, tenantID, sessionID uuid.UUID) error { + if err := s.repos.Session.Revoke(ctx, tenantID, sessionID); err != nil { + return fmt.Errorf("failed to revoke session: %w", err) + } + s.logger.WithFields(logrus.Fields{ + "tenant_id": tenantID, + "session_id": sessionID, + }).Info("session revoked") + return nil +} + +func (s *sessionServiceImpl) RevokeAll(ctx context.Context, tenantID, userID uuid.UUID) error { + if err := s.repos.Session.RevokeAll(ctx, tenantID, userID); err != nil { + return fmt.Errorf("failed to revoke all sessions: %w", err) + } + s.logger.WithFields(logrus.Fields{ + "tenant_id": tenantID, + "user_id": userID, + }).Info("all sessions revoked") + return nil +} + +func (s *sessionServiceImpl) UpdateLastActive(ctx context.Context, sessionID uuid.UUID) error { + session, err := s.repos.Session.GetByTokenID(ctx, sessionID.String()) + if err != nil { + return nil // best-effort, don't fail + } + session.LastActiveAt = time.Now() + return s.repos.Session.Update(ctx, session) +} + +func (s *sessionServiceImpl) CleanupExpired(ctx context.Context) error { + return s.repos.Session.DeleteExpired(ctx) +} diff --git a/internal/services/tests/mfa_service_test.go b/internal/services/tests/mfa_service_test.go new file mode 100644 index 0000000..e4f1b41 --- /dev/null +++ b/internal/services/tests/mfa_service_test.go @@ -0,0 +1,123 @@ +package tests + +import ( + "context" + "testing" + + "shieldgate/internal/crypto" + "shieldgate/internal/models" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// --- in-memory MFASecret store for unit testing --- + +type memMFASecretRepo struct { + records map[string]*models.MFASecret // key: tenantID+":"+userID +} + +func newMemMFASecretRepo() *memMFASecretRepo { + return &memMFASecretRepo{records: make(map[string]*models.MFASecret)} +} + +func (r *memMFASecretRepo) key(t, u uuid.UUID) string { return t.String() + ":" + u.String() } + +func (r *memMFASecretRepo) Create(_ context.Context, s *models.MFASecret) error { + r.records[r.key(s.TenantID, s.UserID)] = s + return nil +} +func (r *memMFASecretRepo) GetByUserID(_ context.Context, t, u uuid.UUID) (*models.MFASecret, error) { + s, ok := r.records[r.key(t, u)] + if !ok { + return nil, models.ErrMFANotSetup + } + return s, nil +} +func (r *memMFASecretRepo) Update(_ context.Context, s *models.MFASecret) error { + r.records[r.key(s.TenantID, s.UserID)] = s + return nil +} +func (r *memMFASecretRepo) Delete(_ context.Context, t, u uuid.UUID) error { + delete(r.records, r.key(t, u)) + return nil +} + +// --- tests --- + +func TestTOTP_GenerateAndValidate(t *testing.T) { + secret, err := crypto.GenerateTOTPSecret() + require.NoError(t, err) + assert.NotEmpty(t, secret) + + code, err := crypto.GenerateTOTP(secret, 30, 6) + require.NoError(t, err) + assert.Len(t, code, 6) + + assert.True(t, crypto.ValidateTOTP(secret, code, 30, 6), "freshly generated code must be valid") +} + +func TestTOTP_WrongCodeRejected(t *testing.T) { + secret, _ := crypto.GenerateTOTPSecret() + assert.False(t, crypto.ValidateTOTP(secret, "000000", 30, 6)) +} + +func TestTOTP_InvalidSecretRejected(t *testing.T) { + assert.False(t, crypto.ValidateTOTP("BAD!!!", "123456", 30, 6)) +} + +func TestBackupCodes_UniqueAndCorrectLength(t *testing.T) { + codes, err := crypto.GenerateBackupCodes(8) + require.NoError(t, err) + assert.Len(t, codes, 8) + seen := make(map[string]bool) + for _, c := range codes { + assert.Len(t, c, 10) + assert.False(t, seen[c]) + seen[c] = true + } +} + +func TestProvisioningURI_Format(t *testing.T) { + uri := crypto.TOTPProvisioningURI("Acme", "alice@acme.com", "JBSWY3DPEHPK3PXP", 30, 6) + assert.Contains(t, uri, "otpauth://totp/") + assert.Contains(t, uri, "issuer=Acme") + assert.Contains(t, uri, "secret=JBSWY3DPEHPK3PXP") +} + +func TestMFASecretRepo_CreateAndGet(t *testing.T) { + repo := newMemMFASecretRepo() + tenantID := uuid.New() + userID := uuid.New() + + secret := &models.MFASecret{ + ID: uuid.New(), + TenantID: tenantID, + UserID: userID, + Secret: "TESTSECRET", + Enabled: false, + } + require.NoError(t, repo.Create(context.Background(), secret)) + + got, err := repo.GetByUserID(context.Background(), tenantID, userID) + require.NoError(t, err) + assert.Equal(t, "TESTSECRET", got.Secret) +} + +func TestMFASecretRepo_NotFound(t *testing.T) { + repo := newMemMFASecretRepo() + _, err := repo.GetByUserID(context.Background(), uuid.New(), uuid.New()) + assert.Error(t, err) +} + +func TestMFASecretRepo_Delete(t *testing.T) { + repo := newMemMFASecretRepo() + tenantID, userID := uuid.New(), uuid.New() + _ = repo.Create(context.Background(), &models.MFASecret{ + ID: uuid.New(), TenantID: tenantID, UserID: userID, Secret: "X", + }) + require.NoError(t, repo.Delete(context.Background(), tenantID, userID)) + _, err := repo.GetByUserID(context.Background(), tenantID, userID) + assert.Error(t, err) +} diff --git a/internal/services/tests/session_service_test.go b/internal/services/tests/session_service_test.go new file mode 100644 index 0000000..d407c2e --- /dev/null +++ b/internal/services/tests/session_service_test.go @@ -0,0 +1,156 @@ +package tests + +import ( + "context" + "testing" + "time" + + "shieldgate/internal/models" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// --- in-memory Session repo --- + +type memSessionRepo struct { + records map[uuid.UUID]*models.Session +} + +func newMemSessionRepo() *memSessionRepo { + return &memSessionRepo{records: make(map[uuid.UUID]*models.Session)} +} + +func (r *memSessionRepo) Create(_ context.Context, s *models.Session) error { + r.records[s.ID] = s + return nil +} +func (r *memSessionRepo) GetByID(_ context.Context, _, id uuid.UUID) (*models.Session, error) { + s, ok := r.records[id] + if !ok { + return nil, models.ErrSessionNotFound + } + return s, nil +} +func (r *memSessionRepo) GetByTokenID(_ context.Context, tokenID string) (*models.Session, error) { + for _, s := range r.records { + if s.TokenID == tokenID { + return s, nil + } + } + return nil, models.ErrSessionNotFound +} +func (r *memSessionRepo) ListByUser(_ context.Context, tenantID, userID uuid.UUID, limit, offset int) ([]*models.Session, int64, error) { + var out []*models.Session + for _, s := range r.records { + if s.TenantID == tenantID && s.UserID == userID { + out = append(out, s) + } + } + return out, int64(len(out)), nil +} +func (r *memSessionRepo) Update(_ context.Context, s *models.Session) error { + r.records[s.ID] = s + return nil +} +func (r *memSessionRepo) Revoke(_ context.Context, _, id uuid.UUID) error { + if s, ok := r.records[id]; ok { + now := time.Now() + s.RevokedAt = &now + } + return nil +} +func (r *memSessionRepo) RevokeAll(_ context.Context, tenantID, userID uuid.UUID) error { + for _, s := range r.records { + if s.TenantID == tenantID && s.UserID == userID { + now := time.Now() + s.RevokedAt = &now + } + } + return nil +} +func (r *memSessionRepo) DeleteExpired(_ context.Context) error { return nil } + +// --- tests --- + +func TestSession_IsActive(t *testing.T) { + s := &models.Session{ + ExpiresAt: time.Now().Add(time.Hour), + RevokedAt: nil, + } + assert.True(t, s.IsActive()) +} + +func TestSession_IsExpired(t *testing.T) { + s := &models.Session{ExpiresAt: time.Now().Add(-time.Hour)} + assert.True(t, s.IsExpired()) + assert.False(t, s.IsActive()) +} + +func TestSession_IsRevoked(t *testing.T) { + now := time.Now() + s := &models.Session{ExpiresAt: time.Now().Add(time.Hour), RevokedAt: &now} + assert.True(t, s.IsRevoked()) + assert.False(t, s.IsActive()) +} + +func TestSessionRepo_CreateAndGet(t *testing.T) { + repo := newMemSessionRepo() + tenantID, userID := uuid.New(), uuid.New() + session := &models.Session{ + ID: uuid.New(), + TenantID: tenantID, + UserID: userID, + TokenID: "tok-abc", + ExpiresAt: time.Now().Add(time.Hour), + } + require.NoError(t, repo.Create(context.Background(), session)) + + got, err := repo.GetByID(context.Background(), tenantID, session.ID) + require.NoError(t, err) + assert.Equal(t, "tok-abc", got.TokenID) +} + +func TestSessionRepo_GetByTokenID(t *testing.T) { + repo := newMemSessionRepo() + session := &models.Session{ + ID: uuid.New(), + TenantID: uuid.New(), + UserID: uuid.New(), + TokenID: "unique-token-id", + ExpiresAt: time.Now().Add(time.Hour), + } + _ = repo.Create(context.Background(), session) + got, err := repo.GetByTokenID(context.Background(), "unique-token-id") + require.NoError(t, err) + assert.Equal(t, session.ID, got.ID) +} + +func TestSessionRepo_Revoke(t *testing.T) { + repo := newMemSessionRepo() + tenantID, sessionID := uuid.New(), uuid.New() + _ = repo.Create(context.Background(), &models.Session{ + ID: sessionID, TenantID: tenantID, UserID: uuid.New(), + TokenID: "t1", ExpiresAt: time.Now().Add(time.Hour), + }) + require.NoError(t, repo.Revoke(context.Background(), tenantID, sessionID)) + got, _ := repo.GetByID(context.Background(), tenantID, sessionID) + assert.NotNil(t, got.RevokedAt) +} + +func TestSessionRepo_RevokeAll(t *testing.T) { + repo := newMemSessionRepo() + tenantID, userID := uuid.New(), uuid.New() + for i := 0; i < 3; i++ { + _ = repo.Create(context.Background(), &models.Session{ + ID: uuid.New(), TenantID: tenantID, UserID: userID, + TokenID: uuid.New().String(), ExpiresAt: time.Now().Add(time.Hour), + }) + } + require.NoError(t, repo.RevokeAll(context.Background(), tenantID, userID)) + sessions, _, _ := repo.ListByUser(context.Background(), tenantID, userID, 10, 0) + for _, s := range sessions { + assert.NotNil(t, s.RevokedAt, "session %s should be revoked", s.ID) + } +}