nxdns/src/server/rate_limiter.zig

//! Per-client query rate limiter for the DNS listeners. Pure: the caller passes
//! the timestamp, so this file holds no clock, no `std.Io` operation and no
//! socket. `check` neither allocates nor fails.
//!
//! Not thread-safe. Phase 7 decides the locking when it wires the limiter into
//! the query path.
//!
//! The window is fixed, not sliding (PLAN §10 reserves the token bucket for the
//! API limiter). A fixed window admits at most twice the limit across a window
//! boundary. That is acceptable for abuse protection at household scale and it
//! costs one counter per client instead of a timestamp ring.
//!
//! Each client's window is anchored at that client's first query rather than at
//! an absolute boundary: callers pass `.awake` timestamps, whose origin is
//! arbitrary, so an absolute alignment would carry no meaning.
//!
//! The table is bounded at `max_clients`. When it is full and the key is
//! unknown the query is allowed and counted under `untracked`. Refusing unseen
//! clients instead would let `max_clients` attackers deny service to every new
//! device on the LAN, and a household LAN never holds `max_clients` honest
//! clients. The counter makes the condition visible.

const std = @import("std");
const address = @import("../platform/address.zig");

const Allocator = std.mem.Allocator;

/// Upper bound on tracked clients. The table never grows past it, so `check`
/// never allocates.
pub const max_clients = 4096;

pub const Config = struct {
    limit: u32,
    window_seconds: u32,
};

/// `allowed + refused` equals the number of `check` calls. `untracked` counts
/// the subset of `allowed` that the full table could not attribute to a client.
pub const Stats = struct {
    allowed: u64 = 0,
    refused: u64 = 0,
    untracked: u64 = 0,
};

const Window = struct {
    start_ns: i96,
    count: u32,
};

const Table = std.AutoHashMapUnmanaged(address.NetAddress.Key, Window);

