From bf080925b7a78338e664561d188a776c97df4777 Mon Sep 17 00:00:00 2001
From: Zxilly <zxilly@outlook.com>
Date: Mon, 22 Jul 2024 20:26:39 +0800
Subject: [PATCH] feat: allow disable connmark with config

---
 README.md                 |  3 +++
 openwrt/files/ua2f.config |  3 ++-
 openwrt/files/ua2f.init   | 31 ++++++++++++++++---------------
 src/cli.c                 |  6 ++++++
 src/config.c              | 12 +++++++++---
 src/config.h              |  1 +
 src/handler.c             | 25 ++++++++++++++-----------
 src/handler.h             |  2 +-
 src/ua2f.c                |  6 +++---
 9 files changed, 55 insertions(+), 34 deletions(-)

diff --git a/README.md b/README.md
index ac0f458..9d6cbe9 100644
--- a/README.md
+++ b/README.md
@@ -33,6 +33,9 @@ uci set ua2f.firewall.handle_intranet=1
 # 使用自定义 User-Agent
 uci set ua2f.main.custom_ua="Test UA/1.0"
 
+# 禁用 Conntrack 标记，这会降低性能，但是有助于和其他修改 Connmark 的软件共存
+uci set ua2f.main.disable_connmark=1
+
 # 应用配置
 uci commit ua2f
 
diff --git a/openwrt/files/ua2f.config b/openwrt/files/ua2f.config
index 4fcb471..fadb852 100644
--- a/openwrt/files/ua2f.config
+++ b/openwrt/files/ua2f.config
@@ -8,4 +8,5 @@ config ua2f 'firewall'
 	option handle_intranet '1'
 
 config ua2f 'main'
-	option custom_ua ''
\ No newline at end of file
+	option custom_ua ''
+	option disable_connmark '0'
\ No newline at end of file
diff --git a/openwrt/files/ua2f.init b/openwrt/files/ua2f.init
index 70c7809..ebd8933 100755
--- a/openwrt/files/ua2f.init
+++ b/openwrt/files/ua2f.init
@@ -34,10 +34,11 @@ setup_firewall() {
 		return 1
 	fi
 
-	local handle_tls handle_intranet handle_mmtls
+	local handle_tls handle_intranet handle_mmtls disable_connmark
 	config_get_bool handle_tls "firewall" "handle_tls" "0"
 	config_get_bool handle_intranet "firewall" "handle_intranet" "0"
 	config_get_bool handle_mmtls "firewall" "handle_mmtls" "0"
+	config_get_bool disable_connmark "main" "disable_connmark" "0"
 
 	if [ -n "$HAS_NFT" ]; then
 		nft -f- <<-EOF
@@ -67,23 +68,23 @@ setup_firewall() {
 					};
 				}
 
-				chain prerouting {
-					type filter hook prerouting priority mangle -5; policy accept;
+				chain postrouting {
+					type filter hook postrouting priority mangle -5; policy accept;
 
 					$([ "$handle_intranet" -ne "1" ] || echo 'ip daddr @localaddr_v4 counter return;')
 					$([ "$handle_intranet" -ne "1" ] || echo 'ip6 daddr @localaddr_v6 counter return;')
 
 					tcp dport 22 counter return comment "!ua2f: bypass SSH";
 					$([ "$handle_tls" -eq "1" ] || echo 'tcp dport 443 counter return comment "!ua2f: bypass HTTPS";')
-					tcp dport 80 counter ct mark set 44;
-					ct mark 43 counter return comment "!ua2f: bypass non-http stream";
+					$([ "$disable_connmark" -eq "1" ] || echo 'tcp dport 80 counter ct mark set 44;')
+					$([ "$disable_connmark" -eq "1" ] || echo 'ct mark 43 counter return comment "!ua2f: bypass non-http stream";')
 					meta l4proto tcp ct direction original counter queue num 10010 bypass;
 				}
 			}
 		EOF
 	else
 		# Flush existing rules
