From 48ff5cf38e042380d6a6842b7f449c335d695a92 Mon Sep 17 00:00:00 2001 From: DualFroz Date: Fri, 11 Sep 2026 17:11:42 +0200 Subject: [PATCH] feat: add WorkerCtx and worker to server --- src/server.zig | 159 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 159 insertions(+) diff --git a/src/server.zig b/src/server.zig index 606e780..4b90135 100644 --- a/src/server.zig +++ b/src/server.zig @@ -2171,3 +2171,162 @@ test "server batch applies puts and deletes atomically" { try std.testing.expect(kept == null); } +test "server rejects a malformed frame with bad_request without crashing" { + const allocator = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const path = try tmp.dir.realpathAlloc(allocator, "."); + defer allocator.free(path); + + var db = try testDb(allocator, path); + defer db.close(); + + const address = try std.net.Address.parseIp("127.0.0.1", 0); + var server = try Server.init(allocator, &db, address); + try server.start(); + defer server.stop(); + + const stream = try std.net.tcpConnectToAddress(server.boundAddress()); + defer stream.close(); + + try protocol.writeFrame(stream.writer(), @intFromEnum(protocol.Op.set), "not a valid set payload with no length prefix boundary"); + var frame = try protocol.readFrame(allocator, stream.reader()); + defer frame.deinit(); + try std.testing.expectEqual(@intFromEnum(protocol.Status.bad_request), frame.tag); +} + +test "server serves multiple concurrent clients without corrupting state" { + const allocator = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const path = try tmp.dir.realpathAlloc(allocator, "."); + defer allocator.free(path); + + var db = try testDb(allocator, path); + defer db.close(); + + const address = try std.net.Address.parseIp("127.0.0.1", 0); + var server = try Server.init(allocator, &db, address); + try server.start(); + defer server.stop(); + + const client_mod = @import("client.zig"); + const WorkerCtx = struct { + allocator: Allocator, + addr: std.net.Address, + index: usize, + ok: bool = false, + }; + + const worker = struct { + fn run(ctx: *WorkerCtx) void { + var client = client_mod.Client.connect(ctx.allocator, ctx.addr) catch return; + defer client.close(); + var buf: [32]u8 = undefined; + const key = std.fmt.bufPrint(&buf, "k{d}", .{ctx.index}) catch return; + client.set(key, "v") catch return; + const got = client.get(ctx.allocator, key) catch return; + defer if (got) |g| ctx.allocator.free(g); + ctx.ok = got != null and std.mem.eql(u8, got.?, "v"); + } + }.run; + + var contexts: [8]WorkerCtx = undefined; + var threads: [8]std.Thread = undefined; + for (0..8) |i| { + contexts[i] = .{ .allocator = allocator, .addr = server.boundAddress(), .index = i }; + threads[i] = try std.Thread.spawn(.{}, worker, .{&contexts[i]}); + } + for (0..8) |i| threads[i].join(); + for (0..8) |i| try std.testing.expect(contexts[i].ok); +} + +test "server enforces a token-bucket rate limit and rejects the overflow" { + const allocator = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const path = try tmp.dir.realpathAlloc(allocator, "."); + defer allocator.free(path); + + var db = try testDb(allocator, path); + defer db.close(); + + const address = try std.net.Address.parseIp("127.0.0.1", 0); + var server = try Server.init(allocator, &db, address); + // Slow refill (1 token/sec): the whole burst below completes in well + // under a second of wall-clock time, so refill contributes nowhere + // near a full token and the outcome is deterministic. + server.enableRateLimit(3, 1); + try server.start(); + defer server.stop(); + + const client_mod = @import("client.zig"); + var client = try client_mod.Client.connect(allocator, server.boundAddress()); + defer client.close(); + + var allowed: usize = 0; + var rejected: usize = 0; + var i: usize = 0; + while (i < 6) : (i += 1) { + client.ping() catch |err| { + try std.testing.expectEqual(client_mod.ClientError.ServerRateLimited, err); + rejected += 1; + continue; + }; + allowed += 1; + } + + try std.testing.expectEqual(@as(usize, 3), allowed); + try std.testing.expectEqual(@as(usize, 3), rejected); + try std.testing.expectEqual(@as(u64, 6), server.requestsTotal.get()); + try std.testing.expectEqual(@as(u64, 3), server.rateLimitedTotal.get()); +} + +test "server without an enabled rate limit never rejects" { + const allocator = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const path = try tmp.dir.realpathAlloc(allocator, "."); + defer allocator.free(path); + + var db = try testDb(allocator, path); + defer db.close(); + + const address = try std.net.Address.parseIp("127.0.0.1", 0); + var server = try Server.init(allocator, &db, address); + try server.start(); + defer server.stop(); + + const client_mod = @import("client.zig"); + var client = try client_mod.Client.connect(allocator, server.boundAddress()); + defer client.close(); + + var i: usize = 0; + while (i < 20) : (i += 1) try client.ping(); + try std.testing.expectEqual(@as(u64, 20), server.requestsTotal.get()); + try std.testing.expectEqual(@as(u64, 0), server.rateLimitedTotal.get()); +} + +test "per-connection rate limit gives each connection its own budget" { + const allocator = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const path = try tmp.dir.realpathAlloc(allocator, "."); + defer allocator.free(path); + + var db = try testDb(allocator, path); + defer db.close(); + + const address = try std.net.Address.parseIp("127.0.0.1", 0); + var server = try Server.init(allocator, &db, address); + server.enablePerConnectionRateLimit(2, 1); + try server.start(); + defer server.stop(); + + const client_mod = @import("client.zig"); + + var clientA = try client_mod.Client.connect(allocator, server.boundAddress()); + defer clientA.close(); + var clientB = try client_mod.Client.connect(allocator, server.boundAddress()); + defer clientB.close(); +