Merge branch 'main' into feat/oidc-preserve-consent

internal/service/auth_service.go

@@ -2,8 +2,10 @@ package service
 
 import (
 	"context"
+	"crypto/rand"
 	"errors"
 	"fmt"
+	"math/big"
 	"net/http"
 	"strings"
 	"sync"
@@ -14,6 +16,7 @@ import (
 	"github.com/tinyauthapp/tinyauth/internal/repository"
 	"github.com/tinyauthapp/tinyauth/internal/utils"
 	"github.com/tinyauthapp/tinyauth/internal/utils/logger"
+	"go.uber.org/dig"
 
 	"github.com/google/uuid"
 	"golang.org/x/crypto/bcrypt"
@@ -24,7 +27,6 @@ import (
 // but for now these are just safety limits to prevent unbounded memory usage
 const MaxOAuthPendingSessions = 256
 const OAuthCleanupCount = 16
-const MaxLoginAttemptRecords = 256
 
 var (
 	ErrUserNotFound = errors.New("user not found")
@@ -57,9 +59,8 @@ type LoginAttempt struct {
 
 type AuthService struct {
 	log     *logger.Logger
-	config  model.Config
-	runtime model.RuntimeConfig
-	helpers *model.RuntimeHelpers
+	config  *model.Config
+	runtime *model.RuntimeConfig
 	ctx     context.Context
 
 	ldap         *LdapService
@@ -81,44 +82,57 @@ type AuthService struct {
 		oauth *CacheStore[OAuthPendingSession]
 		ldap  *CacheStore[[]string]
 	}
+
+	maxLoginLimits int
 }
 
-func NewAuthService(
-	log *logger.Logger,
-	config model.Config,
-	runtime model.RuntimeConfig,
-	helpers *model.RuntimeHelpers,
-	ctx context.Context,
-	dg *ding.Ding,
-	ldap *LdapService,
-	queries repository.Store,
-	oauthBroker *OAuthBrokerService,
-	tailscale *TailscaleService,
-	policy *PolicyEngine,
-) *AuthService {
+type AuthServiceInput struct {
+	dig.In
+
+	Log          *logger.Logger
+	Config       *model.Config
+	Runtime      *model.RuntimeConfig
+	Ctx          context.Context
+	Ding         *ding.Ding
+	LDAP         *LdapService `optional:"true"`
+	Queries      repository.Store
+	OAuthBroker  *OAuthBrokerService
+	Tailscale    *TailscaleService `optional:"true"`
+	PolicyEngine *PolicyEngine
+}
+
+func NewAuthService(i AuthServiceInput) *AuthService {
 	service := &AuthService{
-		log:          log,
-		runtime:      runtime,
-		helpers:      helpers,
-		ctx:          ctx,
-		config:       config,
-		ldap:         ldap,
-		queries:      queries,
-		oauthBroker:  oauthBroker,
-		tailscale:    tailscale,
-		policyEngine: policy,
+		log:          i.Log,
+		runtime:      i.Runtime,
+		ctx:          i.Ctx,
+		config:       i.Config,
+		ldap:         i.LDAP,
+		queries:      i.Queries,
+		oauthBroker:  i.OAuthBroker,
+		tailscale:    i.Tailscale,
+		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
 	}
 
 	// caches setup
 	oauthCache := NewCacheStore[OAuthPendingSession](256)
-	loginCache := NewCacheStore[LoginAttempt](1024)
+	loginCache := NewCacheStore[LoginAttempt](loginCacheSize)
 	ldapCache := NewCacheStore[[]string](1024)
 
 	service.caches.oauth = oauthCache
 	service.caches.login = loginCache
 	service.caches.ldap = ldapCache
 
-	dg.Go(func(ctx context.Context) {
+	i.Ding.Go(func(ctx context.Context) {
 		ticker := time.NewTicker(1 * time.Minute)
 		defer ticker.Stop()
 
@@ -257,7 +271,7 @@ func (auth *AuthService) RecordLoginAttempt(identifier string, success bool) {
 		return
 	}
 
-	if auth.caches.login.Size() >= MaxLoginAttemptRecords {
+	if !success && auth.config.Auth.LockdownEnabled && auth.caches.login.Size() >= auth.maxLoginLimits {
 		if locked, _ := auth.IsInLockdown(); locked {
 			return
 		}
@@ -632,16 +646,17 @@ func (auth *AuthService) lockdownMode() {
 		return
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
+	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
-	auth.lockdown.until = time.Now().Add(time.Duration(auth.config.Auth.LoginTimeout) * time.Second)
 
-	timer := time.NewTimer(time.Until(auth.lockdown.until))
+	d := time.Duration(auth.config.Auth.LoginTimeout) * time.Second
+	auth.lockdown.until = time.Now().Add(d)
+	timer := time.NewTimer(d)
 
 	auth.lockdown.mu.Unlock()
 
@@ -653,14 +668,13 @@ func (auth *AuthService) lockdownMode() {
 		// Timer expired, end lockdown
 	case <-ctx.Done():
 		// Context cancelled, end lockdown
-	case <-auth.ctx.Done():
-		// Service is shutting down, 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
@@ -683,3 +697,32 @@ func (auth *AuthService) IsInLockdown() (bool, int) {
 func (auth *AuthService) ClearLoginAttempts() {
 	auth.caches.login.Clear()
 }
+
+func (auth *AuthService) calculateLockdownLimit() int {
+	userCount := len(auth.runtime.LocalUsers)
+
+	if auth.ldap != nil {
+		ldapUsers, err := auth.ldap.GetUserCount()
+		if err != nil {
+			auth.log.App.Warn().Err(err).Msg("Failed to get LDAP user count")
+		} else {
+			userCount += ldapUsers
+		}
+	}
+
+	limit := userCount * auth.config.Auth.LoginMaxRetries
+
+	jitter, err := rand.Int(rand.Reader, big.NewInt(64))
+
+	if err != nil {
+		auth.log.App.Warn().Err(err).Msg("Failed to generate jitter for lockdown limit")
+	} else {
+		limit += int(jitter.Int64())
+	}
+
+	if limit < 256 {
+		limit = 256
+	}
+
+	return limit
+}