//! Token-bucket rate limiter for the REST API (PLAN §10; milestone-8 ruling 19). //! //! The DNS listeners use a fixed window (`server/rate_limiter.zig`); the API //! uses a bucket, because an admin UI legitimately fires a burst of requests //! when a page loads and then goes quiet. A bucket admits that burst up to its //! capacity and still holds the long-run rate to `rate_per_min` per minute. //! //! One bucket per client address, in a table bounded at `max_clients`. When //! the table is full, an unknown address first reclaims the slot of a bucket //! that a fresh one would answer identically to (refills to capacity, holds no //! SSE connection); if no such bucket exists, the address is refused and //! counted under `untracked`. Admitting it instead — the DNS limiter's choice — //! would let a client cycling addresses bypass the limiter entirely, and here //! the DNS limiter's reason does not apply: API clients speak TCP, so a flood //! of spoofed sources cannot fill the table, and the operator on the box stays //! covered by the localhost exemption. //! //! The same table carries each address's live SSE connection count, since both //! limits key on the address and both are taken and released around one //! request. `/metrics` and `/api/health` never reach this file — Prometheus must //! not be told 429 (ruling 19) — and the router is what exempts them. //! //! Timestamps come from the caller, as everywhere else in this codebase. Pass //! the `.awake` clock: elapsed time is what refills a bucket, and a wall-clock //! step must not hand out a minute of tokens. //! //! Not lock-free but self-locking: connection tasks run concurrently, so the //! mutex lives here rather than in every caller. const std = @import("std"); const address = @import("../platform/address.zig"); const Allocator = std.mem.Allocator; /// Upper bound on tracked addresses. The table never grows past it, so `check` /// never allocates. pub const max_clients = 4096; /// A bucket refills its whole capacity over this window (ruling 19). pub const window_seconds = 60; const window_ns: i96 = @as(i96, window_seconds) * std.time.ns_per_s; /// Tokens are counted in millionths so that a fraction of a token earned /// between two requests is not lost to integer division. One whole token is /// `token_scale`. const token_scale: u64 = 1_000_000; pub const Config = struct { /// `web.api_rate_limit_per_min`: both the bucket capacity and the refill per /// minute. `validate.zig` rejects zero. rate_per_min: u32, /// `web.api_localhost_exempt`. The box's own requests — a script on the /// server, a health probe in a container namespace — are usually the /// operator's own and are not what the limiter defends against. localhost_exempt: bool = true, /// `web.sse_max_connections_per_ip`. sse_max_per_ip: u16, }; /// `allowed + refused` equals the number of `check` calls that were not exempt. /// `untracked` counts the subset of `refused` that a full table could not hold /// a bucket for, and `exempt` the calls that never consulted a bucket. pub const Stats = struct { allowed: u64 = 0, refused: u64 = 0, untracked: u64 = 0, exempt: u64 = 0, sse_refused: u64 = 0, }; pub const Result = struct { allowed: bool, /// Seconds until one token is available again, for the `Retry-After` /// header. Zero when the request was allowed. Never zero when it was /// refused: a client told to retry after zero seconds retries immediately. retry_after_s: u32 = 0, pub const ok: Result = .{ .allowed = true }; }; const Bucket = struct { /// Tokens held, scaled by `token_scale`. tokens: u64, /// When `tokens` was last brought up to date. updated_ns: i96, /// Live SSE responses this address holds open. sse: u16, }; const Table = std.AutoHashMapUnmanaged(address.NetAddress.Key, Bucket); pub const ApiLimiter = struct { /// Guards `table` and `stats`; see the file comment. mutex: std.Io.Mutex, gpa: Allocator, config: Config, capacity: u64, table: Table, /// `sweep` collects the keys to drop before removing any, because a removal /// invalidates a live iterator. Owning the buffer keeps `sweep` /// allocation-free. stale_keys: []address.NetAddress.Key, stats: Stats, /// Asserts `config.rate_per_min` is nonzero; `validate.zig` rejects a zero /// rate before a config reaches this far. pub fn init(gpa: Allocator, config: Config) Allocator.Error!ApiLimiter { std.debug.assert(config.rate_per_min > 0); var table: Table = .empty; errdefer table.deinit(gpa); try table.ensureTotalCapacity(gpa, max_clients); const stale_keys = try gpa.alloc(address.NetAddress.Key, max_clients); return .{ .mutex = .init, .gpa = gpa, .config = config, .capacity = @as(u64, config.rate_per_min) * token_scale, .table = table, .stale_keys = stale_keys, .stats = .{}, }; } pub fn deinit(self: *ApiLimiter) void { self.table.deinit(self.gpa); self.gpa.free(self.stale_keys); self.* = undefined; } /// Spends one token for a request from `addr`. Never allocates, never fails. pub fn check(self: *ApiLimiter, io: std.Io, now: std.Io.Timestamp, addr: address.NetAddress) Result { if (self.config.localhost_exempt and isLoopback(addr)) { self.mutex.lockUncancelable(io); defer self.mutex.unlock(io); self.stats.exempt += 1; return .ok; } self.mutex.lockUncancelable(io); defer self.mutex.unlock(io); const bucket = self.bucketLocked(addr.key(), now) orelse { self.stats.untracked += 1; self.stats.refused += 1; return .{ .allowed = false, .retry_after_s = self.retryAfter(0) }; }; if (bucket.tokens < token_scale) { self.stats.refused += 1; return .{ .allowed = false, .retry_after_s = self.retryAfter(bucket.tokens) }; } bucket.tokens -= token_scale; self.stats.allowed += 1; return .ok; } /// Takes an SSE slot for `addr`. A connect also spends a token, which the /// caller does with `check` first (ruling 19); this call is only the /// per-address connection cap. /// /// The cap applies to loopback too: it bounds a fixed resource (subscriber /// slots in the hub), which the rate exemption has no bearing on. pub fn tryAcquireSse(self: *ApiLimiter, io: std.Io, now: std.Io.Timestamp, addr: address.NetAddress) bool { self.mutex.lockUncancelable(io); defer self.mutex.unlock(io); // A full table cannot hold the counter, so it cannot enforce the cap // either. Refusing is consistent with `check`: an uncounted stream // could otherwise reach the hub's global cap past the per-IP one. const bucket = self.bucketLocked(addr.key(), now) orelse { self.stats.sse_refused += 1; return false; }; if (bucket.sse >= self.config.sse_max_per_ip) { self.stats.sse_refused += 1; return false; } bucket.sse += 1; return true; } /// Releases a slot taken by `tryAcquireSse`. A release whose bucket was /// swept finds no counter and does nothing: the alternative is an /// underflow on a path that must not fail. pub fn releaseSse(self: *ApiLimiter, io: std.Io, addr: address.NetAddress) void { self.mutex.lockUncancelable(io); defer self.mutex.unlock(io); const bucket = self.table.getPtr(addr.key()) orelse return; if (bucket.sse > 0) bucket.sse -= 1; } /// Live SSE connections held by `addr`. pub fn sseConnections(self: *ApiLimiter, io: std.Io, addr: address.NetAddress) u16 { self.mutex.lockUncancelable(io); defer self.mutex.unlock(io); const bucket = self.table.getPtr(addr.key()) orelse return 0; return bucket.sse; } /// Drops every bucket that is full, holds no SSE connection and has been /// idle for a full window: such a bucket answers exactly as a fresh one /// would, so forgetting it changes no decision. Returns how many it dropped. pub fn sweep(self: *ApiLimiter, io: std.Io, now: std.Io.Timestamp) u32 { self.mutex.lockUncancelable(io); defer self.mutex.unlock(io); var stale_count: u32 = 0; var it = self.table.iterator(); while (it.next()) |entry| { const bucket = entry.value_ptr; if (bucket.sse != 0) continue; if (now.nanoseconds - bucket.updated_ns < window_ns) continue; if (refilled(bucket.*, now, self.capacity).tokens < self.capacity) continue; self.stale_keys[stale_count] = entry.key_ptr.*; stale_count += 1; } for (self.stale_keys[0..stale_count]) |key| { const removed = self.table.remove(key); std.debug.assert(removed); } return stale_count; } /// Addresses currently holding a bucket. Reaching `max_clients` with no /// reclaimable bucket is what turns unknown addresses into `untracked` /// refusals. pub fn trackedClients(self: *ApiLimiter, io: std.Io) u32 { self.mutex.lockUncancelable(io); defer self.mutex.unlock(io); return self.table.count(); } pub fn snapshotStats(self: *ApiLimiter, io: std.Io) Stats { self.mutex.lockUncancelable(io); defer self.mutex.unlock(io); return self.stats; } /// The address's bucket, refilled to `now`, or null when the table is full /// and no slot can be reclaimed for the unknown address. Caller holds the /// mutex. fn bucketLocked(self: *ApiLimiter, key: address.NetAddress.Key, now: std.Io.Timestamp) ?*Bucket { if (self.table.getPtr(key)) |bucket| { const state = refilled(bucket.*, now, self.capacity); bucket.tokens = state.tokens; bucket.updated_ns = state.updated_ns; return bucket; } if (self.table.count() >= max_clients and !self.evictReclaimableLocked(now)) return null; const gop = self.table.getOrPutAssumeCapacity(key); gop.value_ptr.* = .{ .tokens = self.capacity, .updated_ns = now.nanoseconds, .sse = 0 }; return gop.value_ptr; } /// Removes one bucket a fresh bucket would answer identically to: full /// after refill and holding no SSE connection. Unlike `sweep` it demands /// no idle window — that hysteresis avoids churn in background sweeping /// but forgets nothing here, since a full bucket decides as a fresh one /// does. Returns whether a slot was reclaimed. Caller holds the mutex. fn evictReclaimableLocked(self: *ApiLimiter, now: std.Io.Timestamp) bool { var it = self.table.iterator(); while (it.next()) |entry| { if (entry.value_ptr.sse != 0) continue; if (refilled(entry.value_ptr.*, now, self.capacity).tokens < self.capacity) continue; const removed = self.table.remove(entry.key_ptr.*); std.debug.assert(removed); return true; } return false; } /// Seconds until `tokens` reaches one whole token, rounded up and never /// below one. fn retryAfter(self: *const ApiLimiter, tokens: u64) u32 { const missing = token_scale - tokens; // missing tokens / (rate_per_min tokens per window) seconds, rounded up. const seconds = (missing * window_seconds + self.capacity - 1) / self.capacity; return @intCast(@max(1, seconds)); } }; const Refill = struct { tokens: u64, updated_ns: i96 }; /// `bucket` brought up to `now`. /// /// The time that bought fewer than one microtoken stays on the clock rather /// than being rounded away: `updated_ns` only advances by the span actually /// converted into tokens. Without that, a client polling faster than one /// microtoken per request would never refill at all, and a refill boundary /// would land a microtoken short of where the configured rate puts it. /// /// A backwards timestamp earns nothing and resets the clock, so a clock that /// steps back cannot later be credited for the time it repeated. fn refilled(bucket: Bucket, now: std.Io.Timestamp, capacity: u64) Refill { const elapsed_ns = now.nanoseconds - bucket.updated_ns; if (elapsed_ns <= 0) return .{ .tokens = @min(bucket.tokens, capacity), .updated_ns = now.nanoseconds }; if (bucket.tokens >= capacity) return .{ .tokens = capacity, .updated_ns = now.nanoseconds }; // A whole window refills the bucket whatever it held, and short-circuiting // here also keeps the multiplication below inside i96. if (elapsed_ns >= window_ns) return .{ .tokens = capacity, .updated_ns = now.nanoseconds }; const capacity_96: i96 = @intCast(capacity); const gained = @divTrunc(elapsed_ns * capacity_96, window_ns); if (gained == 0) return .{ .tokens = bucket.tokens, .updated_ns = bucket.updated_ns }; const tokens = bucket.tokens + @as(u64, @intCast(gained)); if (tokens >= capacity) return .{ .tokens = capacity, .updated_ns = now.nanoseconds }; return .{ .tokens = tokens, .updated_ns = bucket.updated_ns + @divTrunc(gained * window_ns, capacity_96) }; } /// 127.0.0.0/8 and ::1, the addresses a request from the box itself carries. /// An IPv4-mapped loopback address has already normalized to `.ip4` by the time /// a `NetAddress` exists (`address.zig:51`). fn isLoopback(addr: address.NetAddress) bool { return switch (addr) { .ip4 => |b| b[0] == 127, .ip6 => |b| std.mem.eql(u8, &b, &[_]u8{0} ** 15 ++ [_]u8{1}), }; } // --------------------------------------------------------------------------- // tests // --------------------------------------------------------------------------- const testing = std.testing; fn at(seconds: i64) std.Io.Timestamp { return .{ .nanoseconds = @as(i96, seconds) * std.time.ns_per_s }; } fn atMillis(millis: i64) std.Io.Timestamp { return .{ .nanoseconds = @as(i96, millis) * std.time.ns_per_ms }; } fn v4(a: u8, b: u8, c: u8, d: u8) address.NetAddress { return .{ .ip4 = .{ a, b, c, d } }; } fn indexed(index: u32) address.NetAddress { var octets: [4]u8 = undefined; std.mem.writeInt(u32, &octets, index, .big); return .{ .ip4 = octets }; } const Fixture = struct { threaded: std.Io.Threaded, limiter: ApiLimiter, fn init(config: Config) !*Fixture { const self = try testing.allocator.create(Fixture); self.* = .{ .threaded = .init(testing.allocator, .{}), .limiter = try ApiLimiter.init(testing.allocator, config), }; return self; } fn deinit(self: *Fixture) void { self.limiter.deinit(); self.threaded.deinit(); testing.allocator.destroy(self); } fn io(self: *Fixture) std.Io { return self.threaded.io(); } }; test "a burst up to the capacity is allowed and the next request is refused" { const fx = try Fixture.init(.{ .rate_per_min = 60, .localhost_exempt = false, .sse_max_per_ip = 3 }); defer fx.deinit(); const client = v4(192, 168, 1, 10); for (0..60) |_| { try testing.expect(fx.limiter.check(fx.io(), at(0), client).allowed); } const refused = fx.limiter.check(fx.io(), at(0), client); try testing.expect(!refused.allowed); // At 60 per minute a token is worth one second. try testing.expectEqual(@as(u32, 1), refused.retry_after_s); const stats = fx.limiter.snapshotStats(fx.io()); try testing.expectEqual(@as(u64, 60), stats.allowed); try testing.expectEqual(@as(u64, 1), stats.refused); } test "an emptied bucket refills at the configured rate" { const fx = try Fixture.init(.{ .rate_per_min = 60, .localhost_exempt = false, .sse_max_per_ip = 3 }); defer fx.deinit(); const client = v4(10, 0, 0, 1); for (0..60) |_| try testing.expect(fx.limiter.check(fx.io(), at(0), client).allowed); // One token is worth exactly one second, so 999 ms is still short. try testing.expect(!fx.limiter.check(fx.io(), atMillis(999), client).allowed); try testing.expect(fx.limiter.check(fx.io(), atMillis(1000), client).allowed); try testing.expect(!fx.limiter.check(fx.io(), atMillis(1000), client).allowed); // The fractions left behind by the refused calls still accumulate. try testing.expect(fx.limiter.check(fx.io(), atMillis(2000), client).allowed); // A long idle period refills no further than the capacity. for (0..60) |_| try testing.expect(fx.limiter.check(fx.io(), at(3600), client).allowed); try testing.expect(!fx.limiter.check(fx.io(), at(3600), client).allowed); } test "retry-after reports the wait for one token and is never zero" { const fx = try Fixture.init(.{ .rate_per_min = 6, .localhost_exempt = false, .sse_max_per_ip = 3 }); defer fx.deinit(); const client = v4(10, 0, 0, 2); for (0..6) |_| try testing.expect(fx.limiter.check(fx.io(), at(0), client).allowed); // At 6 per minute a token takes 10 seconds. try testing.expectEqual(@as(u32, 10), fx.limiter.check(fx.io(), at(0), client).retry_after_s); try testing.expectEqual(@as(u32, 5), fx.limiter.check(fx.io(), at(5), client).retry_after_s); // Under a second of waiting still reports one second. try testing.expectEqual(@as(u32, 1), fx.limiter.check(fx.io(), atMillis(9_500), client).retry_after_s); try testing.expect(fx.limiter.check(fx.io(), at(10), client).allowed); } test "a rate of one still admits one request per minute" { const fx = try Fixture.init(.{ .rate_per_min = 1, .localhost_exempt = false, .sse_max_per_ip = 1 }); defer fx.deinit(); const client = v4(10, 0, 0, 3); try testing.expect(fx.limiter.check(fx.io(), at(0), client).allowed); const refused = fx.limiter.check(fx.io(), at(0), client); try testing.expect(!refused.allowed); try testing.expectEqual(@as(u32, 60), refused.retry_after_s); try testing.expect(!fx.limiter.check(fx.io(), at(59), client).allowed); try testing.expect(fx.limiter.check(fx.io(), at(60), client).allowed); } test "clients hold independent buckets" { const fx = try Fixture.init(.{ .rate_per_min = 1, .localhost_exempt = false, .sse_max_per_ip = 3 }); defer fx.deinit(); const a = v4(192, 168, 1, 20); const b = try address.NetAddress.parse("fd00::20"); try testing.expect(fx.limiter.check(fx.io(), at(0), a).allowed); try testing.expect(!fx.limiter.check(fx.io(), at(0), a).allowed); try testing.expect(fx.limiter.check(fx.io(), at(0), b).allowed); try testing.expect(!fx.limiter.check(fx.io(), at(0), b).allowed); try testing.expectEqual(@as(u32, 2), fx.limiter.trackedClients(fx.io())); } test "loopback is exempt when configured and limited when not" { const exempt = try Fixture.init(.{ .rate_per_min = 1, .localhost_exempt = true, .sse_max_per_ip = 3 }); defer exempt.deinit(); for (0..10) |_| { try testing.expect(exempt.limiter.check(exempt.io(), at(0), v4(127, 0, 0, 1)).allowed); } try testing.expect(exempt.limiter.check(exempt.io(), at(0), v4(127, 1, 2, 3)).allowed); try testing.expect(exempt.limiter.check(exempt.io(), at(0), try address.NetAddress.parse("::1")).allowed); // An exempt request consults no bucket at all. try testing.expectEqual(@as(u32, 0), exempt.limiter.trackedClients(exempt.io())); try testing.expectEqual(@as(u64, 12), exempt.limiter.snapshotStats(exempt.io()).exempt); // A LAN address is limited either way. try testing.expect(exempt.limiter.check(exempt.io(), at(0), v4(192, 168, 1, 5)).allowed); try testing.expect(!exempt.limiter.check(exempt.io(), at(0), v4(192, 168, 1, 5)).allowed); const strict = try Fixture.init(.{ .rate_per_min = 1, .localhost_exempt = false, .sse_max_per_ip = 3 }); defer strict.deinit(); try testing.expect(strict.limiter.check(strict.io(), at(0), v4(127, 0, 0, 1)).allowed); try testing.expect(!strict.limiter.check(strict.io(), at(0), v4(127, 0, 0, 1)).allowed); try testing.expectEqual(@as(u64, 0), strict.limiter.snapshotStats(strict.io()).exempt); } test "isLoopback covers both families and nothing else" { try testing.expect(isLoopback(try address.NetAddress.parse("127.0.0.1"))); try testing.expect(isLoopback(try address.NetAddress.parse("127.255.255.254"))); try testing.expect(isLoopback(try address.NetAddress.parse("::1"))); // An IPv4-mapped loopback literal normalizes to the IPv4 form. try testing.expect(isLoopback(address.NetAddress.fromIp(try std.Io.net.IpAddress.parse("::ffff:127.0.0.1", 0)))); try testing.expect(!isLoopback(try address.NetAddress.parse("128.0.0.1"))); try testing.expect(!isLoopback(try address.NetAddress.parse("0.0.0.0"))); try testing.expect(!isLoopback(try address.NetAddress.parse("::"))); try testing.expect(!isLoopback(try address.NetAddress.parse("fd00::1"))); } test "sse connections are capped per address and released" { const fx = try Fixture.init(.{ .rate_per_min = 300, .localhost_exempt = false, .sse_max_per_ip = 2 }); defer fx.deinit(); const client = v4(192, 168, 1, 30); const other = v4(192, 168, 1, 31); try testing.expect(fx.limiter.tryAcquireSse(fx.io(), at(0), client)); try testing.expect(fx.limiter.tryAcquireSse(fx.io(), at(0), client)); try testing.expect(!fx.limiter.tryAcquireSse(fx.io(), at(0), client)); try testing.expectEqual(@as(u16, 2), fx.limiter.sseConnections(fx.io(), client)); try testing.expectEqual(@as(u64, 1), fx.limiter.snapshotStats(fx.io()).sse_refused); // The cap is per address. try testing.expect(fx.limiter.tryAcquireSse(fx.io(), at(0), other)); fx.limiter.releaseSse(fx.io(), client); try testing.expectEqual(@as(u16, 1), fx.limiter.sseConnections(fx.io(), client)); try testing.expect(fx.limiter.tryAcquireSse(fx.io(), at(0), client)); fx.limiter.releaseSse(fx.io(), client); fx.limiter.releaseSse(fx.io(), client); try testing.expectEqual(@as(u16, 0), fx.limiter.sseConnections(fx.io(), client)); // An unmatched release neither underflows nor invents a bucket. fx.limiter.releaseSse(fx.io(), client); fx.limiter.releaseSse(fx.io(), v4(203, 0, 113, 9)); try testing.expectEqual(@as(u16, 0), fx.limiter.sseConnections(fx.io(), client)); try testing.expectEqual(@as(u32, 2), fx.limiter.trackedClients(fx.io())); } test "the sse cap applies to an exempt loopback client too" { const fx = try Fixture.init(.{ .rate_per_min = 300, .localhost_exempt = true, .sse_max_per_ip = 1 }); defer fx.deinit(); const local = v4(127, 0, 0, 1); try testing.expect(fx.limiter.tryAcquireSse(fx.io(), at(0), local)); try testing.expect(!fx.limiter.tryAcquireSse(fx.io(), at(0), local)); fx.limiter.releaseSse(fx.io(), local); try testing.expect(fx.limiter.tryAcquireSse(fx.io(), at(0), local)); } test "sweep drops only idle full buckets and keeps sse holders" { const fx = try Fixture.init(.{ .rate_per_min = 60, .localhost_exempt = false, .sse_max_per_ip = 3 }); defer fx.deinit(); const idle = v4(10, 0, 0, 1); const busy = v4(10, 0, 0, 2); const streaming = v4(10, 0, 0, 3); try testing.expect(fx.limiter.check(fx.io(), at(0), idle).allowed); for (0..60) |_| try testing.expect(fx.limiter.check(fx.io(), at(0), busy).allowed); try testing.expect(fx.limiter.tryAcquireSse(fx.io(), at(0), streaming)); try testing.expect(fx.limiter.check(fx.io(), at(0), streaming).allowed); try testing.expectEqual(@as(u32, 3), fx.limiter.trackedClients(fx.io())); // Before a full window nothing is stale, even though `idle` is full again. try testing.expectEqual(@as(u32, 0), fx.limiter.sweep(fx.io(), at(59))); // At 60 s both quiet buckets are full again and go; `streaming` stays // however long it idles, because its counter is still in use. try testing.expectEqual(@as(u32, 2), fx.limiter.sweep(fx.io(), at(60))); try testing.expectEqual(@as(u32, 1), fx.limiter.trackedClients(fx.io())); try testing.expectEqual(@as(u32, 0), fx.limiter.sweep(fx.io(), at(3600))); try testing.expectEqual(@as(u16, 1), fx.limiter.sseConnections(fx.io(), streaming)); // Releasing the stream makes its bucket collectable. fx.limiter.releaseSse(fx.io(), streaming); try testing.expectEqual(@as(u32, 1), fx.limiter.sweep(fx.io(), at(3601))); try testing.expectEqual(@as(u32, 0), fx.limiter.trackedClients(fx.io())); // A swept client starts from a full bucket rather than inheriting a count. for (0..60) |_| try testing.expect(fx.limiter.check(fx.io(), at(3601), busy).allowed); } test "a full table with no reclaimable bucket refuses unknown clients" { const fx = try Fixture.init(.{ .rate_per_min = 1, .localhost_exempt = false, .sse_max_per_ip = 1 }); defer fx.deinit(); // Every bucket spends its only token, so at t=0 none refills to capacity. for (0..max_clients) |i| { try testing.expect(fx.limiter.check(fx.io(), at(0), indexed(@intCast(i))).allowed); } try testing.expectEqual(@as(u32, max_clients), fx.limiter.trackedClients(fx.io())); try testing.expectEqual(@as(u64, 0), fx.limiter.snapshotStats(fx.io()).untracked); const newcomer = indexed(max_clients); const refused = fx.limiter.check(fx.io(), at(0), newcomer); try testing.expect(!refused.allowed); try testing.expectEqual(@as(u32, 60), refused.retry_after_s); try testing.expectEqual(@as(u64, 1), fx.limiter.snapshotStats(fx.io()).untracked); try testing.expectEqual(@as(u64, 1), fx.limiter.snapshotStats(fx.io()).refused); // The refusal did not displace a tracked client. try testing.expectEqual(@as(u32, max_clients), fx.limiter.trackedClients(fx.io())); // No slot means no SSE counter, so the stream is refused too. try testing.expect(!fx.limiter.tryAcquireSse(fx.io(), at(0), newcomer)); try testing.expectEqual(@as(u64, 1), fx.limiter.snapshotStats(fx.io()).sse_refused); // A tracked client is still limited while the table is full. try testing.expect(!fx.limiter.check(fx.io(), at(0), indexed(0)).allowed); // Sweeping frees room and the newcomer becomes tracked. try testing.expectEqual(@as(u32, max_clients), fx.limiter.sweep(fx.io(), at(3600))); try testing.expect(fx.limiter.check(fx.io(), at(3600), newcomer).allowed); try testing.expectEqual(@as(u32, 1), fx.limiter.trackedClients(fx.io())); } test "a full table evicts a refilled bucket to admit a newcomer" { const fx = try Fixture.init(.{ .rate_per_min = 1, .localhost_exempt = false, .sse_max_per_ip = 1 }); defer fx.deinit(); for (0..max_clients) |i| { try testing.expect(fx.limiter.check(fx.io(), at(0), indexed(@intCast(i))).allowed); } // At t=60 every drained bucket has refilled to capacity and is fair game. const newcomer = indexed(max_clients); try testing.expect(fx.limiter.check(fx.io(), at(60), newcomer).allowed); try testing.expectEqual(@as(u32, max_clients), fx.limiter.trackedClients(fx.io())); try testing.expectEqual(@as(u64, 0), fx.limiter.snapshotStats(fx.io()).untracked); // The newcomer got a real bucket: its second request is rate-limited. try testing.expect(!fx.limiter.check(fx.io(), at(60), newcomer).allowed); // An SSE acquire can reclaim a slot the same way. try testing.expect(fx.limiter.tryAcquireSse(fx.io(), at(60), indexed(max_clients + 1))); try testing.expectEqual(@as(u32, max_clients), fx.limiter.trackedClients(fx.io())); } test "buckets holding sse connections are never evicted" { const fx = try Fixture.init(.{ .rate_per_min = 1, .localhost_exempt = false, .sse_max_per_ip = 1 }); defer fx.deinit(); // Each bucket keeps its full token balance but holds a live stream. for (0..max_clients) |i| { try testing.expect(fx.limiter.tryAcquireSse(fx.io(), at(0), indexed(@intCast(i)))); } const newcomer = indexed(max_clients); try testing.expect(!fx.limiter.check(fx.io(), at(3600), newcomer).allowed); try testing.expect(!fx.limiter.tryAcquireSse(fx.io(), at(3600), newcomer)); try testing.expectEqual(@as(u32, max_clients), fx.limiter.trackedClients(fx.io())); // Releasing one stream makes exactly one slot reclaimable. fx.limiter.releaseSse(fx.io(), indexed(0)); try testing.expect(fx.limiter.check(fx.io(), at(3600), newcomer).allowed); try testing.expectEqual(@as(u32, max_clients), fx.limiter.trackedClients(fx.io())); } test "bucket arithmetic holds far from the timestamp origin" { const fx = try Fixture.init(.{ .rate_per_min = 2, .localhost_exempt = false, .sse_max_per_ip = 1 }); defer fx.deinit(); // Beyond the range of i64 nanoseconds, so only the i96 arithmetic works. const base: i96 = 1 << 80; const client = v4(10, 1, 2, 3); try testing.expect(fx.limiter.check(fx.io(), .{ .nanoseconds = base }, client).allowed); try testing.expect(fx.limiter.check(fx.io(), .{ .nanoseconds = base + 1 }, client).allowed); try testing.expect(!fx.limiter.check(fx.io(), .{ .nanoseconds = base + 2 }, client).allowed); try testing.expect(fx.limiter.check(fx.io(), .{ .nanoseconds = base + window_ns / 2 }, client).allowed); try testing.expectEqual(@as(u32, 1), fx.limiter.sweep(fx.io(), .{ .nanoseconds = base + 4 * window_ns })); } test "a timestamp that goes backwards neither refills nor underflows" { const fx = try Fixture.init(.{ .rate_per_min = 2, .localhost_exempt = false, .sse_max_per_ip = 1 }); defer fx.deinit(); const client = v4(10, 4, 5, 6); try testing.expect(fx.limiter.check(fx.io(), at(100), client).allowed); try testing.expect(fx.limiter.check(fx.io(), at(100), client).allowed); try testing.expect(!fx.limiter.check(fx.io(), at(90), client).allowed); try testing.expect(!fx.limiter.check(fx.io(), at(100), client).allowed); } fn initCheckDeinit(allocator: Allocator) !void { var threaded: std.Io.Threaded = .init(allocator, .{}); defer threaded.deinit(); var limiter = try ApiLimiter.init(allocator, .{ .rate_per_min = 10, .localhost_exempt = false, .sse_max_per_ip = 3, }); defer limiter.deinit(); try testing.expect(limiter.check(threaded.io(), at(0), v4(10, 0, 0, 1)).allowed); } test "init surfaces allocation failure without leaking" { try testing.checkAllAllocationFailures(testing.allocator, initCheckDeinit, .{}); }