From 219c612399bef8c0ec3bb2007c00b4e67325e850 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Tue, 8 Apr 2025 15:56:14 +0800
Subject: [PATCH] redirect: Fix UDP rules

---
 redirect_nftables.go       | 60 +++++++++++++++++++++++++++++++++-----
 redirect_nftables_rules.go | 21 +++++++++++++
 2 files changed, 74 insertions(+), 7 deletions(-)

diff --git a/redirect_nftables.go b/redirect_nftables.go
index 2f3adea..e89a246 100644
--- a/redirect_nftables.go
+++ b/redirect_nftables.go
@@ -113,12 +113,6 @@ func (r *autoRedirect) setupNFTables() error {
 			Priority: nftables.ChainPriorityRef(*nftables.ChainPriorityNATDest + 2),
 			Type:     nftables.ChainTypeFilter,
 		})
-		if r.enableIPv4 {
-			nftablesCreateExcludeDestinationIPSet(nft, table, chainPreRoutingUDP, 5, "inet4_local_address_set", nftables.TableFamilyIPv4, false)
-		}
-		if r.enableIPv6 {
-			nftablesCreateExcludeDestinationIPSet(nft, table, chainPreRoutingUDP, 6, "inet6_local_address_set", nftables.TableFamilyIPv6, false)
-		}
 		nft.AddRule(&nftables.Rule{
 			Table: table,
 			Chain: chainPreRoutingUDP,
@@ -128,10 +122,28 @@ func (r *autoRedirect) setupNFTables() error {
 					Register: 1,
 				},
 				&expr.Cmp{
-					Op:       expr.CmpOpEq,
+					Op:       expr.CmpOpNeq,
 					Register: 1,
 					Data:     []byte{unix.IPPROTO_UDP},
 				},
+				&expr.Verdict{
+					Kind: expr.VerdictReturn,
+				},
+			},
+		})
+		nft.AddRule(&nftables.Rule{
+			Table: table,
+			Chain: chainPreRoutingUDP,
+			Exprs: []expr.Any{
+				&expr.Meta{
+					Key:      expr.MetaKeyIIFNAME,
+					Register: 1,
+				},
+				&expr.Cmp{
+					Op:       expr.CmpOpNeq,
+					Register: 1,
+					Data:     nftablesIfname(r.tunOptions.Name),
+				},
 				&expr.Ct{
 					Key:      expr.CtKeyMARK,
 					Register: 1,
@@ -149,6 +161,40 @@ func (r *autoRedirect) setupNFTables() error {
 				&expr.Counter{},
 			},
 		})
+		nft.AddRule(&nftables.Rule{
+			Table: table,
+			Chain: chainPreRoutingUDP,
+			Exprs: []expr.Any{
+				&expr.Ct{
+					Key:      expr.CtKeyMARK,
+					Register: 1,
+				},
+				&expr.Cmp{
+					Op:       expr.CmpOpNeq,
+					Register: 1,
+					Data:     binaryutil.NativeEndian.PutUint32(r.tunOptions.AutoRedirectInputMark),
+				},
+				&expr.Immediate{
+					Register: 1,
+					Data:     binaryutil.NativeEndian.PutUint32(r.tunOptions.AutoRedirectOutputMark),
+				},
+				&expr.Meta{
+					Key:            expr.MetaKeyMARK,
+					Register:       1,
+					SourceRegister: true,
+				},
+				&expr.Meta{
+					Key:      expr.MetaKeyMARK,
+					Register: 1,
+				},
+				&expr.Ct{
+					Key:            expr.CtKeyMARK,
+					Register:       1,
+					SourceRegister: true,
+				},
+				&expr.Counter{},
+			},
+		})
 	}
 
 	err = r.configureOpenWRTFirewall4(nft, false)
diff --git a/redirect_nftables_rules.go b/redirect_nftables_rules.go
index cc694c6..8361917 100644
--- a/redirect_nftables_rules.go
+++ b/redirect_nftables_rules.go
@@ -138,6 +138,27 @@ func (r *autoRedirect) nftablesCreateExcludeRules(nft *nftables.Conn, table *nft
 				},
 			},
 		})
+		if chain.Type == nftables.ChainTypeRoute {
+			nft.AddRule(&nftables.Rule{
+				Table: table,
+				Chain: chain,
+				Exprs: []expr.Any{
+					&expr.Ct{
+						Key:      expr.CtKeyMARK,
+						Register: 1,
+					},
+					&expr.Cmp{
+						Op:       expr.CmpOpEq,
+						Register: 1,
+						Data:     binaryutil.NativeEndian.PutUint32(r.tunOptions.AutoRedirectOutputMark),
+					},
+					&expr.Counter{},
+					&expr.Verdict{
+						Kind: expr.VerdictReturn,
+					},
+				},
+			})
+		}
 	}
 	if chain.Hooknum == nftables.ChainHookPrerouting {
 		if len(r.tunOptions.IncludeInterface) > 0 {