feat: add brute force protection (#59)

* feat: add brute force protection

* fix: bind flags to env

---------

Co-authored-by: Stavros <steveiliop56@gmail.com>

.env.example

@@ -26,5 +26,7 @@ DISABLE_CONTINUE=false
 OAUTH_WHITELIST=
 GENERIC_NAME=My OAuth
 SESSION_EXPIRY=7200
+LOGIN_TIMEOUT=300
+LOGIN_MAX_RETRIES=5
 LOG_LEVEL=0
 APP_TITLE=Tinyauth SSO

cmd/root.go

@@ -121,7 +121,7 @@ var rootCmd = &cobra.Command{
 		HandleError(err, "Failed to initialize docker")
 
 		// Create auth service
-		auth := auth.NewAuth(docker, users, oauthWhitelist, config.SessionExpiry)
+		auth := auth.NewAuth(docker, users, oauthWhitelist, config.SessionExpiry, config.LoginTimeout, config.LoginMaxRetries)
 
 		// Create OAuth providers service
 		providers := providers.NewProviders(oauthConfig)
@@ -198,6 +198,8 @@ func init() {
 	rootCmd.Flags().Bool("disable-continue", false, "Disable continue screen and redirect to app directly.")
 	rootCmd.Flags().String("oauth-whitelist", "", "Comma separated list of email addresses to whitelist when using OAuth.")
 	rootCmd.Flags().Int("session-expiry", 86400, "Session (cookie) expiration time in seconds.")
+	rootCmd.Flags().Int("login-timeout", 300, "Login timeout in seconds after max retries reached (0 to disable).")
+	rootCmd.Flags().Int("login-max-retries", 5, "Maximum login attempts before timeout (0 to disable).")
 	rootCmd.Flags().Int("log-level", 1, "Log level.")
 	rootCmd.Flags().String("app-title", "Tinyauth", "Title of the app.")
 
@@ -232,6 +234,8 @@ func init() {
 	viper.BindEnv("session-expiry", "SESSION_EXPIRY")
 	viper.BindEnv("log-level", "LOG_LEVEL")
 	viper.BindEnv("app-title", "APP_TITLE")
+	viper.BindEnv("login-timeout", "LOGIN_TIMEOUT")
+	viper.BindEnv("login-max-retries", "LOGIN_MAX_RETRIES")
 
 	// Bind flags to viper
 	viper.BindPFlags(rootCmd.Flags())

internal/api/api_test.go

@@ -66,7 +66,7 @@ func getAPI(t *testing.T) *api.API {
 			Username: user.Username,
 			Password: user.Password,
 		},
-	}, nil, apiConfig.SessionExpiry)
+	}, nil, apiConfig.SessionExpiry, 300, 5)
 
 	// Create providers service
 	providers := providers.NewProviders(types.OAuthConfig{})

internal/auth/auth.go

