diff --git a/ratelimit.go b/ratelimit.go index 543fb85..fb5b6e7 100644 --- a/ratelimit.go +++ b/ratelimit.go @@ -15,17 +15,21 @@ import ( // RateLimitConfig holds the rate limit configuration for a role type RateLimitConfig struct { RateCounterTicker *goratecounter.RateCounter - Endpoints []string `json:"endpoints,omitempty"` - Interval time.Duration `json:"interval"` - Req int `json:"req"` - Burst int `json:"burst,omitempty"` + // burstCounter enforces Burst over a short trailing window, independent + // of RateCounterTicker's sustained-rate window. Nil when Burst <= 0 + // (burst disabled). See burstWindow and checkRateLimit for semantics. + burstCounter *burstWindow + Endpoints []string `json:"endpoints,omitempty"` + Interval time.Duration `json:"interval"` + Req int `json:"req"` + Burst int `json:"burst,omitempty"` } // UnmarshalJSON implements custom JSON unmarshaling for RateLimitConfig func (r *RateLimitConfig) UnmarshalJSON(data []byte) error { // Use a temporary struct to unmarshal the JSON data. // It must carry every configurable field: Burst (used by the burst - // allowance at the top of checkRateLimit) and Endpoints would otherwise be + // allowance in checkRateLimit) and Endpoints would otherwise be // silently dropped from the loaded config and left at their zero values. type RateLimitConfigTemp struct { Interval any `json:"interval"` @@ -82,6 +86,72 @@ var ( rateLimitConfigAtomic atomic.Value ) +// burstWindowFraction sets the short burst window as a fraction of the +// sustained Interval: window = Interval / burstWindowFraction. Not +// specified by any upstream spec; chosen so the burst window is clearly +// shorter than the sustained window while remaining simple to reason about. +const burstWindowFraction = 10 + +// minBurstWindow floors the short burst window so a very small Interval +// (for example a few milliseconds) still leaves the burst window a usable, +// non-degenerate duration. +const minBurstWindow = 10 * time.Millisecond + +// burstWindowDuration derives the short burst-window length from the +// sustained Interval. See burstWindowFraction and minBurstWindow. +func burstWindowDuration(interval time.Duration) time.Duration { + w := interval / burstWindowFraction + if w < minBurstWindow { + w = minBurstWindow + } + return w +} + +// burstWindow is a minimal, clock-injectable, trailing-window request +// counter used to enforce RateLimitConfig.Burst independently of the +// sustained-rate RateCounterTicker. +// +// go-ratecounter@v0.1.12 cannot provide this second, shorter window +// cleanly: the window length lives on *RateCounter itself, not per named +// Counter (helpers.go: GetRate -> c.parent.interval.Seconds()), so two +// different window lengths need two separate RateCounter instances. Each +// instance drives its own background goroutine off wall-clock time.Now with +// no seam to inject a fake clock (helpers.go: start(), cleanUpOldValues(), +// addValue() all call time.Now() directly). That would force any test of +// window rollover to sleep on real time. burstWindow is implemented here +// instead, with a swappable clock, so tests can advance time deterministically. +type burstWindow struct { + mu sync.Mutex + now func() time.Time + window time.Duration + hits []time.Time +} + +// newBurstWindow returns a burstWindow that tracks hits over the trailing +// duration window, using the real wall clock. +func newBurstWindow(window time.Duration) *burstWindow { + return &burstWindow{now: time.Now, window: window} +} + +// recordAndCount records a hit at the current time, drops hits older than +// the trailing window, and returns the resulting hit count within the +// window. Safe for concurrent use. +func (b *burstWindow) recordAndCount() int { + b.mu.Lock() + defer b.mu.Unlock() + now := b.now() + cutoff := now.Add(-b.window) + live := b.hits[:0] + for _, t := range b.hits { + if t.After(cutoff) { + live = append(live, t) + } + } + live = append(live, now) + b.hits = live + return len(b.hits) +} + // Variable to hold the current load config function - allows for testing var loadConfigFunc = loadConfigFromPath @@ -171,6 +241,12 @@ func loadConfigFromPath(path string) error { Interval: value.Interval, }) + // Burst disabled (<=0) leaves burstCounter nil, so checkRateLimit + // performs the Req-only check exactly as before burst existed. + if value.Burst > 0 { + value.burstCounter = newBurstWindow(burstWindowDuration(value.Interval)) + } + if cfg.LogLevel == "DEBUG" { cfg.Logger.Debug(&libpack_logger.LogMessage{ Message: "Setting ratelimit config for role", @@ -226,7 +302,34 @@ func rateLimitedRequest(userID, userRole string) bool { return checkRateLimit(userID, userRole, roleConfig, "") } -// checkRateLimit performs the actual rate limit check +// checkRateLimit performs the actual rate limit check. +// +// Two independent gates apply, both must pass: +// +// 1. Sustained: tickerRate (requests measured over the full Interval) must +// stay at or below Req. This is the only gate when Burst <= 0, and its +// comparison is unchanged from before the burst fix. +// 2. Burst (only when roleConfig.burstCounter != nil, i.e. Burst > 0): the +// raw request count within a short trailing window (see +// burstWindowDuration), much shorter than Interval, must stay at or +// below Burst. +// +// Because the burst window is short relative to Interval, a client can +// legally land up to Burst requests in a tight cluster - more than a +// steady Req-paced stream would deliver in that same short span - as long +// as the sustained gate still holds over the full Interval. That is what +// makes burst > req a real spike allowance instead of a no-op. Burst is a +// short-window ceiling layered ON TOP of the sustained Req rate, not a +// replacement for it. burst < req is an unusual config: the burst gate +// becomes the tighter, binding constraint whenever requests cluster in +// time, but the sustained cap for evenly-spread traffic is still Req +// alone. The pre-fix code compared the SAME sustained-window tickerRate +// against both Burst and Req, so when burst < req its Burst check tripped +// first and Burst was effectively the sustained cap too. This gate's +// Req-only sustained check can therefore allow MORE through than the +// pre-fix code did for evenly-spread traffic under a burst < req config. +// In both cases the burst gate never allows more than Burst requests +// within its own window (fail-safe). func checkRateLimit(userID, userRole string, roleConfig RateLimitConfig, endpoint string) bool { roleConfig.RateCounterTicker.Incr(1) tickerRate := roleConfig.RateCounterTicker.GetRate() @@ -245,15 +348,6 @@ func checkRateLimit(userID, userRole string, roleConfig RateLimitConfig, endpoin Pairs: map[string]any{"log_details": logDetails}, }) - // Check burst limit if configured - if roleConfig.Burst > 0 && tickerRate > float64(roleConfig.Burst) { - cfg.Logger.Debug(&libpack_logger.LogMessage{ - Message: "Burst limit exceeded", - Pairs: map[string]any{"log_details": logDetails}, - }) - return false - } - if tickerRate > float64(roleConfig.Req) { cfg.Logger.Debug(&libpack_logger.LogMessage{ Message: "Rate limit exceeded", @@ -262,5 +356,23 @@ func checkRateLimit(userID, userRole string, roleConfig RateLimitConfig, endpoin return false } + if roleConfig.burstCounter != nil { + burstCount := roleConfig.burstCounter.recordAndCount() + burstLogDetails := map[string]any{ + "user_role": userRole, + "user_id": userID, + "burst_count": burstCount, + "config_burst": roleConfig.Burst, + "endpoint": endpoint, + } + if burstCount > roleConfig.Burst { + cfg.Logger.Debug(&libpack_logger.LogMessage{ + Message: "Burst limit exceeded", + Pairs: map[string]any{"log_details": burstLogDetails}, + }) + return false + } + } + return true } diff --git a/ratelimit_test.go b/ratelimit_test.go index 22b707b..563a401 100644 --- a/ratelimit_test.go +++ b/ratelimit_test.go @@ -4,6 +4,8 @@ import ( "fmt" "os" "path/filepath" + "sync" + "sync/atomic" "testing" "time" @@ -299,3 +301,210 @@ func TestRateLimitConfigUnmarshalKeepsBurstAndEndpoints(t *testing.T) { t.Errorf("endpoints = %v, want [\"/v1/graphql\"]", rc.Endpoints) } } + +// fakeClock is a manually-advanced clock injected into a burstWindow so +// tests can move the burst window forward deterministically, without +// sleeping on real time. +type fakeClock struct { + mu sync.Mutex + now time.Time +} + +func newFakeClock(start time.Time) *fakeClock { + return &fakeClock{now: start} +} + +func (f *fakeClock) Now() time.Time { + f.mu.Lock() + defer f.mu.Unlock() + return f.now +} + +func (f *fakeClock) Advance(d time.Duration) { + f.mu.Lock() + defer f.mu.Unlock() + f.now = f.now.Add(d) +} + +// newBurstTestRoleConfig mirrors the RateLimitConfig construction performed +// by loadConfigFromPath: a real RateCounterTicker for the sustained check, +// and (only when burst > 0) a burstCounter sized by burstWindowDuration. +func newBurstTestRoleConfig(interval time.Duration, req, burst int) RateLimitConfig { + rc := RateLimitConfig{ + RateCounterTicker: goratecounter.NewRateCounter().WithConfig(goratecounter.RateCounterConfig{ + Interval: interval, + }), + Interval: interval, + Req: req, + Burst: burst, + } + if burst > 0 { + rc.burstCounter = newBurstWindow(burstWindowDuration(interval)) + } + return rc +} + +// Test_checkRateLimit_BurstSemantics is the table test for the C10 burst +// fix. Each case is self-contained and deterministic: sustained-rate +// timing relies only on test execution being far faster than the +// configured Interval (the same assumption Test_rateLimitedRequest already +// makes), and burst-window timing is controlled directly via fakeClock, so +// nothing here depends on a wall-clock sleep as the pass/fail gate. +func Test_checkRateLimit_BurstSemantics(t *testing.T) { + cfg = &config{} + parseConfig() + cfg.Logger = libpack_logger.New() + + t.Run("a_burst_unset_only_req_enforced", func(t *testing.T) { + rc := newBurstTestRoleConfig(1*time.Second, 2, 0) + if rc.burstCounter != nil { + t.Fatalf("burstCounter = %v, want nil when Burst <= 0", rc.burstCounter) + } + + if !checkRateLimit("u", "role-a", rc, "") { + t.Errorf("request 1: got denied, want allowed (within Req)") + } + if !checkRateLimit("u", "role-a", rc, "") { + t.Errorf("request 2: got denied, want allowed (at Req)") + } + if checkRateLimit("u", "role-a", rc, "") { + t.Errorf("request 3: got allowed, want denied (exceeds Req, burst disabled)") + } + }) + + t.Run("b_burst_above_req_allows_short_spike_sustained_still_capped", func(t *testing.T) { + const ( + interval = 10 * time.Second + req = 3 + burst = 6 // burst > req + ) + rc := newBurstTestRoleConfig(interval, req, burst) + fc := newFakeClock(time.Now()) + rc.burstCounter.now = fc.Now + + // A rapid cluster of `burst` requests must all be allowed, even + // though burst (6) > req (3): the sustained gate only bounds the + // total over the full 10s Interval (tickerRate <= 3 needs count <= + // 30), so a small cluster like this never trips it. + for i := 1; i <= burst; i++ { + if !checkRateLimit("u", "role-b", rc, "") { + t.Fatalf("request %d in first burst: got denied, want allowed (<= Burst=%d)", i, burst) + } + } + + // The (burst+1)th request in the SAME short window must be denied + // by the burst gate specifically - not by the sustained gate, + // which alone would still allow it (tickerRate = 7/10 = 0.7 <= 3). + // This is the proof burst > req is not a no-op: it is the gate + // that actually fires here. + if checkRateLimit("u", "role-b", rc, "") { + t.Fatalf("request %d: got allowed, want denied by burst gate (> Burst=%d)", burst+1, burst) + } + + // Roll the burst window forward (fake clock only - the real + // sustained ticker has not aged, since no real time passed) and + // repeat clusters of `burst` requests. Each cluster alone passes + // burst, but the cumulative sustained count keeps climbing until + // it exceeds req*interval.Seconds() = 30, at which point the + // sustained gate must deny regardless of how the burst window + // resets. This proves "sustained still capped at Req". + sustainedTripped := false + for batch := 0; batch < 10 && !sustainedTripped; batch++ { + fc.Advance(2 * burstWindowDuration(interval)) // clears the burst window + for i := 1; i <= burst; i++ { + allowed := checkRateLimit("u", "role-b", rc, "") + rate := rc.RateCounterTicker.GetRate() + if rate > float64(req) { + if allowed { + t.Fatalf("batch %d req %d: allowed with sustained rate %.2f > Req=%d", batch, i, rate, req) + } + sustainedTripped = true + break + } + } + } + if !sustainedTripped { + t.Fatalf("sustained gate never tripped after repeated bursts; expected it to cap total requests at Req over the Interval") + } + }) + + t.Run("c_burst_below_req_documented_behaviour", func(t *testing.T) { + t.Run("clustered_traffic_capped_below_req_by_burst", func(t *testing.T) { + const ( + interval = 10 * time.Second + req = 6 + burst = 3 // burst < req + ) + rc := newBurstTestRoleConfig(interval, req, burst) + fc := newFakeClock(time.Now()) + rc.burstCounter.now = fc.Now + + for i := 1; i <= burst; i++ { + if !checkRateLimit("u", "role-c1", rc, "") { + t.Fatalf("request %d: got denied, want allowed (<= Burst=%d)", i, burst) + } + } + // The 4th rapid request is within Req (tickerRate = 4/10 = 0.4 + // <= 6) but must still be denied: with burst < req, the burst + // gate is the tighter, binding constraint for clustered + // traffic - it is not a no-op the way the pre-fix code made + // burst > req. + if checkRateLimit("u", "role-c1", rc, "") { + t.Fatalf("request %d: got allowed, want denied by burst gate (> Burst=%d) though within Req", burst+1, burst) + } + }) + + t.Run("evenly_spread_traffic_still_governed_by_req_alone", func(t *testing.T) { + const ( + interval = 1 * time.Second + req = 3 + burst = 1 // burst < req + ) + rc := newBurstTestRoleConfig(interval, req, burst) + fc := newFakeClock(time.Now()) + rc.burstCounter.now = fc.Now + window := burstWindowDuration(interval) + + // One request per (window*2): each lands in its own fresh + // burst window (count resets to 1 <= Burst=1 every time), so + // the burst gate never fires. Only the sustained gate, over + // req=3, eventually denies. + for i := 1; i <= req; i++ { + if !checkRateLimit("u", "role-c2", rc, "") { + t.Fatalf("spread request %d: got denied, want allowed (within Req=%d, burst resets every time)", i, req) + } + fc.Advance(2 * window) + } + if checkRateLimit("u", "role-c2", rc, "") { + t.Fatalf("spread request %d: got allowed, want denied by sustained gate (> Req=%d)", req+1, req) + } + }) + }) + + t.Run("d_fail_safe_never_allows_more_than_burst_in_window", func(t *testing.T) { + const ( + interval = 10 * time.Second + req = 1000 // high enough it never binds; isolates the burst gate + burst = 5 + goroutines = 50 + ) + rc := newBurstTestRoleConfig(interval, req, burst) + + var allowed int64 + var wg sync.WaitGroup + wg.Add(goroutines) + for i := 0; i < goroutines; i++ { + go func() { + defer wg.Done() + if checkRateLimit("u", "role-d", rc, "") { + atomic.AddInt64(&allowed, 1) + } + }() + } + wg.Wait() + + if allowed != burst { + t.Fatalf("allowed = %d concurrent requests, want exactly Burst=%d (fail-safe: never more than Burst allowed in the burst window)", allowed, burst) + } + }) +}