fix(http): rewrite all User-Agent headers and preserve active sessions

CMakeLists.txt

@@ -271,6 +271,8 @@ if (UA2F_BUILD_TESTS)
             test/cli_test.cc
             test/http_session_test.cc
             test/http_parser_ua_test.cc
+            test/http_parser_ua_alloc_test.cc
+            test/alloc_tracker.c
             src/util.c
             src/cache.c
             src/http_session.c
@@ -287,6 +289,9 @@ if (UA2F_BUILD_TESTS)
             pthread
     )
     ua2f_link_atomic_if_needed(ua2f_test)
+    # Count and inject allocation failures only inside explicitly scoped tests.
+    target_link_options(ua2f_test PRIVATE
+            -Wl,--wrap=malloc -Wl,--wrap=calloc -Wl,--wrap=realloc -Wl,--wrap=free)
     target_include_directories(ua2f_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src)
     target_include_directories(ua2f_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src/third/llhttp)
     if (UA2F_ENABLE_UCI)

src/handler.c

@@ -277,6 +277,8 @@ void handle_packet(const struct packet_io *io, void *io_ctx, const struct nf_pac
     assert(pkt->payload != NULL && "Packet payload cannot be NULL");
     assert(pkt->payload_len > 0 && "Packet payload length must be positive");
     struct pkt_buff *pkt_buff = NULL;
+    struct ua_mangle_entry ua_entries_inline[UA_INLINE_ENTRIES];
+    struct ua_mangle_entry *ua_entries_copy = ua_entries_inline;
     bool ct_ok = use_conntrack && pkt->has_conntrack;
     bool verdict_sent = false;
 
@@ -469,16 +471,34 @@ void handle_packet(const struct packet_io *io, void *io_ctx, const struct nf_pac
     // Level 3: feed to llhttp (session is valid, protected by its own state lock)
     session_state_lock(session);
     session_reset_per_packet(session, tcp_payload);
-    const int parse_ret = http_parser_feed(session, (const char *)tcp_payload, tcp_payload_len);
+    int parse_ret = http_parser_feed(session, (const char *)tcp_payload, tcp_payload_len);
 
     // Copy results out before releasing lock
-    const int ua_count = session->ua_entry_count;
-    struct ua_mangle_entry ua_entries_copy[UA_MAX_ENTRIES];
-    if (ua_count > 0) {
-        memcpy(ua_entries_copy, session->ua_entries, ua_count * sizeof(struct ua_mangle_entry));
+    const size_t ua_count = session->ua_entry_count;
+    if (parse_ret == 0 && ua_count > UA_INLINE_ENTRIES) {
+        ua_entries_copy = malloc(ua_count * sizeof(*ua_entries_copy));
+        if (ua_entries_copy == NULL) {
+            session->ua_allocation_failed = true;
+            parse_ret = HTTP_PARSER_NO_MEMORY;
+        }
+    }
+    if (parse_ret == 0 && ua_count > 0) {
+        for (size_t i = 0; i < ua_count; i++) {
+            ua_entries_copy[i] = *session_ua_entry_const(session, i);
+        }
     }
     session_state_unlock(session);
 
+    if (parse_ret == HTTP_PARSER_NO_MEMORY) {
+        // Keep the failed session so a retransmitted continuation cannot be
+        // mistaken for a new, non-HTTP stream and bypass rewriting.
+        session_release(session);
+        session = NULL;
+        syslog(LOG_ERR, "Failed to allocate User-Agent entries, dropping packet");
+        SEND_VERDICT(NF_DROP, MARK_NONE, NULL);
+        goto end;
+    }
+
     if (parse_ret != 0) {
         session_wrlock();
         session_delete(session);
@@ -499,7 +519,7 @@ void handle_packet(const struct packet_io *io, void *io_ctx, const struct nf_pac
     session = NULL;
 
     // Mangle UA entries (using copied data, session lock released)
-    for (int i = 0; i < ua_count; i++) {
+    for (size_t i = 0; i < ua_count; i++) {
         const size_t ua_offset = ua_entries_copy[i].offset;
         const size_t ua_len = ua_entries_copy[i].len;
         const size_t replacement_offset = ua_entries_copy[i].replacement_offset;
@@ -539,6 +559,9 @@ void handle_packet(const struct packet_io *io, void *io_ctx, const struct nf_pac
     SEND_VERDICT(NF_ACCEPT, (ct_ok && new_session) ? MARK_HTTP : MARK_NONE, pkt_buff);
 
 end:
+    if (ua_entries_copy != ua_entries_inline) {
+        free(ua_entries_copy);
+    }
     if (!verdict_sent) {
         SEND_VERDICT(NF_ACCEPT, MARK_NONE, NULL);
     }

src/http_parser_ua.c

@@ -1,6 +1,8 @@
 #include "http_parser_ua.h"
 
+#include <limits.h>
 #include <pthread.h>
+#include <stdlib.h>
 #include <string.h>
 #include <strings.h>
 #include <sys/syslog.h>
@@ -9,6 +11,38 @@
 #include "statistics.h"
 #include "third/llhttp/llhttp.h"
 
+// Override only at this expansion site: the upstream header remains unchanged.
+#undef utarray_oom
+#define utarray_oom() goto allocation_failed
+
+static bool append_ua_entry(struct http_session *session, const struct ua_mangle_entry *entry) {
+    if (session->ua_entry_count < UA_INLINE_ENTRIES) {
+        session->ua_entries_inline[session->ua_entry_count++] = *entry;
+        return true;
+    }
+
+    UT_array *overflow = &session->ua_entries_overflow;
+    const unsigned old_capacity = overflow->n;
+    // utarray uses unsigned counts and doubles capacity. Guard both its count
+    // arithmetic and byte-size multiplication, including the inline prefix.
+    const size_t max_capacity = SIZE_MAX / sizeof(*entry) - UA_INLINE_ENTRIES;
+    if (overflow->i == UINT_MAX ||
+        (overflow->i == overflow->n &&
+         (overflow->n > UINT_MAX / 2 || overflow->n > max_capacity / 2))) {
+        return false;
+    }
+    utarray_push_back(overflow, entry);
+    session->ua_entry_count++;
+    return true;
+
+allocation_failed:
+    // reserve changes n before realloc; i and d still describe the old buffer.
+    overflow->n = old_capacity;
+    return false;
+}
+
+#undef utarray_oom
+
 static int on_header_field(llhttp_t *parser, const char *data, size_t len) {
     struct http_session *session = (struct http_session *)parser->data;
 
@@ -82,19 +116,19 @@ static int on_header_value(llhttp_t *parser, const char *data, size_t len) {
 
     if (session->in_ua_value && session->ua_entry_count > 0) {
         // Continuation of the same UA value — extend current entry
-        size_t *entry_len = &session->ua_entries[session->ua_entry_count - 1].len;
+        size_t *entry_len = &session_ua_entry(session, session->ua_entry_count - 1)->len;
         if (SIZE_MAX - *entry_len < len) {
             *entry_len = SIZE_MAX;
         } else {
             *entry_len += len;
         }
     } else {
         // New UA entry
-        if (session->ua_entry_count < UA_MAX_ENTRIES) {
-            session->ua_entries[session->ua_entry_count].offset = offset;
-            session->ua_entries[session->ua_entry_count].len = len;
-            session->ua_entries[session->ua_entry_count].replacement_offset = session->ua_value_seen_len;
-            session->ua_entry_count++;
+        const struct ua_mangle_entry entry = {offset, len, session->ua_value_seen_len};
+        if (!append_ua_entry(session, &entry)) {
+            session->ua_allocation_failed = true;
+            llhttp_set_error_reason(parser, "Failed to allocate User-Agent entries");
+            return HPE_USER;
         }
         session->in_ua_value = true;
     }
@@ -118,7 +152,6 @@ static int on_headers_complete(llhttp_t *parser) {
 static int on_message_complete(llhttp_t *parser) {
     struct http_session *session = (struct http_session *)parser->data;
     session_reset_per_message(session);
-    session->last_active = time(NULL);
     return 0;
 }
 
@@ -138,16 +171,33 @@ void http_parser_init_session(struct http_session *session) {
 
     llhttp_init(&session->parser, HTTP_REQUEST, &shared_settings);
     session->parser.data = session;
+    session->ua_allocation_failed = false;
+
+    if (session->ua_entries_overflow.icd.sz == 0) {
+        const UT_icd entry_icd = {sizeof(struct ua_mangle_entry), NULL, NULL, NULL};
+        utarray_init(&session->ua_entries_overflow, &entry_icd);
+    }
 
     session_reset_per_message(session);
 }
 
 int http_parser_feed(struct http_session *session, const char *data, size_t len) {
+    // TTL measures idle time, including unfinished headers/bodies and retries
+    // on a session whose rewriting failed. Empty feeds are not activity.
+    if (len > 0) {
+        session->last_active = time(NULL);
+    }
+    // A failed allocation may leave llhttp partway through a payload. Keep
+    // rejecting this stream until it is idle or closed, rather than allowing
+    // later fragments to be reclassified as non-HTTP traffic.
+    if (session->ua_allocation_failed) {
+        return HTTP_PARSER_NO_MEMORY;
+    }
     llhttp_errno_t err = llhttp_execute(&session->parser, data, len);
     if (err != HPE_OK) {
         syslog(LOG_DEBUG, "llhttp parse error: %s (%s)", llhttp_errno_name(err),
                llhttp_get_error_reason(&session->parser));
-        return -1;
+        return session->ua_allocation_failed ? HTTP_PARSER_NO_MEMORY : -1;
     }
     return 0;
 }

src/http_parser_ua.h

@@ -6,8 +6,11 @@
 // Initialize llhttp parser and callbacks on a session.
 void http_parser_init_session(struct http_session *session);
 
-// Feed TCP payload to llhttp parser. Updates session->ua_entries.
-// Returns: 0 on success, -1 on parse error.
+// Allocation failures must not be treated as non-HTTP traffic and forwarded.
+#define HTTP_PARSER_NO_MEMORY (-2)
+
+// Feed TCP payload to llhttp parser. Records entries accessible through session_ua_entry().
+// Returns: 0 on success, -1 on parse error, HTTP_PARSER_NO_MEMORY on allocation failure.
 int http_parser_feed(struct http_session *session, const char *data, size_t len);
 
 #endif /* UA2F_HTTP_PARSER_UA_H */

src/http_session.c

@@ -188,11 +188,16 @@ bool session_state_init(struct http_session *session) {
 }
 
 void session_state_destroy(struct http_session *session) {
-    if (session == NULL || !session->state_lock_initialized) {
+    if (session == NULL) {
         return;
     }
-    pthread_mutex_destroy(&session->state_lock);
-    session->state_lock_initialized = false;
+    utarray_done(&session->ua_entries_overflow);
+    memset(&session->ua_entries_overflow, 0, sizeof(session->ua_entries_overflow));
+    session->ua_entry_count = 0;
+    if (session->state_lock_initialized) {
+        pthread_mutex_destroy(&session->state_lock);
+        session->state_lock_initialized = false;
+    }
 }
 
 void session_state_lock(struct http_session *session) {
@@ -209,9 +214,8 @@ void session_state_unlock(struct http_session *session) {
 
 void session_reset_per_packet(struct http_session *session, const void *tcp_payload_base) {
     session->ua_entry_count = 0;
+    utarray_clear(&session->ua_entries_overflow);
     session->tcp_payload_base = tcp_payload_base;
-    // last_active is updated in session_create and by the cleaner's TTL check.
-    // Avoid time() syscall on every packet — the TTL is coarse (300s default).
 }
 
 void session_reset_per_message(struct http_session *session) {

src/http_session.h

@@ -10,8 +10,9 @@
 #include "third/llhttp/llhttp.h"
 #include "third/nfqueue-mnl/nfqueue-mnl.h"
 #include "third/uthash/uthash.h"
+#include "third/uthash/utarray.h"
 
-#define UA_MAX_ENTRIES 8
+#define UA_INLINE_ENTRIES 8
 #define FIELD_BUF_SIZE 32
 
 struct ua_mangle_entry {
@@ -50,15 +51,33 @@ struct http_session {
     bool in_ua_value;
     size_t ua_value_seen_len;
 
-    struct ua_mangle_entry ua_entries[UA_MAX_ENTRIES];
-    int ua_entry_count;
+    // Keep the common case allocation-free; grow for large batches of headers.
+    struct ua_mangle_entry ua_entries_inline[UA_INLINE_ENTRIES];
+    UT_array ua_entries_overflow;
+    size_t ua_entry_count;
+    bool ua_allocation_failed; // fail closed until this session is destroyed/reinitialized
 
     const void *tcp_payload_base;
 
     time_t last_active;
     UT_hash_handle hh;
 };
 
+// Entry storage is split: the first eight slots are inline, the rest are utarray-owned.
+static inline struct ua_mangle_entry *session_ua_entry(struct http_session *session, size_t index) {
+    if (index < UA_INLINE_ENTRIES) {
+        return &session->ua_entries_inline[index];
+    }
+    return (struct ua_mangle_entry *)utarray_eltptr(&session->ua_entries_overflow, index - UA_INLINE_ENTRIES);
+}
+
+static inline const struct ua_mangle_entry *session_ua_entry_const(const struct http_session *session, size_t index) {
+    if (index < UA_INLINE_ENTRIES) {
+        return &session->ua_entries_inline[index];
+    }
+    return (const struct ua_mangle_entry *)utarray_eltptr(&session->ua_entries_overflow, index - UA_INLINE_ENTRIES);
+}
+
 void init_http_sessions(int max_sessions);
 struct session_key session_key_from_connid(uint32_t conn_id);
 struct session_key session_key_from_tuple(const struct ip_tuple *tuple);
@@ -73,6 +92,7 @@ int session_cleanup_expired(int ttl_seconds);
 void session_wrlock(void);
 void session_wrunlock(void);
 bool session_state_init(struct http_session *session);
+// Release parser entry storage and, if initialized, the state mutex.
 void session_state_destroy(struct http_session *session);
 void session_state_lock(struct http_session *session);
 void session_state_unlock(struct http_session *session);

src/proxy.c

@@ -416,6 +416,7 @@ static void free_closed_connections(struct proxy_context *ctx) {
         struct proxy_connection *conn = ctx->closing;
         ctx->closing = conn->close_next;
         connection_unlink(ctx, conn);
+        session_state_destroy(&conn->session);
         free(conn);
     }
 }
@@ -516,10 +517,11 @@ static void rewrite_user_agent_entries(uint8_t *buf, size_t len, const struct ht
     }
     const size_t replacement_len = UA2F_MAX_USER_AGENT_LENGTH;
 
-    for (int i = 0; i < session->ua_entry_count; i++) {
-        const size_t offset = session->ua_entries[i].offset;
-        const size_t ua_len = session->ua_entries[i].len;
-        const size_t replacement_offset = session->ua_entries[i].replacement_offset;
+    for (size_t i = 0; i < session->ua_entry_count; i++) {
+        const struct ua_mangle_entry *entry = session_ua_entry_const(session, i);
+        const size_t offset = entry->offset;
+        const size_t ua_len = entry->len;
+        const size_t replacement_offset = entry->replacement_offset;
         if (offset > len || ua_len > len - offset) {
             continue;
         }
@@ -535,14 +537,18 @@ static void rewrite_user_agent_entries(uint8_t *buf, size_t len, const struct ht
     }
 }
 
-static void process_client_payload(struct proxy_connection *conn, uint8_t *buf, size_t len) {
+static int process_client_payload(struct proxy_connection *conn, uint8_t *buf, size_t len) {
     if (conn->rewrite_disabled) {
-        return;
+        return 0;
     }
 
     count_tcp_packet();
     session_reset_per_packet(&conn->session, buf);
     const int parse_ret = http_parser_feed(&conn->session, (const char *)buf, len);
+    if (parse_ret == HTTP_PARSER_NO_MEMORY) {
+        syslog(LOG_ERR, "Failed to allocate User-Agent entries, closing connection");
+        return -1;
+    }
     if (conn->session.ua_entry_count > 0) {
         rewrite_user_agent_entries(buf, len, &conn->session);
         count_user_agent_packet();
@@ -553,6 +559,7 @@ static void process_client_payload(struct proxy_connection *conn, uint8_t *buf,
     if (parse_ret != 0) {
         conn->rewrite_disabled = true;
     }
+    return 0;
 }
 
 static int flush_buffer(int fd, struct proxy_buffer *buf) {
@@ -607,8 +614,9 @@ static int read_into_buffer(struct proxy_connection *conn, enum proxy_side side)
             return -1;
         }
 
-        if (side == PROXY_SIDE_CLIENT) {
-            process_client_payload(conn, out->data + out->len, (size_t)n);
+        if (side == PROXY_SIDE_CLIENT &&
+            process_client_payload(conn, out->data + out->len, (size_t)n) != 0) {
+            return -1;
         }
         out->len += (size_t)n;
     }

src/third/uthash/README.md

@@ -0,0 +1,22 @@
+# uthash
+
+The existing `uthash.h` declares version 2.3.0. The added `utarray.h` is
+vendored from the upstream v2.3.0 release:
+
+- Upstream: https://github.com/troydhanson/uthash
+- Pinned commit: `e493aa90a2833b4655927598f169c31cfcdf7861`
+- `utarray.h` is copied unchanged from `src/utarray.h` at that commit
+- The upstream BSD license notice is retained in each header
+
+UA2F embeds a `UT_array` for User-Agent entries beyond the eight inline slots.
+`utarray_init` allocates nothing, `utarray_clear` retains overflow capacity, and
+`utarray_done` frees only the overflow buffer. The inline storage is never given
+to utarray. The parser's append helper overrides utarray's default fatal OOM
+handling locally and restores capacity when a reserve fails.
+
+The zero-allocation guarantee is limited to UA entry storage for at most eight
+entries per parser feed. Session objects and network buffers are separate.
+Overflow parser capacity is reused across feeds; NFQUEUE still allocates a
+separate copy-out buffer for each payload with more than eight entries.
+Allocation-count and failure-injection tests use GNU-compatible linker `--wrap`
+options; production builds do not replace the allocator.