@@ -4,6 +4,7 @@ import (
 	"regexp"
 	"slices"
 	"strings"
+	"sync"
 	"time"
 	"tinyauth/internal/docker"
 	"tinyauth/internal/types"
@@ -14,20 +15,34 @@ import (
 	"golang.org/x/crypto/bcrypt"
 )
 
-func NewAuth(docker *docker.Docker, userList types.Users, oauthWhitelist []string, sessionExpiry int) *Auth {
+func NewAuth(docker *docker.Docker, userList types.Users, oauthWhitelist []string, sessionExpiry int, loginTimeout int, loginMaxRetries int) *Auth {
 	return &Auth{
 		Docker:         docker,
 		Users:          userList,
 		OAuthWhitelist: oauthWhitelist,
 		SessionExpiry:  sessionExpiry,
+		LoginTimeout:   loginTimeout,
+		LoginMaxRetries: loginMaxRetries,
+		LoginAttempts:  make(map[string]*LoginAttempt),
 	}
 }
 
+// LoginAttempt tracks information about login attempts for rate limiting
+type LoginAttempt struct {
+	FailedAttempts int
+	LastAttempt    time.Time
+	LockedUntil    time.Time
+}
+
 type Auth struct {
 	Users          types.Users
 	Docker         *docker.Docker
 	OAuthWhitelist []string
 	SessionExpiry  int
+	LoginTimeout   int
+	LoginMaxRetries int
+	LoginAttempts  map[string]*LoginAttempt // Map of username/IP to login attempts
+	LoginMutex     sync.RWMutex            // Mutex to protect the LoginAttempts map
 }
 
 func (auth *Auth) GetUser(username string) *types.User {
@@ -45,6 +60,70 @@ func (auth *Auth) CheckPassword(user types.User, password string) bool {
 	return bcrypt.CompareHashAndPassword([]byte(user.Password), []byte(password)) == nil
 }
 
+// IsAccountLocked checks if a username or IP is locked due to too many failed login attempts
+func (auth *Auth) IsAccountLocked(identifier string) (bool, int) {
+	auth.LoginMutex.RLock()
+	defer auth.LoginMutex.RUnlock()
+	
+	// Return false if rate limiting is not configured
+	if auth.LoginMaxRetries <= 0 || auth.LoginTimeout <= 0 {
+		return false, 0
+	}
+	
+	// Check if the identifier exists in the map
+	attempt, exists := auth.LoginAttempts[identifier]
+	if !exists {
+		return false, 0
+	}
+	
+	// If account is locked, check if lock time has expired
+	if attempt.LockedUntil.After(time.Now()) {
+		// Calculate remaining lockout time in seconds
+		remaining := int(time.Until(attempt.LockedUntil).Seconds())
+		return true, remaining
+	}
+	
+	// Lock has expired
+	return false, 0
+}
+
+// RecordLoginAttempt records a login attempt for rate limiting
+func (auth *Auth) RecordLoginAttempt(identifier string, success bool) {
+	// Skip if rate limiting is not configured
+	if auth.LoginMaxRetries <= 0 || auth.LoginTimeout <= 0 {
+		return
+	}
+	
+	auth.LoginMutex.Lock()
+	defer auth.LoginMutex.Unlock()
+	
+	// Get current attempt record or create a new one
+	attempt, exists := auth.LoginAttempts[identifier]
+	if !exists {
+		attempt = &LoginAttempt{}
+		auth.LoginAttempts[identifier] = attempt
+	}
+	
+	// Update last attempt time
+	attempt.LastAttempt = time.Now()
+	
+	// If successful login, reset failed attempts
+	if success {
+		attempt.FailedAttempts = 0
+		attempt.LockedUntil = time.Time{} // Reset lock time
+		return
+	}
+	
+	// Increment failed attempts
+	attempt.FailedAttempts++
+	
+	// If max retries reached, lock the account
+	if attempt.FailedAttempts >= auth.LoginMaxRetries {
+		attempt.LockedUntil = time.Now().Add(time.Duration(auth.LoginTimeout) * time.Second)
+		log.Warn().Str("identifier", identifier).Int("timeout", auth.LoginTimeout).Msg("Account locked due to too many failed login attempts")
+	}
+}
+
 func (auth *Auth) EmailWhitelisted(emailSrc string) bool {
 	// If the whitelist is empty, allow all emails
 	if len(auth.OAuthWhitelist) == 0 {

internal/auth/auth_test.go

@@ -0,0 +1,120 @@
+package auth_test
+
+import (
+	"testing"
+	"time"
+	"tinyauth/internal/auth"
+	"tinyauth/internal/docker"
+	"tinyauth/internal/types"
+)
+
+func TestLoginRateLimiting(t *testing.T) {
+	// Initialize a new auth service with 3 max retries and 5 seconds timeout
+	authService := auth.NewAuth(&docker.Docker{}, types.Users{}, []string{}, 3600, 5, 3)
+
+	// Test identifier
+	identifier := "test_user"
+
+	// Test successful login - should not lock account
+	t.Log("Testing successful login")
+	authService.RecordLoginAttempt(identifier, true)
+	locked, _ := authService.IsAccountLocked(identifier)
+	if locked {
+		t.Fatalf("Account should not be locked after successful login")
+	}
+
+	// Test 2 failed attempts - should not lock account yet
+	t.Log("Testing 2 failed login attempts")
+	authService.RecordLoginAttempt(identifier, false)
+	authService.RecordLoginAttempt(identifier, false)
+	locked, _ = authService.IsAccountLocked(identifier)
+	if locked {
+		t.Fatalf("Account should not be locked after only 2 failed attempts")
+	}
+
+	// Add one more failed attempt (total 3) - should lock account with maxRetries=3
+	t.Log("Testing 3 failed login attempts")
+	authService.RecordLoginAttempt(identifier, false)
+	locked, remainingTime := authService.IsAccountLocked(identifier)
+	if !locked {
+		t.Fatalf("Account should be locked after reaching max retries")
+	}
+	if remainingTime <= 0 || remainingTime > 5 {
+		t.Fatalf("Expected remaining time between 1-5 seconds, got %d", remainingTime)
+	}
+
+	// Test reset after waiting for timeout - use 1 second timeout for fast testing
+	t.Log("Testing unlocking after timeout")
+	// Create a new service for this test with very short timeout (1 second)
+	authService2 := auth.NewAuth(&docker.Docker{}, types.Users{}, []string{}, 3600, 1, 3)
+	// Add enough failed attempts to lock the account
+	for i := 0; i < 3; i++ {
+		authService2.RecordLoginAttempt(identifier, false)
+	}
+	// Verify it's locked
+	locked, _ = authService2.IsAccountLocked(identifier)
+	if !locked {
+		t.Fatalf("Account should be locked initially")
+	}
+	
+	// Wait a bit and verify it gets unlocked after timeout
+	time.Sleep(1500 * time.Millisecond) // Wait longer than the timeout
+	locked, _ = authService2.IsAccountLocked(identifier)
+	if locked {
+		t.Fatalf("Account should be unlocked after timeout period")
+	}
+
+	// Test disabled rate limiting
+	t.Log("Testing disabled rate limiting")
+	authDisabled := auth.NewAuth(&docker.Docker{}, types.Users{}, []string{}, 3600, 0, 0)
+	for i := 0; i < 10; i++ {
+		authDisabled.RecordLoginAttempt(identifier, false)
+	}
+	locked, _ = authDisabled.IsAccountLocked(identifier)
+	if locked {
+		t.Fatalf("Account should not be locked when rate limiting is disabled")
+	}
+}
+
+func TestConcurrentLoginAttempts(t *testing.T) {
+	// Initialize a new auth service with 2 max retries and 5 seconds timeout
+	authService := auth.NewAuth(&docker.Docker{}, types.Users{}, []string{}, 3600, 5, 2)
+
+	// Test multiple identifiers
+	identifiers := []string{"user1", "user2", "user3"}
+
+	// Test that locking one identifier doesn't affect others
+	t.Log("Testing multiple identifiers")
+	
+	// Add enough failed attempts to lock first user (2 attempts with maxRetries=2)
+	authService.RecordLoginAttempt(identifiers[0], false)
+	authService.RecordLoginAttempt(identifiers[0], false)
+	
+	// Check if first user is locked
+	locked, _ := authService.IsAccountLocked(identifiers[0])
+	if !locked {
+		t.Fatalf("User1 should be locked after reaching max retries")
+	}
+	
+	// Check that other users are not affected
+	for i := 1; i < len(identifiers); i++ {
+		locked, _ := authService.IsAccountLocked(identifiers[i])
+		if locked {
+			t.Fatalf("User%d should not be locked", i+1)
+		}
+	}
+
+	// Test successful login after failed attempts (but before lock)
+	t.Log("Testing successful login after failed attempts but before lock")
+	// One failed attempt for user2
+	authService.RecordLoginAttempt(identifiers[1], false)
+	// Successful login should reset the counter
+	authService.RecordLoginAttempt(identifiers[1], true)
+	
+	// Now try a failed login again - should not be locked as counter was reset
+	authService.RecordLoginAttempt(identifiers[1], false)
+	locked, _ = authService.IsAccountLocked(identifiers[1])
+	if locked {
+		t.Fatalf("User2 should not be locked after successful login reset")
+	}
+}

internal/handlers/handlers.go

@@ -248,6 +248,26 @@ func (h *Handlers) LoginHandler(c *gin.Context) {
 	}
 
 	log.Debug().Msg("Got login request")
+	
+	// Get client IP for rate limiting
+	clientIP := c.ClientIP()
+	
+	// Create an identifier for rate limiting (username or IP if username doesn't exist yet)
+	rateIdentifier := login.Username
+	if rateIdentifier == "" {
+		rateIdentifier = clientIP
+	}
+	
+	// Check if the account is locked due to too many failed attempts
+	locked, remainingTime := h.Auth.IsAccountLocked(rateIdentifier)
+	if locked {
+		log.Warn().Str("identifier", rateIdentifier).Int("remaining_seconds", remainingTime).Msg("Account is locked due to too many failed login attempts")
+		c.JSON(429, gin.H{
+			"status":  429,
+			"message": fmt.Sprintf("Too many failed login attempts. Try again in %d seconds", remainingTime),
+		})
+		return
+	}
 
 	// Get user based on username
 	user := h.Auth.GetUser(login.Username)
@@ -255,6 +275,8 @@ func (h *Handlers) LoginHandler(c *gin.Context) {
 	// User does not exist
 	if user == nil {
 		log.Debug().Str("username", login.Username).Msg("User not found")
+		// Record failed login attempt
+		h.Auth.RecordLoginAttempt(rateIdentifier, false)
 		c.JSON(401, gin.H{
 			"status":  401,
 			"message": "Unauthorized",
@@ -267,6 +289,8 @@ func (h *Handlers) LoginHandler(c *gin.Context) {
 	// Check if password is correct
 	if !h.Auth.CheckPassword(*user, login.Password) {
 		log.Debug().Str("username", login.Username).Msg("Password incorrect")
+		// Record failed login attempt
+		h.Auth.RecordLoginAttempt(rateIdentifier, false)
 		c.JSON(401, gin.H{
 			"status":  401,
 			"message": "Unauthorized",
@@ -275,6 +299,9 @@ func (h *Handlers) LoginHandler(c *gin.Context) {
 	}
 
 	log.Debug().Msg("Password correct, checking totp")
+	
+	// Record successful login attempt (will reset failed attempt counter)
+	h.Auth.RecordLoginAttempt(rateIdentifier, true)
 
 	// Check if user has totp enabled
 	if user.TotpSecret != "" {

internal/types/types.go

@@ -56,6 +56,8 @@ type Config struct {
 	LogLevel                  int8   `mapstructure:"log-level" validate:"min=-1,max=5"`
 	Title                     string `mapstructure:"app-title"`
 	EnvFile                   string `mapstructure:"env-file"`
+	LoginTimeout              int    `mapstructure:"login-timeout"`
+	LoginMaxRetries           int    `mapstructure:"login-max-retries"`
 }
 
 // UserContext is the context for the user