nxdns/tests/integration_tests.zig

const std = @import("std");
const testing = std.testing;

const handler_mod = @import("handler");
const rate_limiter = @import("rate_limiter");
const cache = @import("cache");
const blocklist_mod = @import("blocklist");
const packet = @import("packet");
const types = @import("types");
const Name = @import("name").Name;

// ============================================================================
// Test DNS Query/Response Packets
// ============================================================================

/// Standard A query for example.com
fn createTestQuery() [29]u8 {
    return [_]u8{
        // Header
        0x00, 0x01, // ID
        0x01, 0x00, // Flags: standard query, RD=1
        0x00, 0x01, // QDCOUNT: 1
        0x00, 0x00, // ANCOUNT: 0
        0x00, 0x00, // NSCOUNT: 0
        0x00, 0x00, // ARCOUNT: 0
        // Question: example.com A IN
        0x07, 'e', 'x', 'a', 'm', 'p', 'l', 'e',
        0x03, 'c', 'o', 'm',
        0x00, // null
        0x00, 0x01, // TYPE = A
        0x00, 0x01, // CLASS = IN
    };
}

/// Create a query for a specific domain
fn createQueryForDomain(domain: []const u8, buf: *[512]u8) usize {
    // Header
    buf[0] = 0x00;
    buf[1] = 0x02; // ID = 2
    buf[2] = 0x01;
    buf[3] = 0x00; // RD=1
    buf[4] = 0x00;
    buf[5] = 0x01; // QDCOUNT=1
    buf[6] = 0x00;
    buf[7] = 0x00;
    buf[8] = 0x00;
    buf[9] = 0x00;
    buf[10] = 0x00;
    buf[11] = 0x00;

    // Question section - encode domain name
    var pos: usize = 12;

    // Split domain by dots and encode labels
    var iter = std.mem.splitScalar(u8, domain, '.');
    while (iter.next()) |label| {
        if (label.len > 63 or label.len == 0) continue;
        buf[pos] = @intCast(label.len);
        pos += 1;
        @memcpy(buf[pos..][0..label.len], label);
        pos += label.len;
    }
    buf[pos] = 0; // Null terminator
    pos += 1;

    // QTYPE = A
    buf[pos] = 0x00;
    buf[pos + 1] = 0x01;
    // QCLASS = IN
    buf[pos + 2] = 0x00;
    buf[pos + 3] = 0x01;
    pos += 4;

    return pos;
}

// ============================================================================
// Mock Upstream for Testing
// ============================================================================

const MockUpstream = struct {
    response: ?[]const u8,
    allocator: std.mem.Allocator,
    call_count: usize = 0,

    pub fn init(response: ?[]const u8, allocator: std.mem.Allocator) MockUpstream {
        return .{
            .response = response,
            .allocator = allocator,
            .call_count = 0,
        };
    }

    pub fn toHandlerUpstream(self: *MockUpstream) handler_mod.Upstream {
        return .{
            .context = self,
            .queryFn = queryWrapper,
        };
    }

    fn queryWrapper(ctx: *anyopaque, _: []const u8, allocator: std.mem.Allocator) ?[]const u8 {
        const self: *MockUpstream = @ptrCast(@alignCast(ctx));
        self.call_count += 1;
        if (self.response) |r| {
            return allocator.dupe(u8, r) catch null;
        }
        return null;
    }
};

// ============================================================================
// Handler Tests
// ============================================================================

test "Handler - returns SERVFAIL without upstream" {
    const allocator = testing.allocator;

    var handler = handler_mod.Handler.init(allocator);
    const query = createTestQuery();
    const addr = std.net.Address.initIp4([4]u8{ 127, 0, 0, 1 }, 12345);

    const response = handler.handle(&query, addr, allocator);

    if (response) |r| {
        defer allocator.free(r);
        // Should be at least header size
        try testing.expect(r.len >= types.DNS_HEADER_SIZE);
        // Check RCODE is SERVFAIL (2) - in flags byte
        const rcode = r[3] & 0x0F;
        try testing.expectEqual(@as(u8, 2), rcode);
    }
}

