From 2d17f0c1e39fe43aed792dbab3c7a4e4dbc646e1 Mon Sep 17 00:00:00 2001
From: Zxilly <zxilly@outlook.com>
Date: Tue, 3 Dec 2024 18:57:21 +0800
Subject: [PATCH] fix: read ip packet protocol

---
 src/cache.c      | 14 ++++------
 src/handler.c    | 66 ++++++++++++++++++++++--------------------------
 src/statistics.c |  4 +--
 src/ua2f.c       |  1 +
 4 files changed, 38 insertions(+), 47 deletions(-)

diff --git a/src/cache.c b/src/cache.c
index a97af3c..383b783 100644
--- a/src/cache.c
+++ b/src/cache.c
@@ -55,21 +55,17 @@ void init_not_http_cache(const int interval) {
 }
 
 bool cache_contains(struct addr_port target) {
-    pthread_rwlock_rdlock(&cacheLock);
+    pthread_rwlock_wrlock(&cacheLock);
 
     struct cache *s;
     HASH_FIND(hh, not_http_dst_cache, &target, sizeof(struct addr_port), s);
+    if (s != NULL) {
+        s->last_time = time(NULL);
+    }
 
     pthread_rwlock_unlock(&cacheLock);
 
-    if (s != NULL) {
-        pthread_rwlock_wrlock(&cacheLock);
-        s->last_time = time(NULL);
-        pthread_rwlock_unlock(&cacheLock);
-        return true;
-    }
-
-    return false;
+    return s != NULL;
 }
 
 void cache_add(struct addr_port addr_port) {
diff --git a/src/handler.c b/src/handler.c
index 199e6eb..e6649fd 100644
--- a/src/handler.c
+++ b/src/handler.c
@@ -10,6 +10,7 @@
 #endif
 
 #include <assert.h>
+#include <linux/if_ether.h>
 #include <libnetfilter_queue/libnetfilter_queue_ipv4.h>
 #include <libnetfilter_queue/libnetfilter_queue_ipv6.h>
 #include <libnetfilter_queue/libnetfilter_queue_tcp.h>
@@ -84,7 +85,7 @@ static void send_verdict(const struct nf_queue *queue, const struct nf_packet *p
         syslog(LOG_ERR, "failed to put nfqueue header");
         goto end;
     }
-    nfq_nlmsg_verdict_put(nlh, pkt->packet_id, NF_ACCEPT);
+    nfq_nlmsg_verdict_put(nlh, (int)pkt->packet_id, NF_ACCEPT);
 
     if (mark.should_set) {
         struct nlattr *nest = mnl_attr_nest_start_check(nlh, SEND_BUF_LEN, NFQA_CT);
@@ -202,32 +203,19 @@ static bool ipv6_set_transport_header(struct pkt_buff *pkt_buff) {
     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;
+int get_pkt_ip_version(const struct nf_packet *pkt) {
+    if (pkt->has_conntrack) {
+        return pkt->orig.ip_version;
     }
 
-    // unknown ip type
-    if (ipv4_set_transport_header(pkt_buff)) {
-        count_ipv4_packet();
-        return IPV4;
+    switch (pkt->hw_protocol) {
+        case ETH_P_IP:
+            return IPV4;
+        case ETH_P_IPV6:
+            return IPV6;
+        default:
+            return IP_UNK;
     }
-    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) {
@@ -244,18 +232,22 @@ 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);
+    if (pkt_buff == NULL) {
+        syslog(LOG_ERR, "Failed to allocate packet buffer");
+        goto end;
+    }
 
-    int type;
-    if (pkt->has_conntrack) {
-        type = pkt->orig.ip_version;
-        set_transport_header(pkt_buff, type);
-    } else {
-        type = set_transport_header(pkt_buff, IP_UNK);
-        if (type == IP_UNK) {
-            syslog(LOG_ERR, "Failed to set transport header");
-            goto end;
-        }
+    int type = get_pkt_ip_version(pkt);
+    if (type == IP_UNK) {
+        // will this happen?
+        send_verdict(queue, pkt, get_next_mark(pkt, false), NULL);
+        syslog(LOG_WARNING, "Received unknown ip packet %x. You may set wrong firewall rules.", pkt->hw_protocol);
+    }
+
+    if (type == IPV4) {
+        assert(ipv4_set_transport_header(pkt_buff));
+    } else if (type == IPV6) {
+        assert(ipv6_set_transport_header(pkt_buff));
     }
 
     const __auto_type tcp_hdr = nfq_tcp_get_hdr(pkt_buff);
@@ -347,7 +339,9 @@ void handle_packet(const struct nf_queue *queue, const struct nf_packet *pkt) {
 
 end:
     free(pkt->payload);
-    pktb_free(pkt_buff);
+    if (pkt_buff != NULL) {
+        pktb_free(pkt_buff);
+    }
 
     try_print_statistics();
 }
diff --git a/src/statistics.c b/src/statistics.c
index 07a10cb..28e373f 100644
--- a/src/statistics.c
+++ b/src/statistics.c
@@ -10,7 +10,7 @@ static long long tcp_packet_count = 0;
 
 static long long ipv4_packet_count = 0;
 static long long ipv6_packet_count = 0;
-static long long last_report_count = 4;
+static long long last_report_count = 1;
 
 static time_t start_t;
 
@@ -29,7 +29,7 @@ void count_ipv4_packet() { ipv4_packet_count++; }
 
 void count_ipv6_packet() { ipv6_packet_count++; }
 
-static char time_string_buffer[100];
+static char time_string_buffer[512];
 
 char *fill_time_string(const double sec) {
     const int s = (int)sec;
diff --git a/src/ua2f.c b/src/ua2f.c
index e3d59b7..9bf8afc 100644
--- a/src/ua2f.c
+++ b/src/ua2f.c
@@ -72,6 +72,7 @@ void main_loop(struct nf_queue *queue) {
                     break;
                 }
             } else {
+                should_exit = true;
                 break;
             }
         }