From 891f6ee7c2d9c66199ab52695ff12cc811f8e2e3 Mon Sep 17 00:00:00 2001
From: Vladimir Avtsenov <vladimir.lsk.cool@gmail.com>
Date: Fri, 14 Feb 2025 20:55:37 +0300
Subject: [PATCH] fix dns routing

---
 netfilter-helper/ipset-to-link.go |  1 +
 netfilter-helper/port-remap.go    | 72 ++++++++++++++++++++++++++-----
 2 files changed, 63 insertions(+), 10 deletions(-)

diff --git a/netfilter-helper/ipset-to-link.go b/netfilter-helper/ipset-to-link.go
index 6866cbf..5935601 100644
--- a/netfilter-helper/ipset-to-link.go
+++ b/netfilter-helper/ipset-to-link.go
@@ -37,6 +37,7 @@ func (r *IPSetToLink) insertIPTablesRules(table string) error {
 		}
 
 		for _, iptablesArgs := range [][]string{
+			{"-j", "CONNMARK", "--restore-mark"},
 			{"-j", "MARK", "--set-mark", strconv.Itoa(int(r.mark))},
 			{"-j", "CONNMARK", "--save-mark"},
 		} {
diff --git a/netfilter-helper/port-remap.go b/netfilter-helper/port-remap.go
index e8512f7..32e52c9 100644
--- a/netfilter-helper/port-remap.go
+++ b/netfilter-helper/port-remap.go
@@ -21,7 +21,8 @@ type PortRemap struct {
 
 func (r *PortRemap) insertIPTablesRules(table string) error {
 	if table == "" || table == "nat" {
-		err := r.IPTables.NewChain("nat", r.ChainName)
+		preroutingChain := r.ChainName + "_PRR"
+		err := r.IPTables.NewChain("nat", preroutingChain)
 		if err != nil {
 			// If not "AlreadyExists"
 			if eerr, eok := err.(*iptables.Error); !(eok && eerr.ExitStatus() == 1) {
@@ -34,18 +35,62 @@ func (r *PortRemap) insertIPTablesRules(table string) error {
 				continue
 			}
 
-			for _, iptablesArgs := range [][]string{
-				{"-p", "tcp", "-d", addr.IP.String(), "--dport", strconv.Itoa(int(r.From)), "-j", "DNAT", "--to-destination", fmt.Sprintf(":%d", r.To)},
-				{"-p", "udp", "-d", addr.IP.String(), "--dport", strconv.Itoa(int(r.From)), "-j", "DNAT", "--to-destination", fmt.Sprintf(":%d", r.To)},
-			} {
-				err = r.IPTables.AppendUnique("nat", r.ChainName, iptablesArgs...)
-				if err != nil {
-					return fmt.Errorf("failed to append rule: %w", err)
+			if r.IPTables.Proto() != iptables.ProtocolIPv6 {
+				for _, iptablesArgs := range [][]string{
+					{"-p", "tcp", "-d", addr.IP.String(), "--dport", fmt.Sprintf("%d", r.From), "-j", "REDIRECT", "--to-port", fmt.Sprintf("%d", r.To)},
+					{"-p", "udp", "-d", addr.IP.String(), "--dport", fmt.Sprintf("%d", r.From), "-j", "REDIRECT", "--to-port", fmt.Sprintf("%d", r.To)},
+				} {
+					err = r.IPTables.AppendUnique("nat", preroutingChain, iptablesArgs...)
+					if err != nil {
+						return fmt.Errorf("failed to append rule: %w", err)
+					}
+				}
+			} else {
+				for _, iptablesArgs := range [][]string{
+					{"-p", "tcp", "-d", addr.IP.String(), "--dport", strconv.Itoa(int(r.From)), "-j", "DNAT", "--to-destination", fmt.Sprintf(":%d", r.To)},
+					{"-p", "udp", "-d", addr.IP.String(), "--dport", strconv.Itoa(int(r.From)), "-j", "DNAT", "--to-destination", fmt.Sprintf(":%d", r.To)},
+				} {
+					err = r.IPTables.AppendUnique("nat", preroutingChain, iptablesArgs...)
+					if err != nil {
+						return fmt.Errorf("failed to append rule: %w", err)
+					}
 				}
 			}
 		}
 
-		err = r.IPTables.InsertUnique("nat", "PREROUTING", 1, "-j", r.ChainName)
+		err = r.IPTables.InsertUnique("nat", "PREROUTING", 1, "-j", preroutingChain)
+		if err != nil {
+			return fmt.Errorf("failed to linking chain: %w", err)
+		}
+
+		postroutingChain := r.ChainName + "_POR"
+		err = r.IPTables.NewChain("nat", postroutingChain)
+		if err != nil {
+			// If not "AlreadyExists"
+			if eerr, eok := err.(*iptables.Error); !(eok && eerr.ExitStatus() == 1) {
+				return fmt.Errorf("failed to create chain: %w", err)
+			}
+		}
+
+		for _, addr := range r.Addresses {
+			if !((r.IPTables.Proto() == iptables.ProtocolIPv4 && len(addr.IP) == net.IPv4len) || (r.IPTables.Proto() == iptables.ProtocolIPv6 && len(addr.IP) == net.IPv6len)) {
+				continue
+			}
+
+			if r.IPTables.Proto() == iptables.ProtocolIPv4 {
+				for _, iptablesArgs := range [][]string{
+					{"-p", "tcp", "-d", addr.IP.String(), "--sport", strconv.Itoa(int(r.To)), "-j", "SNAT", "--to-source", addr.IP.String()},
+					{"-p", "udp", "-d", addr.IP.String(), "--sport", strconv.Itoa(int(r.To)), "-j", "SNAT", "--to-source", addr.IP.String()},
+				} {
+					err = r.IPTables.AppendUnique("nat", postroutingChain, iptablesArgs...)
+					if err != nil {
+						return fmt.Errorf("failed to append rule: %w", err)
+					}
+				}
+			}
+		}
+
+		err = r.IPTables.InsertUnique("nat", "POSTROUTING", 1, "-j", postroutingChain)
 		if err != nil {
 			return fmt.Errorf("failed to linking chain: %w", err)
 		}
@@ -57,7 +102,14 @@ func (r *PortRemap) insertIPTablesRules(table string) error {
 func (r *PortRemap) deleteIPTablesRules() []error {
 	var errs []error
 
-	err := r.IPTables.DeleteIfExists("nat", "PREROUTING", "-j", r.ChainName)
+	preroutingChain := r.ChainName + "_PRR"
+	err := r.IPTables.DeleteIfExists("nat", "PREROUTING", "-j", preroutingChain)
+	if err != nil {
+		errs = append(errs, fmt.Errorf("failed to unlinking chain: %w", err))
+	}
+
+	postroutingChain := r.ChainName + "_POR"
+	err = r.IPTables.DeleteIfExists("nat", "POSTROUTING", "-j", postroutingChain)
 	if err != nil {
 		errs = append(errs, fmt.Errorf("failed to unlinking chain: %w", err))
 	}