test "Handler - cache integration" {
    const allocator = testing.allocator;

    var dns_cache = cache.DnsCache.init(allocator);
    defer dns_cache.deinit();

    var handler = handler_mod.Handler.init(allocator);
    handler.setCache(dns_cache.toHandlerCache());

    // Pre-populate cache
    const cached_response = [_]u8{
        0x00, 0x01, // ID (will be overwritten)
        0x81, 0x80, // Flags: response
        0x00, 0x01, // QDCOUNT
        0x00, 0x01, // ANCOUNT
        0x00, 0x00, // NSCOUNT
        0x00, 0x00, // ARCOUNT
        // Question
        0x07, 'e', 'x', 'a', 'm', 'p', 'l', 'e',
        0x03, 'c', 'o', 'm', 0x00,
        0x00, 0x01, 0x00, 0x01,
        // Answer
        0xC0, 0x0C, // Compression pointer
        0x00, 0x01, // TYPE A
        0x00, 0x01, // CLASS IN
        0x00, 0x00, 0x01, 0x2C, // TTL 300
        0x00, 0x04, // RDLENGTH
        0x01, 0x02, 0x03, 0x04, // IP: 1.2.3.4
    };

    dns_cache.put("example.com", types.QType.A, &cached_response, 300);

    const query = createTestQuery();
    const addr = std.net.Address.initIp4([4]u8{ 127, 0, 0, 1 }, 12345);

    const response = handler.handle(&query, addr, allocator);
    try testing.expect(response != null);

    if (response) |r| {
        defer allocator.free(r);
        // ID should be updated to match query
        try testing.expectEqual(@as(u8, 0x00), r[0]);
        try testing.expectEqual(@as(u8, 0x01), r[1]);
    }
}

// ============================================================================
// Rate Limiter Tests
// ============================================================================

test "RateLimiter - allows requests under limit" {
    const allocator = testing.allocator;

    var limiter = rate_limiter.RateLimiter.initWithConfig(allocator, .{
        .max_qps = 10,
        .window_ms = 1000,
    });
    defer limiter.deinit();

    // Should allow first 10 requests
    for (0..10) |_| {
        try testing.expect(limiter.checkRequest("192.168.1.1"));
    }
}

test "RateLimiter - blocks requests over limit" {
    const allocator = testing.allocator;

    var limiter = rate_limiter.RateLimiter.initWithConfig(allocator, .{
        .max_qps = 5,
        .window_ms = 1000,
    });
    defer limiter.deinit();

    // First 5 should be allowed
    for (0..5) |_| {
        try testing.expect(limiter.checkRequest("10.0.0.1"));
    }

    // 6th should be blocked
    try testing.expect(!limiter.checkRequest("10.0.0.1"));
}

test "RateLimiter - tracks clients independently" {
    const allocator = testing.allocator;

    var limiter = rate_limiter.RateLimiter.initWithConfig(allocator, .{
        .max_qps = 2,
        .window_ms = 1000,
    });
    defer limiter.deinit();

    // Client A uses quota
    try testing.expect(limiter.checkRequest("192.168.1.1"));
    try testing.expect(limiter.checkRequest("192.168.1.1"));
    try testing.expect(!limiter.checkRequest("192.168.1.1")); // Blocked

    // Client B still has quota
    try testing.expect(limiter.checkRequest("192.168.1.2"));
    try testing.expect(limiter.checkRequest("192.168.1.2"));
    try testing.expect(!limiter.checkRequest("192.168.1.2")); // Blocked
}

test "RateLimiter - can be disabled" {
    const allocator = testing.allocator;

    var limiter = rate_limiter.RateLimiter.initWithConfig(allocator, .{
        .max_qps = 1,
        .enabled = false,
    });
    defer limiter.deinit();

    // All requests allowed when disabled
    for (0..100) |_| {
        try testing.expect(limiter.checkRequest("any.ip"));
    }
}

