Optimize lock usage in proxy connection handling

- Change activeTunnel.conns from slice to map for O(1) add/remove
- Improve lock scoping in UpdateLocalSNIs: use read lock for diff
  computation, minimize write lock hold time
- Move cache invalidation outside lock (go-cache is thread-safe)

main.go

@@ -355,8 +355,21 @@ func main() {
 		return periodicBandwidthCheck(groupCtx, remoteConfigURL+"/gerbil/receive-bandwidth")
 	})
 
-	// Parse local overrides and trusted upstreams early so that both the relay
+	// Start the UDP proxy server
-	// and the SNI proxy share the same configuration values.
+	relayPort := wgconfig.RelayPort
+	if relayPort == 0 {
+		relayPort = 21820 // in case there is no relay port set, use 21820
+	}
+	proxyRelay = relay.NewUDPProxyServer(groupCtx, fmt.Sprintf(":%d", relayPort), remoteConfigURL, key, reachableAt)
+	err = proxyRelay.Start()
+	if err != nil {
+		logger.Fatal("Failed to start UDP proxy server: %v", err)
+	}
+	defer proxyRelay.Stop()
+
+	// TODO: WE SHOULD PULL THIS OUT OF THE CONFIG OR SOMETHING
+	// 		 SO YOU DON'T NEED TO SET THIS SEPARATELY
+	// Parse local overrides
 	var localOverrides []string
 	if localOverridesStr != "" {
 		localOverrides = strings.Split(localOverridesStr, ",")
@@ -375,21 +388,6 @@ func main() {
 		logger.Info("Trusted upstreams configured: %v", trustedUpstreams)
 	}
 
-	// Start the UDP proxy server.
-	// proxyProtocol and trustedUpstreams are forwarded so the relay can strip
-	// PROXY protocol v2 headers from load-balancer traffic and recover the
-	// original client IP for hole-punch registration.
-	relayPort := wgconfig.RelayPort
-	if relayPort == 0 {
-		relayPort = 21820 // in case there is no relay port set, use 21820
-	}
-	proxyRelay = relay.NewUDPProxyServer(groupCtx, fmt.Sprintf(":%d", relayPort), remoteConfigURL, key, reachableAt, proxyProtocol, trustedUpstreams)
-	err = proxyRelay.Start()
-	if err != nil {
-		logger.Fatal("Failed to start UDP proxy server: %v", err)
-	}
-	defer proxyRelay.Stop()
-
 	proxySNI, err = proxy.NewSNIProxy(sniProxyPort, remoteConfigURL, key.PublicKey().String(), localProxyAddr, localProxyPort, localOverrides, proxyProtocol, trustedUpstreams)
 	if err != nil {
 		logger.Fatal("Failed to create proxy: %v", err)

proxy/proxy.go

@@ -11,12 +11,12 @@ import (
 	"log"
 	"net"
 	"net/http"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
 
 	"github.com/fosrl/gerbil/logger"
-	"github.com/fosrl/gerbil/proxyproto"
 	"github.com/patrickmn/go-cache"
 )
 
@@ -32,6 +32,16 @@ type RouteAPIResponse struct {
 	Endpoints []string `json:"endpoints"`
 }
 
+// ProxyProtocolInfo holds information parsed from incoming PROXY protocol header
+type ProxyProtocolInfo struct {
+	Protocol     string // TCP4 or TCP6
+	SrcIP        string
+	DestIP       string
+	SrcPort      int
+	DestPort     int
+	OriginalConn net.Conn // The original connection after PROXY protocol parsing
+}
+
 // SNIProxy represents the main proxy server
 type SNIProxy struct {
 	port            int
@@ -62,7 +72,7 @@ type SNIProxy struct {
 }
 
 type activeTunnel struct {
-	conns []net.Conn
+	conns map[net.Conn]struct{}
 }
 
 // readOnlyConn is a wrapper for io.Reader that implements net.Conn
@@ -79,6 +89,249 @@ func (conn readOnlyConn) SetDeadline(t time.Time) error      { return nil }
 func (conn readOnlyConn) SetReadDeadline(t time.Time) error  { return nil }
 func (conn readOnlyConn) SetWriteDeadline(t time.Time) error { return nil }
 
+// parseProxyProtocolHeader parses a PROXY protocol v1 header from the connection
+func (p *SNIProxy) parseProxyProtocolHeader(conn net.Conn) (*ProxyProtocolInfo, net.Conn, error) {
+	// Check if the connection comes from a trusted upstream
+	remoteHost, _, err := net.SplitHostPort(conn.RemoteAddr().String())
+	if err != nil {
+		return nil, conn, fmt.Errorf("failed to parse remote address: %w", err)
+	}
+
+	// Resolve the remote IP to hostname to check if it's trusted
+	// For simplicity, we'll check the IP directly in trusted upstreams
+	// In production, you might want to do reverse DNS lookup
+	if _, isTrusted := p.trustedUpstreams[remoteHost]; !isTrusted {
+		// Not from trusted upstream, return original connection
+		return nil, conn, nil
+	}
+
+	// Set read timeout for PROXY protocol parsing
+	if err := conn.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil {
+		return nil, conn, fmt.Errorf("failed to set read deadline: %w", err)
+	}
+
+	// Read the first line (PROXY protocol header)
+	buffer := make([]byte, 512) // PROXY protocol header should be much smaller
+	n, err := conn.Read(buffer)
+	if err != nil {
+		// If we can't read from trusted upstream, treat as regular connection
+		logger.Debug("Could not read from trusted upstream %s, treating as regular connection: %v", remoteHost, err)
+		// Clear read timeout before returning
+		if clearErr := conn.SetReadDeadline(time.Time{}); clearErr != nil {
+			logger.Debug("Failed to clear read deadline: %v", clearErr)
+		}
+		return nil, conn, nil
+	}
+
+	// Find the end of the first line (CRLF)
+	headerEnd := bytes.Index(buffer[:n], []byte("\r\n"))
+	if headerEnd == -1 {
+		// No PROXY protocol header found, treat as regular TLS connection
+		// Return the connection with the buffered data prepended
+		logger.Debug("No PROXY protocol header from trusted upstream %s, treating as regular TLS connection", remoteHost)
+
+		// Clear read timeout
+		if err := conn.SetReadDeadline(time.Time{}); err != nil {
+			logger.Debug("Failed to clear read deadline: %v", err)
+		}
+
+		// Create a reader that includes the buffered data + original connection
+		newReader := io.MultiReader(bytes.NewReader(buffer[:n]), conn)
+		wrappedConn := &proxyProtocolConn{
+			Conn:   conn,
+			reader: newReader,
+		}
+		return nil, wrappedConn, nil
+	}
+
+	headerLine := string(buffer[:headerEnd])
+	remainingData := buffer[headerEnd+2 : n]
+
+	// Parse PROXY protocol line: "PROXY TCP4/TCP6 srcIP destIP srcPort destPort"
+	parts := strings.Fields(headerLine)
+	if len(parts) != 6 || parts[0] != "PROXY" {
+		// Check for PROXY UNKNOWN
+		if len(parts) == 2 && parts[0] == "PROXY" && parts[1] == "UNKNOWN" {
+			// PROXY UNKNOWN - use original connection info
+			return nil, conn, nil
+		}
+		// Invalid PROXY protocol, but might be regular TLS - treat as such
+		logger.Debug("Invalid PROXY protocol from trusted upstream %s, treating as regular TLS connection: %s", remoteHost, headerLine)
+
+		// Clear read timeout
+		if err := conn.SetReadDeadline(time.Time{}); err != nil {
+			logger.Debug("Failed to clear read deadline: %v", err)
+		}
+
+		// Return the connection with all buffered data prepended
+		newReader := io.MultiReader(bytes.NewReader(buffer[:n]), conn)
+		wrappedConn := &proxyProtocolConn{
+			Conn:   conn,
+			reader: newReader,
+		}
+		return nil, wrappedConn, nil
+	}
+
+	protocol := parts[1]
+	srcIP := parts[2]
+	destIP := parts[3]
+	srcPort, err := strconv.Atoi(parts[4])
+	if err != nil {
+		return nil, conn, fmt.Errorf("invalid source port in PROXY header: %s", parts[4])
+	}
+	destPort, err := strconv.Atoi(parts[5])
+	if err != nil {
+		return nil, conn, fmt.Errorf("invalid destination port in PROXY header: %s", parts[5])
+	}
+
+	// Create a new reader that includes remaining data + original connection
+	var newReader io.Reader
+	if len(remainingData) > 0 {
+		newReader = io.MultiReader(bytes.NewReader(remainingData), conn)
+	} else {
+		newReader = conn
+	}
+
+	// Create a wrapper connection that reads from the combined reader
+	wrappedConn := &proxyProtocolConn{
+		Conn:   conn,
+		reader: newReader,
+	}
+
+	proxyInfo := &ProxyProtocolInfo{
+		Protocol:     protocol,
+		SrcIP:        srcIP,
+		DestIP:       destIP,
+		SrcPort:      srcPort,
+		DestPort:     destPort,
+		OriginalConn: wrappedConn,
+	}
+
+	// Clear read timeout
+	if err := conn.SetReadDeadline(time.Time{}); err != nil {
+		return nil, conn, fmt.Errorf("failed to clear read deadline: %w", err)
+	}
+
+	return proxyInfo, wrappedConn, nil
+}
+
+// proxyProtocolConn wraps a connection to read from a custom reader
+type proxyProtocolConn struct {
+	net.Conn
+	reader io.Reader
+}
+
+func (c *proxyProtocolConn) Read(b []byte) (int, error) {
+	return c.reader.Read(b)
+}
+
+// buildProxyProtocolHeaderFromInfo creates a PROXY protocol v1 header using ProxyProtocolInfo
+func (p *SNIProxy) buildProxyProtocolHeaderFromInfo(proxyInfo *ProxyProtocolInfo, targetAddr net.Addr) string {
+	targetTCP, ok := targetAddr.(*net.TCPAddr)
+	if !ok {
+		// Fallback for unknown address types
+		return "PROXY UNKNOWN\r\n"
+	}
+
+	// Use the original client information from the PROXY protocol
+	var targetIP string
+	var protocol string
+
+	// Parse source IP to determine protocol family
+	srcIP := net.ParseIP(proxyInfo.SrcIP)
+	if srcIP == nil {
+		return "PROXY UNKNOWN\r\n"
+	}
+
+	if srcIP.To4() != nil {
+		// Source is IPv4, use TCP4 protocol
+		protocol = "TCP4"
+		if targetTCP.IP.To4() != nil {
+			// Target is also IPv4, use as-is
+			targetIP = targetTCP.IP.String()
+		} else {
+			// Target is IPv6, but we need IPv4 for consistent protocol family
+			if targetTCP.IP.IsLoopback() {
+				targetIP = "127.0.0.1"
+			} else {
+				targetIP = "127.0.0.1" // Safe fallback
+			}
+		}
+	} else {
+		// Source is IPv6, use TCP6 protocol
+		protocol = "TCP6"
+		if targetTCP.IP.To4() != nil {
+			// Target is IPv4, convert to IPv6 representation
+			targetIP = "::ffff:" + targetTCP.IP.String()
+		} else {
+			// Target is also IPv6, use as-is
+			targetIP = targetTCP.IP.String()
+		}
+	}
+
+	return fmt.Sprintf("PROXY %s %s %s %d %d\r\n",
+		protocol,
+		proxyInfo.SrcIP,
+		targetIP,
+		proxyInfo.SrcPort,
+		targetTCP.Port)
+}
+
+// buildProxyProtocolHeader creates a PROXY protocol v1 header
+func buildProxyProtocolHeader(clientAddr, targetAddr net.Addr) string {
+	clientTCP, ok := clientAddr.(*net.TCPAddr)
+	if !ok {
+		// Fallback for unknown address types
+		return "PROXY UNKNOWN\r\n"
+	}
+
+	targetTCP, ok := targetAddr.(*net.TCPAddr)
+	if !ok {
+		// Fallback for unknown address types
+		return "PROXY UNKNOWN\r\n"
+	}
+
+	// Determine protocol family based on client IP and normalize target IP accordingly
+	var protocol string
+	var targetIP string
+
+	if clientTCP.IP.To4() != nil {
+		// Client is IPv4, use TCP4 protocol
+		protocol = "TCP4"
+		if targetTCP.IP.To4() != nil {
+			// Target is also IPv4, use as-is
+			targetIP = targetTCP.IP.String()
+		} else {
+			// Target is IPv6, but we need IPv4 for consistent protocol family
+			// Use the IPv4 loopback if target is IPv6 loopback, otherwise use 127.0.0.1
+			if targetTCP.IP.IsLoopback() {
+				targetIP = "127.0.0.1"
+			} else {
+				// For non-loopback IPv6 targets, we could try to extract embedded IPv4
+				// or fall back to a sensible IPv4 address based on the target
+				targetIP = "127.0.0.1" // Safe fallback
+			}
+		}
+	} else {
+		// Client is IPv6, use TCP6 protocol
+		protocol = "TCP6"
+		if targetTCP.IP.To4() != nil {
+			// Target is IPv4, convert to IPv6 representation
+			targetIP = "::ffff:" + targetTCP.IP.String()
+		} else {
+			// Target is also IPv6, use as-is
+			targetIP = targetTCP.IP.String()
+		}
+	}
+
+	return fmt.Sprintf("PROXY %s %s %s %d %d\r\n",
+		protocol,
+		clientTCP.IP.String(),
+		targetIP,
+		clientTCP.Port,
+		targetTCP.Port)
+}
+
 // NewSNIProxy creates a new SNI proxy instance
 func NewSNIProxy(port int, remoteConfigURL, publicKey, localProxyAddr string, localProxyPort int, localOverrides []string, proxyProtocol bool, trustedUpstreams []string) (*SNIProxy, error) {
 	ctx, cancel := context.WithCancel(context.Background())
@@ -237,12 +490,12 @@ func (p *SNIProxy) handleConnection(clientConn net.Conn) {
 	logger.Debug("Accepted connection from %s", clientConn.RemoteAddr())
 
 	// Check for PROXY protocol from trusted upstream
-	var proxyInfo *proxyproto.Info
+	var proxyInfo *ProxyProtocolInfo
 	var actualClientConn net.Conn = clientConn
 
 	if len(p.trustedUpstreams) > 0 {
 		var err error
-		proxyInfo, actualClientConn, err = proxyproto.ParseV1Header(clientConn, p.trustedUpstreams)
+		proxyInfo, actualClientConn, err = p.parseProxyProtocolHeader(clientConn)
 		if err != nil {
 			logger.Debug("Failed to parse PROXY protocol: %v", err)
 			return
@@ -322,10 +575,10 @@ func (p *SNIProxy) handleConnection(clientConn net.Conn) {
 		var proxyHeader string
 		if proxyInfo != nil {
 			// Use original client info from PROXY protocol
-			proxyHeader = proxyproto.BuildV1HeaderFromInfo(proxyInfo, targetConn.LocalAddr())
+			proxyHeader = p.buildProxyProtocolHeaderFromInfo(proxyInfo, targetConn.LocalAddr())
 		} else {
 			// Use direct client connection info
-			proxyHeader = proxyproto.BuildV1Header(clientConn.RemoteAddr(), targetConn.LocalAddr())
+			proxyHeader = buildProxyProtocolHeader(clientConn.RemoteAddr(), targetConn.LocalAddr())
 		}
 		logger.Debug("Sending PROXY protocol header: %s", strings.TrimSpace(proxyHeader))
 
@@ -339,26 +592,19 @@ func (p *SNIProxy) handleConnection(clientConn net.Conn) {
 	p.activeTunnelsLock.Lock()
 	tunnel, ok := p.activeTunnels[hostname]
 	if !ok {
-		tunnel = &activeTunnel{}
+		tunnel = &activeTunnel{conns: make(map[net.Conn]struct{})}
 		p.activeTunnels[hostname] = tunnel
 	}
-	tunnel.conns = append(tunnel.conns, actualClientConn)
+	tunnel.conns[actualClientConn] = struct{}{}
 	p.activeTunnelsLock.Unlock()
 
 	defer func() {
-		// Remove this conn from active tunnels
+		// Remove this conn from active tunnels - O(1) with map
 		p.activeTunnelsLock.Lock()
 		if tunnel, ok := p.activeTunnels[hostname]; ok {
-			newConns := make([]net.Conn, 0, len(tunnel.conns))
+			delete(tunnel.conns, actualClientConn)
-			for _, c := range tunnel.conns {
+			if len(tunnel.conns) == 0 {
-				if c != actualClientConn {
-					newConns = append(newConns, c)
-				}
-			}
-			if len(newConns) == 0 {
 				delete(p.activeTunnels, hostname)
-			} else {
-				tunnel.conns = newConns
 			}
 		}
 		p.activeTunnelsLock.Unlock()
@@ -557,32 +803,42 @@ func (p *SNIProxy) ClearCache() {
 
 // UpdateLocalSNIs updates the local SNIs and invalidates cache for changed domains
 func (p *SNIProxy) UpdateLocalSNIs(fullDomains []string) {
-	newSNIs := make(map[string]struct{})
+	newSNIs := make(map[string]struct{}, len(fullDomains))
 	for _, domain := range fullDomains {
 		newSNIs[domain] = struct{}{}
-		// Invalidate any cached route for this domain
-		p.cache.Delete(domain)
 	}
 
-	// Update localSNIs
+	// Get old SNIs with read lock to compute diff outside write lock
-	p.localSNIsLock.Lock()
+	p.localSNIsLock.RLock()
+	oldSNIs := p.localSNIs
+	p.localSNIsLock.RUnlock()
+
+	// Compute removed SNIs outside the lock
 	removed := make([]string, 0)
-	for sni := range p.localSNIs {
+	for sni := range oldSNIs {
 		if _, stillLocal := newSNIs[sni]; !stillLocal {
 			removed = append(removed, sni)
 		}
 	}
+
+	// Swap with minimal write lock hold time
+	p.localSNIsLock.Lock()
 	p.localSNIs = newSNIs
 	p.localSNIsLock.Unlock()
 
+	// Invalidate cache for new domains (cache is thread-safe)
+	for domain := range newSNIs {
+		p.cache.Delete(domain)
+	}
+
 	logger.Debug("Updated local SNIs, added %d, removed %d", len(newSNIs), len(removed))
 
 	// Terminate tunnels for removed SNIs
 	if len(removed) > 0 {
 		p.activeTunnelsLock.Lock()
 		for _, sni := range removed {
-			if tunnels, ok := p.activeTunnels[sni]; ok {
+			if tunnel, ok := p.activeTunnels[sni]; ok {
-				for _, conn := range tunnels.conns {
+				for conn := range tunnel.conns {
 					conn.Close()
 				}
 				delete(p.activeTunnels, sni)

proxy/proxy_test.go

@@ -3,8 +3,6 @@ package proxy
 import (
 	"net"
 	"testing"
-
-	"github.com/fosrl/gerbil/proxyproto"
 )
 
 func TestBuildProxyProtocolHeader(t *testing.T) {
@@ -58,7 +56,7 @@ func TestBuildProxyProtocolHeader(t *testing.T) {
 				t.Fatalf("Failed to resolve target address: %v", err)
 			}
 
-			result := proxyproto.BuildV1Header(clientTCP, targetTCP)
+			result := buildProxyProtocolHeader(clientTCP, targetTCP)
 			if result != tt.expected {
 				t.Errorf("Expected %q, got %q", tt.expected, result)
 			}
@@ -71,7 +69,7 @@ func TestBuildProxyProtocolHeaderUnknownType(t *testing.T) {
 	clientAddr := &net.UDPAddr{IP: net.ParseIP("192.168.1.100"), Port: 12345}
 	targetAddr := &net.UDPAddr{IP: net.ParseIP("10.0.0.1"), Port: 443}
 
-	result := proxyproto.BuildV1Header(clientAddr, targetAddr)
+	result := buildProxyProtocolHeader(clientAddr, targetAddr)
 	expected := "PROXY UNKNOWN\r\n"
 
 	if result != expected {
@@ -80,8 +78,13 @@ func TestBuildProxyProtocolHeaderUnknownType(t *testing.T) {
 }
 
 func TestBuildProxyProtocolHeaderFromInfo(t *testing.T) {
+	proxy, err := NewSNIProxy(8443, "", "", "127.0.0.1", 443, nil, true, nil)
+	if err != nil {
+		t.Fatalf("Failed to create SNI proxy: %v", err)
+	}
+
 	// Test IPv4 case
-	info := &proxyproto.Info{
+	proxyInfo := &ProxyProtocolInfo{
 		Protocol: "TCP4",
 		SrcIP:    "10.0.0.1",
 		DestIP:   "192.168.1.100",
@@ -90,7 +93,7 @@ func TestBuildProxyProtocolHeaderFromInfo(t *testing.T) {
 	}
 
 	targetAddr, _ := net.ResolveTCPAddr("tcp", "127.0.0.1:8080")
-	header := proxyproto.BuildV1HeaderFromInfo(info, targetAddr)
+	header := proxy.buildProxyProtocolHeaderFromInfo(proxyInfo, targetAddr)
 
 	expected := "PROXY TCP4 10.0.0.1 127.0.0.1 12345 8080\r\n"
 	if header != expected {
@@ -98,7 +101,7 @@ func TestBuildProxyProtocolHeaderFromInfo(t *testing.T) {
 	}
 
 	// Test IPv6 case
-	info = &proxyproto.Info{
+	proxyInfo = &ProxyProtocolInfo{
 		Protocol: "TCP6",
 		SrcIP:    "2001:db8::1",
 		DestIP:   "2001:db8::2",
@@ -107,99 +110,10 @@ func TestBuildProxyProtocolHeaderFromInfo(t *testing.T) {
 	}
 
 	targetAddr, _ = net.ResolveTCPAddr("tcp6", "[::1]:8080")
-	header = proxyproto.BuildV1HeaderFromInfo(info, targetAddr)
+	header = proxy.buildProxyProtocolHeaderFromInfo(proxyInfo, targetAddr)
 
 	expected = "PROXY TCP6 2001:db8::1 ::1 12345 8080\r\n"
 	if header != expected {
 		t.Errorf("Expected header '%s', got '%s'", expected, header)
 	}
 }
-
-func TestParseV2UDPHeader(t *testing.T) {
-	// Build a minimal PROXY v2 header for IPv4 UDP
-	// Magic (12) + ver/cmd (1) + fam/proto (1) + len (2) + src IP (4) + dst IP (4) + src port (2) + dst port (2) = 28 bytes
-	header := []byte{
-		// Magic signature
-		0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A,
-		// Version 2 (0x2x), PROXY command (0x01)
-		0x21,
-		// AF_INET (0x1x), DGRAM/UDP (0x02)
-		0x12,
-		// Address block length: 12 bytes (4+4+2+2)
-		0x00, 0x0C,
-		// Source IP: 192.168.1.100
-		192, 168, 1, 100,
-		// Destination IP: 10.0.0.1
-		10, 0, 0, 1,
-		// Source port: 4500
-		0x11, 0x94,
-		// Destination port: 21820
-		0x55, 0x3C,
-	}
-
-	// Append a fake application payload
-	payload := []byte{0x01, 0x02, 0x03}
-	data := append(header, payload...)
-
-	info, remaining, ok := proxyproto.ParseV2UDPHeader(data)
-	if !ok {
-		t.Fatal("Expected ParseV2UDPHeader to return ok=true")
-	}
-	if info == nil {
-		t.Fatal("Expected non-nil Info")
-	}
-	if info.Protocol != "UDP4" {
-		t.Errorf("Expected protocol UDP4, got %s", info.Protocol)
-	}
-	if info.SrcIP != "192.168.1.100" {
-		t.Errorf("Expected SrcIP 192.168.1.100, got %s", info.SrcIP)
-	}
-	if info.DestIP != "10.0.0.1" {
-		t.Errorf("Expected DestIP 10.0.0.1, got %s", info.DestIP)
-	}
-	if info.SrcPort != 4500 {
-		t.Errorf("Expected SrcPort 4500, got %d", info.SrcPort)
-	}
-	if info.DestPort != 21820 {
-		t.Errorf("Expected DestPort 21820, got %d", info.DestPort)
-	}
-	if len(remaining) != len(payload) {
-		t.Errorf("Expected %d remaining bytes, got %d", len(payload), len(remaining))
-	}
-}
-
-func TestParseV2UDPHeaderNoHeader(t *testing.T) {
-	// Data that does NOT start with v2 magic should be returned as-is
-	data := []byte{0x01, 0x02, 0x03}
-	info, remaining, ok := proxyproto.ParseV2UDPHeader(data)
-	if ok {
-		t.Error("Expected ok=false for non-v2 data")
-	}
-	if info != nil {
-		t.Error("Expected nil Info for non-v2 data")
-	}
-	if len(remaining) != len(data) {
-		t.Errorf("Expected remaining to equal original data length %d, got %d", len(data), len(remaining))
-	}
-}
-
-func TestIsV2Header(t *testing.T) {
-	valid := []byte{
-		0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A,
-		// extra bytes beyond the magic
-		0x21, 0x12,
-	}
-	if !proxyproto.IsV2Header(valid) {
-		t.Error("Expected IsV2Header=true for valid magic")
-	}
-
-	invalid := []byte{0x01, 0x02, 0x03}
-	if proxyproto.IsV2Header(invalid) {
-		t.Error("Expected IsV2Header=false for non-magic data")
-	}
-
-	tooShort := []byte{0x0D, 0x0A}
-	if proxyproto.IsV2Header(tooShort) {
-		t.Error("Expected IsV2Header=false for too-short data")
-	}
-}

proxyproto/proxyproto.go

@@ -1,370 +0,0 @@
-// Package proxyproto provides shared PROXY protocol v1 (TCP) and v2 (UDP) parsing
-// and header building utilities used by both the SNI proxy and UDP relay components.
-package proxyproto
-
-import (
-	"bytes"
-	"encoding/binary"
-	"fmt"
-	"io"
-	"net"
-	"strconv"
-	"strings"
-	"time"
-
-	"github.com/fosrl/gerbil/logger"
-)
-
-// v2Signature is the 12-byte magic prefix for PROXY protocol v2 headers.
-var v2Signature = []byte{
-	0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A,
-}
-
-// Info holds information parsed from an incoming PROXY protocol header (v1 or v2).
-type Info struct {
-	Protocol string // e.g. "TCP4", "TCP6", "UDP4", "UDP6"
-	SrcIP    string
-	DestIP   string
-	SrcPort  int
-	DestPort int
-}
-
-// Conn wraps a net.Conn so that reads are satisfied from a pre-pended buffered
-// reader first (remaining bytes after PROXY header parsing) and then from the
-// underlying connection. All other net.Conn methods are forwarded unchanged.
-type Conn struct {
-	net.Conn
-	Reader io.Reader
-}
-
-// Read satisfies net.Conn, draining the buffered reader before falling through
-// to the underlying connection.
-func (c *Conn) Read(b []byte) (int, error) {
-	return c.Reader.Read(b)
-}
-
-// IsV2Header returns true when data begins with the 12-byte PROXY protocol v2
-// magic signature.
-func IsV2Header(data []byte) bool {
-	if len(data) < 12 {
-		return false
-	}
-	return bytes.Equal(data[:12], v2Signature)
-}
-
-// ParseV2UDPHeader tries to parse a PROXY protocol v2 header from the front of
-// a UDP datagram payload.
-//
-// Three return values are provided:
-//   - *Info  – filled when a PROXY command header was parsed successfully; nil
-//     for a LOCAL command or unrecognised address family.
-//   - []byte – the remaining payload that follows the header (the actual
-//     application data).
-//   - bool   – true when a v2 header was detected (and consumed), false when
-//     no v2 magic is present and data should be treated as-is.
-func ParseV2UDPHeader(data []byte) (*Info, []byte, bool) {
-	if !IsV2Header(data) {
-		return nil, data, false
-	}
-
-	// Minimum fixed header size: 12 (magic) + 1 (ver/cmd) + 1 (fam/proto) + 2 (len) = 16
-	if len(data) < 16 {
-		return nil, data, false
-	}
-
-	// Byte 12: version (high nibble) + command (low nibble)
-	versionCmd := data[12]
-	version := (versionCmd >> 4) & 0x0F
-	command := versionCmd & 0x0F
-
-	if version != 2 {
-		return nil, data, false
-	}
-
-	// Byte 13: address family (high nibble) + transport protocol (low nibble)
-	familyProto := data[13]
-	family := (familyProto >> 4) & 0x0F
-	protocol := familyProto & 0x0F
-
-	// Bytes 14-15: length of the address block that follows, big-endian
-	addrLen := int(binary.BigEndian.Uint16(data[14:16]))
-	totalHeaderLen := 16 + addrLen
-
-	if len(data) < totalHeaderLen {
-		// Truncated packet – signal that a header was detected but is malformed
-		return nil, data, false
-	}
-
-	payload := data[totalHeaderLen:]
-
-	// LOCAL command (0) carries no address information.
-	if command == 0 {
-		return nil, payload, true
-	}
-
-	if command != 1 {
-		// Unknown command – consume the header and return no info
-		return nil, payload, true
-	}
-
-	addrBlock := data[16:totalHeaderLen]
-
-	var (
-		srcIP, destIP net.IP
-		srcPort       uint16
-		destPort      uint16
-		protocolStr   string
-	)
-
-	switch {
-	case family == 1 && protocol == 1: // AF_INET / STREAM  (TCP over IPv4)
-		if len(addrBlock) < 12 {
-			return nil, payload, false
-		}
-		srcIP = net.IP(addrBlock[0:4])
-		destIP = net.IP(addrBlock[4:8])
-		srcPort = binary.BigEndian.Uint16(addrBlock[8:10])
-		destPort = binary.BigEndian.Uint16(addrBlock[10:12])
-		protocolStr = "TCP4"
-
-	case family == 1 && protocol == 2: // AF_INET / DGRAM   (UDP over IPv4)
-		if len(addrBlock) < 12 {
-			return nil, payload, false
-		}
-		srcIP = net.IP(addrBlock[0:4])
-		destIP = net.IP(addrBlock[4:8])
-		srcPort = binary.BigEndian.Uint16(addrBlock[8:10])
-		destPort = binary.BigEndian.Uint16(addrBlock[10:12])
-		protocolStr = "UDP4"
-
-	case family == 2 && protocol == 1: // AF_INET6 / STREAM (TCP over IPv6)
-		if len(addrBlock) < 36 {
-			return nil, payload, false
-		}
-		srcIP = net.IP(addrBlock[0:16])
-		destIP = net.IP(addrBlock[16:32])
-		srcPort = binary.BigEndian.Uint16(addrBlock[32:34])
-		destPort = binary.BigEndian.Uint16(addrBlock[34:36])
-		protocolStr = "TCP6"
-
-	case family == 2 && protocol == 2: // AF_INET6 / DGRAM  (UDP over IPv6)
-		if len(addrBlock) < 36 {
-			return nil, payload, false
-		}
-		srcIP = net.IP(addrBlock[0:16])
-		destIP = net.IP(addrBlock[16:32])
-		srcPort = binary.BigEndian.Uint16(addrBlock[32:34])
-		destPort = binary.BigEndian.Uint16(addrBlock[34:36])
-		protocolStr = "UDP6"
-
-	default:
-		// UNSPEC or AF_UNIX – consume the header, no address info available
-		return nil, payload, true
-	}
-
-	info := &Info{
-		Protocol: protocolStr,
-		SrcIP:    srcIP.String(),
-		DestIP:   destIP.String(),
-		SrcPort:  int(srcPort),
-		DestPort: int(destPort),
-	}
-	return info, payload, true
-}
-
-// ParseV1Header attempts to parse a PROXY protocol v1 (text) header from the
-// given TCP connection.
-//
-// The function first checks whether the remote address appears in
-// trustedUpstreams. If it does not, it returns (nil, conn, nil) and the caller
-// should treat the connection as a plain (non-proxied) connection.
-//
-// When a trusted upstream is detected the function reads up to 512 bytes,
-// locates the CRLF-terminated header line, and parses the proxy information.
-// Whatever bytes were consumed (including any data beyond the header line) are
-// re-prepended via a *Conn wrapper so that subsequent reads by the caller are
-// transparent.
-//
-// Return values:
-//   - *Info    – non-nil when a valid PROXY header was parsed.
-//   - net.Conn – always a valid connection (possibly a *Conn wrapper).
-//   - error    – non-nil only on hard failures (e.g. bad port numbers).
-func ParseV1Header(conn net.Conn, trustedUpstreams map[string]struct{}) (*Info, net.Conn, error) {
-	remoteHost, _, err := net.SplitHostPort(conn.RemoteAddr().String())
-	if err != nil {
-		return nil, conn, fmt.Errorf("failed to parse remote address: %w", err)
-	}
-
-	if _, isTrusted := trustedUpstreams[remoteHost]; !isTrusted {
-		return nil, conn, nil
-	}
-
-	// Give the upstream 5 s to deliver the PROXY header before timing out.
-	if err := conn.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil {
-		return nil, conn, fmt.Errorf("failed to set read deadline: %w", err)
-	}
-
-	// The PROXY v1 spec mandates the header fits in 108 bytes; 512 is generous.
-	buffer := make([]byte, 512)
-	n, err := conn.Read(buffer)
-	if err != nil {
-		logger.Debug("Could not read from trusted upstream %s, treating as regular connection: %v", remoteHost, err)
-		if clearErr := conn.SetReadDeadline(time.Time{}); clearErr != nil {
-			logger.Debug("Failed to clear read deadline: %v", clearErr)
-		}
-		return nil, conn, nil
-	}
-
-	// Locate the CRLF that terminates the PROXY header line.
-	headerEnd := bytes.Index(buffer[:n], []byte("\r\n"))
-	if headerEnd == -1 {
-		logger.Debug("No PROXY protocol header from trusted upstream %s, treating as regular TLS connection", remoteHost)
-		if clearErr := conn.SetReadDeadline(time.Time{}); clearErr != nil {
-			logger.Debug("Failed to clear read deadline: %v", clearErr)
-		}
-		newReader := io.MultiReader(bytes.NewReader(buffer[:n]), conn)
-		return nil, &Conn{Conn: conn, Reader: newReader}, nil
-	}
-
-	headerLine := string(buffer[:headerEnd])
-	remainingData := buffer[headerEnd+2 : n]
-
-	parts := strings.Fields(headerLine)
-
-	// Handle "PROXY UNKNOWN" – upstream knows the real source but we don't need it.
-	if len(parts) == 2 && parts[0] == "PROXY" && parts[1] == "UNKNOWN" {
-		if clearErr := conn.SetReadDeadline(time.Time{}); clearErr != nil {
-			logger.Debug("Failed to clear read deadline: %v", clearErr)
-		}
-		var newConn net.Conn
-		if len(remainingData) > 0 {
-			newConn = &Conn{Conn: conn, Reader: io.MultiReader(bytes.NewReader(remainingData), conn)}
-		} else {
-			newConn = conn
-		}
-		return nil, newConn, nil
-	}
-
-	if len(parts) != 6 || parts[0] != "PROXY" {
-		// Malformed line from a trusted upstream – re-prepend everything and
-		// let the caller deal with it as a plain TLS connection.
-		logger.Debug("Invalid PROXY protocol from trusted upstream %s, treating as regular TLS connection: %s", remoteHost, headerLine)
-		if clearErr := conn.SetReadDeadline(time.Time{}); clearErr != nil {
-			logger.Debug("Failed to clear read deadline: %v", clearErr)
-		}
-		newReader := io.MultiReader(bytes.NewReader(buffer[:n]), conn)
-		return nil, &Conn{Conn: conn, Reader: newReader}, nil
-	}
-
-	protocol := parts[1]
-	srcIP := parts[2]
-	destIP := parts[3]
-
-	srcPort, err := strconv.Atoi(parts[4])
-	if err != nil {
-		return nil, conn, fmt.Errorf("invalid source port in PROXY header: %s", parts[4])
-	}
-	destPort, err := strconv.Atoi(parts[5])
-	if err != nil {
-		return nil, conn, fmt.Errorf("invalid destination port in PROXY header: %s", parts[5])
-	}
-
-	// Re-assemble a reader that returns any bytes read beyond the header first.
-	var newReader io.Reader
-	if len(remainingData) > 0 {
-		newReader = io.MultiReader(bytes.NewReader(remainingData), conn)
-	} else {
-		newReader = conn
-	}
-	wrappedConn := &Conn{Conn: conn, Reader: newReader}
-
-	if clearErr := conn.SetReadDeadline(time.Time{}); clearErr != nil {
-		return nil, conn, fmt.Errorf("failed to clear read deadline: %w", clearErr)
-	}
-
-	info := &Info{
-		Protocol: protocol,
-		SrcIP:    srcIP,
-		DestIP:   destIP,
-		SrcPort:  srcPort,
-		DestPort: destPort,
-	}
-	return info, wrappedConn, nil
-}
-
-// BuildV1Header constructs a PROXY protocol v1 header string from two TCP
-// addresses, normalising the protocol family so that v1's constraint of a
-// single family per header is satisfied.
-func BuildV1Header(clientAddr, targetAddr net.Addr) string {
-	clientTCP, ok := clientAddr.(*net.TCPAddr)
-	if !ok {
-		return "PROXY UNKNOWN\r\n"
-	}
-	targetTCP, ok := targetAddr.(*net.TCPAddr)
-	if !ok {
-		return "PROXY UNKNOWN\r\n"
-	}
-
-	var protocol, targetIP string
-
-	if clientTCP.IP.To4() != nil {
-		// IPv4 client
-		protocol = "TCP4"
-		if targetTCP.IP.To4() != nil {
-			targetIP = targetTCP.IP.String()
-		} else if targetTCP.IP.IsLoopback() {
-			targetIP = "127.0.0.1"
-		} else {
-			targetIP = "127.0.0.1" // safe fallback for mixed-family
-		}
-	} else {
-		// IPv6 client
-		protocol = "TCP6"
-		if targetTCP.IP.To4() != nil {
-			targetIP = "::ffff:" + targetTCP.IP.String()
-		} else {
-			targetIP = targetTCP.IP.String()
-		}
-	}
-
-	return fmt.Sprintf("PROXY %s %s %s %d %d\r\n",
-		protocol, clientTCP.IP.String(), targetIP, clientTCP.Port, targetTCP.Port)
-}
-
-// BuildV1HeaderFromInfo constructs a PROXY protocol v1 header string using a
-// previously-parsed *Info (i.e. when this server itself sits behind an
-// upstream proxy) and the target TCP address.
-func BuildV1HeaderFromInfo(info *Info, targetAddr net.Addr) string {
-	targetTCP, ok := targetAddr.(*net.TCPAddr)
-	if !ok {
-		return "PROXY UNKNOWN\r\n"
-	}
-
-	srcIP := net.ParseIP(info.SrcIP)
-	if srcIP == nil {
-		return "PROXY UNKNOWN\r\n"
-	}
-
-	var protocol, targetIP string
-
-	if srcIP.To4() != nil {
-		protocol = "TCP4"
-		if targetTCP.IP.To4() != nil {
-			targetIP = targetTCP.IP.String()
-		} else if targetTCP.IP.IsLoopback() {
-			targetIP = "127.0.0.1"
-		} else {
-			targetIP = "127.0.0.1"
-		}
-	} else {
-		protocol = "TCP6"
-		if targetTCP.IP.To4() != nil {
-			targetIP = "::ffff:" + targetTCP.IP.String()
-		} else {
-			targetIP = targetTCP.IP.String()
-		}
-	}
-
-	return fmt.Sprintf("PROXY %s %s %s %d %d\r\n",
-		protocol, info.SrcIP, targetIP, info.SrcPort, targetTCP.Port)
-}

relay/relay.go

@@ -9,13 +9,10 @@ import (
 	"io"
 	"net"
 	"net/http"
-	"runtime"
-	"strings"
 	"sync"
 	"time"
 
 	"github.com/fosrl/gerbil/logger"
-	"github.com/fosrl/gerbil/proxyproto"
 	"golang.org/x/crypto/chacha20poly1305"
 	"golang.org/x/crypto/curve25519"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
@@ -121,13 +118,6 @@ type Packet struct {
 	n          int
 }
 
-// holePunchRateLimitEntry tracks hole punch message counts within a sliding 1-second window.
-type holePunchRateLimitEntry struct {
-	mu          sync.Mutex
-	count       int
-	windowStart time.Time
-}
-
 // WireGuard message types
 const (
 	WireGuardMessageTypeHandshakeInitiation = 1
@@ -163,57 +153,21 @@ type UDPProxyServer struct {
 	// Communication pattern tracking for rebuilding sessions
 	// Key format: "clientIP:clientPort-destIP:destPort"
 	commPatterns sync.Map
-	// Rate limiter for encrypted hole punch messages, keyed by "ip:port"
-	holePunchRateLimiter sync.Map
-	// Cache for resolved UDP addresses to avoid per-packet DNS lookups
-	// Key: "ip:port" string, Value: *net.UDPAddr
-	addrCache sync.Map
 	// ReachableAt is the URL where this server can be reached
 	ReachableAt string
-
-	// proxyProtocol enables PROXY protocol v2 header parsing for incoming UDP packets.
-	// When enabled, packets from trustedUpstreams that carry a v2 header will have
-	// their source address overridden with the address reported in the header.
-	proxyProtocol    bool
-	trustedUpstreams map[string]struct{}
 }
 
 // NewUDPProxyServer initializes the server with a buffered packet channel and derived context.
-//
+func NewUDPProxyServer(parentCtx context.Context, addr, serverURL string, privateKey wgtypes.Key, reachableAt string) *UDPProxyServer {
-// proxyProtocol enables PROXY protocol v2 parsing for datagrams arriving from
-// any address listed in trustedUpstreams (plain IPs or resolvable hostnames).
-// When a trusted datagram carries a v2 header its source address is replaced
-// with the address carried inside the header before further processing, so that
-// hole-punch endpoints reflect the original client IP rather than the load
-// balancer's address.
-func NewUDPProxyServer(parentCtx context.Context, addr, serverURL string, privateKey wgtypes.Key, reachableAt string, proxyProtocol bool, trustedUpstreams []string) *UDPProxyServer {
 	ctx, cancel := context.WithCancel(parentCtx)
-
-	trustedMap := make(map[string]struct{})
-	for _, upstream := range trustedUpstreams {
-		upstream = strings.TrimSpace(upstream)
-		if upstream == "" {
-			continue
-		}
-		trustedMap[upstream] = struct{}{}
-		// Also resolve any hostnames to their current IPs so we can match by IP.
-		if ips, err := net.LookupIP(upstream); err == nil {
-			for _, ip := range ips {
-				trustedMap[ip.String()] = struct{}{}
-			}
-		}
-	}
-
 	return &UDPProxyServer{
-		addr:             addr,
+		addr:        addr,
-		serverURL:        serverURL,
+		serverURL:   serverURL,
-		privateKey:       privateKey,
+		privateKey:  privateKey,
-		packetChan:       make(chan Packet, 50000), // Increased from 1000 to handle high throughput
+		packetChan:  make(chan Packet, 1000),
-		ReachableAt:      reachableAt,
+		ReachableAt: reachableAt,
-		ctx:              ctx,
+		ctx:         ctx,
-		cancel:           cancel,
+		cancel:      cancel,
-		proxyProtocol:    proxyProtocol,
-		trustedUpstreams: trustedMap,
 	}
 }
 
@@ -235,13 +189,8 @@ func (s *UDPProxyServer) Start() error {
 	s.conn = conn
 	logger.Info("UDP server listening on %s", s.addr)
 
-	// Start worker goroutines based on CPU cores for better parallelism
+	// Start a fixed number of worker goroutines.
-	// At high throughput (160+ Mbps), we need many workers to avoid bottlenecks
+	workerCount := 10 // TODO: Make this configurable or pick it better!
-	workerCount := runtime.NumCPU() * 10
-	if workerCount < 20 {
-		workerCount = 20 // Minimum 20 workers
-	}
-	logger.Info("Starting %d packet workers (CPUs: %d)", workerCount, runtime.NumCPU())
 	for i := 0; i < workerCount; i++ {
 		go s.packetWorker()
 	}
@@ -261,9 +210,6 @@ func (s *UDPProxyServer) Start() error {
 	// Start the communication pattern cleanup routine
 	go s.cleanupIdleCommunicationPatterns()
 
-	// Start the hole punch rate limiter cleanup routine
-	go s.cleanupHolePunchRateLimiter()
-
 	return nil
 }
 
@@ -321,70 +267,14 @@ func (s *UDPProxyServer) readPackets() {
 // packetWorker processes incoming packets from the channel.
 func (s *UDPProxyServer) packetWorker() {
 	for packet := range s.packetChan {
-		// effectiveData and effectiveAddr represent the application-layer payload
+		// Determine packet type by inspecting the first byte.
-		// and the true originating address.  They start as the raw UDP values and
+		if packet.n > 0 && packet.data[0] >= 1 && packet.data[0] <= 4 {
-		// may be updated below when a PROXY protocol v2 header is present.
-		effectiveData := packet.data[:packet.n]
-		effectiveAddr := packet.remoteAddr
-
-		// ---------- PROXY protocol v2 (UDP) ------------------------------------
-		// If proxy protocol is enabled and this datagram arrives from a trusted
-		// upstream (e.g. a load balancer), attempt to parse the v2 header so
-		// that we use the original client address for hole-punch registration and
-		// WireGuard session tracking rather than the load balancer's address.
-		if s.proxyProtocol && len(s.trustedUpstreams) > 0 {
-			remoteHost := packet.remoteAddr.IP.String()
-			if _, trusted := s.trustedUpstreams[remoteHost]; trusted {
-				if info, payload, ok := proxyproto.ParseV2UDPHeader(effectiveData); ok {
-					if info != nil {
-						// Override source address with what the proxy reported.
-						if srcIP := net.ParseIP(info.SrcIP); srcIP != nil {
-							effectiveAddr = &net.UDPAddr{
-								IP:   srcIP,
-								Port: info.SrcPort,
-							}
-							logger.Debug("PROXY protocol v2: overriding source %s → %s:%d",
-								packet.remoteAddr, info.SrcIP, info.SrcPort)
-						}
-					}
-					// Always advance past the header so the remainder is treated
-					// as the real application payload.
-					effectiveData = payload
-				}
-			}
-		}
-		// -----------------------------------------------------------------------
-
-		// Determine packet type by inspecting the first byte of the (possibly
-		// stripped) application payload.
-		if len(effectiveData) > 0 && effectiveData[0] >= 1 && effectiveData[0] <= 4 {
 			// Process as a WireGuard packet.
-			s.handleWireGuardPacket(effectiveData, effectiveAddr)
+			s.handleWireGuardPacket(packet.data, packet.remoteAddr)
 		} else {
-			// Rate limit: allow at most 2 hole punch messages per IP:Port per second
-			rateLimitKey := effectiveAddr.String()
-			entryVal, _ := s.holePunchRateLimiter.LoadOrStore(rateLimitKey, &holePunchRateLimitEntry{
-				windowStart: time.Now(),
-			})
-			rlEntry := entryVal.(*holePunchRateLimitEntry)
-			rlEntry.mu.Lock()
-			now := time.Now()
-			if now.Sub(rlEntry.windowStart) >= time.Second {
-				rlEntry.count = 0
-				rlEntry.windowStart = now
-			}
-			rlEntry.count++
-			allowed := rlEntry.count <= 2
-			rlEntry.mu.Unlock()
-			if !allowed {
-				// logger.Debug("Rate limiting hole punch message from %s", rateLimitKey)
-				bufferPool.Put(packet.data[:1500])
-				continue
-			}
-
 			// Process as an encrypted hole punch message
 			var encMsg EncryptedHolePunchMessage
-			if err := json.Unmarshal(effectiveData, &encMsg); err != nil {
+			if err := json.Unmarshal(packet.data, &encMsg); err != nil {
 				logger.Error("Error unmarshaling encrypted message: %v", err)
 				// Return the buffer to the pool for reuse and continue with next packet
 				bufferPool.Put(packet.data[:1500])
@@ -401,7 +291,7 @@ func (s *UDPProxyServer) packetWorker() {
 			// This appears to be an encrypted message
 			decryptedData, err := s.decryptMessage(encMsg)
 			if err != nil {
-				// logger.Error("Failed to decrypt message: %v", err)
+				logger.Error("Failed to decrypt message: %v", err)
 				// Return the buffer to the pool for reuse and continue with next packet
 				bufferPool.Put(packet.data[:1500])
 				continue
@@ -420,14 +310,14 @@ func (s *UDPProxyServer) packetWorker() {
 				NewtID:            msg.NewtID,
 				OlmID:             msg.OlmID,
 				Token:             msg.Token,
-				IP:                effectiveAddr.IP.String(),
+				IP:                packet.remoteAddr.IP.String(),
-				Port:              effectiveAddr.Port,
+				Port:              packet.remoteAddr.Port,
 				Timestamp:         time.Now().Unix(),
 				ReachableAt:       s.ReachableAt,
 				ExitNodePublicKey: s.privateKey.PublicKey().String(),
 				ClientPublicKey:   msg.PublicKey,
 			}
-			logger.Debug("Created endpoint from packet remoteAddr %s: IP=%s, Port=%d", effectiveAddr.String(), endpoint.IP, endpoint.Port)
+			logger.Debug("Created endpoint from packet remoteAddr %s: IP=%s, Port=%d", packet.remoteAddr.String(), endpoint.IP, endpoint.Port)
 			s.notifyServer(endpoint)
 			s.clearSessionsForIP(endpoint.IP) // Clear sessions for this IP to allow re-establishment
 		}
@@ -526,43 +416,6 @@ func extractWireGuardIndices(packet []byte) (uint32, uint32, bool) {
 	return 0, 0, false
 }
 
-// cachedAddr holds a resolved UDP address with TTL
-type cachedAddr struct {
-	addr      *net.UDPAddr
-	expiresAt time.Time
-}
-
-// addrCacheTTL is how long resolved addresses are cached before re-resolving
-const addrCacheTTL = 5 * time.Minute
-
-// getCachedAddr returns a cached UDP address or resolves and caches it.
-// This avoids per-packet DNS lookups which are a major throughput bottleneck.
-func (s *UDPProxyServer) getCachedAddr(ip string, port int) (*net.UDPAddr, error) {
-	key := fmt.Sprintf("%s:%d", ip, port)
-
-	// Check cache first
-	if cached, ok := s.addrCache.Load(key); ok {
-		entry := cached.(*cachedAddr)
-		if time.Now().Before(entry.expiresAt) {
-			return entry.addr, nil
-		}
-		// Cache expired, delete and re-resolve
-		s.addrCache.Delete(key)
-	}
-
-	// Resolve and cache
-	addr, err := net.ResolveUDPAddr("udp", key)
-	if err != nil {
-		return nil, err
-	}
-
-	s.addrCache.Store(key, &cachedAddr{
-		addr:      addr,
-		expiresAt: time.Now().Add(addrCacheTTL),
-	})
-	return addr, nil
-}
-
 // Updated to handle multi-peer WireGuard communication
 func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UDPAddr) {
 	if len(packet) == 0 {
@@ -597,7 +450,7 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD
 		logger.Debug("Forwarding handshake initiation from %s (sender index: %d) to peers %v", remoteAddr, senderIndex, proxyMapping.Destinations)
 
 		for _, dest := range proxyMapping.Destinations {
-			destAddr, err := s.getCachedAddr(dest.DestinationIP, dest.DestinationPort)
+			destAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", dest.DestinationIP, dest.DestinationPort))
 			if err != nil {
 				logger.Error("Failed to resolve destination address: %v", err)
 				continue
@@ -633,7 +486,7 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD
 
 		// Forward the response to the original sender
 		for _, dest := range proxyMapping.Destinations {
-			destAddr, err := s.getCachedAddr(dest.DestinationIP, dest.DestinationPort)
+			destAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", dest.DestinationIP, dest.DestinationPort))
 			if err != nil {
 				logger.Error("Failed to resolve destination address: %v", err)
 				continue
@@ -690,7 +543,7 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD
 			// No known session, fall back to forwarding to all peers
 			logger.Debug("No session found for receiver index %d, forwarding to all destinations", receiverIndex)
 			for _, dest := range proxyMapping.Destinations {
-				destAddr, err := s.getCachedAddr(dest.DestinationIP, dest.DestinationPort)
+				destAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", dest.DestinationIP, dest.DestinationPort))
 				if err != nil {
 					logger.Error("Failed to resolve destination address: %v", err)
 					continue
@@ -718,7 +571,7 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD
 
 		// Forward to all peers
 		for _, dest := range proxyMapping.Destinations {
-			destAddr, err := s.getCachedAddr(dest.DestinationIP, dest.DestinationPort)
+			destAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", dest.DestinationIP, dest.DestinationPort))
 			if err != nil {
 				logger.Error("Failed to resolve destination address: %v", err)
 				continue
@@ -1177,30 +1030,6 @@ func (s *UDPProxyServer) tryRebuildSession(pattern *CommunicationPattern) {
 }
 
 // cleanupIdleCommunicationPatterns periodically removes idle communication patterns
-// cleanupHolePunchRateLimiter periodically evicts stale rate limit entries to prevent unbounded growth.
-func (s *UDPProxyServer) cleanupHolePunchRateLimiter() {
-	ticker := time.NewTicker(30 * time.Second)
-	defer ticker.Stop()
-	for {
-		select {
-		case <-ticker.C:
-			now := time.Now()
-			s.holePunchRateLimiter.Range(func(key, value interface{}) bool {
-				rlEntry := value.(*holePunchRateLimitEntry)
-				rlEntry.mu.Lock()
-				stale := now.Sub(rlEntry.windowStart) > 10*time.Second
-				rlEntry.mu.Unlock()
-				if stale {
-					s.holePunchRateLimiter.Delete(key)
-				}
-				return true
-			})
-		case <-s.ctx.Done():
-			return
-		}
-	}
-}
-
 func (s *UDPProxyServer) cleanupIdleCommunicationPatterns() {
 	ticker := time.NewTicker(10 * time.Minute)
 	defer ticker.Stop()