-		$IPT_M -D FORWARD -p tcp -m conntrack --ctdir ORIGINAL -j ua2f 2>"/dev/null"
+		$IPT_M -D POSTROUTING -p tcp -m conntrack --ctdir ORIGINAL -j ua2f 2>"/dev/null"
 
 		if $IPT_M -N ua2f; then
 			if [ "$handle_intranet" -ne "1" ]; then
@@ -98,16 +99,16 @@ setup_firewall() {
 			fi
 
 			[ "$handle_tls" -eq "1" ] || $IPT_M -A ua2f -p tcp --dport 443 -j RETURN # 不处理 HTTPS
-			$IPT_M -A ua2f -p tcp --dport 80 -j CONNMARK --set-mark 44
-			$IPT_M -A ua2f -m connmark --mark 43 -j RETURN # 不处理标记为非 http 的流
+			[ "$disable_connmark" -eq "1" ] || $IPT_M -A ua2f -p tcp --dport 80 -j CONNMARK --set-mark 44
+      [ "$disable_connmark" -eq "1" ] || $IPT_M -A ua2f -m connmark --mark 43 -j RETURN # 不处理标记为非 http 的流
 			[ "$handle_mmtls" -eq "1" ] || $IPT_M -A ua2f -p tcp --dport 80 -m string --string "/mmtls/" --algo bm -j RETURN # 不处理微信的mmtls
 			$IPT_M -A ua2f -j NFQUEUE --queue-num 10010 --queue-bypass
 		fi
-		$IPT_M -A FORWARD -p tcp -m conntrack --ctdir ORIGINAL -j ua2f
+		$IPT_M -A POSTROUTING -p tcp -m conntrack --ctdir ORIGINAL -j ua2f
 
 		if [ -n "$HAS_IPT6" ]; then
 			# Flush existing rules
-			$IPT6_M -D FORWARD -p tcp -m conntrack --ctdir ORIGINAL -j ua2f 2>"/dev/null"
+			$IPT6_M -D POSTROUTING -p tcp -m conntrack --ctdir ORIGINAL -j ua2f 2>"/dev/null"
 
 			if $IPT6_M -N ua2f; then
 				if [ "$handle_intranet" -ne "1" ]; then
@@ -127,12 +128,12 @@ setup_firewall() {
 				fi
 
 				[ "$handle_tls" -eq "1" ] || $IPT6_M -A ua2f -p tcp --dport 443 -j RETURN # 不处理 HTTPS
-				$IPT6_M -A ua2f -p tcp --dport 80 -j CONNMARK --set-mark 44
-				$IPT6_M -A ua2f -m connmark --mark 43 -j RETURN # 不处理标记为非 http 的流
+				[ "$disable_connmark" -eq "1" ] || $IPT6_M -A ua2f -p tcp --dport 80 -j CONNMARK --set-mark 44
+        [ "$disable_connmark" -eq "1" ] || $IPT6_M -A ua2f -m connmark --mark 43 -j RETURN # 不处理标记为非 http 的流
 				[ "$handle_mmtls" -eq "1" ] || $IPT6_M -A ua2f -p tcp --dport 80 -m string --string "/mmtls/" --algo bm -j RETURN # 不处理微信的mmtls
 				$IPT6_M -A ua2f -j NFQUEUE --queue-num 10010 --queue-bypass
 			fi
-			$IPT6_M -A FORWARD -p tcp -m conntrack --ctdir ORIGINAL -j ua2f
+			$IPT6_M -A POSTROUTING -p tcp -m conntrack --ctdir ORIGINAL -j ua2f
 		fi
 	fi
 }
@@ -173,12 +174,12 @@ stop_service() {
 		nft flush table inet ua2f
 		nft delete table inet ua2f
 	else
-		$IPT_M -D FORWARD -p tcp -m conntrack --ctdir ORIGINAL -j ua2f 2>"/dev/null"
+		$IPT_M -D POSTROUTING -p tcp -m conntrack --ctdir ORIGINAL -j ua2f 2>"/dev/null"
 		$IPT_M -F ua2f 2>"/dev/null"
 		$IPT_M -X ua2f 2>"/dev/null"
 
 		if [ -n "$HAS_IPT6" ]; then
-			$IPT6_M -D FORWARD -p tcp -m conntrack --ctdir ORIGINAL -j ua2f 2>"/dev/null"
+			$IPT6_M -D POSTROUTING -p tcp -m conntrack --ctdir ORIGINAL -j ua2f 2>"/dev/null"
 			$IPT6_M -F ua2f 2>"/dev/null"
 			$IPT6_M -X ua2f 2>"/dev/null"
 		fi
diff --git a/src/cli.c b/src/cli.c
index be92266..0b94dbf 100644
--- a/src/cli.c
+++ b/src/cli.c
@@ -29,6 +29,12 @@ void try_print_info(const int argc, char *argv[]) {
         } else {
             printf("Config UA: not set\n");
         }
+
+        if (config.disable_connmark) {
+            printf("Conntrack cache: disabled\n");
+        } else {
+            printf("Conntrack cache: auto\n");
+        }
 #else
         printf("UCI support disabled\n");
 #endif
diff --git a/src/config.c b/src/config.c
index 467d454..e05988b 100644
--- a/src/config.c
+++ b/src/config.c
@@ -8,11 +8,12 @@
 struct ua2f_config config = {
     .use_custom_ua = false,
     .custom_ua = NULL,
+    .disable_connmark = false,
 };
 
 void load_config() {
     const __auto_type ctx = uci_alloc_context();
-    if (!ctx) {
+    if (ctx == NULL) {
         syslog(LOG_ERR, "Failed to allocate uci context");
         return;
     }
@@ -24,12 +25,12 @@ void load_config() {
 
     // find ua2f.main.custom_ua
     const __auto_type section = uci_lookup_section(ctx, package, "main");
-    if (!section) {
+    if (section == NULL) {
         goto cleanup;
     }
 
     const __auto_type custom_ua = uci_lookup_option_string(ctx, section, "custom_ua");
-    if (!custom_ua) {
+    if (custom_ua == NULL) {
         goto cleanup;
     }
     if (strlen(custom_ua) > 0) {
@@ -37,6 +38,11 @@ void load_config() {
         config.custom_ua = strdup(custom_ua);
     }
 
+    const __auto_type disable_connmark = uci_lookup_option_string(ctx, section, "disable_connmark");
+    if (disable_connmark != NULL && strcmp(disable_connmark, "1") == 0) {
+        config.disable_connmark = true;
+    }
+
 cleanup:
     uci_free_context(ctx);
 }
diff --git a/src/config.h b/src/config.h
index 08bc43c..cce482e 100644
--- a/src/config.h
+++ b/src/config.h
@@ -9,6 +9,7 @@
 struct ua2f_config {
     bool use_custom_ua;
     char *custom_ua;
+    bool disable_connmark;
 };
 
 void load_config();
diff --git a/src/handler.c b/src/handler.c
index 04db694..0bf2c67 100644
--- a/src/handler.c
+++ b/src/handler.c
@@ -27,6 +27,9 @@ static char *replacement_user_agent_string = NULL;
 #define CONNMARK_NOT_HTTP 43
 #define CONNMARK_HTTP 44
 
+bool use_conntrack = true;
+static bool cache_initialized = false;
+
 void init_handler() {
     replacement_user_agent_string = malloc(MAX_USER_AGENT_LENGTH);
 
@@ -39,6 +42,11 @@ void init_handler() {
         syslog(LOG_INFO, "Using config user agent string: %s", replacement_user_agent_string);
         ua_set = true;
     }
+
+    if (config.disable_connmark) {
+        use_conntrack = false;
+        syslog(LOG_INFO, "Conntrack cache disabled by config.");
+    }
 #endif
 
 #ifdef UA2F_CUSTOM_UA
@@ -100,9 +108,6 @@ end:
     }
 }
 
-bool conntrack_info_available = true;
-static bool cache_initialized = false;
-
 static void add_to_cache(const struct nf_packet *pkt) {
     struct addr_port target = {
         .addr = pkt->orig.dst,
@@ -113,7 +118,7 @@ static void add_to_cache(const struct nf_packet *pkt) {
 }
 
 static struct mark_op get_next_mark(const struct nf_packet *pkt, const bool has_ua) {
-    if (!conntrack_info_available) {
+    if (!use_conntrack) {
         return (struct mark_op){false, 0};
     }
 
@@ -162,9 +167,9 @@ bool should_ignore(const struct nf_packet *pkt) {
 }
 
 void handle_packet(const struct nf_queue *queue, const struct nf_packet *pkt) {
-    if (conntrack_info_available) {
+    if (use_conntrack) {
         if (!pkt->has_conntrack) {
-            conntrack_info_available = false;
+            use_conntrack = false;
             syslog(LOG_WARNING, "Packet has no conntrack. Switching to no cache mode.");
             syslog(LOG_WARNING, "Note that this may lead to performance degradation. Especially on low-end routers.");
         } else {
@@ -175,19 +180,17 @@ void handle_packet(const struct nf_queue *queue, const struct nf_packet *pkt) {
         }
     }
 
-    struct pkt_buff *pkt_buff = NULL;
-    if (conntrack_info_available && should_ignore(pkt)) {
+    if (use_conntrack && should_ignore(pkt)) {
         send_verdict(queue, pkt, (struct mark_op){true, CONNMARK_NOT_HTTP}, NULL);
         goto end;
     }
 
-    pkt_buff = pktb_alloc(AF_INET, pkt->payload, pkt->payload_len, 0);
-
+    struct pkt_buff *pkt_buff = pktb_alloc(AF_INET, pkt->payload, pkt->payload_len, 0);
     ASSERT(pkt_buff != NULL);
 
     int type;
 
-    if (conntrack_info_available) {
+    if (use_conntrack) {
         type = pkt->orig.ip_version;
     } else {
         const __auto_type ip_hdr = nfq_ip_get_hdr(pkt_buff);
diff --git a/src/handler.h b/src/handler.h
index 7be7341..e5df06c 100644
--- a/src/handler.h
+++ b/src/handler.h
@@ -3,7 +3,7 @@
 
 #include "third/nfqueue-mnl.h"
 
-extern bool conntrack_info_available;
+extern bool use_conntrack;
 
 void init_handler();
 
diff --git a/src/ua2f.c b/src/ua2f.c
index 2426d0b..e38e26f 100644
--- a/src/ua2f.c
+++ b/src/ua2f.c
@@ -50,10 +50,10 @@ int read_buffer(struct nf_queue *queue, struct nf_buffer *buf) {
     }
 }
 
-bool retry_disable_conntrack(struct nf_queue *queue) {
+bool retry_without_conntrack(struct nf_queue *queue) {
     nfqueue_close(queue);
 
-    syslog(LOG_INFO, "Retrying to disable conntrack");
+    syslog(LOG_INFO, "Retry without conntrack");
     const __auto_type ret = nfqueue_open(queue, QUEUE_NUM, 0, true);
     if (!ret) {
         syslog(LOG_ERR, "Failed to open nfqueue with conntrack disabled");
@@ -70,7 +70,7 @@ void main_loop(struct nf_queue *queue) {
         if (read_buffer(queue, buf) == IO_ERROR) {
             if (!retried) {
                 retried = true;
-                if (!retry_disable_conntrack(queue)) {
+                if (!retry_without_conntrack(queue)) {
                     break;
                 }
             } else {