Recover DNS detection across split TCP responses

.github/workflows/test.yml

@@ -72,12 +72,8 @@ jobs:
           # the C ABI the NetGuard engine calls (wgbridge-rs/tests/).
           cargo test --manifest-path wgbridge-rs/Cargo.toml --workspace --locked --offline
 
-      - name: Run DNS-over-TCP framing host tests
-        run: |
-          cc -Wall -Wextra -Werror -Iapp/src/main/jni/netguard \
-            -o /tmp/dns_frame_test \
-            app/src/test/native/dns_frame_test.c app/src/main/jni/netguard/dns_frame.c
-          /tmp/dns_frame_test
+      - name: Run DNS-over-TCP framing and detection regression tests
+        run: bash app/src/test/native/run_dns_frame_tests.sh
 
       - name: Run IPv6 extension header walk host tests
         run: |

app/src/main/jni/netguard/dns.c

@@ -22,6 +22,7 @@
 struct tcdns_ctx {
     const struct arguments *args;
     const struct ng_session *s;
+    int detection_only;
 };
 
 static void tcdns_record_answer(void *opaque, const char *qname, const char *aname,
@@ -32,7 +33,7 @@ static void tcdns_record_answer(void *opaque, const char *qname, const char *ana
 
 static int tcdns_is_domain_blocked(void *opaque, const char *qname) {
     const struct tcdns_ctx *ctx = (const struct tcdns_ctx *) opaque;
-    return is_domain_blocked(ctx->args, qname) != 0;
+    return !ctx->detection_only && is_domain_blocked(ctx->args, qname) != 0;
 }
 
 static uint8_t tcdns_blocked_rcode(void *opaque) {
@@ -43,6 +44,8 @@ static uint8_t tcdns_blocked_rcode(void *opaque) {
 static void tcdns_on_blanked(void *opaque, const char *qname,
                              uint16_t qtype, uint8_t rcode) {
     const struct tcdns_ctx *ctx = (const struct tcdns_ctx *) opaque;
+    if (ctx->detection_only)
+        return;
     const struct arguments *args = ctx->args;
     const struct ng_session *s = ctx->s;
 
@@ -115,3 +118,11 @@ void parse_dns_partial_response(const struct arguments *args, const struct ng_se
     if (result != TCDNS_UNCHANGED)
         *blanked = 1;
 }
+
+// This is a completed private copy: recording is useful, but policy/logging
+// cannot claim to have rewritten bytes already forwarded to the client.
+void record_dns_response(const struct arguments *args, uint8_t *data, size_t datalen) {
+    struct tcdns_ctx ctx = { .args = args, .s = NULL, .detection_only = 1 };
+    tcdns_callbacks cb = TCDNS_CALLBACKS_INIT;
+    (void) tcdns_process_response(data, datalen, &cb, &ctx);
+}

app/src/main/jni/netguard/dns_frame.c

@@ -17,10 +17,68 @@
     Copyright 2015-2019 by Marcel Bokhorst (M66B)
 */
 
+#include <stdlib.h>
 #include <string.h>
 
 #include "dns_frame.h"
 
+static void clear_frame_buffer(struct dns_stream_state *state) {
+    free(state->frame_buffer);
+    state->frame_buffer = NULL;
+    state->frame_length = 0;
+    state->frame_received = 0;
+}
+
+void dns_frame_reset(struct dns_stream_state *state) {
+    if (state == NULL)
+        return;
+
+    free(state->frame_buffer);
+    memset(state, 0, sizeof(*state));
+}
+
+static void start_frame_buffer(struct dns_stream_state *state, size_t frame_len,
+                               const uint8_t *data, size_t bytes) {
+    clear_frame_buffer(state);
+    if (frame_len == 0)
+        return;
+
+    state->frame_buffer = malloc(frame_len);
+    if (state->frame_buffer == NULL)
+        return;
+
+    state->frame_length = (uint32_t) frame_len;
+    if (bytes > frame_len)
+        bytes = frame_len;
+    if (bytes > 0)
+        memcpy(state->frame_buffer, data, bytes);
+    state->frame_received = (uint32_t) bytes;
+}
+
+static void append_frame_buffer(struct dns_stream_state *state,
+                                const uint8_t *data, size_t bytes) {
+    if (state->frame_buffer == NULL || bytes == 0)
+        return;
+
+    size_t available = state->frame_length - state->frame_received;
+    if (bytes > available)
+        bytes = available;
+    memcpy(state->frame_buffer + state->frame_received, data, bytes);
+    state->frame_received += (uint32_t) bytes;
+}
+
+static void parse_completed_buffer(struct dns_stream_state *state,
+                                   dns_frame_parse_fn parse, void *ctx) {
+    if (state->frame_buffer != NULL &&
+        state->frame_received == state->frame_length) {
+        int blank_rest = 0;
+        (void) parse(ctx, state->frame_buffer, state->frame_length,
+                     DNS_FRAME_REPLAY,
+                     &blank_rest);
+    }
+    clear_frame_buffer(state);
+}
+
 size_t dns_frame_process_stream(uint8_t *buffer, size_t bytes,
                                 struct dns_stream_state *state,
                                 dns_frame_parse_fn parse, void *ctx) {
@@ -32,18 +90,21 @@ size_t dns_frame_process_stream(uint8_t *buffer, size_t bytes,
     size_t end = bytes; // bytes to forward; only ever shrinks, never below cursor
 
     // 1. Continuation of a frame whose earlier bytes were already forwarded in
-    //    a previous recv(). Those bytes are neither parsed (their DNS header is
-    //    gone) nor rewritten (their prefix is already on the wire).
+    //    a previous recv(). Preserve the original bytes for a complete parse
+    //    before applying any carried blanking decision to the forwarded copy.
     if (state->frame_remaining > 0) {
         size_t remaining = end - cursor;
         size_t skip = (state->frame_remaining < remaining
                        ? (size_t) state->frame_remaining : remaining);
+        append_frame_buffer(state, buffer + cursor, skip);
         if (state->blank_remaining != 0 && skip > 0)
             memset(buffer + cursor, 0, skip);
         state->frame_remaining -= (uint32_t) skip;
         cursor += skip;
-        if (state->frame_remaining == 0)
+        if (state->frame_remaining == 0) {
             state->blank_remaining = 0;
+            parse_completed_buffer(state, parse, ctx);
+        }
         if (cursor >= end)
             return end;
     }
@@ -61,15 +122,18 @@ size_t dns_frame_process_stream(uint8_t *buffer, size_t bytes,
         size_t avail = end - cursor;
         if (frame_len > avail) {
             int blank_rest = 0;
+            start_frame_buffer(state, frame_len, buffer + cursor, avail);
             if (avail > 0)
-                (void) parse(ctx, buffer + cursor, avail, 1, &blank_rest);
+                (void) parse(ctx, buffer + cursor, avail, DNS_FRAME_PARTIAL,
+                             &blank_rest);
             state->frame_remaining = (uint32_t) (frame_len - avail);
             state->blank_remaining = (uint8_t) (blank_rest != 0);
             return end;
         }
         if (frame_len > 0) {
             int blank_rest = 0;
-            (void) parse(ctx, buffer + cursor, frame_len, 1, &blank_rest); // shrink ignored
+            (void) parse(ctx, buffer + cursor, frame_len, DNS_FRAME_PARTIAL,
+                         &blank_rest); // shrink ignored
             cursor += frame_len;
         }
     }
@@ -98,10 +162,13 @@ size_t dns_frame_process_stream(uint8_t *buffer, size_t bytes,
 
         if (frame_len > avail) {
             // Frame runs past this read: parse what is visible (blanking only,
-            // any shrink is ignored) and remember the overflow.
+            // any shrink is ignored), retain its original bytes for detection
+            // when complete, and remember the overflow.
             int blank_rest = 0;
+            start_frame_buffer(state, frame_len, buffer + cursor, avail);
             if (avail > 0)
-                (void) parse(ctx, buffer + cursor, avail, 1, &blank_rest);
+                (void) parse(ctx, buffer + cursor, avail, DNS_FRAME_PARTIAL,
+                             &blank_rest);
             state->frame_remaining = (uint32_t) (frame_len - avail);
             state->blank_remaining = (uint8_t) (blank_rest != 0);
             return end;
@@ -110,7 +177,8 @@ size_t dns_frame_process_stream(uint8_t *buffer, size_t bytes,
         // Complete frame, prefix and payload both inside this buffer: nothing
         // here has been forwarded yet, so it may be shortened.
         int blank_rest = 0;
-        size_t new_dlen = parse(ctx, buffer + cursor, frame_len, 0, &blank_rest);
+        size_t new_dlen = parse(ctx, buffer + cursor, frame_len,
+                                DNS_FRAME_COMPLETE, &blank_rest);
         if (new_dlen > frame_len)
             new_dlen = frame_len; // defensive: a parser must never grow a frame
 

app/src/main/jni/netguard/dns_frame.h

@@ -21,7 +21,7 @@
 #define DNS_FRAME_H
 
 /*
- * Bufferless cursor over a DNS-over-TCP (port 53) byte stream, extracted out
+ * Cursor over a DNS-over-TCP (port 53) byte stream, extracted out
  * of check_tcp_socket() (tcp.c) so it can be unit-tested on the host without
  * pulling in JNI/session dependencies. This header and its implementation
  * (dns_frame.c) must only depend on libc: no netguard.h, no JNI.
@@ -35,19 +35,11 @@
  *   - a lone byte that is the first half of a length prefix.
  *
  * dns_frame_process_stream() walks every frame boundary inside one recv()
- * buffer and carries enough state to stay aligned into the next call without
- * buffering stream bytes. A frame is offered to the DNS parser only in the
- * recv() where its prefix is completed. If its payload continues into later
- * reads, only the bytes visible in that first call are parsed; the parser is
- * told that the frame is partial, and a blocking result causes the visible
- * answer tail and all later continuation bytes to be zeroed in place. Policy
- * enforcement for such split frames therefore remains best-effort: once the
- * answer section is cut short, per-answer detection stops at the truncation
- * point and SVCB-triggered blanking is unavailable, leaving the domain of the
- * question as the only signal the block decision can use. That signal needs the
- * header and the whole question inside the visible bytes, so a read split
- * within the first few bytes of a response blocks nothing at all and the frame
- * passes through intact.
+ * buffer and carries alignment into the next call. A split payload is also
+ * copied into a bounded buffer (the DNS-over-TCP length field limits it to
+ * 65535 bytes). The visible first part is parsed immediately for enforcement;
+ * once the payload is complete, the copy is parsed again so answers that were
+ * beyond the first recv() still reach tracker detection.
  *
  * A blanked split frame keeps its original 2-byte length prefix, so what
  * reaches the client is a DNS message with all three counts cleared followed by
@@ -90,19 +82,35 @@ struct dns_stream_state {
     uint8_t prefix_hi;        /* stashed first byte of a length prefix split
                                  across recv()s */
     uint8_t have_prefix_hi;   /* nonzero when prefix_hi is valid */
+    uint8_t *frame_buffer;    /* original bytes of a split frame */
+    uint32_t frame_length;
+    uint32_t frame_received;
+};
+
+/* Releases any split-frame buffer and restores the all-zero initial state. */
+void dns_frame_reset(struct dns_stream_state *state);
+
+enum dns_frame_parse_mode {
+    DNS_FRAME_COMPLETE = 0,
+    DNS_FRAME_PARTIAL = 1,
+    DNS_FRAME_REPLAY = 2,
 };
 
 /*
  * Called for each DNS payload (or the visible part of one) found in the
- * buffer; stands in for parse_dns_response(). partial != 0 means the frame is
- * not wholly inside this buffer, so the callback must not shrink it. On a
- * partial call, *blank_rest is set nonzero when the callback blanked the
- * visible part and the caller must blank the frame's later continuation bytes.
+ * buffer; stands in for parse_dns_response(). mode is DNS_FRAME_COMPLETE for
+ * a complete frame, DNS_FRAME_PARTIAL for the first visible part of a split
+ * frame, and DNS_FRAME_REPLAY when replaying the completed buffered copy of a
+ * split frame for detection. Partial and replay calls must not shrink the
+ * forwarded stream. On a partial call, *blank_rest is set nonzero when the
+ * callback blanked the visible part and the caller must blank the frame's later
+ * continuation bytes.
  * Returns the possibly-shrunk payload length for a complete frame; a return
  * > dlen must be treated by the caller as "unchanged" (defensive clamp).
  */
 typedef size_t (*dns_frame_parse_fn)(void *ctx, uint8_t *data, size_t dlen,
-                                     int partial, int *blank_rest);
+                                     enum dns_frame_parse_mode mode,
+                                     int *blank_rest);
 
 /*
  * Processes one recv() buffer of a DNS-over-TCP stream in place.

app/src/main/jni/netguard/netguard.h

@@ -414,6 +414,8 @@ void check_udp_socket(const struct arguments *args, const struct epoll_event *ev
 void parse_dns_response(const struct arguments *args, const struct ng_session *session,
                         uint8_t *data, size_t *datalen);
 
+void record_dns_response(const struct arguments *args, uint8_t *data, size_t datalen);
+
 void parse_dns_partial_response(const struct arguments *args,
                                 const struct ng_session *session,
                                 uint8_t *data, size_t *datalen, int *blanked);

app/src/main/jni/netguard/tcp.c

@@ -28,6 +28,7 @@ extern char socks5_password[127 + 1];
 extern FILE *pcap_file;
 
 void clear_tcp_data(struct tcp_session *cur) {
+    dns_frame_reset(&cur->dns_stream);
     struct segment *s = cur->forward;
     while (s != NULL) {
         struct segment *p = s;
@@ -103,6 +104,7 @@ int check_tcp_session(const struct arguments *args, struct ng_session *s,
 
         s->tcp.time = time(NULL);
         s->tcp.state = TCP_CLOSE;
+        dns_frame_reset(&s->tcp.dns_stream);
     }
 
     if ((s->tcp.state == TCP_CLOSING || s->tcp.state == TCP_CLOSE) &&
@@ -249,10 +251,15 @@ struct dns_stream_parse_ctx {
 // Adapter matching dns_frame_parse_fn: complete frames use
 // parse_dns_response() and may shrink; partial frames use the in-place path
 // and report whether their later continuation bytes must be blanked.
+// Replay parses only a private copy; its result cannot rewrite forwarded bytes.
 static size_t tcp_dns_parse_frame(void *ctx, uint8_t *data, size_t dlen,
-                                  int partial, int *blank_rest) {
+                                  enum dns_frame_parse_mode mode, int *blank_rest) {
     struct dns_stream_parse_ctx *pctx = (struct dns_stream_parse_ctx *) ctx;
-    if (partial != 0) {
+    if (mode == DNS_FRAME_REPLAY) {
+        record_dns_response(pctx->args, data, dlen);
+        return dlen;
+    }
+    if (mode == DNS_FRAME_PARTIAL) {
         int blanked = 0;
         parse_dns_partial_response(pctx->args, pctx->s, data, &dlen, &blanked);
         *blank_rest = blanked;

app/src/test/native/dns_frame_allocation_test.c

@@ -0,0 +1,59 @@
+/* Deterministic allocation failure and lifetime tests for the stream buffer. */
+#include <assert.h>
+#include <stdlib.h>
+#include "dns_frame.h"
+
+static int fail_allocation;
+static size_t live_allocations, last_allocation, replays;
+static void *frame_alloc(size_t size) {
+    last_allocation = size;
+    if (fail_allocation) return NULL;
+    void *p = malloc(size);
+    assert(p != NULL);
+    live_allocations++;
+    return p;
+}
+static void frame_free(void *p) {
+    if (p != NULL) { assert(live_allocations > 0); live_allocations--; }
+    free(p);
+}
+#define malloc frame_alloc
+#define free frame_free
+#include "dns_frame.c"
+#undef malloc
+#undef free
+
+static size_t parse(void *ctx, uint8_t *data, size_t len,
+                    enum dns_frame_parse_mode mode, int *blank_rest) {
+    (void)ctx; (void)data;
+    *blank_rest = 0;
+    if (mode == DNS_FRAME_REPLAY) replays++;
+    return len;
+}
+int main(void) {
+    struct dns_stream_state state = {0};
+    uint8_t first[] = {0, 5, 1}, rest[] = {2, 3, 4, 5};
+    assert(dns_frame_process_stream(first, sizeof(first), &state, parse, NULL) == sizeof(first));
+    assert(live_allocations == 1 && last_allocation == 5);
+    assert(dns_frame_process_stream(rest, sizeof(rest), &state, parse, NULL) == sizeof(rest));
+    assert(replays == 1 && live_allocations == 0 && state.frame_buffer == NULL);
+    dns_frame_reset(&state);
+
+    /* A connection closes before the announced maximum frame arrives. */
+    uint8_t maximum[] = {255, 255};
+    dns_frame_process_stream(maximum, sizeof(maximum), &state, parse, NULL);
+    assert(last_allocation == 65535 && live_allocations == 1);
+    dns_frame_reset(&state);
+    dns_frame_reset(&state);
+    assert(live_allocations == 0 && state.frame_remaining == 0);
+
+    /* OOM must preserve forwarding/alignment and the previous partial parser. */
+    fail_allocation = 1;
+    assert(dns_frame_process_stream(first, sizeof(first), &state, parse, NULL) == sizeof(first));
+    assert(state.frame_buffer == NULL && state.frame_remaining == 4);
+    assert(dns_frame_process_stream(rest, sizeof(rest), &state, parse, NULL) == sizeof(rest));
+    assert(replays == 1 && state.frame_remaining == 0 && live_allocations == 0);
+    assert(first[2] == 1 && rest[0] == 2 && rest[3] == 5);
+    dns_frame_reset(&state);
+    return 0;
+}