test "RateLimiter - statistics tracking" {
    const allocator = testing.allocator;

    var limiter = rate_limiter.RateLimiter.initWithConfig(allocator, .{
        .max_qps = 3,
        .window_ms = 1000,
    });
    defer limiter.deinit();

    _ = limiter.checkRequest("1.1.1.1"); // allowed
    _ = limiter.checkRequest("1.1.1.1"); // allowed
    _ = limiter.checkRequest("1.1.1.1"); // allowed
    _ = limiter.checkRequest("1.1.1.1"); // blocked
    _ = limiter.checkRequest("2.2.2.2"); // allowed (different client)

    const stats = limiter.getStats();
    try testing.expectEqual(@as(u64, 5), stats.total_requests);
    try testing.expectEqual(@as(u64, 1), stats.rate_limited);
}

// ============================================================================
// Cache Tests
// ============================================================================

test "DnsCache - stores and retrieves entries" {
    const allocator = testing.allocator;

    var dns_cache = cache.DnsCache.init(allocator);
    defer dns_cache.deinit();

    const response = "test response data";
    dns_cache.put("test.com", types.QType.A, response, 300);

    const cached = try dns_cache.getCopy("test.com", types.QType.A);
    defer if (cached) |c| allocator.free(c);

    try testing.expect(cached != null);
    try testing.expectEqualStrings(response, cached.?);
}

test "DnsCache - returns null for missing entries" {
    const allocator = testing.allocator;

    var dns_cache = cache.DnsCache.init(allocator);
    defer dns_cache.deinit();

    const cached = try dns_cache.getCopy("nonexistent.com", types.QType.A);
    try testing.expect(cached == null);
}

test "DnsCache - separates entries by qtype" {
    const allocator = testing.allocator;

    var dns_cache = cache.DnsCache.init(allocator);
    defer dns_cache.deinit();

    dns_cache.put("example.com", types.QType.A, "A record", 300);
    dns_cache.put("example.com", types.QType.AAAA, "AAAA record", 300);

    const a_cached = try dns_cache.getCopy("example.com", types.QType.A);
    defer if (a_cached) |c| allocator.free(c);

    const aaaa_cached = try dns_cache.getCopy("example.com", types.QType.AAAA);
    defer if (aaaa_cached) |c| allocator.free(c);

    try testing.expectEqualStrings("A record", a_cached.?);
    try testing.expectEqualStrings("AAAA record", aaaa_cached.?);
}

test "DnsCache - respects max entries limit" {
    const allocator = testing.allocator;

    var dns_cache = cache.DnsCache.initWithConfig(allocator, 3, 60, 86400);
    defer dns_cache.deinit();

    // Add 4 entries, should evict oldest
    dns_cache.put("one.com", types.QType.A, "1", 300);
    dns_cache.put("two.com", types.QType.A, "2", 300);
    dns_cache.put("three.com", types.QType.A, "3", 300);
    dns_cache.put("four.com", types.QType.A, "4", 300);

    const stats = dns_cache.getStats();
    try testing.expect(stats.entry_count <= 3);
}

test "DnsCache - updates existing entries" {
    const allocator = testing.allocator;

    var dns_cache = cache.DnsCache.init(allocator);
    defer dns_cache.deinit();

    dns_cache.put("update.com", types.QType.A, "original", 300);
    dns_cache.put("update.com", types.QType.A, "updated", 300);

    const cached = try dns_cache.getCopy("update.com", types.QType.A);
    defer if (cached) |c| allocator.free(c);

    try testing.expectEqualStrings("updated", cached.?);
}

test "DnsCache - does not cache TTL=0 responses" {
    const allocator = testing.allocator;

    var dns_cache = cache.DnsCache.init(allocator);
    defer dns_cache.deinit();

    // Try to cache with TTL=0 - should NOT be cached per RFC 2308
    dns_cache.put("nocache.com", types.QType.A, "should not cache", 0);

    const cached = try dns_cache.getCopy("nocache.com", types.QType.A);
    try testing.expect(cached == null);
}

