diff --git a/src/main.zig b/src/main.zig index c466896..1a2a3fe 100644 --- a/src/main.zig +++ b/src/main.zig @@ -432,4 +432,5 @@ test { _ = @import("rate_limiter.zig"); _ = @import("server.zig"); _ = @import("relay_metrics.zig"); + _ = @import("subscriptions.zig"); } diff --git a/src/rate_limiter.zig b/src/rate_limiter.zig index f7fcdf4..2ce1b5c 100644 --- a/src/rate_limiter.zig +++ b/src/rate_limiter.zig @@ -1,10 +1,33 @@ const std = @import("std"); const nostr = @import("nostr.zig"); +// Per-IP limiter state is sharded across independent mutex+map buckets so +// worker threads keying on different IPs do not serialize on one process-global +// lock on every event/query/connection. The shard is chosen by hashing the +// final bucket key, so a given IP always maps to the same shard and per-IP +// semantics are preserved exactly. +const SHARD_COUNT = 16; + +// The hash is seeded with a per-process random value chosen at init. The seed is +// public-input-independent, so an attacker who controls the bucket key (e.g. via +// a spoofable X-Forwarded-For under trust_proxy) cannot craft keys that all fall +// on one shard and collapse the 16-way split back to a single mutex. The seed +// only needs to be stable for the process lifetime (per-IP -> same shard), not +// across restarts. +fn shardIndex(seed: u64, key: []const u8) usize { + return std.hash.Wyhash.hash(seed, key) % SHARD_COUNT; +} + +fn randomSeed() u64 { + var buf: [8]u8 = undefined; + nostr.io.randomBytes(&buf); + return std.mem.readInt(u64, &buf, .little); +} + pub const ConnectionLimiter = struct { allocator: std.mem.Allocator, - mutex: std.Io.Mutex, - ip_buckets: std.StringHashMap(IpBucket), + shards: [SHARD_COUNT]Shard, + shard_seed: u64, max_connections_per_ip: u32, cleanup_interval_seconds: i64, @@ -13,32 +36,45 @@ pub const ConnectionLimiter = struct { last_activity: i64, }; + const Shard = struct { + mutex: std.Io.Mutex, + ip_buckets: std.StringHashMap(IpBucket), + }; + pub fn init(allocator: std.mem.Allocator, max_connections_per_ip: u32) ConnectionLimiter { - return .{ + var self: ConnectionLimiter = .{ .allocator = allocator, - .mutex = .init, - .ip_buckets = std.StringHashMap(IpBucket).init(allocator), + .shards = undefined, + .shard_seed = randomSeed(), .max_connections_per_ip = max_connections_per_ip, .cleanup_interval_seconds = 300, }; + for (&self.shards) |*shard| { + shard.* = .{ .mutex = .init, .ip_buckets = std.StringHashMap(IpBucket).init(allocator) }; + } + return self; } pub fn deinit(self: *ConnectionLimiter) void { - var iter = self.ip_buckets.keyIterator(); - while (iter.next()) |key| { - self.allocator.free(key.*); + for (&self.shards) |*shard| { + var iter = shard.ip_buckets.keyIterator(); + while (iter.next()) |key| { + self.allocator.free(key.*); + } + shard.ip_buckets.deinit(); } - self.ip_buckets.deinit(); } pub fn canConnect(self: *ConnectionLimiter, raw_ip: []const u8) bool { const io = nostr.io.io(); - self.mutex.lockUncancelable(io); - defer self.mutex.unlock(io); - var key_buf: [19]u8 = undefined; const ip = bucketKey(raw_ip, &key_buf); - if (self.ip_buckets.get(ip)) |bucket| { + const shard = &self.shards[shardIndex(self.shard_seed, ip)]; + + shard.mutex.lockUncancelable(io); + defer shard.mutex.unlock(io); + + if (shard.ip_buckets.get(ip)) |bucket| { return bucket.connection_count < self.max_connections_per_ip; } return true; @@ -49,14 +85,16 @@ pub const ConnectionLimiter = struct { /// closing the canConnect/addConnection check-then-act race. pub fn tryAcquireConnection(self: *ConnectionLimiter, raw_ip: []const u8) bool { const io = nostr.io.io(); - self.mutex.lockUncancelable(io); - defer self.mutex.unlock(io); + var key_buf: [19]u8 = undefined; + const ip = bucketKey(raw_ip, &key_buf); + const shard = &self.shards[shardIndex(self.shard_seed, ip)]; + + shard.mutex.lockUncancelable(io); + defer shard.mutex.unlock(io); const now = nostr.io.timestamp(); - var key_buf: [19]u8 = undefined; - const ip = bucketKey(raw_ip, &key_buf); - if (self.ip_buckets.getPtr(ip)) |bucket| { + if (shard.ip_buckets.getPtr(ip)) |bucket| { if (bucket.connection_count >= self.max_connections_per_ip) return false; bucket.connection_count += 1; bucket.last_activity = now; @@ -64,7 +102,7 @@ pub const ConnectionLimiter = struct { } const ip_copy = self.allocator.dupe(u8, ip) catch return false; - self.ip_buckets.put(ip_copy, .{ + shard.ip_buckets.put(ip_copy, .{ .connection_count = 1, .last_activity = now, }) catch { @@ -76,12 +114,14 @@ pub const ConnectionLimiter = struct { pub fn removeConnection(self: *ConnectionLimiter, raw_ip: []const u8) void { const io = nostr.io.io(); - self.mutex.lockUncancelable(io); - defer self.mutex.unlock(io); - var key_buf: [19]u8 = undefined; const ip = bucketKey(raw_ip, &key_buf); - if (self.ip_buckets.getPtr(ip)) |bucket| { + const shard = &self.shards[shardIndex(self.shard_seed, ip)]; + + shard.mutex.lockUncancelable(io); + defer shard.mutex.unlock(io); + + if (shard.ip_buckets.getPtr(ip)) |bucket| { if (bucket.connection_count > 0) { bucket.connection_count -= 1; } @@ -90,43 +130,90 @@ pub const ConnectionLimiter = struct { pub fn cleanup(self: *ConnectionLimiter) void { const io = nostr.io.io(); - self.mutex.lockUncancelable(io); - defer self.mutex.unlock(io); - const now = nostr.io.timestamp(); - var to_remove: std.ArrayListUnmanaged([]const u8) = .empty; - defer to_remove.deinit(self.allocator); - var iter = self.ip_buckets.iterator(); - while (iter.next()) |entry| { - if (entry.value_ptr.connection_count == 0 and - now - entry.value_ptr.last_activity > self.cleanup_interval_seconds) - { - to_remove.append(self.allocator, entry.key_ptr.*) catch continue; + for (&self.shards) |*shard| { + shard.mutex.lockUncancelable(io); + defer shard.mutex.unlock(io); + + var to_remove: std.ArrayListUnmanaged([]const u8) = .empty; + defer to_remove.deinit(self.allocator); + + var iter = shard.ip_buckets.iterator(); + while (iter.next()) |entry| { + if (entry.value_ptr.connection_count == 0 and + now - entry.value_ptr.last_activity > self.cleanup_interval_seconds) + { + to_remove.append(self.allocator, entry.key_ptr.*) catch continue; + } } - } - for (to_remove.items) |key| { - _ = self.ip_buckets.remove(key); - self.allocator.free(key); + for (to_remove.items) |key| { + _ = shard.ip_buckets.remove(key); + self.allocator.free(key); + } } } pub fn getStats(self: *ConnectionLimiter) Stats { const io = nostr.io.io(); - self.mutex.lockUncancelable(io); - defer self.mutex.unlock(io); - - return .{ - .tracked_ips = self.ip_buckets.count(), - }; + var tracked: usize = 0; + for (&self.shards) |*shard| { + shard.mutex.lockUncancelable(io); + defer shard.mutex.unlock(io); + tracked += shard.ip_buckets.count(); + } + return .{ .tracked_ips = tracked }; } pub const Stats = struct { tracked_ips: usize, }; + + // Locate the bucket for an IP so a test can backdate last_activity. Mirrors + // EventRateLimiter.tracks: test-only, no lock (single-threaded tests). + fn bucketPtrForTest(self: *ConnectionLimiter, raw_ip: []const u8) ?*IpBucket { + var key_buf: [19]u8 = undefined; + const ip = bucketKey(raw_ip, &key_buf); + return self.shards[shardIndex(self.shard_seed, ip)].ip_buckets.getPtr(ip); + } }; +test "ConnectionLimiter cleanup reclaims only idle buckets across shards" { + var limiter = ConnectionLimiter.init(std.testing.allocator, 5); + defer limiter.deinit(); + + // Track four IPs. With a per-process seed they spread across shards, and + // cleanup() sweeps every shard, so the assertion holds regardless of layout. + try std.testing.expect(limiter.tryAcquireConnection("1.1.1.1")); + try std.testing.expect(limiter.tryAcquireConnection("2.2.2.2")); + try std.testing.expect(limiter.tryAcquireConnection("3.3.3.3")); + try std.testing.expect(limiter.tryAcquireConnection("4.4.4.4")); + try std.testing.expectEqual(@as(usize, 4), limiter.getStats().tracked_ips); + + const now = nostr.io.timestamp(); + + // Two IPs go idle (count 0) and stale (last_activity older than the 300s + // cleanup interval): these must be reclaimed. + limiter.removeConnection("1.1.1.1"); + limiter.removeConnection("2.2.2.2"); + limiter.bucketPtrForTest("1.1.1.1").?.last_activity = now - 400; + limiter.bucketPtrForTest("2.2.2.2").?.last_activity = now - 400; + + // 3.3.3.3 stays active (count > 0). 4.4.4.4 is idle but recent. Both remain. + limiter.removeConnection("4.4.4.4"); + limiter.bucketPtrForTest("4.4.4.4").?.last_activity = now; + + limiter.cleanup(); + + try std.testing.expectEqual(@as(usize, 2), limiter.getStats().tracked_ips); + + // A fresh limiter tracks nothing. + var empty = ConnectionLimiter.init(std.testing.allocator, 5); + defer empty.deinit(); + try std.testing.expectEqual(@as(usize, 0), empty.getStats().tracked_ips); +} + pub const IpFilter = struct { allocator: std.mem.Allocator, whitelist: std.StringHashMap(void), @@ -375,8 +462,8 @@ test "normalizeIp" { pub const EventRateLimiter = struct { allocator: std.mem.Allocator, - mutex: std.Io.Mutex, - ip_buckets: std.StringHashMap(EventBucket), + shards: [SHARD_COUNT]Shard, + shard_seed: u64, events_per_minute: u32, // Idle buckets are reclaimed after this many seconds; a bucket idle for a @@ -392,21 +479,32 @@ pub const EventRateLimiter = struct { last_refill: i64, }; + const Shard = struct { + mutex: std.Io.Mutex, + ip_buckets: std.StringHashMap(EventBucket), + }; + pub fn init(allocator: std.mem.Allocator, events_per_minute: u32) EventRateLimiter { - return .{ + var self: EventRateLimiter = .{ .allocator = allocator, - .mutex = .init, - .ip_buckets = std.StringHashMap(EventBucket).init(allocator), + .shards = undefined, + .shard_seed = randomSeed(), .events_per_minute = events_per_minute, }; + for (&self.shards) |*shard| { + shard.* = .{ .mutex = .init, .ip_buckets = std.StringHashMap(EventBucket).init(allocator) }; + } + return self; } pub fn deinit(self: *EventRateLimiter) void { - var iter = self.ip_buckets.keyIterator(); - while (iter.next()) |key| { - self.allocator.free(key.*); + for (&self.shards) |*shard| { + var iter = shard.ip_buckets.keyIterator(); + while (iter.next()) |key| { + self.allocator.free(key.*); + } + shard.ip_buckets.deinit(); } - self.ip_buckets.deinit(); } pub fn checkAndRecord(self: *EventRateLimiter, ip: []const u8) bool { @@ -417,13 +515,14 @@ pub const EventRateLimiter = struct { fn checkAndRecordAt(self: *EventRateLimiter, ip: []const u8, now: i64) bool { const io = nostr.io.io(); - self.mutex.lockUncancelable(io); - defer self.mutex.unlock(io); + const shard = &self.shards[shardIndex(self.shard_seed, ip)]; + shard.mutex.lockUncancelable(io); + defer shard.mutex.unlock(io); const capacity: f64 = @floatFromInt(self.events_per_minute); const refill_per_sec: f64 = capacity / 60.0; - if (self.ip_buckets.getPtr(ip)) |bucket| { + if (shard.ip_buckets.getPtr(ip)) |bucket| { const elapsed: f64 = @floatFromInt(@max(@as(i64, 0), now - bucket.last_refill)); bucket.tokens = @min(capacity, bucket.tokens + elapsed * refill_per_sec); bucket.last_refill = now; @@ -436,7 +535,7 @@ pub const EventRateLimiter = struct { // First event from this IP: start with a full bucket, then consume one. const ip_copy = self.allocator.dupe(u8, ip) catch return false; - self.ip_buckets.put(ip_copy, .{ + shard.ip_buckets.put(ip_copy, .{ .tokens = capacity - 1.0, .last_refill = now, }) catch { @@ -452,24 +551,39 @@ pub const EventRateLimiter = struct { fn cleanupAt(self: *EventRateLimiter, now: i64) void { const io = nostr.io.io(); - self.mutex.lockUncancelable(io); - defer self.mutex.unlock(io); - var to_remove: std.ArrayListUnmanaged([]const u8) = .empty; - defer to_remove.deinit(self.allocator); + for (&self.shards) |*shard| { + shard.mutex.lockUncancelable(io); + defer shard.mutex.unlock(io); - var iter = self.ip_buckets.iterator(); - while (iter.next()) |entry| { - if (now - entry.value_ptr.last_refill >= IDLE_SECONDS) { - to_remove.append(self.allocator, entry.key_ptr.*) catch continue; + var to_remove: std.ArrayListUnmanaged([]const u8) = .empty; + defer to_remove.deinit(self.allocator); + + var iter = shard.ip_buckets.iterator(); + while (iter.next()) |entry| { + if (now - entry.value_ptr.last_refill >= IDLE_SECONDS) { + to_remove.append(self.allocator, entry.key_ptr.*) catch continue; + } } - } - for (to_remove.items) |key| { - _ = self.ip_buckets.remove(key); - self.allocator.free(key); + for (to_remove.items) |key| { + _ = shard.ip_buckets.remove(key); + self.allocator.free(key); + } } } + + // Total tracked IPs across all shards. Test-only helper; the sharded map has + // no single count() to assert against. + fn trackedCount(self: *EventRateLimiter) usize { + var total: usize = 0; + for (&self.shards) |*shard| total += shard.ip_buckets.count(); + return total; + } + + fn tracks(self: *EventRateLimiter, ip: []const u8) bool { + return self.shards[shardIndex(self.shard_seed, ip)].ip_buckets.contains(ip); + } }; test "EventRateLimiter enforces a limit above 255" { @@ -548,11 +662,11 @@ test "EventRateLimiter cleanup reclaims only idle buckets" { // At t=1059, the first bucket is idle for 59s (<60), so nothing is reclaimed. limiter.cleanupAt(1059); - try std.testing.expectEqual(@as(u32, 2), limiter.ip_buckets.count()); + try std.testing.expectEqual(@as(usize, 2), limiter.trackedCount()); // At t=1060, the first bucket hits IDLE_SECONDS and is removed; the second // (idle 10s) is retained. limiter.cleanupAt(1060); - try std.testing.expectEqual(@as(u32, 1), limiter.ip_buckets.count()); - try std.testing.expect(limiter.ip_buckets.contains("2.2.2.2")); + try std.testing.expectEqual(@as(usize, 1), limiter.trackedCount()); + try std.testing.expect(limiter.tracks("2.2.2.2")); } diff --git a/src/subscriptions.zig b/src/subscriptions.zig index 364338a..c61d22c 100644 --- a/src/subscriptions.zig +++ b/src/subscriptions.zig @@ -153,6 +153,28 @@ pub const Subscriptions = struct { sub_id: []const u8, }; + // Reused per-broadcast scratch. forEachMatching runs on the single writer + // thread in durable sync modes but on multiple worker threads concurrently + // in sync=none, so the scratch is threadlocal (no shared-mutable state) and + // cleared per call rather than reallocated, avoiding per-event glibc arena + // churn in the fan-out hot path. Reusing one buffer per thread makes + // forEachMatching non-reentrant on a thread: callees (the post-unlock socket + // writes) must never re-enter it, or the inner clear would corrupt the outer + // drain and unbalance write_guard. + // + // These are struct-scoped threadlocals shared across every Subscriptions + // instance on a thread, and tl_seen is bound to whichever instance's + // allocator first initializes it. That is correct ONLY because there is + // exactly one Subscriptions instance, backed by one process-global allocator; + // a second instance with a different allocator would alias this state. + // + // The buffers are never freed or shrunk: each thread retains its peak + // fan-out capacity for the process lifetime. This is bounded and intentional, + // and it assumes the worker-pool threads live for the whole process. A + // respawned pool thread would orphan (leak) its buffers. + threadlocal var tl_pending: std.ArrayListUnmanaged(PendingWrite) = .empty; + threadlocal var tl_seen: ?std.AutoHashMap(u64, void) = null; + pub fn forEachMatching( self: *Subscriptions, event: *const nostr.Event, @@ -166,23 +188,24 @@ pub const Subscriptions = struct { // send timeout. Each snapshotted connection has its write_guard bumped // while still in the registry; removeConnection plus close() drain the // guard before freeing, so writing post-unlock is use-after-free safe. - var pending: std.ArrayListUnmanaged(PendingWrite) = .empty; - defer pending.deinit(self.allocator); + const pending = &tl_pending; + pending.clearRetainingCapacity(); + + if (tl_seen == null) tl_seen = std.AutoHashMap(u64, void).init(self.allocator); + const seen = &tl_seen.?; + seen.clearRetainingCapacity(); { self.rwlock.lockSharedUncancelable(io); defer self.rwlock.unlockShared(io); - var seen = std.AutoHashMap(u64, void).init(self.allocator); - defer seen.deinit(); - // Kind-indexed candidates: only connections subscribed to this kind. if (self.kind_index.get(event.kind())) |conn_ids| { for (conn_ids.items) |conn_id| { if (seen.contains(conn_id)) continue; seen.put(conn_id, {}) catch continue; const conn = self.connections.get(conn_id) orelse continue; - snapshotMatch(self, &pending, conn, event); + snapshotMatch(self, pending, conn, event); } } @@ -194,7 +217,7 @@ pub const Subscriptions = struct { if (seen.contains(conn.id)) continue; if (!connHasWildcard(conn)) continue; seen.put(conn.id, {}) catch continue; - snapshotMatch(self, &pending, conn, event); + snapshotMatch(self, pending, conn, event); } } @@ -243,3 +266,71 @@ pub const Subscriptions = struct { }; } }; + +const testing = std.testing; + +// 64- and 128-hex-char fields so Event.parse accepts the structure. +const HID = "0" ** 63 ++ "1"; +const HPK = "0" ** 63 ++ "2"; +const HSIG = "0" ** 127 ++ "3"; + +fn eventJson(comptime kind: []const u8) []const u8 { + return "[\"EVENT\",{\"id\":\"" ++ HID ++ "\",\"pubkey\":\"" ++ HPK ++ + "\",\"created_at\":1700000000,\"kind\":" ++ kind ++ + ",\"tags\":[],\"content\":\"hi\",\"sig\":\"" ++ HSIG ++ "\"}]"; +} + +test "forEachMatching reuses threadlocal scratch cleanly across calls" { + // The Subscriptions instance owns the threadlocal scratch (tl_pending/tl_seen) + // which is intentionally never freed for the process lifetime, so back it with + // a non-leak-checking allocator. Connections use testing.allocator and are + // deinitialized normally. + var subs = Subscriptions.init(std.heap.page_allocator); + defer subs.deinit(); + + var conn1: Connection = undefined; + conn1.init(testing.allocator, 1); + defer conn1.deinit(); + var conn7: Connection = undefined; + conn7.init(testing.allocator, 7); + defer conn7.deinit(); + + try subs.addConnection(&conn1); + try subs.addConnection(&conn7); + + var req1 = try nostr.ClientMsg.parse("[\"REQ\",\"k1\",{\"kinds\":[1]}]"); + defer req1.deinit(); + try subs.subscribe(&conn1, "k1", try req1.getFilters(conn1.allocator()), 5); + + var req7 = try nostr.ClientMsg.parse("[\"REQ\",\"k7\",{\"kinds\":[7]}]"); + defer req7.deinit(); + try subs.subscribe(&conn7, "k7", try req7.getFilters(conn7.allocator()), 5); + + const msg_buf = try testing.allocator.create([65536]u8); + defer testing.allocator.destroy(msg_buf); + + // First broadcast: a kind-1 event reaches only conn1. + var ev1 = try nostr.ClientMsg.parse(eventJson("1")); + defer ev1.deinit(); + const event1 = try ev1.getEvent(); + subs.forEachMatching(&event1, msg_buf); + + try testing.expectEqual(@as(u64, 1), conn1.events_sent.load(.monotonic)); + try testing.expectEqual(@as(u64, 0), conn7.events_sent.load(.monotonic)); + // The post-unlock drain balances every guard it took. + try testing.expectEqual(@as(u32, 0), conn1.write_guard.load(.acquire)); + try testing.expectEqual(@as(u32, 0), conn7.write_guard.load(.acquire)); + + // Second broadcast on the same thread: a kind-7 event reaches only conn7. If + // the threadlocal scratch carried stale entries from call 1, conn1 would be + // written again or its guard would leak. + var ev7 = try nostr.ClientMsg.parse(eventJson("7")); + defer ev7.deinit(); + const event7 = try ev7.getEvent(); + subs.forEachMatching(&event7, msg_buf); + + try testing.expectEqual(@as(u64, 1), conn1.events_sent.load(.monotonic)); + try testing.expectEqual(@as(u64, 1), conn7.events_sent.load(.monotonic)); + try testing.expectEqual(@as(u32, 0), conn1.write_guard.load(.acquire)); + try testing.expectEqual(@as(u32, 0), conn7.write_guard.load(.acquire)); +}