From 7ab58fe4e41d670c14452ecb568fe0e34476b557 Mon Sep 17 00:00:00 2001 From: DualFroz Date: Thu, 10 Sep 2026 21:35:17 +0200 Subject: [PATCH] feat: add enablePerConnectionRateLimit and enableAuth to server --- src/server.zig | 156 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 156 insertions(+) diff --git a/src/server.zig b/src/server.zig index 54a8c04..9e677cb 100644 --- a/src/server.zig +++ b/src/server.zig @@ -158,3 +158,159 @@ pub const Server = struct { self.limiter = metrics.TokenBucket.init(capacity, refillPerSecond, @intCast(std.time.nanoTimestamp())); } + /// Gives every accepted connection its own private token bucket, + /// independent of every other connection's budget and of the shared + /// `limiter` above. A fresh bucket is created per connection in + /// `handleConnection` (not shared state, so no locking is needed beyond + /// what TokenBucket already does internally) - one noisy client cannot + /// starve another client's share of the server, which a single global + /// bucket cannot express. Both limits can be enabled together: a + /// request must clear the per-connection budget before it is even + /// checked against the global one. Disabled by default; call once + /// before `start()`. + pub fn enablePerConnectionRateLimit(self: *Server, capacity: f64, refillPerSecond: f64) void { + self.perConnLimit = .{ .capacity = capacity, .refillPerSecond = refillPerSecond }; + } + + /// Turns on username/password authentication and role enforcement for + /// every connection: after this call, every opcode except `ping` and + /// `auth` itself requires a prior successful `auth` on that connection + /// (`unauthorized` otherwise), and the authenticated role must meet + /// `requiredRole(op)` (`forbidden` otherwise). `store` is not owned by + /// the server - the caller keeps managing the underlying `Db` user + /// records (creating the first admin, etc.) through `store` directly. + /// Disabled (no auth required, matching every prior version of this + /// server) by default; call once before `start()`. + pub fn enableAuth(self: *Server, store: authmod.Store) void { + self.authStore = store; + } + + /// Layers expiring session tokens on top of password auth (requires + /// `enableAuth` to also be called - sessions are only ever minted from + /// a successful password check, never a substitute for one). Once + /// enabled: a successful `auth` response carries a fresh token in its + /// payload (empty payload otherwise, unchanged from before this + /// existed) valid for `ttlNanos`; a NEW connection can skip the + /// password entirely by sending that token as `auth_token` instead; + /// and - the actual behavior change from plain password auth - a + /// connection's authenticated role stops being honored the moment its + /// session's expiry passes, even if the TCP connection itself has + /// stayed open the whole time (`dispatch` re-checks the wall clock on + /// every request, not just at `auth` time). `store` is not owned by + /// the server, exactly like `enableAuth`'s `authmod.Store` - the + /// caller keeps it alive and can inspect/revoke sessions directly. + pub fn enableSessions(self: *Server, store: *sessionmod.SessionStore, ttlNanos: i128) void { + self.sessions = store; + self.sessionTtlNanos = ttlNanos; + } + + /// Runs `SessionStore.sweepExpired` on a background thread every + /// `intervalNanos`, starting once `start()` is called. Without this, a + /// session nobody ever presents again (the client that held it simply + /// vanished) sits in `sessions` forever - `validate()` only ever cleans + /// up a token someone actually tries to use. Requires `enableSessions` + /// to have been called first (checked in `start()`, not here, since + /// callers are free to call these two in either order before `start`). + pub fn enableSessionSweep(self: *Server, intervalNanos: u64) void { + self.sweepIntervalNanos = intervalNanos; + } + + /// Runs `ttl.sweepExpired` on a background thread every `intervalNanos`, + /// starting once `start()` is called - independent of `enableAuth`/ + /// `enableSessions`/`enableSessionSweep` above, since key TTLs (set via + /// the `ttl_op` opcode or the CLI's local `put --ttl`/`expire`) are a + /// property of the `Db` itself, not of authentication. Without this, a + /// key whose TTL elapsed but that nobody ever reads again (so + /// `ttl.getLive`'s lazy expiry never fires) sits in the store forever - + /// exactly the same gap `enableSessionSweep` closes for session tokens. + /// Call once before `start()`. + pub fn enableTtlSweep(self: *Server, intervalNanos: u64) void { + self.ttlSweepIntervalNanos = intervalNanos; + } + + /// Runs `collttl.sweepExpired` on a background thread every + /// `intervalNanos`, starting once `start()` is called - the whole- + /// collection counterpart of `enableTtlSweep` above, independent of it + /// (a server can enable one, both, or neither). Without this, a + /// collection whose TTL elapsed but that nobody ever touches again (so + /// `coll_ttl_op`'s lazy `expireIfDue` never fires, matching the CLI's + /// local `coll` dispatcher) sits in the store forever. Call once before + /// `start()`. + pub fn enableCollTtlSweep(self: *Server, intervalNanos: u64) void { + self.collTtlSweepIntervalNanos = intervalNanos; + } + + /// Registers this server's request counters into `reg` under fixed + /// names, so they show up alongside any other Prometheus metrics an + /// embedding application already exports. + pub fn registerMetrics(self: *Server, reg: *metrics.Registry) !void { + try reg.registerCounter("cobalt_server_requests_total", &self.requestsTotal); + try reg.registerCounter("cobalt_server_rate_limited_total", &self.rateLimitedTotal); + try reg.registerCounter("cobalt_server_conn_rate_limited_total", &self.connRateLimitedTotal); + try reg.registerCounter("cobalt_server_unauthorized_total", &self.unauthorizedTotal); + try reg.registerCounter("cobalt_server_forbidden_total", &self.forbiddenTotal); + try reg.registerCounter("cobalt_server_session_expired_total", &self.sessionExpiredTotal); + try reg.registerCounter("cobalt_server_sessions_swept_total", &self.sessionsSweptTotal); + try reg.registerCounter("cobalt_server_ttl_swept_total", &self.ttlSweptTotal); + try reg.registerCounter("cobalt_server_coll_ttl_swept_total", &self.collTtlSweptTotal); + try reg.registerCounter("cobalt_server_messages_published_total", &self.messagesPublishedTotal); + } + + pub fn start(self: *Server) !void { + self.acceptThread = try std.Thread.spawn(.{}, acceptLoop, .{self}); + if (self.sessions != null and self.sweepIntervalNanos > 0) { + self.sweeperThread = try std.Thread.spawn(.{}, sweepLoop, .{self}); + } + if (self.ttlSweepIntervalNanos > 0) { + self.ttlSweeperThread = try std.Thread.spawn(.{}, ttlSweepLoop, .{self}); + } + if (self.collTtlSweepIntervalNanos > 0) { + self.collTtlSweeperThread = try std.Thread.spawn(.{}, collTtlSweepLoop, .{self}); + } + } + + /// Blocks until the accept loop exits (only happens via `stop()` or + /// an unrecoverable accept error). Used by long-running foreground + /// servers such as the CLI's `serve` command. + pub fn wait(self: *Server) void { + if (self.acceptThread) |t| t.join(); + } + + /// Background loop started by `start()` when `enableSessionSweep` set a + /// nonzero interval and sessions are enabled. Sleeps the full interval + /// between sweeps (a sweep is pure housekeeping, never load-bearing for + /// correctness - `validate()` already refuses an expired token on its + /// own), and checks `stopping` both before and after each sleep so + /// `stop()` never has to wait more than one interval to join it. + fn sweepLoop(self: *Server) void { + const store = self.sessions orelse return; + while (!self.stopping.load(.acquire)) { + std.time.sleep(self.sweepIntervalNanos); + if (self.stopping.load(.acquire)) return; + const dropped = store.sweepExpired(std.time.nanoTimestamp()); + if (dropped > 0) self.sessionsSweptTotal.incrBy(@intCast(dropped)); + } + } + + /// Background loop started by `start()` when `enableTtlSweep` set a + /// nonzero interval. Unlike `sweepLoop` above, `ttl.sweepExpired` reads + /// and mutates the same `Db` every request handler touches, so this + /// loop must take `self.mutex` around the call exactly like `dispatch` + /// does for every opcode - `Db` itself keeps no internal locking (see + /// the doc comment on `Server`). A failed sweep (e.g. a transient I/O + /// error) is skipped rather than crashing the server; the next tick + /// tries again. + fn ttlSweepLoop(self: *Server) void { + while (!self.stopping.load(.acquire)) { + std.time.sleep(self.ttlSweepIntervalNanos); + if (self.stopping.load(.acquire)) return; + self.mutex.lock(); + const reaped = ttlmod.sweepExpired(self.db, self.allocator, std.time.nanoTimestamp()) catch { + self.mutex.unlock(); + continue; + }; + self.mutex.unlock(); + if (reaped > 0) self.ttlSweptTotal.incrBy(@intCast(reaped)); + } + } +