// ============================================================================
// Handler with Rate Limiting Integration
// ============================================================================

test "Handler - respects rate limiter" {
    const allocator = testing.allocator;

    var limiter = rate_limiter.RateLimiter.initWithConfig(allocator, .{
        .max_qps = 2,
        .window_ms = 1000,
    });
    defer limiter.deinit();

    var handler = handler_mod.Handler.init(allocator);
    handler.setRateLimiter(&limiter);

    const query = createTestQuery();
    const addr = std.net.Address.initIp4([4]u8{ 10, 0, 0, 1 }, 12345);

    // First 2 requests succeed (return SERVFAIL due to no upstream)
    const r1 = handler.handle(&query, addr, allocator);
    const r2 = handler.handle(&query, addr, allocator);

    if (r1) |r| allocator.free(r);
    if (r2) |r| allocator.free(r);

    // 3rd request should be rate limited (REFUSED)
    const r3 = handler.handle(&query, addr, allocator);
    if (r3) |r| {
        defer allocator.free(r);
        // REFUSED = RCODE 5
        const rcode = r[3] & 0x0F;
        try testing.expectEqual(@as(u8, 5), rcode);
    }
}

// ============================================================================
// Blocklist Integration Tests
// ============================================================================

test "Handler with blocklist - blocks matching domains" {
    const allocator = testing.allocator;

    var blocklist = blocklist_mod.Blocklist.init(allocator);
    defer blocklist.deinit();

    try blocklist.addBlockedDomain("ads.example.com", 0);

    var handler = handler_mod.Handler.init(allocator);
    handler.setBlocklist(blocklist.toHandlerBlocklist());

    // Query for blocked domain
    var query_buf: [512]u8 = undefined;
    const query_len = createQueryForDomain("ads.example.com", &query_buf);
    const addr = std.net.Address.initIp4([4]u8{ 127, 0, 0, 1 }, 12345);

    const response = handler.handle(query_buf[0..query_len], addr, allocator);
    try testing.expect(response != null);

    if (response) |r| {
        defer allocator.free(r);
        // Should be a valid response with blocked content (0.0.0.0 or NXDOMAIN)
        try testing.expect(r.len >= types.DNS_HEADER_SIZE);
        // QR bit should be set (response)
        try testing.expect((r[2] & 0x80) != 0);
    }
}

test "Handler with blocklist - allows non-blocked domains" {
    const allocator = testing.allocator;

    var blocklist = blocklist_mod.Blocklist.init(allocator);
    defer blocklist.deinit();

    try blocklist.addBlockedDomain("blocked.com", 0);

    // Create mock upstream that returns a valid response
    const mock_response = [_]u8{
        0x00, 0x02, 0x81, 0x80, // Header (response)
        0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
        // Question: allowed.com A IN
        0x07, 'a', 'l', 'l', 'o', 'w', 'e', 'd',
        0x03, 'c', 'o', 'm', 0x00,
        0x00, 0x01, 0x00, 0x01,
        // Answer
        0xC0, 0x0C, 0x00, 0x01, 0x00, 0x01,
        0x00, 0x00, 0x01, 0x2C, // TTL 300
        0x00, 0x04, 0x08, 0x08, 0x08, 0x08, // IP 8.8.8.8
    };

    var mock_upstream = MockUpstream.init(&mock_response, allocator);

    var handler = handler_mod.Handler.init(allocator);
    handler.setBlocklist(blocklist.toHandlerBlocklist());
    handler.setUpstream(mock_upstream.toHandlerUpstream());

    // Query for non-blocked domain
    var query_buf: [512]u8 = undefined;
    const query_len = createQueryForDomain("allowed.com", &query_buf);
    const addr = std.net.Address.initIp4([4]u8{ 127, 0, 0, 1 }, 12345);

    const response = handler.handle(query_buf[0..query_len], addr, allocator);
    try testing.expect(response != null);

    if (response) |r| {
        defer allocator.free(r);
        // Should forward to upstream and return its response
        try testing.expect(mock_upstream.call_count == 1);
    }
}

