fix: preserve LDAP STS rate limits without penalizing success

cmd/sts-handlers.go

@@ -43,7 +43,6 @@ import (
 	"github.com/minio/mux"
 	"github.com/minio/pkg/v3/policy"
 	"github.com/minio/pkg/v3/wildcard"
-	"golang.org/x/time/rate"
 )
 
 const (
@@ -272,6 +271,11 @@ type stsLDAPLoginRateLimiter struct {
 	user   *stsLDAPLoginKeyLimiterSet
 }
 
+type stsLDAPLoginReservation struct {
+	source *stsLDAPLoginKeyReservation
+	user   *stsLDAPLoginKeyReservation
+}
+
 type stsLDAPLoginKeyLimiterSet struct {
 	mu          sync.Mutex
 	refillEvery time.Duration
@@ -282,8 +286,16 @@ type stsLDAPLoginKeyLimiterSet struct {
 }
 
 type stsLDAPLoginKeyLimiter struct {
-	limiter  *rate.Limiter
-	lastSeen time.Time
+	tokens     float64
+	lastRefill time.Time
+	lastSeen   time.Time
+	inFlight   int
+}
+
+type stsLDAPLoginKeyReservation struct {
+	set       *stsLDAPLoginKeyLimiterSet
+	entry     *stsLDAPLoginKeyLimiter
+	finalized bool
 }
 
 func newSTSLDAPLoginRateLimiter(refillEvery time.Duration, burst int, ttl time.Duration) *stsLDAPLoginRateLimiter {
@@ -302,21 +314,80 @@ func newSTSLDAPLoginKeyLimiterSet(refillEvery time.Duration, burst int, ttl time
 	}
 }
 
+func normalizeSTSLDAPUsername(username string) string {
+	return strings.ToLower(strings.TrimSpace(username))
+}
+
 func (l *stsLDAPLoginRateLimiter) Allow(sourceIP, username string) bool {
-	now := UTCNow()
-	if sourceIP != "" && !l.source.Allow(now, sourceIP) {
+	reservation := l.Reserve(sourceIP, username)
+	if reservation == nil {
 		return false
 	}
-
-	username = strings.ToLower(strings.TrimSpace(username))
-	if username != "" && !l.user.Allow(now, username) {
-		return false
-	}
-
+	reservation.Commit()
 	return true
 }
 
+func (l *stsLDAPLoginRateLimiter) Reserve(sourceIP, username string) *stsLDAPLoginReservation {
+	now := UTCNow()
+	reservation := &stsLDAPLoginReservation{}
+
+	if sourceIP != "" {
+		reservation.source = l.source.Reserve(now, sourceIP)
+		if reservation.source == nil {
+			return nil
+		}
+	}
+
+	username = normalizeSTSLDAPUsername(username)
+	if username != "" {
+		reservation.user = l.user.Reserve(now, username)
+		if reservation.user == nil {
+			reservation.Cancel()
+			return nil
+		}
+	}
+
+	return reservation
+}
+
+func (r *stsLDAPLoginReservation) Commit() {
+	if r == nil {
+		return
+	}
+	if r.source != nil {
+		r.source.CommitAt(UTCNow())
+		r.source = nil
+	}
+	if r.user != nil {
+		r.user.CommitAt(UTCNow())
+		r.user = nil
+	}
+}
+
+func (r *stsLDAPLoginReservation) Cancel() {
+	if r == nil {
+		return
+	}
+	if r.source != nil {
+		r.source.CancelAt(UTCNow())
+		r.source = nil
+	}
+	if r.user != nil {
+		r.user.CancelAt(UTCNow())
+		r.user = nil
+	}
+}
+
 func (l *stsLDAPLoginKeyLimiterSet) Allow(now time.Time, key string) bool {
+	reservation := l.Reserve(now, key)
+	if reservation == nil {
+		return false
+	}
+	reservation.CommitAt(now)
+	return true
+}
+
+func (l *stsLDAPLoginKeyLimiterSet) Reserve(now time.Time, key string) *stsLDAPLoginKeyReservation {
 	l.mu.Lock()
 	defer l.mu.Unlock()
 
@@ -325,30 +396,102 @@ func (l *stsLDAPLoginKeyLimiterSet) Allow(now time.Time, key string) bool {
 		l.lastCleanup = now
 	}
 
-	entry, ok := l.entries[key]
-	if !ok {
-		entry = &stsLDAPLoginKeyLimiter{
-			limiter:  rate.NewLimiter(rate.Every(l.refillEvery), l.burst),
-			lastSeen: now,
-		}
-		l.entries[key] = entry
-	} else {
+	entry := l.getOrCreateLocked(now, key)
+	l.refillLocked(now, entry)
+	if entry.tokens < 1 {
 		entry.lastSeen = now
+		return nil
 	}
 
-	return entry.limiter.AllowN(now, 1)
+	entry.tokens--
+	entry.inFlight++
+	entry.lastSeen = now
+	return &stsLDAPLoginKeyReservation{set: l, entry: entry}
 }
 
 func (l *stsLDAPLoginKeyLimiterSet) cleanup(now time.Time) {
 	for key, entry := range l.entries {
-		if now.Sub(entry.lastSeen) > l.ttl {
+		if entry.inFlight == 0 && now.Sub(entry.lastSeen) > l.ttl {
 			delete(l.entries, key)
 		}
 	}
 }
 
-func allowSTSLDAPLogin(r *http.Request) bool {
-	return globalSTSLDAPLoginRateLimiter.Allow(handlers.GetSourceIPRaw(r), r.Form.Get(stsLDAPUsername))
+func (l *stsLDAPLoginKeyLimiterSet) getOrCreateLocked(now time.Time, key string) *stsLDAPLoginKeyLimiter {
+	entry, ok := l.entries[key]
+	if !ok {
+		entry = &stsLDAPLoginKeyLimiter{
+			tokens:     float64(l.burst),
+			lastRefill: now,
+			lastSeen:   now,
+		}
+		l.entries[key] = entry
+	}
+	return entry
+}
+
+func (l *stsLDAPLoginKeyLimiterSet) refillLocked(now time.Time, entry *stsLDAPLoginKeyLimiter) {
+	if entry.lastRefill.IsZero() {
+		entry.lastRefill = now
+	}
+
+	lastRefill := entry.lastRefill
+	if now.Before(lastRefill) {
+		lastRefill = now
+	}
+	if l.refillEvery <= 0 {
+		entry.lastRefill = now
+		entry.tokens = float64(l.burst)
+		return
+	}
+
+	entry.tokens += float64(now.Sub(lastRefill)) / float64(l.refillEvery)
+	if maxTokens := float64(l.burst); entry.tokens > maxTokens {
+		entry.tokens = maxTokens
+	}
+	entry.lastRefill = now
+}
+
+func (r *stsLDAPLoginKeyReservation) CommitAt(now time.Time) {
+	r.finalize(now, false)
+}
+
+func (r *stsLDAPLoginKeyReservation) CancelAt(now time.Time) {
+	r.finalize(now, true)
+}
+
+func (r *stsLDAPLoginKeyReservation) finalize(now time.Time, refund bool) {
+	if r == nil {
+		return
+	}
+
+	r.set.mu.Lock()
+	defer r.set.mu.Unlock()
+
+	if r.finalized {
+		return
+	}
+
+	r.set.refillLocked(now, r.entry)
+	r.entry.lastSeen = now
+	if r.entry.inFlight > 0 {
+		r.entry.inFlight--
+	}
+	if refund {
+		r.entry.tokens++
+		if maxTokens := float64(r.set.burst); r.entry.tokens > maxTokens {
+			r.entry.tokens = maxTokens
+		}
+	}
+
+	r.finalized = true
+}
+
+// reserveSTSLDAPLogin acquires immediate tokens from the per-source and
+// per-username limiters before contacting LDAP. Call Commit on auth failures
+// and Cancel when the attempt should not count as an authentication failure.
+func reserveSTSLDAPLogin(r *http.Request) *stsLDAPLoginReservation {
+	return globalSTSLDAPLoginRateLimiter.Reserve(handlers.GetSourceIPRaw(r), r.Form.Get(stsLDAPUsername))
 }
 
 func ldapBindErrorToSTS(err error) (STSErrorCode, error) {
@@ -806,20 +949,26 @@ func (sts *stsAPIHandlers) AssumeRoleWithLDAPIdentity(w http.ResponseWriter, r *
 		return
 	}
 
-	if !allowSTSLDAPLogin(r) {
+	loginReservation := reserveSTSLDAPLogin(r)
+	if loginReservation == nil {
 		writeSTSThrottledResponse(w)
 		return
 	}
+	defer loginReservation.Cancel()
 
 	lookupResult, groupDistNames, err := globalIAMSys.LDAPConfig.Bind(ldapUsername, ldapPassword)
 	if err != nil {
 		errCode, errResp := ldapBindErrorToSTS(err)
 		if errCode == ErrSTSUpstreamError {
+			loginReservation.Cancel()
 			stsLogIf(ctx, err, logger.ErrorKind)
+		} else {
+			loginReservation.Commit()
 		}
 		writeSTSErrorResponse(ctx, w, errCode, errResp)
 		return
 	}
+	loginReservation.Cancel()
 	ldapUserDN := lookupResult.NormDN
 	ldapActualUserDN := lookupResult.ActualDN