From cf471eb9f7229f91d79d762db3eee21d91a3ef6e Mon Sep 17 00:00:00 2001 From: DualFroz Date: Fri, 11 Sep 2026 12:16:59 +0200 Subject: [PATCH] feat: add handleMetrics and handleAuth to server --- src/server.zig | 157 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 157 insertions(+) diff --git a/src/server.zig b/src/server.zig index 76e847f..a4264ae 100644 --- a/src/server.zig +++ b/src/server.zig @@ -929,3 +929,160 @@ pub const Server = struct { } } + /// Answers the `metrics` opcode: the network counterpart of the CLI's + /// `stats --prom` (which only ever renders `Db.stats()` locally) plus + /// this server's own request/rate-limit/auth/session counters from + /// `registerMetrics` - until now those counters were only ever + /// observable in-process, by a test holding a direct `*Server` + /// pointer. The registry is built fresh per request and thrown away + /// immediately after rendering: there is nothing to keep alive between + /// calls, since every counter it reads lives on `self` regardless. + fn handleMetrics(self: *Server, writer: anytype) !void { + var reg = metrics.Registry.init(self.allocator); + defer reg.deinit(); + self.registerMetrics(®) catch |err| return writeErr(writer, .internal, @errorName(err)); + + var buf = std.ArrayList(u8).init(self.allocator); + defer buf.deinit(); + reg.exportText(self.allocator, buf.writer()) catch |err| return writeErr(writer, .internal, @errorName(err)); + metrics.formatDbStats(buf.writer(), "cobalt", self.db.stats()) catch |err| return writeErr(writer, .internal, @errorName(err)); + + try protocol.writeFrame(writer, @intFromEnum(protocol.Status.ok), buf.items); + } + + /// Verifies credentials against `store` and, on success, stamps the + /// resulting role into this connection's `connRole` slot so every + /// subsequent request on the same connection is dispatched under it. + /// Locks `self.mutex` only around the actual `Db` read `verify` + /// performs, same as every other handler here. + fn handleAuth(self: *Server, store: authmod.Store, payload: []const u8, writer: anytype, connRole: *?authmod.Role, connExpiresAtNanos: *?i128) !void { + const args = protocol.decodeAuth(payload) catch { + return writeErr(writer, .bad_request, "malformed auth payload"); + }; + + self.mutex.lock(); + const result = store.verify(self.allocator, args.username, args.password); + self.mutex.unlock(); + + const role = result catch |err| { + self.unauthorizedTotal.incr(); + return switch (err) { + error.InvalidCredentials, error.Malformed => writeErr(writer, .unauthorized, "invalid credentials"), + else => writeErr(writer, .internal, @errorName(err)), + }; + }; + connRole.* = role; + connExpiresAtNanos.* = null; + + if (self.sessions) |store_| { + const now = std.time.nanoTimestamp(); + const token = store_.issue(role, self.sessionTtlNanos, now) catch |err| { + return writeErr(writer, .internal, @errorName(err)); + }; + defer self.allocator.free(token); + connExpiresAtNanos.* = now + self.sessionTtlNanos; + const responsePayload = protocol.encodeAuthOk(self.allocator, token) catch |err| { + return writeErr(writer, .internal, @errorName(err)); + }; + defer self.allocator.free(responsePayload); + return protocol.writeFrame(writer, @intFromEnum(protocol.Status.ok), responsePayload); + } + try protocol.writeFrame(writer, @intFromEnum(protocol.Status.ok), ""); + } + + /// Authenticates a connection from a previously issued session token + /// instead of a username/password pair (see `enableSessions`). Only + /// reachable when `self.sessions` is set, since `dispatch` only routes + /// `auth_token` here while `self.authStore` is also set. + fn handleAuthToken(self: *Server, payload: []const u8, writer: anytype, connRole: *?authmod.Role, connExpiresAtNanos: *?i128) !void { + const store = self.sessions orelse { + self.unauthorizedTotal.incr(); + return writeErr(writer, .unauthorized, "session tokens are not enabled on this server"); + }; + const token = protocol.decodeAuthToken(payload); + const session = store.validate(token, std.time.nanoTimestamp()) orelse { + self.unauthorizedTotal.incr(); + return writeErr(writer, .unauthorized, "invalid or expired session token"); + }; + connRole.* = session.role; + connExpiresAtNanos.* = session.expiresAtNanos; + try protocol.writeFrame(writer, @intFromEnum(protocol.Status.ok), ""); + } + + fn handleScan(self: *Server, payload: []const u8, writer: anytype) !void { + const args = protocol.decodeScan(payload) catch { + return writeErr(writer, .bad_request, "malformed scan payload"); + }; + var entries = std.ArrayList(dbmod.Entry).init(self.allocator); + defer entries.deinit(); + + if (args.reverse) { + var it = self.db.scanReverse(self.allocator, .{ .start = args.start, .end = args.end }) catch |err| { + return writeErr(writer, .internal, @errorName(err)); + }; + defer it.deinit(); + while (try it.next()) |e| { + try entries.append(e); + if (args.limit != 0 and entries.items.len >= args.limit) break; + } + } else { + var it = self.db.scanRange(self.allocator, .{ .start = args.start, .end = args.end }) catch |err| { + return writeErr(writer, .internal, @errorName(err)); + }; + defer it.deinit(); + while (try it.next()) |e| { + try entries.append(e); + if (args.limit != 0 and entries.items.len >= args.limit) break; + } + } + + const out = protocol.encodeEntries(self.allocator, entries.items) catch |err| { + return writeErr(writer, .internal, @errorName(err)); + }; + defer self.allocator.free(out); + try protocol.writeFrame(writer, @intFromEnum(protocol.Status.ok), out); + } + + fn handlePrefix(self: *Server, payload: []const u8, writer: anytype) !void { + const args = protocol.decodePrefix(payload) catch { + return writeErr(writer, .bad_request, "malformed prefix payload"); + }; + var it = self.db.scanPrefix(self.allocator, args.prefix) catch |err| { + return writeErr(writer, .internal, @errorName(err)); + }; + defer it.deinit(); + + var entries = std.ArrayList(dbmod.Entry).init(self.allocator); + defer entries.deinit(); + while (try it.next()) |e| { + try entries.append(e); + if (args.limit != 0 and entries.items.len >= args.limit) break; + } + + const out = protocol.encodeEntries(self.allocator, entries.items) catch |err| { + return writeErr(writer, .internal, @errorName(err)); + }; + defer self.allocator.free(out); + try protocol.writeFrame(writer, @intFromEnum(protocol.Status.ok), out); + } + + fn handleBatch(self: *Server, payload: []const u8, writer: anytype) !void { + const ops = protocol.decodeBatch(self.allocator, payload) catch { + return writeErr(writer, .bad_request, "malformed batch payload"); + }; + defer self.allocator.free(ops); + + var wb = dbmod.WriteBatch.init(self.allocator); + defer wb.deinit(); + for (ops) |op| { + switch (op) { + .put => |kv| try wb.put(kv.key, kv.value), + .del => |key| try wb.delete(key), + } + } + self.db.apply(&wb) catch |err| { + return writeErr(writer, .internal, @errorName(err)); + }; + try protocol.writeFrame(writer, @intFromEnum(protocol.Status.ok), ""); + } +