update 2024-11-27 10:57:37

UA2F/scripts/test.py

@@ -50,10 +50,12 @@ def start_ua2f(u: str):
 # iptables 设置函数
 def setup_iptables():
     os.system(f"sudo iptables -A OUTPUT -p tcp --dport {PORT} -j NFQUEUE --queue-num 10010")
+    # os.system(f"sudo ip6tables -A OUTPUT -p tcp --dport {PORT} -j NFQUEUE --queue-num 10010")
 
 
 def cleanup_iptables():
     os.system(f"sudo iptables -D OUTPUT -p tcp --dport {PORT} -j NFQUEUE --queue-num 10010")
+    # os.system(f"sudo ip6tables -D OUTPUT -p tcp --dport {PORT} -j NFQUEUE --queue-num 10010")
 
 
 if __name__ == "__main__":
@@ -81,13 +83,21 @@ if __name__ == "__main__":
 
     time.sleep(3)
 
-    for i in tqdm(range(10000)):
+    for i in tqdm(range(2000)):
         nxt = ua.random
-        response = requests.get(f"http://localhost:{PORT}", headers={
+        response = requests.get(f"http://127.0.0.1:{PORT}", headers={
             "User-Agent": nxt
         })
         assert response.ok
         assert response.text == str(len(nxt))
 
+    # for i in tqdm(range(2000)):
+    #     nxt = ua.random
+    #     response = requests.get(f"http://[::1]:{PORT}", headers={
+    #         "User-Agent": nxt
+    #     })
+    #     assert response.ok
+    #     assert response.text == str(len(nxt))
+
     # clean
     cleanup_iptables()

UA2F/src/cache.c

@@ -63,18 +63,10 @@ bool cache_contains(struct addr_port target) {
     pthread_rwlock_unlock(&cacheLock);
 
     if (s != NULL) {
-        bool ret;
         pthread_rwlock_wrlock(&cacheLock);
-        if (difftime(time(NULL), s->last_time) > check_interval * 2) {
+        s->last_time = time(NULL);
-            HASH_DEL(not_http_dst_cache, s);
-            free(s);
-            ret = false;
-        } else {
-            s->last_time = time(NULL);
-            ret = true;
-        }
         pthread_rwlock_unlock(&cacheLock);
-        return ret;
+        return true;
     }
 
     return false;

UA2F/src/handler.c

@@ -9,6 +9,7 @@
 #include "config.h"
 #endif
 
+#include <assert.h>
 #include <libnetfilter_queue/libnetfilter_queue_ipv4.h>
 #include <libnetfilter_queue/libnetfilter_queue_ipv6.h>
 #include <libnetfilter_queue/libnetfilter_queue_tcp.h>
@@ -114,7 +115,7 @@ end:
 }
 
 static void add_to_cache(const struct nf_packet *pkt) {
-    struct addr_port target = {
+    const struct addr_port target = {
         .addr = pkt->orig.dst,
         .port = pkt->orig.dst_port,
     };
@@ -171,6 +172,64 @@ bool should_ignore(const struct nf_packet *pkt) {
     return retval;
 }
 
+enum {
+    IP_UNK = 0,
+};
+
+static bool ipv4_set_transport_header(struct pkt_buff *pkt_buff) {
+    struct iphdr *ip_hdr = nfq_ip_get_hdr(pkt_buff);
+    if (ip_hdr == NULL) {
+        return false;
+    }
+
+    if (nfq_ip_set_transport_header(pkt_buff, ip_hdr) < 0) {
+        syslog(LOG_ERR, "Failed to set ipv4 transport header");
+        return false;
+    }
+    return true;
+}
+
+static bool ipv6_set_transport_header(struct pkt_buff *pkt_buff) {
+    struct ip6_hdr *ip_hdr = nfq_ip6_get_hdr(pkt_buff);
+    if (ip_hdr == NULL) {
+        return false;
+    }
+
+    if (nfq_ip6_set_transport_header(pkt_buff, ip_hdr, IPPROTO_TCP) < 0) {
+        syslog(LOG_ERR, "Failed to set ipv6 transport header");
+        return false;
+    }
+    return true;
+}
+
+static int set_transport_header(struct pkt_buff* pkt_buff, const int ip_type) {
+    if (ip_type == IPV4) {
+        if (ipv4_set_transport_header(pkt_buff)) {
+            count_ipv4_packet();
+            return IPV4;
+        }
+        return IP_UNK;
+    }
+    if (ip_type == IPV6) {
+        if (ipv6_set_transport_header(pkt_buff)) {
+            count_ipv6_packet();
+            return IPV6;
+        }
+        return IP_UNK;
+    }
+
+    // unknown ip type
+    if (ipv4_set_transport_header(pkt_buff)) {
+        count_ipv4_packet();
+        return IPV4;
+    }
+    if (ipv6_set_transport_header(pkt_buff)) {
+        count_ipv6_packet();
+        return IPV6;
+    }
+    return IP_UNK;
+}
+
 void handle_packet(const struct nf_queue *queue, const struct nf_packet *pkt) {
     if (use_conntrack) {
         if (!pkt->has_conntrack) {
@@ -191,36 +250,16 @@ void handle_packet(const struct nf_queue *queue, const struct nf_packet *pkt) {
     }
 
     struct pkt_buff *pkt_buff = pktb_alloc(AF_INET, pkt->payload, pkt->payload_len, 0);
-    ASSERT(pkt_buff != NULL);
+    assert(pkt_buff != NULL);
 
     int type;
     if (use_conntrack) {
         type = pkt->orig.ip_version;
+        set_transport_header(pkt_buff, type);
     } else {
-        const __auto_type ip_hdr = nfq_ip_get_hdr(pkt_buff);
+        type = set_transport_header(pkt_buff, IP_UNK);
-        if (ip_hdr == NULL) {
+        if (type == IP_UNK) {
-            type = IPV6;
+            syslog(LOG_ERR, "Failed to set transport header");
-        } else {
-            type = IPV4;
-        }
-    }
-
-    if (type == IPV4) {
-        count_ipv4_packet();
-    } else {
-        count_ipv6_packet();
-    }
-
-    if (type == IPV4) {
-        const __auto_type ip_hdr = nfq_ip_get_hdr(pkt_buff);
-        if (nfq_ip_set_transport_header(pkt_buff, ip_hdr) < 0) {
-            syslog(LOG_ERR, "Failed to set ipv4 transport header");
-            goto end;
-        }
-    } else {
-        const __auto_type ip_hdr = nfq_ip6_get_hdr(pkt_buff);
-        if (nfq_ip6_set_transport_header(pkt_buff, ip_hdr, IPPROTO_TCP) < 0) {
-            syslog(LOG_ERR, "Failed to set ipv6 transport header");
             goto end;
         }
     }
@@ -234,9 +273,13 @@ void handle_packet(const struct nf_queue *queue, const struct nf_packet *pkt) {
     }
 
     const __auto_type tcp_payload = nfq_tcp_get_payload(tcp_hdr, pkt_buff);
-    const __auto_type tcp_payload_len = nfq_tcp_get_payload_len(tcp_hdr, pkt_buff);
+    if (tcp_payload == NULL) {
+        send_verdict(queue, pkt, get_next_mark(pkt, false), NULL);
+        goto end;
+    }
 
-    if (tcp_payload == NULL || tcp_payload_len < USER_AGENT_MATCH_LENGTH) {
+    const __auto_type tcp_payload_len = nfq_tcp_get_payload_len(tcp_hdr, pkt_buff);
+    if (tcp_payload_len < USER_AGENT_MATCH_LENGTH) {
         send_verdict(queue, pkt, get_next_mark(pkt, false), NULL);
         goto end;
     }

UA2F/src/ua2f.c

@@ -8,10 +8,12 @@
 #include "config.h"
 #endif
 
+#include <assert.h>
 #include <signal.h>
 #include <stdlib.h>
 #include <string.h>
 #include <syslog.h>
+#include <stdbool.h>
 
 #pragma clang diagnostic push
 #pragma ide diagnostic ignored "EndlessLoop"
@@ -28,11 +30,9 @@ int parse_packet(const struct nf_queue *queue, struct nf_buffer *buf) {
 
     while (!should_exit) {
         const __auto_type status = nfqueue_next(buf, packet);
-        switch (status) {
+        if (status == IO_READY) {
-        case IO_READY:
             handle_packet(queue, packet);
-            break;
+        } else {
-        default:
             return status;
         }
     }
@@ -42,12 +42,10 @@ int parse_packet(const struct nf_queue *queue, struct nf_buffer *buf) {
 
 int read_buffer(struct nf_queue *queue, struct nf_buffer *buf) {
     const __auto_type buf_status = nfqueue_receive(queue, buf, 0);
-    switch (buf_status) {
+    if (buf_status == IO_READY) {
-    case IO_READY:
         return parse_packet(queue, buf);
-    default:
-        return buf_status;
     }
+    return buf_status;
 }
 
 bool retry_without_conntrack(struct nf_queue *queue) {
@@ -109,6 +107,8 @@ int main(const int argc, char *argv[]) {
         syslog(LOG_ERR, "Failed to open nfqueue");
         return EXIT_FAILURE;
     }
+    assert(queue->queue_num == QUEUE_NUM);
+    assert(queue->nl_socket != NULL);
 
     main_loop(queue);
 
@@ -117,4 +117,4 @@ int main(const int argc, char *argv[]) {
     return EXIT_SUCCESS;
 }
 
-#pragma clang diagnostic pop
+#pragma clang diagnostic pop

UA2F/test/cache_test.cc

@@ -6,7 +6,7 @@ extern "C" {
 
 class CacheTest : public ::testing::Test {
 protected:
-    struct addr_port test_addr{};
+    addr_port test_addr{};
 
     void SetUp() override {
         test_addr.addr.ip4 = 12345;
@@ -39,11 +39,11 @@ TEST_F(CacheTest, AddAndRemoveFromCache) {
     cache_add(test_addr);
     EXPECT_TRUE(cache_contains(test_addr));
     sleep(5);
-    EXPECT_FALSE(cache_contains(test_addr));
+    EXPECT_TRUE(cache_contains(test_addr));
 }
 
 TEST_F(CacheTest, CacheDoesNotContainNonexistentEntry) {
-    struct addr_port nonexistent_addr{};
+    addr_port nonexistent_addr{};
     nonexistent_addr.addr.ip4 = 54321;
     EXPECT_FALSE(cache_contains(nonexistent_addr));
 }