Merge branch 'fosrl:main' into optimize-reverse-nat-lookup

2026-03-06 18:56:41 +00:00 · 2026-02-22 12:09:01 +00:00
parent 1cf75b00ff b7af49d759
commit 5977667291
41 changed files with 2263 additions and 411 deletions
--- a/netstack2/proxy.go
+++ b/netstack2/proxy.go
@@ -22,10 +22,12 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip/transport/udp"
 )

-// PortRange represents an allowed range of ports (inclusive)
+// PortRange represents an allowed range of ports (inclusive) with optional protocol filtering
+// Protocol can be "tcp", "udp", or "" (empty string means both protocols)
 type PortRange struct {
-	Min uint16
-	Max uint16
+	Min      uint16
+	Max      uint16
+	Protocol string // "tcp", "udp", or "" for both
 }

 // SubnetRule represents a subnet with optional port restrictions and source address
@@ -41,6 +43,7 @@ type PortRange struct {
 type SubnetRule struct {
 	SourcePrefix netip.Prefix // Source IP prefix (who is sending)
 	DestPrefix   netip.Prefix // Destination IP prefix (where it's going)
+	DisableIcmp  bool         // If true, ICMP traffic is blocked for this subnet
 	RewriteTo    string       // Optional rewrite address for DNAT - can be IP/CIDR or domain name
 	PortRanges   []PortRange  // empty slice means all ports allowed
 }
@@ -67,7 +70,7 @@ func NewSubnetLookup() *SubnetLookup {
 // AddSubnet adds a subnet rule with source and destination prefixes and optional port restrictions
 // If portRanges is nil or empty, all ports are allowed for this subnet
 // rewriteTo can be either an IP/CIDR (e.g., "192.168.1.1/32") or a domain name (e.g., "example.com")
-func (sl *SubnetLookup) AddSubnet(sourcePrefix, destPrefix netip.Prefix, rewriteTo string, portRanges []PortRange) {
+func (sl *SubnetLookup) AddSubnet(sourcePrefix, destPrefix netip.Prefix, rewriteTo string, portRanges []PortRange, disableIcmp bool) {
 	sl.mu.Lock()
 	defer sl.mu.Unlock()

@@ -79,6 +82,7 @@ func (sl *SubnetLookup) AddSubnet(sourcePrefix, destPrefix netip.Prefix, rewrite
 	sl.rules[key] = &SubnetRule{
 		SourcePrefix: sourcePrefix,
 		DestPrefix:   destPrefix,
+		DisableIcmp:  disableIcmp,
 		RewriteTo:    rewriteTo,
 		PortRanges:   portRanges,
 	}
@@ -97,14 +101,16 @@ func (sl *SubnetLookup) RemoveSubnet(sourcePrefix, destPrefix netip.Prefix) {
 	delete(sl.rules, key)
 }

-// Match checks if a source IP, destination IP, and port match any subnet rule
-// Returns the matched rule if BOTH:
+// Match checks if a source IP, destination IP, port, and protocol match any subnet rule
+// Returns the matched rule if ALL of these conditions are met:
 //   - The source IP is in the rule's source prefix
 //   - The destination IP is in the rule's destination prefix
 //   - The port is in an allowed range (or no port restrictions exist)
+//   - The protocol matches (or the port range allows both protocols)
 //
+// proto should be header.TCPProtocolNumber or header.UDPProtocolNumber
 // Returns nil if no rule matches
-func (sl *SubnetLookup) Match(srcIP, dstIP netip.Addr, port uint16) *SubnetRule {
+func (sl *SubnetLookup) Match(srcIP, dstIP netip.Addr, port uint16, proto tcpip.TransportProtocolNumber) *SubnetRule {
 	sl.mu.RLock()
 	defer sl.mu.RUnlock()

@@ -119,16 +125,31 @@ func (sl *SubnetLookup) Match(srcIP, dstIP netip.Addr, port uint16) *SubnetRule
 			continue
 		}

+		if rule.DisableIcmp && (proto == header.ICMPv4ProtocolNumber || proto == header.ICMPv6ProtocolNumber) {
+		    // ICMP is disabled for this subnet
+			return nil
+		}
+
 		// Both IPs match - now check port restrictions
 		// If no port ranges specified, all ports are allowed
 		if len(rule.PortRanges) == 0 {
 			return rule
 		}

-		// Check if port is in any of the allowed ranges
+		// Check if port and protocol are in any of the allowed ranges
 		for _, pr := range rule.PortRanges {
 			if port >= pr.Min && port <= pr.Max {
-				return rule
+				// Check protocol compatibility
+				if pr.Protocol == "" {
+					// Empty protocol means allow both TCP and UDP
+					return rule
+				}
+				// Check if the packet protocol matches the port range protocol
+				if (pr.Protocol == "tcp" && proto == header.TCPProtocolNumber) ||
+					(pr.Protocol == "udp" && proto == header.UDPProtocolNumber) {
+					return rule
+				}
+				// Port matches but protocol doesn't - continue checking other ranges
 			}
 		}
 	}
@@ -177,24 +198,28 @@ type ProxyHandler struct {
 	proxyNotifyHandle *channel.NotificationHandle
 	tcpHandler        *TCPHandler
 	udpHandler        *UDPHandler
+	icmpHandler       *ICMPHandler
 	subnetLookup      *SubnetLookup
 	natTable          map[connKey]*natState
 	reverseNatTable   map[reverseConnKey]*natState // Reverse lookup map for O(1) reply packet NAT
 	destRewriteTable  map[destKey]netip.Addr       // Maps original dest to rewritten dest for handler lookups
 	natMu             sync.RWMutex
 	enabled           bool
+	icmpReplies       chan []byte          // Channel for ICMP reply packets to be sent back through the tunnel
+	notifiable        channel.Notification // Notification handler for triggering reads
 }

 // ProxyHandlerOptions configures the proxy handler
 type ProxyHandlerOptions struct {
-	EnableTCP bool
-	EnableUDP bool
-	MTU       int
+	EnableTCP  bool
+	EnableUDP  bool
+	EnableICMP bool
+	MTU        int
 }

 // NewProxyHandler creates a new proxy handler for promiscuous mode
 func NewProxyHandler(options ProxyHandlerOptions) (*ProxyHandler, error) {
-	if !options.EnableTCP && !options.EnableUDP {
+	if !options.EnableTCP && !options.EnableUDP && !options.EnableICMP {
 		return nil, nil // No proxy needed
 	}

@@ -204,6 +229,7 @@ func NewProxyHandler(options ProxyHandlerOptions) (*ProxyHandler, error) {
 		natTable:         make(map[connKey]*natState),
 		reverseNatTable:  make(map[reverseConnKey]*natState),
 		destRewriteTable: make(map[destKey]netip.Addr),
+		icmpReplies:      make(chan []byte, 256), // Buffer for ICMP reply packets
 		proxyEp:          channel.New(1024, uint32(options.MTU), ""),
 		proxyStack: stack.New(stack.Options{
 			NetworkProtocols: []stack.NetworkProtocolFactory{
@@ -235,6 +261,15 @@ func NewProxyHandler(options ProxyHandlerOptions) (*ProxyHandler, error) {
 		}
 	}

+	// Initialize ICMP handler if enabled
+	if options.EnableICMP {
+		handler.icmpHandler = NewICMPHandler(handler.proxyStack, handler)
+		if err := handler.icmpHandler.InstallICMPHandler(); err != nil {
+			return nil, fmt.Errorf("failed to install ICMP handler: %v", err)
+		}
+		logger.Debug("ProxyHandler: ICMP handler enabled")
+	}
+
 	// // Example 1: Add a rule with no port restrictions (all ports allowed)
 	// // This accepts all traffic FROM 10.0.0.0/24 TO 10.20.20.0/24
 	// sourceSubnet := netip.MustParsePrefix("10.0.0.0/24")
@@ -259,11 +294,11 @@ func NewProxyHandler(options ProxyHandlerOptions) (*ProxyHandler, error) {
 // destPrefix: The IP prefix of the destination
 // rewriteTo: Optional address to rewrite destination to - can be IP/CIDR or domain name
 // If portRanges is nil or empty, all ports are allowed for this subnet
-func (p *ProxyHandler) AddSubnetRule(sourcePrefix, destPrefix netip.Prefix, rewriteTo string, portRanges []PortRange) {
+func (p *ProxyHandler) AddSubnetRule(sourcePrefix, destPrefix netip.Prefix, rewriteTo string, portRanges []PortRange, disableIcmp bool) {
 	if p == nil || !p.enabled {
 		return
 	}
-	p.subnetLookup.AddSubnet(sourcePrefix, destPrefix, rewriteTo, portRanges)
+	p.subnetLookup.AddSubnet(sourcePrefix, destPrefix, rewriteTo, portRanges, disableIcmp)
 }

 // RemoveSubnetRule removes a subnet from the proxy handler
@@ -342,6 +377,9 @@ func (p *ProxyHandler) Initialize(notifiable channel.Notification) error {
 		return nil
 	}

+	// Store notifiable for triggering notifications on ICMP replies
+	p.notifiable = notifiable
+
 	// Add notification handler
 	p.proxyNotifyHandle = p.proxyEp.AddNotify(notifiable)

@@ -420,14 +458,21 @@ func (p *ProxyHandler) HandleIncomingPacket(packet []byte) bool {
 		}
 		udpHeader := header.UDP(packet[headerLen:])
 		dstPort = udpHeader.DestinationPort()
-	default:
-		// For other protocols (ICMP, etc.), use port 0 (must match rules with no port restrictions)
+	case header.ICMPv4ProtocolNumber:
+		// ICMP doesn't have ports, use port 0 (must match rules with no port restrictions)
 		dstPort = 0
+		logger.Debug("HandleIncomingPacket: ICMP packet from %s to %s", srcAddr, dstAddr)
+	default:
+		// For other protocols, use port 0 (must match rules with no port restrictions)
+		dstPort = 0
+		logger.Debug("HandleIncomingPacket: Unknown protocol %d from %s to %s", protocol, srcAddr, dstAddr)
 	}

-	// Check if the source IP, destination IP, and port match any subnet rule
-	matchedRule := p.subnetLookup.Match(srcAddr, dstAddr, dstPort)
+	// Check if the source IP, destination IP, port, and protocol match any subnet rule
+	matchedRule := p.subnetLookup.Match(srcAddr, dstAddr, dstPort, protocol)
 	if matchedRule != nil {
+		logger.Debug("HandleIncomingPacket: Matched rule for %s -> %s (proto=%d, port=%d)",
+			srcAddr, dstAddr, protocol, dstPort)
 		// Check if we need to perform DNAT
 		if matchedRule.RewriteTo != "" {
 			// Create connection tracking key using original destination
@@ -527,9 +572,12 @@ func (p *ProxyHandler) HandleIncomingPacket(packet []byte) bool {
 			Payload: buffer.MakeWithData(packet),
 		})
 		p.proxyEp.InjectInbound(header.IPv4ProtocolNumber, pkb)
+		logger.Debug("HandleIncomingPacket: Injected packet into proxy stack (proto=%d)", protocol)
 		return true
 	}

+	// logger.Debug("HandleIncomingPacket: No matching rule for %s -> %s (proto=%d, port=%d)",
+		// srcAddr, dstAddr, protocol, dstPort)
 	return false
 }

@@ -652,6 +700,15 @@ func (p *ProxyHandler) ReadOutgoingPacket() *buffer.View {
 		return nil
 	}

+	// First check for ICMP reply packets (non-blocking)
+	select {
+	case icmpReply := <-p.icmpReplies:
+		logger.Debug("ReadOutgoingPacket: Returning ICMP reply packet (%d bytes)", len(icmpReply))
+		return buffer.NewViewWithData(icmpReply)
+	default:
+		// No ICMP reply available, continue to check proxy endpoint
+	}
+
 	pkt := p.proxyEp.Read()
 	if pkt != nil {
 		view := pkt.ToView()
@@ -681,6 +738,11 @@ func (p *ProxyHandler) ReadOutgoingPacket() *buffer.View {
 					srcPort = udpHeader.SourcePort()
 					dstPort = udpHeader.DestinationPort()
 				}
+			case header.ICMPv4ProtocolNumber:
+				// ICMP packets don't need NAT translation in our implementation
+				// since we construct reply packets with the correct addresses
+				logger.Debug("ReadOutgoingPacket: ICMP packet from %s to %s", srcIP, dstIP)
+				return view
 			}

 			// Look up NAT state for reverse translation using O(1) reverse lookup map
@@ -716,12 +778,37 @@ func (p *ProxyHandler) ReadOutgoingPacket() *buffer.View {
 	return nil
 }

+// QueueICMPReply queues an ICMP reply packet to be sent back through the tunnel
+func (p *ProxyHandler) QueueICMPReply(packet []byte) bool {
+	if p == nil || !p.enabled {
+		return false
+	}
+
+	select {
+	case p.icmpReplies <- packet:
+		logger.Debug("QueueICMPReply: Queued ICMP reply packet (%d bytes)", len(packet))
+		// Trigger notification so WriteNotify picks up the packet
+		if p.notifiable != nil {
+			p.notifiable.WriteNotify()
+		}
+		return true
+	default:
+		logger.Info("QueueICMPReply: ICMP reply channel full, dropping packet")
+		return false
+	}
+}
+
 // Close cleans up the proxy handler resources
 func (p *ProxyHandler) Close() error {
 	if p == nil || !p.enabled {
 		return nil
 	}

+	// Close ICMP replies channel
+	if p.icmpReplies != nil {
+		close(p.icmpReplies)
+	}
+
 	if p.proxyStack != nil {
 		p.proxyStack.RemoveNIC(1)
 		p.proxyStack.Close()