test "Handler with blocklist - blocks subdomains" {
    const allocator = testing.allocator;

    var blocklist = blocklist_mod.Blocklist.init(allocator);
    defer blocklist.deinit();

    // Block parent domain
    try blocklist.addBlockedDomain("doubleclick.net", 0);

    var handler = handler_mod.Handler.init(allocator);
    handler.setBlocklist(blocklist.toHandlerBlocklist());

    // Query for subdomain - should also be blocked
    var query_buf: [512]u8 = undefined;
    const query_len = createQueryForDomain("ads.doubleclick.net", &query_buf);
    const addr = std.net.Address.initIp4([4]u8{ 127, 0, 0, 1 }, 12345);

    const response = handler.handle(query_buf[0..query_len], addr, allocator);
    try testing.expect(response != null);

    if (response) |r| {
        defer allocator.free(r);
        // Should be blocked
        try testing.expect(r.len >= types.DNS_HEADER_SIZE);
    }
}

test "Blocklist - allow rules override block rules" {
    const allocator = testing.allocator;

    var blocklist = blocklist_mod.Blocklist.init(allocator);
    defer blocklist.deinit();

    try blocklist.addBlockedDomain("example.com", 0);
    try blocklist.addAllowRule("allowed.example.com", 0);

    // Subdomain blocked
    try testing.expect(blocklist.isBlocked("blocked.example.com", 0));
    // But allowed.example.com is explicitly allowed
    try testing.expect(!blocklist.isBlocked("allowed.example.com", 0));
}

test "Blocklist - group isolation" {
    const allocator = testing.allocator;

    var blocklist = blocklist_mod.Blocklist.init(allocator);
    defer blocklist.deinit();

    // Block domain only for group 1
    try blocklist.addBlockedDomain("group1only.com", 1);

    // Group 1 sees it blocked
    try testing.expect(blocklist.isBlocked("group1only.com", 1));
    // Group 0 and 2 do not see it blocked
    try testing.expect(!blocklist.isBlocked("group1only.com", 0));
    try testing.expect(!blocklist.isBlocked("group1only.com", 2));
}

// ============================================================================
// Compression Loop Detection Tests
// ============================================================================

test "Name parsing - detects compression loop" {
    const allocator = testing.allocator;

    // Packet with self-referencing compression pointer
    const bad_packet = [_]u8{
        0x00, 0x01, 0x01, 0x00,
        0x00, 0x01, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00,
        // Name with compression pointer pointing to itself (offset 12)
        0xC0, 0x0C,
        0x00, 0x01, 0x00, 0x01,
    };

    // Parsing should detect the loop and fail
    const result = Name.parse(bad_packet[12..], &bad_packet, allocator);
    try testing.expectError(error.CompressionLoop, result);
}

test "Name parsing - detects indirect compression loop" {
    const allocator = testing.allocator;

    // Packet where pointer A -> pointer B -> pointer A
    const bad_packet = [_]u8{
        0x00, 0x01, 0x01, 0x00,
        0x00, 0x01, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00,
        // Offset 12: pointer to offset 14
        0xC0, 0x0E,
        // Offset 14: pointer to offset 12
        0xC0, 0x0C,
        0x00, 0x01, 0x00, 0x01,
    };

    const result = Name.parse(bad_packet[12..], &bad_packet, allocator);
    try testing.expectError(error.CompressionLoop, result);
}

// ============================================================================
// Malformed Packet Handling Tests
// ============================================================================

test "Handler - handles truncated packet" {
    const allocator = testing.allocator;
    var handler = handler_mod.Handler.init(allocator);

    // Only 2 bytes - too short for DNS header
    const truncated = [_]u8{ 0x00, 0x01 };
    const addr = std.net.Address.initIp4([4]u8{ 127, 0, 0, 1 }, 12345);

    const response = handler.handle(&truncated, addr, allocator);

    if (response) |r| {
        defer allocator.free(r);
        // Should return FORMERR
        const rcode = r[3] & 0x0F;
        try testing.expectEqual(@as(u8, 1), rcode); // FORMERR
    }
}