pub const RateLimiter = struct {
    gpa: Allocator,
    config: Config,
    window_ns: i96,
    table: Table,
    /// `sweep` collects the keys to drop before it removes any of them, because
    /// a removal invalidates a live iterator. The buffer is owned so that
    /// `sweep` allocates nothing either.
    stale_keys: []address.NetAddress.Key,
    stats: Stats,

    /// Asserts `config.window_seconds` is nonzero; `validate.zig` rejects a zero
    /// window before a config reaches this far.
    pub fn init(gpa: Allocator, config: Config) Allocator.Error!RateLimiter {
        std.debug.assert(config.window_seconds > 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 .{
            .gpa = gpa,
            .config = config,
            .window_ns = @as(i96, config.window_seconds) * std.time.ns_per_s,
            .table = table,
            .stale_keys = stale_keys,
            .stats = .{},
        };
    }

    pub fn deinit(self: *RateLimiter) void {
        self.table.deinit(self.gpa);
        self.gpa.free(self.stale_keys);
        self.* = undefined;
    }

    /// True = process the query; false = answer REFUSED. Never errors, never
    /// allocates.
    pub fn check(self: *RateLimiter, now: std.Io.Timestamp, key: address.NetAddress.Key) bool {
        const window = self.table.getPtr(key) orelse unknown: {
            if (self.table.count() >= max_clients) {
                self.stats.untracked += 1;
                self.stats.allowed += 1;
                return true;
            }
            const gop = self.table.getOrPutAssumeCapacity(key);
            gop.value_ptr.* = .{ .start_ns = now.nanoseconds, .count = 0 };
            break :unknown gop.value_ptr;
        };

        if (now.nanoseconds - window.start_ns >= self.window_ns) {
            window.start_ns = now.nanoseconds;
            window.count = 0;
        }
        if (window.count >= self.config.limit) {
            self.stats.refused += 1;
            return false;
        }
        window.count += 1;
        self.stats.allowed += 1;
        return true;
    }

    /// Drops every entry whose window ended more than one full window before
    /// `now`, that is `now - start_ns > 2 * window_ns`. Returns how many it
    /// dropped. Phase 7 schedules it.
    pub fn sweep(self: *RateLimiter, now: std.Io.Timestamp) u32 {
        const stale_after = 2 * self.window_ns;
        var stale_count: u32 = 0;

        var it = self.table.iterator();
        while (it.next()) |entry| {
            if (now.nanoseconds - entry.value_ptr.start_ns > stale_after) {
                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;
    }

    /// Clients currently holding a window. Reaching `max_clients` is what turns
    /// unknown clients into `untracked` allowances.
    pub fn trackedClients(self: *const RateLimiter) u32 {
        return self.table.count();
    }
};

const testing = std.testing;

fn at(seconds: i64) std.Io.Timestamp {
    return .{ .nanoseconds = @as(i96, seconds) * std.time.ns_per_s };
}

fn v4Key(a: u8, b: u8, c: u8, d: u8) address.NetAddress.Key {
    const addr: address.NetAddress = .{ .ip4 = .{ a, b, c, d } };
    return addr.key();
}

fn indexedKey(index: u32) address.NetAddress.Key {
    var octets: [4]u8 = undefined;
    std.mem.writeInt(u32, &octets, index, .big);
    const addr: address.NetAddress = .{ .ip4 = octets };
    return addr.key();
}

test "allows up to the limit and refuses beyond it" {
    var limiter = try RateLimiter.init(testing.allocator, .{ .limit = 3, .window_seconds = 60 });
    defer limiter.deinit();

    const client = v4Key(192, 168, 1, 10);
    for (0..3) |_| try testing.expect(limiter.check(at(0), client));
    try testing.expect(!limiter.check(at(0), client));
    try testing.expect(!limiter.check(at(59), client));

    try testing.expectEqual(@as(u64, 3), limiter.stats.allowed);
    try testing.expectEqual(@as(u64, 2), limiter.stats.refused);
}

test "a new window resets the count" {
    var limiter = try RateLimiter.init(testing.allocator, .{ .limit = 2, .window_seconds = 60 });
    defer limiter.deinit();

    const client = v4Key(10, 0, 0, 1);
    try testing.expect(limiter.check(at(0), client));
    try testing.expect(limiter.check(at(0), client));
    try testing.expect(!limiter.check(at(0), client));

    // The window is anchored at the first query, so it ends at t = 60.
    try testing.expect(!limiter.check(at(59), client));
    try testing.expect(limiter.check(at(60), client));
    try testing.expect(limiter.check(at(119), client));
    try testing.expect(!limiter.check(at(119), client));
}

test "ipv4 and ipv6 clients hold independent windows" {
    var limiter = try RateLimiter.init(testing.allocator, .{ .limit = 1, .window_seconds = 60 });
    defer limiter.deinit();

    const v4 = (try address.NetAddress.parse("192.168.1.20")).key();
    const v6 = (try address.NetAddress.parse("fd00::20")).key();

    try testing.expect(limiter.check(at(0), v4));
    try testing.expect(!limiter.check(at(0), v4));
    try testing.expect(limiter.check(at(0), v6));
    try testing.expect(!limiter.check(at(0), v6));
    try testing.expectEqual(@as(u32, 2), limiter.trackedClients());
}

test "an ipv4-mapped ipv6 client shares the ipv4 bucket" {
    var limiter = try RateLimiter.init(testing.allocator, .{ .limit = 2, .window_seconds = 60 });
    defer limiter.deinit();

    const mapped = address.NetAddress.fromIp(try std.Io.net.IpAddress.parse("::ffff:192.168.1.30", 53)).key();
    const plain = v4Key(192, 168, 1, 30);
    try testing.expectEqualSlices(u8, &plain, &mapped);

    try testing.expect(limiter.check(at(0), plain));
    try testing.expect(limiter.check(at(0), mapped));
    try testing.expect(!limiter.check(at(0), plain));
    try testing.expectEqual(@as(u32, 1), limiter.trackedClients());
}

test "sweep removes only entries stale by more than one full window" {
    var limiter = try RateLimiter.init(testing.allocator, .{ .limit = 5, .window_seconds = 60 });
    defer limiter.deinit();

    const old = v4Key(10, 0, 0, 1);
    const boundary = v4Key(10, 0, 0, 2);
    const fresh = v4Key(10, 0, 0, 3);

    try testing.expect(limiter.check(at(0), old));
    try testing.expect(limiter.check(at(0), boundary));
    try testing.expect(limiter.check(at(100), fresh));
    try testing.expectEqual(@as(u32, 3), limiter.trackedClients());

    // At t = 120 the boundary entry is exactly 2 windows old and survives.
    try testing.expectEqual(@as(u32, 0), limiter.sweep(at(120)));
    try testing.expectEqual(@as(u32, 3), limiter.trackedClients());

    try testing.expectEqual(@as(u32, 2), limiter.sweep(at(121)));
    try testing.expectEqual(@as(u32, 1), limiter.trackedClients());
    try testing.expect(!limiter.table.contains(old));
    try testing.expect(!limiter.table.contains(boundary));
    try testing.expect(limiter.table.contains(fresh));

    // A swept client starts a fresh window rather than inheriting the old count.
    try testing.expect(limiter.check(at(121), old));
    try testing.expectEqual(@as(u32, 2), limiter.trackedClients());
}

test "a full table allows unknown clients and counts them untracked" {
    var limiter = try RateLimiter.init(testing.allocator, .{ .limit = 1, .window_seconds = 60 });
    defer limiter.deinit();

    for (0..max_clients) |i| {
        try testing.expect(limiter.check(at(0), indexedKey(@intCast(i))));
    }
    try testing.expectEqual(@as(u32, max_clients), limiter.trackedClients());
    try testing.expectEqual(@as(u64, 0), limiter.stats.untracked);

    const newcomer = indexedKey(max_clients);
    try testing.expect(limiter.check(at(0), newcomer));
    try testing.expect(limiter.check(at(0), newcomer));
    try testing.expectEqual(@as(u64, 2), limiter.stats.untracked);
    try testing.expectEqual(@as(u32, max_clients), limiter.trackedClients());

    // A tracked client is still limited while the table is full.
    try testing.expect(!limiter.check(at(0), indexedKey(0)));

    // Sweeping frees room, and the newcomer becomes tracked.
    try testing.expectEqual(@as(u32, max_clients), limiter.sweep(at(200)));
    try testing.expect(limiter.check(at(200), newcomer));
    try testing.expectEqual(@as(u32, 1), limiter.trackedClients());
    try testing.expectEqual(@as(u64, 2), limiter.stats.untracked);
}

test "stats account for every check" {
    var limiter = try RateLimiter.init(testing.allocator, .{ .limit = 4, .window_seconds = 30 });
    defer limiter.deinit();

    var checks: u64 = 0;
    for (0..10) |i| {
        for (0..3) |_| {
            _ = limiter.check(at(@intCast(i)), v4Key(172, 16, 0, @intCast(i)));
            checks += 1;
        }
    }
    try testing.expectEqual(checks, limiter.stats.allowed + limiter.stats.refused);
    try testing.expect(limiter.stats.untracked <= limiter.stats.allowed);
}

test "window arithmetic holds far from the timestamp origin" {
    var limiter = try RateLimiter.init(testing.allocator, .{ .limit = 2, .window_seconds = 60 });
    defer limiter.deinit();

    // Beyond the range of i64 nanoseconds, so only the i96 arithmetic works.
    const base: i96 = 1 << 80;
    const window_ns: i96 = 60 * std.time.ns_per_s;
    const client = v4Key(10, 1, 2, 3);

    try testing.expect(limiter.check(.{ .nanoseconds = base }, client));
    try testing.expect(limiter.check(.{ .nanoseconds = base + 1 }, client));
    try testing.expect(!limiter.check(.{ .nanoseconds = base + window_ns - 1 }, client));
    try testing.expect(limiter.check(.{ .nanoseconds = base + window_ns }, client));
    try testing.expectEqual(@as(u32, 0), limiter.sweep(.{ .nanoseconds = base + 3 * window_ns }));
    try testing.expectEqual(@as(u32, 1), limiter.sweep(.{ .nanoseconds = base + 4 * window_ns }));
}

test "a limit of zero refuses every query" {
    var limiter = try RateLimiter.init(testing.allocator, .{ .limit = 0, .window_seconds = 60 });
    defer limiter.deinit();

    try testing.expect(!limiter.check(at(0), v4Key(10, 0, 0, 1)));
    try testing.expect(!limiter.check(at(0), v4Key(10, 0, 0, 2)));
    try testing.expectEqual(@as(u32, 2), limiter.trackedClients());
    try testing.expectEqual(@as(u64, 2), limiter.stats.refused);
    try testing.expectEqual(@as(u64, 0), limiter.stats.allowed);
}

fn initCheckDeinit(allocator: Allocator) !void {
    var limiter = try RateLimiter.init(allocator, .{ .limit = 10, .window_seconds = 60 });
    defer limiter.deinit();
    try testing.expect(limiter.check(at(0), v4Key(10, 0, 0, 1)));
}

test "init surfaces allocation failure without leaking" {
    try testing.checkAllAllocationFailures(testing.allocator, initCheckDeinit, .{});
}