refactor: rework rate limit logic (#1008)

.env.example

@@ -101,8 +101,6 @@ TINYAUTH_AUTH_SESSIONMAXLIFETIME=0
 TINYAUTH_AUTH_LOGINTIMEOUT=300
 # Maximum login retries.
 TINYAUTH_AUTH_LOGINMAXRETRIES=3
-# Enable lockdown mode after maximum login retries. Lockdown mode limit is calculated automatically.
-TINYAUTH_AUTH_LOCKDOWNENABLED=true
 # Comma-separated list of trusted proxy addresses.
 TINYAUTH_AUTH_TRUSTEDPROXIES=
 # ACL policy for allow-by-default or deny-by-default, available options are allow and deny, default is allow.

internal/controller/user_controller.go

@@ -72,27 +72,12 @@ func (controller *UserController) loginHandler(c *gin.Context) {
 
 	controller.log.App.Debug().Str("username", req.Username).Msg("Login attempt")
 
-	isLocked, remaining := controller.auth.IsAccountLocked(req.Username)
-
-	if isLocked {
-		controller.log.App.Warn().Str("username", req.Username).Msg("Account is locked due to too many failed login attempts")
-		controller.log.AuditLoginFailure(req.Username, "local", c.ClientIP(), "account locked")
-		c.Writer.Header().Add("x-tinyauth-lock-locked", "true")
-		c.Writer.Header().Add("x-tinyauth-lock-reset", time.Now().Add(time.Duration(remaining)*time.Second).Format(time.RFC3339))
-		c.JSON(429, gin.H{
-			"status":  429,
-			"message": fmt.Sprintf("Too many failed login attempts. Try again in %d seconds", remaining),
-		})
-		return
-	}
-
 	search, err := controller.auth.SearchUser(req.Username)
 
 	if err != nil {
 		if errors.Is(err, service.ErrUserNotFound) {
 			controller.auth.DummyPasswordCheck(req.Password)
 			controller.log.App.Warn().Str("username", req.Username).Msg("User not found during login attempt")
-			controller.auth.RecordLoginAttempt(req.Username, false)
 			controller.log.AuditLoginFailure(req.Username, "unknown", c.ClientIP(), "user not found")
 			c.JSON(401, gin.H{
 				"status":  401,
@@ -108,14 +93,24 @@ func (controller *UserController) loginHandler(c *gin.Context) {
 		return
 	}
 
+	isLocked, remaining := controller.auth.IsAccountLocked(req.Username)
+
+	if isLocked {
+		controller.log.App.Warn().Str("username", req.Username).Msg("Account is locked due to too many failed login attempts")
+		controller.log.AuditLoginFailure(req.Username, search.Type.String(), c.ClientIP(), "account locked")
+		c.Writer.Header().Add("x-tinyauth-lock-locked", "true")
+		c.Writer.Header().Add("x-tinyauth-lock-reset", time.Now().Add(time.Duration(remaining)*time.Second).Format(time.RFC3339))
+		c.JSON(429, gin.H{
+			"status":  429,
+			"message": fmt.Sprintf("Too many failed login attempts. Try again in %d seconds", remaining),
+		})
+		return
+	}
+
 	if err := controller.auth.CheckUserPassword(*search, req.Password); err != nil {
 		controller.log.App.Warn().Str("username", req.Username).Msg("Invalid password during login attempt")
 		controller.auth.RecordLoginAttempt(req.Username, false)
-		if search.Type == model.UserLocal {
+		controller.log.AuditLoginFailure(req.Username, search.Type.String(), c.ClientIP(), "invalid password")
-			controller.log.AuditLoginFailure(req.Username, "local", c.ClientIP(), "invalid password")
-		} else {
-			controller.log.AuditLoginFailure(req.Username, "ldap", c.ClientIP(), "invalid password")
-		}
 		c.JSON(401, gin.H{
 			"status":  401,
 			"message": "Unauthorized",
@@ -216,11 +211,7 @@ func (controller *UserController) loginHandler(c *gin.Context) {
 
 	controller.log.App.Info().Str("username", req.Username).Msg("Login successful")
 
-	if search.Type == model.UserLocal {
+	controller.log.AuditLoginSuccess(req.Username, search.Type.String(), c.ClientIP())
-		controller.log.AuditLoginSuccess(req.Username, "local", c.ClientIP())
-	} else {
-		controller.log.AuditLoginSuccess(req.Username, "ldap", c.ClientIP())
-	}
 
 	controller.auth.RecordLoginAttempt(req.Username, true)
 

internal/model/config.go

@@ -49,7 +49,6 @@ func NewDefaultConfiguration(runtimeEnv RuntimeEnv) *Config {
 			ACLs: ACLsConfig{
 				Policy: "allow",
 			},
-			LockdownEnabled: true,
 		},
 		UI: UIConfig{
 			Title:                 "Tinyauth",
@@ -151,7 +150,6 @@ type AuthConfig struct {
 	SessionMaxLifetime int                       `description:"Maximum session lifetime in seconds." yaml:"sessionMaxLifetime,omitempty"`
 	LoginTimeout       int                       `description:"Login timeout in seconds." yaml:"loginTimeout,omitempty"`
 	LoginMaxRetries    int                       `description:"Maximum login retries." yaml:"loginMaxRetries,omitempty"`
-	LockdownEnabled    bool                      `description:"Enable lockdown mode after maximum login retries. Lockdown mode limit is calculated automatically." yaml:"lockdownEnabled,omitempty"`
 	TrustedProxies     []string                  `description:"Comma-separated list of trusted proxy addresses." yaml:"trustedProxies,omitempty"`
 	ACLs               ACLsConfig                `description:"ACLs configuration." yaml:"acls,omitempty"`
 }

internal/model/users.go

@@ -7,6 +7,16 @@ const (
 	UserLDAP
 )
 
+func (t UserSearchType) String() string {
+	switch t {
+	case UserLocal:
+		return "local"
+	case UserLDAP:
+		return "ldap"
+	}
+	return "unknown"
+}
+
 type LDAPUser struct {
 	DN     string
 	Groups []string

internal/service/auth_service.go

@@ -8,7 +8,6 @@ import (
 	"math/big"
 	"net/http"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/steveiliop56/ding"
@@ -71,21 +70,11 @@ type AuthService struct {
 
 	dummyHash string
 
-	lockdown struct {
-		active     bool
-		until      time.Time
-		ctx        context.Context
-		cancelFunc context.CancelFunc
-		mu         sync.RWMutex
-	}
-
 	caches struct {
 		login *CacheStore[LoginAttempt]
 		oauth *CacheStore[OAuthPendingSession]
 		ldap  *CacheStore[[]string]
 	}
-
-	maxLoginLimits int
 }
 
 type AuthServiceInput struct {
@@ -116,15 +105,6 @@ func NewAuthService(i AuthServiceInput) (*AuthService, error) {
 		policyEngine: i.PolicyEngine,
 	}
 
-	// get the max login limits based on the number of users and the configured max retries
-	service.maxLoginLimits = service.calculateLockdownLimit()
-
-	loginCacheSize := 0
-
-	if !service.config.Auth.LockdownEnabled {
-		loginCacheSize = service.maxLoginLimits
-	}
-
 	// dummy hash
 	dummyHash, err := bcrypt.GenerateFromPassword([]byte(utils.GenerateString(8)), bcrypt.DefaultCost)
 
@@ -136,7 +116,7 @@ func NewAuthService(i AuthServiceInput) (*AuthService, error) {
 
 	// caches setup
 	oauthCache := NewCacheStore[OAuthPendingSession](256)
-	loginCache := NewCacheStore[LoginAttempt](loginCacheSize)
+	loginCache := NewCacheStore[LoginAttempt](service.calculateLockdownLimit())
 	ldapCache := NewCacheStore[[]string](1024)
 
 	service.caches.oauth = oauthCache
@@ -159,6 +139,23 @@ func NewAuthService(i AuthServiceInput) (*AuthService, error) {
 		}
 	}, ding.RingMinor)
 
+	i.Ding.Go(func(ctx context.Context) {
+		ticker := time.NewTicker(15 * time.Minute)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ticker.C:
+				service.log.App.Debug().Msg("Updating login cache limits")
+				service.caches.login.SetMaxSize(service.calculateLockdownLimit())
+				service.log.App.Debug().Msg("Login cache limits updated")
+			case <-ctx.Done():
+				return
+			}
+		}
+
+	}, ding.RingMinor)
+
 	return service, nil
 }
 
@@ -260,10 +257,6 @@ func (auth *AuthService) GetLDAPUser(userDN string) (*model.LDAPUser, error) {
 }
 
 func (auth *AuthService) IsAccountLocked(identifier string) (bool, int) {
-	if locked, remaining := auth.IsInLockdown(); locked {
-		return true, remaining
-	}
-
 	if auth.config.Auth.LoginMaxRetries <= 0 || auth.config.Auth.LoginTimeout <= 0 {
 		return false, 0
 	}
@@ -286,14 +279,6 @@ func (auth *AuthService) RecordLoginAttempt(identifier string, success bool) {
 		return
 	}
 
-	if !success && auth.config.Auth.LockdownEnabled && auth.caches.login.Size() >= auth.maxLoginLimits {
-		if locked, _ := auth.IsInLockdown(); locked {
-			return
-		}
-		go auth.lockdownMode()
-		return
-	}
-
 	auth.caches.login.WithLock(func(actions CacheStoreActions[LoginAttempt]) {
 		entry, ok := actions.Get(identifier)
 
@@ -359,7 +344,7 @@ func (auth *AuthService) CreateSession(ctx context.Context, data repository.Sess
 		return nil, fmt.Errorf("tailscale service not configured, cannot create session for tailscale user")
 	}
 
-	uuid, err := uuid.NewRandom()
+	u, err := uuid.NewRandom()
 
 	if err != nil {
 		return nil, fmt.Errorf("failed to generate session uuid: %w", err)
@@ -376,7 +361,7 @@ func (auth *AuthService) CreateSession(ctx context.Context, data repository.Sess
 	expiresAt := time.Now().Add(time.Duration(expiry) * time.Second)
 
 	session := repository.CreateSessionParams{
-		UUID:        uuid.String(),
+		UUID:        u.String(),
 		Username:    data.Username,
 		Email:       data.Email,
 		Name:        data.Name,
@@ -631,62 +616,7 @@ func (auth *AuthService) GetOAuthPendingSession(sessionId string) (*OAuthPending
 	return &session, nil
 }
 
-func (auth *AuthService) lockdownMode() {
+// ClearLoginAttempts is a testing function, not useful for anything else
-	auth.lockdown.mu.Lock()
-
-	if auth.lockdown.active {
-		auth.lockdown.mu.Unlock()
-		return
-	}
-
-	ctx, cancel := context.WithCancel(auth.ctx)
-
-	auth.log.App.Warn().Msg("Too many failed login attempts, entering lockdown mode")
-
-	auth.lockdown.active = true
-	auth.lockdown.ctx = ctx
-	auth.lockdown.cancelFunc = cancel
-
-	d := time.Duration(auth.config.Auth.LoginTimeout) * time.Second
-	auth.lockdown.until = time.Now().Add(d)
-	timer := time.NewTimer(d)
-
-	auth.lockdown.mu.Unlock()
-
-	defer cancel()
-	defer timer.Stop()
-
-	select {
-	case <-timer.C:
-		// Timer expired, end lockdown
-	case <-ctx.Done():
-		// Context cancelled, end lockdown
-	}
-
-	auth.lockdown.mu.Lock()
-
-	auth.log.App.Info().Msg("Exiting lockdown mode")
-
-	auth.caches.login.Clear()
-	auth.lockdown.active = false
-	auth.lockdown.until = time.Time{}
-	auth.lockdown.ctx = nil
-	auth.lockdown.cancelFunc = nil
-
-	auth.lockdown.mu.Unlock()
-}
-
-func (auth *AuthService) IsInLockdown() (bool, int) {
-	auth.lockdown.mu.RLock()
-	defer auth.lockdown.mu.RUnlock()
-	if auth.lockdown.active {
-		remaining := int(time.Until(auth.lockdown.until).Seconds())
-		return true, remaining
-	}
-	return false, 0
-}
-
-// mostly a testing function, not useful for anything else
 func (auth *AuthService) ClearLoginAttempts() {
 	auth.caches.login.Clear()
 }

internal/service/cache_store.go

@@ -195,3 +195,19 @@ func (cs *CacheStore[T]) Clear() {
 	cs.cache = make(map[string]cacheEntry[T])
 	cs.order = make([]string, 0)
 }
+
+func (cs *CacheStore[T]) SetMaxSize(maxSize int) {
+	cs.mu.Lock()
+	defer cs.mu.Unlock()
+	cs.maxSize = maxSize
+	for len(cs.cache) > maxSize {
+		if !cs.evictOne() {
+			break
+		}
+	}
+}
+func (cs *CacheStore[T]) GetMaxSize() int {
+	cs.mu.Lock()
+	defer cs.mu.Unlock()
+	return cs.maxSize
+}

internal/service/cache_store_test.go

@@ -316,6 +316,21 @@ func TestCacheStoreSizeAndClear(t *testing.T) {
 	assert.False(t, ok)
 }
 
+func TestCacheStoreWithMaxSize(t *testing.T) {
+	cs := NewCacheStore[string](0)
+	assert.Equal(t, 0, cs.Size())
+
+	for i := 0; i < 100; i++ {
+		cs.Set(strconv.Itoa(i), strconv.Itoa(i), 0)
+	}
+
+	assert.Equal(t, 100, cs.Size())
+
+	cs.SetMaxSize(10)
+
+	assert.Equal(t, 10, cs.Size())
+}
+
 func TestCacheStoreWithLock(t *testing.T) {
 	cs := NewCacheStore[int](0)
 	cs.Set("counter", 1, 0)