test "Handler - handles empty question section" {
    const allocator = testing.allocator;
    var handler = handler_mod.Handler.init(allocator);

    // Valid header but QDCOUNT = 0
    const no_question = [_]u8{
        0x00, 0x01, 0x01, 0x00,
        0x00, 0x00, // QDCOUNT = 0
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    };
    const addr = std.net.Address.initIp4([4]u8{ 127, 0, 0, 1 }, 12345);

    const response = handler.handle(&no_question, addr, allocator);

    if (response) |r| {
        defer allocator.free(r);
        const rcode = r[3] & 0x0F;
        try testing.expectEqual(@as(u8, 1), rcode); // FORMERR
    }
}

// ============================================================================
// Full Query Flow Integration Tests
// ============================================================================

test "Full flow - query hits cache" {
    const allocator = testing.allocator;

    var dns_cache = cache.DnsCache.init(allocator);
    defer dns_cache.deinit();

    var mock_upstream = MockUpstream.init(null, allocator);

    var handler = handler_mod.Handler.init(allocator);
    handler.setCache(dns_cache.toHandlerCache());
    handler.setUpstream(mock_upstream.toHandlerUpstream());

    // Pre-populate cache
    const cached_response = [_]u8{
        0x00, 0x01, 0x81, 0x80,
        0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
        0x07, 'e', 'x', 'a', 'm', 'p', 'l', 'e',
        0x03, 'c', 'o', 'm', 0x00,
        0x00, 0x01, 0x00, 0x01,
        0xC0, 0x0C, 0x00, 0x01, 0x00, 0x01,
        0x00, 0x00, 0x01, 0x2C,
        0x00, 0x04, 0x01, 0x02, 0x03, 0x04,
    };
    dns_cache.put("example.com", types.QType.A, &cached_response, 300);

    const query = createTestQuery();
    const addr = std.net.Address.initIp4([4]u8{ 127, 0, 0, 1 }, 12345);

    const response = handler.handle(&query, addr, allocator);
    try testing.expect(response != null);

    if (response) |r| {
        defer allocator.free(r);
        // Upstream should NOT have been called
        try testing.expectEqual(@as(usize, 0), mock_upstream.call_count);
    }
}

test "Full flow - cache miss goes to upstream" {
    const allocator = testing.allocator;

    var dns_cache = cache.DnsCache.init(allocator);
    defer dns_cache.deinit();

    const upstream_response = [_]u8{
        0x00, 0x01, 0x81, 0x80,
        0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
        0x07, 'e', 'x', 'a', 'm', 'p', 'l', 'e',
        0x03, 'c', 'o', 'm', 0x00,
        0x00, 0x01, 0x00, 0x01,
        0xC0, 0x0C, 0x00, 0x01, 0x00, 0x01,
        0x00, 0x00, 0x01, 0x2C,
        0x00, 0x04, 0x08, 0x08, 0x08, 0x08,
    };

    var mock_upstream = MockUpstream.init(&upstream_response, allocator);

    var handler = handler_mod.Handler.init(allocator);
    handler.setCache(dns_cache.toHandlerCache());
    handler.setUpstream(mock_upstream.toHandlerUpstream());

    const query = createTestQuery();
    const addr = std.net.Address.initIp4([4]u8{ 127, 0, 0, 1 }, 12345);

    const response = handler.handle(&query, addr, allocator);
    try testing.expect(response != null);

    if (response) |r| {
        defer allocator.free(r);
        // Upstream should have been called
        try testing.expectEqual(@as(usize, 1), mock_upstream.call_count);

        // Result should now be cached
        const cached = try dns_cache.getCopy("example.com", types.QType.A);
        try testing.expect(cached != null);
        if (cached) |c| allocator.free(c);
    }
}

