6852a8f8c3
- Replace split store checks with one mutex-guarded Attempt operation - Use domainless trusted-client keys for anonymous inline throttles - Preserve fixed-window headers, expiry, stacking, and principal isolation
94 lines
2.2 KiB
Go
94 lines
2.2 KiB
Go
package surf
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// Store owns fixed-window admission as one atomic operation. Attempt performs
|
|
// lazy expiry, threshold comparison, and an admitted increment together so
|
|
// concurrent callers cannot pass a split check-then-increment boundary.
|
|
type Store interface {
|
|
Attempt(key string, max int, decay time.Duration) (allowed bool, attempts int, retryAfter time.Duration)
|
|
}
|
|
|
|
type counterEntry struct {
|
|
count int
|
|
resetAt time.Time
|
|
}
|
|
|
|
// MemoryStore is an in-process, mutex-guarded Store.
|
|
type MemoryStore struct {
|
|
mu sync.Mutex
|
|
entries map[string]*counterEntry
|
|
sweep time.Duration
|
|
stop chan struct{}
|
|
}
|
|
|
|
// NewMemoryStore returns an in-process, mutex-guarded Store. sweep controls
|
|
// the background expired-entry cleanup interval (memory hygiene only --
|
|
// correctness does not depend on it, since expiry is checked lazily).
|
|
// A non-positive sweep disables the background goroutine.
|
|
func NewMemoryStore(sweep time.Duration) *MemoryStore {
|
|
s := &MemoryStore{
|
|
entries: make(map[string]*counterEntry),
|
|
sweep: sweep,
|
|
stop: make(chan struct{}),
|
|
}
|
|
if sweep > 0 {
|
|
go s.loop()
|
|
}
|
|
return s
|
|
}
|
|
|
|
func (s *MemoryStore) loop() {
|
|
ticker := time.NewTicker(s.sweep)
|
|
defer ticker.Stop()
|
|
for {
|
|
select {
|
|
case <-ticker.C:
|
|
s.purge()
|
|
case <-s.stop:
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *MemoryStore) purge() {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
now := time.Now()
|
|
for k, e := range s.entries {
|
|
if now.After(e.resetAt) {
|
|
delete(s.entries, k)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Attempt admits and counts one request when key is below max. The first
|
|
// attempt opens the window; later attempts never extend it. A denied attempt
|
|
// leaves the exhausted count unchanged.
|
|
func (s *MemoryStore) Attempt(key string, max int, decay time.Duration) (bool, int, time.Duration) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
now := time.Now()
|
|
e, ok := s.entries[key]
|
|
if ok && now.After(e.resetAt) {
|
|
delete(s.entries, key)
|
|
ok = false
|
|
}
|
|
if !ok {
|
|
e = &counterEntry{count: 0, resetAt: now.Add(decay)}
|
|
s.entries[key] = e
|
|
}
|
|
retryAfter := e.resetAt.Sub(now)
|
|
if retryAfter < 0 {
|
|
retryAfter = 0
|
|
}
|
|
if e.count >= max {
|
|
return false, e.count, retryAfter
|
|
}
|
|
e.count++
|
|
return true, e.count, retryAfter
|
|
}
|