test "Full flow - blocklist takes precedence over cache" {
    const allocator = testing.allocator;

    var blocklist = blocklist_mod.Blocklist.init(allocator);
    defer blocklist.deinit();
    try blocklist.addBlockedDomain("blocked.com", 0);

    var dns_cache = cache.DnsCache.init(allocator);
    defer dns_cache.deinit();

    // Pre-populate cache with a response for blocked.com
    const cached_response = [_]u8{
        0x00, 0x01, 0x81, 0x80,
        0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
        0x07, 'b', 'l', 'o', 'c', 'k', 'e', 'd',
        0x03, 'c', 'o', 'm', 0x00,
        0x00, 0x01, 0x00, 0x01,
        0xC0, 0x0C, 0x00, 0x01, 0x00, 0x01,
        0x00, 0x00, 0x01, 0x2C,
        0x00, 0x04, 0x08, 0x08, 0x08, 0x08, // 8.8.8.8
    };
    dns_cache.put("blocked.com", types.QType.A, &cached_response, 300);

    var handler = handler_mod.Handler.init(allocator);
    handler.setBlocklist(blocklist.toHandlerBlocklist());
    handler.setCache(dns_cache.toHandlerCache());

    var query_buf: [512]u8 = undefined;
    const query_len = createQueryForDomain("blocked.com", &query_buf);
    const addr = std.net.Address.initIp4([4]u8{ 127, 0, 0, 1 }, 12345);

    const response = handler.handle(query_buf[0..query_len], addr, allocator);
    try testing.expect(response != null);

    if (response) |r| {
        defer allocator.free(r);
        // Parse response to check answer
        var pkt = packet.Packet.parse(r, allocator) catch {
            try testing.expect(false);
            return;
        };
        defer pkt.deinit();

        // Should be blocked (0.0.0.0), not cached (8.8.8.8)
        if (pkt.answers.len > 0) {
            const ip = pkt.answers[0].getA();
            if (ip) |addr_bytes| {
                try testing.expectEqual([4]u8{ 0, 0, 0, 0 }, addr_bytes);
            }
        }
    }
}

// ============================================================================
// Packet Encoding/Decoding Roundtrip Tests
// ============================================================================

test "Packet roundtrip - query" {
    const allocator = testing.allocator;

    const original_query = createTestQuery();
    var pkt = try packet.Packet.parse(&original_query, allocator);
    defer pkt.deinit();

    var buf: [512]u8 = undefined;
    const encoded_len = try pkt.encode(&buf);

    var decoded = try packet.Packet.parse(buf[0..encoded_len], allocator);
    defer decoded.deinit();

    try testing.expectEqual(pkt.header.id, decoded.header.id);
    try testing.expectEqual(pkt.header.qr, decoded.header.qr);
    try testing.expectEqual(pkt.questions.len, decoded.questions.len);
}

test "Packet - createBlockedResponse has 0.0.0.0" {
    const allocator = testing.allocator;

    const query = createTestQuery();
    var query_pkt = try packet.Packet.parse(&query, allocator);
    defer query_pkt.deinit();

    var response = try packet.Packet.createBlockedResponse(&query_pkt, allocator);
    defer response.deinit();

    try testing.expect(response.header.qr); // Is response
    try testing.expectEqual(@as(usize, 1), response.answers.len);

    const ip = response.answers[0].getA();
    try testing.expect(ip != null);
    try testing.expectEqual([4]u8{ 0, 0, 0, 0 }, ip.?);
}

test "Packet - createNxdomainResponse has NXDOMAIN RCODE" {
    const allocator = testing.allocator;

    const query = createTestQuery();
    var query_pkt = try packet.Packet.parse(&query, allocator);
    defer query_pkt.deinit();

    var response = try packet.Packet.createNxdomainResponse(&query_pkt, allocator);
    defer response.deinit();

    try testing.expect(response.header.qr);
    try testing.expectEqual(types.RCode.NXDomain, response.header.rcode);
    try testing.expectEqual(@as(usize, 0), response.answers.len);
}

// ============================================================================
// DoH/DoT URL Parsing Tests
// ============================================================================

const dot = @import("dot");
const doh = @import("doh");

test "DoT URL parsing - valid URLs" {
    const allocator = testing.allocator;

    // Standard TLS URL
    const client1 = try dot.DotClient.fromUrl("tls://cloudflare-dns.com", allocator);
    try testing.expectEqualStrings("cloudflare-dns.com", client1.host);
    try testing.expectEqual(@as(u16, 853), client1.port);

    // With explicit port
    const client2 = try dot.DotClient.fromUrl("tls://1.1.1.1:853", allocator);
    try testing.expectEqualStrings("1.1.1.1", client2.host);
    try testing.expectEqual(@as(u16, 853), client2.port);

    // Custom port
    const client3 = try dot.DotClient.fromUrl("tls://dns.google:8853", allocator);
    try testing.expectEqualStrings("dns.google", client3.host);
    try testing.expectEqual(@as(u16, 8853), client3.port);
}

test "DoT URL parsing - invalid URLs" {
    const allocator = testing.allocator;

    // Wrong scheme
    try testing.expectError(error.InvalidHost, dot.DotClient.fromUrl("https://example.com", allocator));
    try testing.expectError(error.InvalidHost, dot.DotClient.fromUrl("not-a-url", allocator));
}

test "DoH URL parsing - valid URLs" {
    const allocator = testing.allocator;

    // Standard DoH URL
    const client1 = try doh.DohClient.init("https://cloudflare-dns.com/dns-query", allocator);
    try testing.expectEqualStrings("cloudflare-dns.com", client1.host);
    try testing.expectEqualStrings("/dns-query", client1.path);
    try testing.expectEqual(@as(u16, 443), client1.port);

    // With custom port
    const client2 = try doh.DohClient.init("https://dns.quad9.net:8443/dns-query", allocator);
    try testing.expectEqualStrings("dns.quad9.net", client2.host);
    try testing.expectEqual(@as(u16, 8443), client2.port);

    // No path (defaults to /dns-query)
    const client3 = try doh.DohClient.init("https://example.com", allocator);
    try testing.expectEqualStrings("example.com", client3.host);
    try testing.expectEqualStrings("/dns-query", client3.path);
}

test "DoH URL parsing - invalid URLs" {
    const allocator = testing.allocator;

    // Wrong scheme
    try testing.expectError(error.InvalidUrl, doh.DohClient.init("http://example.com/dns-query", allocator));
    try testing.expectError(error.InvalidUrl, doh.DohClient.init("not-a-url", allocator));
}

// ============================================================================
// Edge Case Tests
// ============================================================================

test "Handler - IPv6 client address" {
    const allocator = testing.allocator;
    var handler = handler_mod.Handler.init(allocator);

    const query = createTestQuery();

    // Create IPv6 address
    var addr: std.net.Address = undefined;
    addr.in6 = std.net.Ip6Address.init([16]u8{
        0x20, 0x01, 0x0d, 0xb8, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
    }, 12345);
    addr.any.family = std.posix.AF.INET6;

    // Should handle IPv6 without crashing
    const response = handler.handle(&query, addr, allocator);

    if (response) |r| {
        defer allocator.free(r);
        try testing.expect(r.len >= types.DNS_HEADER_SIZE);
    }
}

test "Cache - TTL clamping" {
    const allocator = testing.allocator;

    // Configure with min_ttl=60, max_ttl=3600
    var dns_cache = cache.DnsCache.initWithConfig(allocator, 100, 60, 3600);
    defer dns_cache.deinit();

    // Put with TTL below min - should be clamped to 60
    dns_cache.put("test1.com", types.QType.A, "response", 10);
    const cached1 = try dns_cache.getCopy("test1.com", types.QType.A);
    try testing.expect(cached1 != null);
    if (cached1) |c| allocator.free(c);

    // Put with TTL above max - should be clamped to 3600
    dns_cache.put("test2.com", types.QType.A, "response", 100000);
    const cached2 = try dns_cache.getCopy("test2.com", types.QType.A);
    try testing.expect(cached2 != null);
    if (cached2) |c| allocator.free(c);
}