Add pp to relay

perf(relay): scale packet workers and queue depth for throughput

.github/workflows/cicd.yml

@@ -16,7 +16,8 @@ permissions:
 on:
     push:
         tags:
-      - "*"
+            - "[0-9]+.[0-9]+.[0-9]+"
+            - "[0-9]+.[0-9]+.[0-9]+.rc.[0-9]+"
 
 concurrency:
   group: ${{ github.ref }}
@@ -30,18 +31,18 @@ jobs:
     timeout-minutes: 120
     env:
       # Target images
-      DOCKERHUB_IMAGE: docker.io/${{ secrets.DOCKER_HUB_USERNAME }}/${{ github.event.repository.name }}
+      DOCKERHUB_IMAGE: docker.io/fosrl/${{ github.event.repository.name }}
       GHCR_IMAGE: ghcr.io/${{ github.repository_owner }}/${{ github.event.repository.name }}
 
     steps:
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
 
       - name: Set up QEMU
         uses: docker/setup-qemu-action@c7c53464625b32c7a7e944ae62b3e17d2b600130 # v3.7.0
 
       - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 # v3.11.1
+        uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
 
       - name: Log in to Docker Hub
         uses: docker/login-action@5e57cd118135c172c3672efd75eb46360885c0ef # v3.6.0
@@ -56,7 +57,7 @@ jobs:
         shell: bash
 
       - name: Install Go
-        uses: actions/setup-go@4dc6199c7b1a012772edbd06daecab0f50c9053c # v6.1.0
+        uses: actions/setup-go@7a3fe6cf4cb3a834922a1244abfce67bcef6a0c5 # v6.2.0
         with:
           go-version: 1.25
 
@@ -154,7 +155,7 @@ jobs:
         shell: bash
 
       - name: Upload artifacts from /bin
-        uses: actions/upload-artifact@330a01c490aca151604b8cf639adc76d48f6c5d4 # v5.0.0
+        uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
         with:
           name: binaries
           path: bin/

.github/workflows/test.yml

@@ -14,10 +14,10 @@ jobs:
     runs-on: amd64-runner
 
     steps:
-      - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
+      - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
 
       - name: Set up Go
-        uses: actions/setup-go@4dc6199c7b1a012772edbd06daecab0f50c9053c # v6.1.0
+        uses: actions/setup-go@7a3fe6cf4cb3a834922a1244abfce67bcef6a0c5 # v6.2.0
         with:
           go-version: 1.25
 

README.md

@@ -42,13 +42,12 @@ In single node (self hosted) Pangolin deployments this can be bypassed by using
 
 ## CLI Args
 
+Important:
 - `reachableAt`: How should the remote server reach Gerbil's API?
 - `generateAndSaveKeyTo`: Where to save the generated WireGuard private key to persist across restarts.
-- `remoteConfig` (optional): Remote config location to HTTP get the JSON based config from. See `example_config.json`
+- `remoteConfig`: Remote config location to HTTP get the JSON based config from.
-- `config` (optional): Local JSON file path to load config. Used if remote config is not supplied. See `example_config.json`
-
-Note: You must use either `config` or `remoteConfig` to configure WireGuard.
 
+Others:
 - `reportBandwidthTo` (optional): **DEPRECATED** - Use `remoteConfig` instead. Remote HTTP endpoint to send peer bandwidth data
 - `interface` (optional): Name of the WireGuard interface created by Gerbil. Default: `wg0`
 - `listen` (optional): Port to listen on for HTTP server. Default: `:3004`
@@ -66,7 +65,6 @@ Note: You must use either `config` or `remoteConfig` to configure WireGuard.
 All CLI arguments can also be provided via environment variables:
 
 - `INTERFACE`: Name of the WireGuard interface
-- `CONFIG`: Path to local configuration file
 - `REMOTE_CONFIG`: URL of the remote config server
 - `LISTEN`: Address to listen on for HTTP server
 - `GENERATE_AND_SAVE_KEY_TO`: Path to save generated private key

config_example.json

@@ -1,23 +0,0 @@
-{
-    "privateKey": "kBGTgk7c+zncEEoSnMl+jsLjVh5ZVoL/HwBSQem+d1M=",
-    "listenPort": 51820,
-    "ipAddress": "10.0.0.1/24",
-    "peers": [
-        {
-            "publicKey": "5UzzoeveFVSzuqK3nTMS5bA1jIMs1fQffVQzJ8MXUQM=",
-            "allowedIps": ["10.0.0.0/28"]
-        },
-        {
-            "publicKey": "kYrZpuO2NsrFoBh1GMNgkhd1i9Rgtu1rAjbJ7qsfngU=",
-            "allowedIps": ["10.0.0.16/28"]
-        },
-        {
-            "publicKey": "1YfPUVr9ZF4zehkbI2BQhCxaRLz+Vtwa4vJwH+mpK0A=",
-            "allowedIps": ["10.0.0.32/28"]
-        },
-        {
-            "publicKey": "2/U4oyZ+sai336Dal/yExCphL8AxyqvIxMk4qsUy4iI=",
-            "allowedIps": ["10.0.0.48/28"]
-        }
-    ]
-}

go.mod

@@ -5,7 +5,8 @@ go 1.25
 require (
 	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/vishvananda/netlink v1.3.1
-	golang.org/x/crypto v0.45.0
+	golang.org/x/crypto v0.46.0
+	golang.org/x/sync v0.1.0
 	golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6
 )
 
@@ -17,7 +18,6 @@ require (
 	github.com/mdlayher/socket v0.4.1 // indirect
 	github.com/vishvananda/netns v0.0.5 // indirect
 	golang.org/x/net v0.47.0 // indirect
-	golang.org/x/sync v0.1.0 // indirect
+	golang.org/x/sys v0.39.0 // indirect
-	golang.org/x/sys v0.38.0 // indirect
 	golang.zx2c4.com/wireguard v0.0.0-20230325221338-052af4a8072b // indirect
 )

go.sum

@@ -16,16 +16,16 @@ github.com/vishvananda/netlink v1.3.1 h1:3AEMt62VKqz90r0tmNhog0r/PpWKmrEShJU0wJW
 github.com/vishvananda/netlink v1.3.1/go.mod h1:ARtKouGSTGchR8aMwmkzC0qiNPrrWO5JS/XMVl45+b4=
 github.com/vishvananda/netns v0.0.5 h1:DfiHV+j8bA32MFM7bfEunvT8IAqQ/NzSJHtcmW5zdEY=
 github.com/vishvananda/netns v0.0.5/go.mod h1:SpkAiCQRtJ6TvvxPnOSyH3BMl6unz3xZlaprSwhNNJM=
-golang.org/x/crypto v0.45.0 h1:jMBrvKuj23MTlT0bQEOBcAE0mjg8mK9RXFhRH6nyF3Q=
+golang.org/x/crypto v0.46.0 h1:cKRW/pmt1pKAfetfu+RCEvjvZkA9RimPbh7bhFjGVBU=
-golang.org/x/crypto v0.45.0/go.mod h1:XTGrrkGJve7CYK7J8PEww4aY7gM3qMCElcJQ8n8JdX4=
+golang.org/x/crypto v0.46.0/go.mod h1:Evb/oLKmMraqjZ2iQTwDwvCtJkczlDuTmdJXoZVzqU0=
 golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY=
 golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU=
 golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o=
 golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.2.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.10.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc=
+golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk=
-golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
+golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
 golang.zx2c4.com/wireguard v0.0.0-20230325221338-052af4a8072b h1:J1CaxgLerRR5lgx3wnr6L04cJFbWoceSK9JWBdglINo=
 golang.zx2c4.com/wireguard v0.0.0-20230325221338-052af4a8072b/go.mod h1:tqur9LnfstdR9ep2LaJT4lFUl0EjlHtge+gAjmsHUG4=
 golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6 h1:CawjfCvYQH2OU3/TnxLx97WDSUDRABfT18pCOYwc2GE=

main.go

@@ -8,11 +8,15 @@ import (
 	"flag"
 	"fmt"
 	"io"
+	"log"
 	"net"
 	"net/http"
+	_ "net/http/pprof"
 	"os"
 	"os/exec"
 	"os/signal"
+	"runtime"
+	"runtime/pprof"
 	"strconv"
 	"strings"
 	"sync"
@@ -38,11 +42,13 @@ var (
 	notifyURL        string
 	proxyRelay       *relay.UDPProxyServer
 	proxySNI         *proxy.SNIProxy
+	doTrafficShaping bool
 )
 
 type WgConfig struct {
 	PrivateKey string `json:"privateKey"`
 	ListenPort int    `json:"listenPort"`
+	RelayPort  int    `json:"relayPort"`
 	IpAddress  string `json:"ipAddress"`
 	Peers      []Peer `json:"peers"`
 }
@@ -111,6 +117,8 @@ func parseLogLevel(level string) logger.LogLevel {
 }
 
 func main() {
+	go monitorMemory(1024 * 1024 * 512) // trigger if memory usage exceeds 512MB
+
 	var (
 		err                  error
 		wgconfig             WgConfig
@@ -144,6 +152,7 @@ func main() {
 	localOverridesStr = os.Getenv("LOCAL_OVERRIDES")
 	trustedUpstreamsStr = os.Getenv("TRUSTED_UPSTREAMS")
 	proxyProtocolStr := os.Getenv("PROXY_PROTOCOL")
+	doTrafficShapingStr := os.Getenv("DO_TRAFFIC_SHAPING")
 
 	if interfaceName == "" {
 		flag.StringVar(&interfaceName, "interface", "wg0", "Name of the WireGuard interface")
@@ -215,6 +224,13 @@ func main() {
 		flag.BoolVar(&proxyProtocol, "proxy-protocol", true, "Enable PROXY protocol v1 for preserving client IP")
 	}
 
+	if doTrafficShapingStr != "" {
+		doTrafficShaping = strings.ToLower(doTrafficShapingStr) == "true"
+	}
+	if doTrafficShapingStr == "" {
+		flag.BoolVar(&doTrafficShaping, "do-traffic-shaping", false, "Whether to set up traffic shaping rules for peers (requires tc command and root privileges)")
+	}
+
 	flag.Parse()
 
 	logger.Init()
@@ -339,17 +355,8 @@ func main() {
 		return periodicBandwidthCheck(groupCtx, remoteConfigURL+"/gerbil/receive-bandwidth")
 	})
 
-	// Start the UDP proxy server
+	// Parse local overrides and trusted upstreams early so that both the relay
-	proxyRelay = relay.NewUDPProxyServer(groupCtx, ":21820", remoteConfigURL, key, reachableAt)
+	// and the SNI proxy share the same configuration values.
-	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, ",")
@@ -368,6 +375,21 @@ 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)
@@ -382,6 +404,7 @@ func main() {
 	http.HandleFunc("/update-proxy-mapping", handleUpdateProxyMapping)
 	http.HandleFunc("/update-destinations", handleUpdateDestinations)
 	http.HandleFunc("/update-local-snis", handleUpdateLocalSNIs)
+	http.HandleFunc("/healthz", handleHealthz)
 	logger.Info("Starting HTTP server on %s", listenAddr)
 
 	// HTTP server with graceful shutdown on context cancel
@@ -543,6 +566,10 @@ func ensureWireguardInterface(wgconfig WgConfig) error {
 		logger.Warn("Failed to ensure MSS clamping: %v", err)
 	}
 
+	if err := ensureWireguardFirewall(); err != nil {
+		logger.Warn("Failed to ensure WireGuard firewall rules: %v", err)
+	}
+
 	logger.Info("WireGuard interface %s created and configured", interfaceName)
 
 	return nil
@@ -711,6 +738,113 @@ func ensureMSSClamping() error {
 	return nil
 }
 
+func ensureWireguardFirewall() error {
+	// Rules to enforce:
+	// 1. Allow established/related connections (responses to our outbound traffic)
+	// 2. Allow ICMP ping packets
+	// 3. Drop all other inbound traffic from peers
+
+	// Define the rules we want to ensure exist
+	rules := [][]string{
+		// Allow established and related connections (responses to outbound traffic)
+		{
+			"-A", "INPUT",
+			"-i", interfaceName,
+			"-m", "conntrack",
+			"--ctstate", "ESTABLISHED,RELATED",
+			"-j", "ACCEPT",
+		},
+		// Allow ICMP ping requests
+		{
+			"-A", "INPUT",
+			"-i", interfaceName,
+			"-p", "icmp",
+			"--icmp-type", "8",
+			"-j", "ACCEPT",
+		},
+		// Drop all other inbound traffic from WireGuard interface
+		{
+			"-A", "INPUT",
+			"-i", interfaceName,
+			"-j", "DROP",
+		},
+	}
+
+	// First, try to delete any existing rules for this interface
+	for _, rule := range rules {
+		deleteArgs := make([]string, len(rule))
+		copy(deleteArgs, rule)
+		// Change -A to -D for deletion
+		for i, arg := range deleteArgs {
+			if arg == "-A" {
+				deleteArgs[i] = "-D"
+				break
+			}
+		}
+
+		deleteCmd := exec.Command("/usr/sbin/iptables", deleteArgs...)
+		logger.Debug("Attempting to delete existing firewall rule: %v", deleteArgs)
+
+		// Try deletion multiple times to handle multiple existing rules
+		for i := 0; i < 5; i++ {
+			out, err := deleteCmd.CombinedOutput()
+			if err != nil {
+				if exitErr, ok := err.(*exec.ExitError); ok {
+					logger.Debug("Deletion stopped: %v (output: %s)", exitErr.String(), string(out))
+				}
+				break // No more rules to delete
+			}
+			logger.Info("Deleted existing firewall rule (attempt %d)", i+1)
+		}
+	}
+
+	// Now add the rules
+	var errors []error
+	for i, rule := range rules {
+		addCmd := exec.Command("/usr/sbin/iptables", rule...)
+		logger.Info("Adding WireGuard firewall rule %d: %v", i+1, rule)
+
+		if out, err := addCmd.CombinedOutput(); err != nil {
+			errMsg := fmt.Sprintf("Failed to add firewall rule %d: %v (output: %s)", i+1, err, string(out))
+			logger.Error("%s", errMsg)
+			errors = append(errors, fmt.Errorf("%s", errMsg))
+			continue
+		}
+
+		// Verify the rule was added by checking
+		checkArgs := make([]string, len(rule))
+		copy(checkArgs, rule)
+		// Change -A to -C for check
+		for j, arg := range checkArgs {
+			if arg == "-A" {
+				checkArgs[j] = "-C"
+				break
+			}
+		}
+
+		checkCmd := exec.Command("/usr/sbin/iptables", checkArgs...)
+		if out, err := checkCmd.CombinedOutput(); err != nil {
+			errMsg := fmt.Sprintf("Rule verification failed for rule %d: %v (output: %s)", i+1, err, string(out))
+			logger.Error("%s", errMsg)
+			errors = append(errors, fmt.Errorf("%s", errMsg))
+			continue
+		}
+
+		logger.Info("Successfully added and verified WireGuard firewall rule %d", i+1)
+	}
+
+	if len(errors) > 0 {
+		var errMsgs []string
+		for _, err := range errors {
+			errMsgs = append(errMsgs, err.Error())
+		}
+		return fmt.Errorf("WireGuard firewall setup encountered errors:\n%s", strings.Join(errMsgs, "\n"))
+	}
+
+	logger.Info("WireGuard firewall rules successfully configured for interface %s", interfaceName)
+	return nil
+}
+
 func handlePeer(w http.ResponseWriter, r *http.Request) {
 	switch r.Method {
 	case http.MethodPost:
@@ -722,6 +856,15 @@ func handlePeer(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+func handleHealthz(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodGet {
+		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
+		return
+	}
+	w.WriteHeader(http.StatusOK)
+	_, _ = w.Write([]byte("ok"))
+}
+
 func handleAddPeer(w http.ResponseWriter, r *http.Request) {
 	var peer Peer
 	if err := json.NewDecoder(r.Body).Decode(&peer); err != nil {
@@ -754,17 +897,23 @@ func addPeerInternal(peer Peer) error {
 		return fmt.Errorf("failed to parse public key: %v", err)
 	}
 
+	logger.Debug("Adding peer %s with AllowedIPs: %v", peer.PublicKey, peer.AllowedIPs)
+
 	// parse allowed IPs into array of net.IPNet
 	var allowedIPs []net.IPNet
 	var wgIPs []string
 	for _, ipStr := range peer.AllowedIPs {
+		logger.Debug("Parsing AllowedIP: %s", ipStr)
 		_, ipNet, err := net.ParseCIDR(ipStr)
 		if err != nil {
+			logger.Warn("Failed to parse allowed IP '%s' for peer %s: %v", ipStr, peer.PublicKey, err)
 			return fmt.Errorf("failed to parse allowed IP: %v", err)
 		}
 		allowedIPs = append(allowedIPs, *ipNet)
 		// Extract the IP address from the CIDR for relay cleanup
-		wgIPs = append(wgIPs, ipNet.IP.String())
+		extractedIP := ipNet.IP.String()
+		wgIPs = append(wgIPs, extractedIP)
+		logger.Debug("Extracted IP %s from AllowedIP %s", extractedIP, ipStr)
 	}
 
 	peerConfig := wgtypes.PeerConfig{
@@ -780,6 +929,18 @@ func addPeerInternal(peer Peer) error {
 		return fmt.Errorf("failed to add peer: %v", err)
 	}
 
+	// Setup bandwidth limiting for each peer IP
+	if doTrafficShaping {
+		logger.Debug("doTrafficShaping is true, setting up bandwidth limits for %d IPs", len(wgIPs))
+		for _, wgIP := range wgIPs {
+			if err := setupPeerBandwidthLimit(wgIP); err != nil {
+				logger.Warn("Failed to setup bandwidth limit for peer IP %s: %v", wgIP, err)
+			}
+		}
+	} else {
+		logger.Debug("doTrafficShaping is false, skipping bandwidth limit setup")
+	}
+
 	// Clear relay connections for the peer's WireGuard IPs
 	if proxyRelay != nil {
 		for _, wgIP := range wgIPs {
@@ -824,9 +985,8 @@ func removePeerInternal(publicKey string) error {
 		return fmt.Errorf("failed to parse public key: %v", err)
 	}
 
-	// Get current peer info before removing to clear relay connections
+	// Get current peer info before removing to clear relay connections and bandwidth limits
 	var wgIPs []string
-	if proxyRelay != nil {
 	device, err := wgClient.Device(interfaceName)
 	if err == nil {
 		for _, peer := range device.Peers {
@@ -839,7 +999,6 @@ func removePeerInternal(publicKey string) error {
 			}
 		}
 	}
-	}
 
 	peerConfig := wgtypes.PeerConfig{
 		PublicKey: pubKey,
@@ -854,6 +1013,15 @@ func removePeerInternal(publicKey string) error {
 		return fmt.Errorf("failed to remove peer: %v", err)
 	}
 
+	// Remove bandwidth limits for each peer IP
+	if doTrafficShaping {
+		for _, wgIP := range wgIPs {
+			if err := removePeerBandwidthLimit(wgIP); err != nil {
+				logger.Warn("Failed to remove bandwidth limit for peer IP %s: %v", wgIP, err)
+			}
+		}
+	}
+
 	// Clear relay connections for the peer's WireGuard IPs
 	if proxyRelay != nil {
 		for _, wgIP := range wgIPs {
@@ -1161,3 +1329,177 @@ func notifyPeerChange(action, publicKey string) {
 		logger.Warn("Notify server returned non-OK: %s", resp.Status)
 	}
 }
+
+func monitorMemory(limit uint64) {
+	var m runtime.MemStats
+	for {
+		runtime.ReadMemStats(&m)
+		if m.Alloc > limit {
+			fmt.Printf("Memory spike detected (%d bytes). Dumping profile...\n", m.Alloc)
+
+			f, err := os.Create(fmt.Sprintf("/var/config/heap/heap-spike-%d.pprof", time.Now().Unix()))
+			if err != nil {
+				log.Println("could not create profile:", err)
+			} else {
+				pprof.WriteHeapProfile(f)
+				f.Close()
+			}
+
+			// Wait a while before checking again to avoid spamming profiles
+			time.Sleep(5 * time.Minute)
+		}
+		time.Sleep(5 * time.Second)
+	}
+}
+
+// setupPeerBandwidthLimit sets up TC (Traffic Control) to limit bandwidth for a specific peer IP
+// Currently hardcoded to 20 Mbps per peer
+func setupPeerBandwidthLimit(peerIP string) error {
+	logger.Debug("setupPeerBandwidthLimit called for peer IP: %s", peerIP)
+	const bandwidthLimit = "50mbit" // 50 Mbps limit per peer
+
+	// Parse the IP to get just the IP address (strip any CIDR notation if present)
+	ip := peerIP
+	if strings.Contains(peerIP, "/") {
+		parsedIP, _, err := net.ParseCIDR(peerIP)
+		if err != nil {
+			return fmt.Errorf("failed to parse peer IP: %v", err)
+		}
+		ip = parsedIP.String()
+	}
+
+	// First, ensure we have a root qdisc on the interface (HTB - Hierarchical Token Bucket)
+	// Check if qdisc already exists
+	cmd := exec.Command("tc", "qdisc", "show", "dev", interfaceName)
+	output, err := cmd.CombinedOutput()
+	if err != nil {
+		return fmt.Errorf("failed to check qdisc: %v, output: %s", err, string(output))
+	}
+
+	// If no HTB qdisc exists, create one
+	if !strings.Contains(string(output), "htb") {
+		cmd = exec.Command("tc", "qdisc", "add", "dev", interfaceName, "root", "handle", "1:", "htb", "default", "9999")
+		if output, err := cmd.CombinedOutput(); err != nil {
+			return fmt.Errorf("failed to add root qdisc: %v, output: %s", err, string(output))
+		}
+		logger.Info("Created HTB root qdisc on %s", interfaceName)
+	}
+
+	// Generate a unique class ID based on the IP address
+	// We'll use the last octet of the IP as part of the class ID
+	ipParts := strings.Split(ip, ".")
+	if len(ipParts) != 4 {
+		return fmt.Errorf("invalid IPv4 address: %s", ip)
+	}
+	lastOctet := ipParts[3]
+	classID := fmt.Sprintf("1:%s", lastOctet)
+	logger.Debug("Generated class ID %s for peer IP %s", classID, ip)
+
+	// Create a class for this peer with bandwidth limit
+	cmd = exec.Command("tc", "class", "add", "dev", interfaceName, "parent", "1:", "classid", classID,
+		"htb", "rate", bandwidthLimit, "ceil", bandwidthLimit)
+	if output, err := cmd.CombinedOutput(); err != nil {
+		logger.Debug("tc class add failed for %s: %v, output: %s", ip, err, string(output))
+		// If class already exists, try to replace it
+		if strings.Contains(string(output), "File exists") {
+			cmd = exec.Command("tc", "class", "replace", "dev", interfaceName, "parent", "1:", "classid", classID,
+				"htb", "rate", bandwidthLimit, "ceil", bandwidthLimit)
+			if output, err := cmd.CombinedOutput(); err != nil {
+				return fmt.Errorf("failed to replace class: %v, output: %s", err, string(output))
+			}
+			logger.Debug("Successfully replaced existing class %s for peer IP %s", classID, ip)
+		} else {
+			return fmt.Errorf("failed to add class: %v, output: %s", err, string(output))
+		}
+	} else {
+		logger.Debug("Successfully added new class %s for peer IP %s", classID, ip)
+	}
+
+	// Add a filter to match traffic from this peer IP (ingress)
+	cmd = exec.Command("tc", "filter", "add", "dev", interfaceName, "protocol", "ip", "parent", "1:",
+		"prio", "1", "u32", "match", "ip", "src", ip, "flowid", classID)
+	if output, err := cmd.CombinedOutput(); err != nil {
+		// If filter fails, log but don't fail the peer addition
+		logger.Warn("Failed to add ingress filter for peer IP %s: %v, output: %s", ip, err, string(output))
+	}
+
+	// Add a filter to match traffic to this peer IP (egress)
+	cmd = exec.Command("tc", "filter", "add", "dev", interfaceName, "protocol", "ip", "parent", "1:",
+		"prio", "1", "u32", "match", "ip", "dst", ip, "flowid", classID)
+	if output, err := cmd.CombinedOutput(); err != nil {
+		// If filter fails, log but don't fail the peer addition
+		logger.Warn("Failed to add egress filter for peer IP %s: %v, output: %s", ip, err, string(output))
+	}
+
+	logger.Info("Setup bandwidth limit of %s for peer IP %s (class %s)", bandwidthLimit, ip, classID)
+	return nil
+}
+
+// removePeerBandwidthLimit removes TC rules for a specific peer IP
+func removePeerBandwidthLimit(peerIP string) error {
+	// Parse the IP to get just the IP address
+	ip := peerIP
+	if strings.Contains(peerIP, "/") {
+		parsedIP, _, err := net.ParseCIDR(peerIP)
+		if err != nil {
+			return fmt.Errorf("failed to parse peer IP: %v", err)
+		}
+		ip = parsedIP.String()
+	}
+
+	// Generate the class ID based on the IP
+	ipParts := strings.Split(ip, ".")
+	if len(ipParts) != 4 {
+		return fmt.Errorf("invalid IPv4 address: %s", ip)
+	}
+	lastOctet := ipParts[3]
+	classID := fmt.Sprintf("1:%s", lastOctet)
+
+	// Remove filters for this IP
+	// List all filters to find the ones for this class
+	cmd := exec.Command("tc", "filter", "show", "dev", interfaceName, "parent", "1:")
+	output, err := cmd.CombinedOutput()
+	if err != nil {
+		logger.Warn("Failed to list filters for peer IP %s: %v, output: %s", ip, err, string(output))
+	} else {
+		// Parse the output to find filter handles that match this classID
+		// The output format includes lines like:
+		// filter parent 1: protocol ip pref 1 u32 chain 0 fh 800::800 order 2048 key ht 800 bkt 0 flowid 1:4
+		lines := strings.Split(string(output), "\n")
+		for _, line := range lines {
+			// Look for lines containing our flowid (classID)
+			if strings.Contains(line, "flowid "+classID) && strings.Contains(line, "fh ") {
+				// Extract handle (format: fh 800::800)
+				parts := strings.Fields(line)
+				var handle string
+				for j, part := range parts {
+					if part == "fh" && j+1 < len(parts) {
+						handle = parts[j+1]
+						break
+					}
+				}
+				if handle != "" {
+					// Delete this filter using the handle
+					delCmd := exec.Command("tc", "filter", "del", "dev", interfaceName, "parent", "1:", "handle", handle, "prio", "1", "u32")
+					if delOutput, delErr := delCmd.CombinedOutput(); delErr != nil {
+						logger.Debug("Failed to delete filter handle %s for peer IP %s: %v, output: %s", handle, ip, delErr, string(delOutput))
+					} else {
+						logger.Debug("Deleted filter handle %s for peer IP %s", handle, ip)
+					}
+				}
+			}
+		}
+	}
+
+	// Remove the class
+	cmd = exec.Command("tc", "class", "del", "dev", interfaceName, "classid", classID)
+	if output, err := cmd.CombinedOutput(); err != nil {
+		// It's okay if the class doesn't exist
+		if !strings.Contains(string(output), "No such file or directory") && !strings.Contains(string(output), "Cannot find") {
+			logger.Warn("Failed to remove class for peer IP %s: %v, output: %s", ip, err, string(output))
+		}
+	}
+
+	logger.Info("Removed bandwidth limit for peer IP %s (class %s)", ip, classID)
+	return nil
+}

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,16 +32,6 @@ 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
@@ -89,249 +79,6 @@ 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())
@@ -490,12 +237,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 *ProxyProtocolInfo
+	var proxyInfo *proxyproto.Info
 	var actualClientConn net.Conn = clientConn
 
 	if len(p.trustedUpstreams) > 0 {
 		var err error
-		proxyInfo, actualClientConn, err = p.parseProxyProtocolHeader(clientConn)
+		proxyInfo, actualClientConn, err = proxyproto.ParseV1Header(clientConn, p.trustedUpstreams)
 		if err != nil {
 			logger.Debug("Failed to parse PROXY protocol: %v", err)
 			return
@@ -575,10 +322,10 @@ func (p *SNIProxy) handleConnection(clientConn net.Conn) {
 		var proxyHeader string
 		if proxyInfo != nil {
 			// Use original client info from PROXY protocol
-			proxyHeader = p.buildProxyProtocolHeaderFromInfo(proxyInfo, targetConn.LocalAddr())
+			proxyHeader = proxyproto.BuildV1HeaderFromInfo(proxyInfo, targetConn.LocalAddr())
 		} else {
 			// Use direct client connection info
-			proxyHeader = buildProxyProtocolHeader(clientConn.RemoteAddr(), targetConn.LocalAddr())
+			proxyHeader = proxyproto.BuildV1Header(clientConn.RemoteAddr(), targetConn.LocalAddr())
 		}
 		logger.Debug("Sending PROXY protocol header: %s", strings.TrimSpace(proxyHeader))
 
@@ -758,14 +505,20 @@ func (p *SNIProxy) pipe(clientConn, targetConn net.Conn, clientReader io.Reader)
 	var wg sync.WaitGroup
 	wg.Add(2)
 
+	// closeOnce ensures we only close connections once
+	var closeOnce sync.Once
+	closeConns := func() {
+		closeOnce.Do(func() {
+			// Close both connections to unblock any pending reads
+			clientConn.Close()
+			targetConn.Close()
+		})
+	}
+
 	// Copy data from client to target (using the buffered reader)
 	go func() {
 		defer wg.Done()
-		defer func() {
+		defer closeConns()
-			if tcpConn, ok := targetConn.(*net.TCPConn); ok {
-				tcpConn.CloseWrite()
-			}
-		}()
 
 		// Use a large buffer for better performance
 		buf := make([]byte, 32*1024)
@@ -778,11 +531,7 @@ func (p *SNIProxy) pipe(clientConn, targetConn net.Conn, clientReader io.Reader)
 	// Copy data from target to client
 	go func() {
 		defer wg.Done()
-		defer func() {
+		defer closeConns()
-			if tcpConn, ok := clientConn.(*net.TCPConn); ok {
-				tcpConn.CloseWrite()
-			}
-		}()
 
 		// Use a large buffer for better performance
 		buf := make([]byte, 32*1024)

proxy/proxy_test.go

@@ -3,6 +3,8 @@ package proxy
 import (
 	"net"
 	"testing"
+
+	"github.com/fosrl/gerbil/proxyproto"
 )
 
 func TestBuildProxyProtocolHeader(t *testing.T) {
@@ -56,7 +58,7 @@ func TestBuildProxyProtocolHeader(t *testing.T) {
 				t.Fatalf("Failed to resolve target address: %v", err)
 			}
 
-			result := buildProxyProtocolHeader(clientTCP, targetTCP)
+			result := proxyproto.BuildV1Header(clientTCP, targetTCP)
 			if result != tt.expected {
 				t.Errorf("Expected %q, got %q", tt.expected, result)
 			}
@@ -69,7 +71,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 := buildProxyProtocolHeader(clientAddr, targetAddr)
+	result := proxyproto.BuildV1Header(clientAddr, targetAddr)
 	expected := "PROXY UNKNOWN\r\n"
 
 	if result != expected {
@@ -78,13 +80,8 @@ 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
-	proxyInfo := &ProxyProtocolInfo{
+	info := &proxyproto.Info{
 		Protocol: "TCP4",
 		SrcIP:    "10.0.0.1",
 		DestIP:   "192.168.1.100",
@@ -93,7 +90,7 @@ func TestBuildProxyProtocolHeaderFromInfo(t *testing.T) {
 	}
 
 	targetAddr, _ := net.ResolveTCPAddr("tcp", "127.0.0.1:8080")
-	header := proxy.buildProxyProtocolHeaderFromInfo(proxyInfo, targetAddr)
+	header := proxyproto.BuildV1HeaderFromInfo(info, targetAddr)
 
 	expected := "PROXY TCP4 10.0.0.1 127.0.0.1 12345 8080\r\n"
 	if header != expected {
@@ -101,7 +98,7 @@ func TestBuildProxyProtocolHeaderFromInfo(t *testing.T) {
 	}
 
 	// Test IPv6 case
-	proxyInfo = &ProxyProtocolInfo{
+	info = &proxyproto.Info{
 		Protocol: "TCP6",
 		SrcIP:    "2001:db8::1",
 		DestIP:   "2001:db8::2",
@@ -110,10 +107,99 @@ func TestBuildProxyProtocolHeaderFromInfo(t *testing.T) {
 	}
 
 	targetAddr, _ = net.ResolveTCPAddr("tcp6", "[::1]:8080")
-	header = proxy.buildProxyProtocolHeaderFromInfo(proxyInfo, targetAddr)
+	header = proxyproto.BuildV1HeaderFromInfo(info, 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

@@ -0,0 +1,370 @@
+// 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,10 +9,13 @@ 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"
@@ -118,6 +121,13 @@ 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
@@ -153,21 +163,57 @@ 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,
 		serverURL:        serverURL,
 		privateKey:       privateKey,
-		packetChan:  make(chan Packet, 1000),
+		packetChan:       make(chan Packet, 50000), // Increased from 1000 to handle high throughput
 		ReachableAt:      reachableAt,
 		ctx:              ctx,
 		cancel:           cancel,
+		proxyProtocol:    proxyProtocol,
+		trustedUpstreams: trustedMap,
 	}
 }
 
@@ -189,8 +235,13 @@ func (s *UDPProxyServer) Start() error {
 	s.conn = conn
 	logger.Info("UDP server listening on %s", s.addr)
 
-	// Start a fixed number of worker goroutines.
+	// Start worker goroutines based on CPU cores for better parallelism
-	workerCount := 10 // TODO: Make this configurable or pick it better!
+	// At high throughput (160+ Mbps), we need many workers to avoid bottlenecks
+	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()
 	}
@@ -210,6 +261,9 @@ 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
 }
 
@@ -267,14 +321,70 @@ func (s *UDPProxyServer) readPackets() {
 // packetWorker processes incoming packets from the channel.
 func (s *UDPProxyServer) packetWorker() {
 	for packet := range s.packetChan {
-		// Determine packet type by inspecting the first byte.
+		// effectiveData and effectiveAddr represent the application-layer payload
-		if packet.n > 0 && packet.data[0] >= 1 && packet.data[0] <= 4 {
+		// and the true originating address.  They start as the raw UDP values and
+		// 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(packet.data, packet.remoteAddr)
+			s.handleWireGuardPacket(effectiveData, effectiveAddr)
 		} 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(packet.data, &encMsg); err != nil {
+			if err := json.Unmarshal(effectiveData, &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])
@@ -291,7 +401,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
@@ -310,14 +420,14 @@ func (s *UDPProxyServer) packetWorker() {
 				NewtID:            msg.NewtID,
 				OlmID:             msg.OlmID,
 				Token:             msg.Token,
-				IP:                packet.remoteAddr.IP.String(),
+				IP:                effectiveAddr.IP.String(),
-				Port:              packet.remoteAddr.Port,
+				Port:              effectiveAddr.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", packet.remoteAddr.String(), endpoint.IP, endpoint.Port)
+			logger.Debug("Created endpoint from packet remoteAddr %s: IP=%s, Port=%d", effectiveAddr.String(), endpoint.IP, endpoint.Port)
 			s.notifyServer(endpoint)
 			s.clearSessionsForIP(endpoint.IP) // Clear sessions for this IP to allow re-establishment
 		}
@@ -416,6 +526,43 @@ 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 {
@@ -450,7 +597,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 := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", dest.DestinationIP, dest.DestinationPort))
+			destAddr, err := s.getCachedAddr(dest.DestinationIP, dest.DestinationPort)
 			if err != nil {
 				logger.Error("Failed to resolve destination address: %v", err)
 				continue
@@ -464,7 +611,7 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD
 
 			_, err = conn.Write(packet)
 			if err != nil {
-				logger.Error("Failed to forward handshake initiation: %v", err)
+				logger.Debug("Failed to forward handshake initiation: %v", err)
 			}
 		}
 
@@ -486,7 +633,7 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD
 
 		// Forward the response to the original sender
 		for _, dest := range proxyMapping.Destinations {
-			destAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", dest.DestinationIP, dest.DestinationPort))
+			destAddr, err := s.getCachedAddr(dest.DestinationIP, dest.DestinationPort)
 			if err != nil {
 				logger.Error("Failed to resolve destination address: %v", err)
 				continue
@@ -543,7 +690,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 := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", dest.DestinationIP, dest.DestinationPort))
+				destAddr, err := s.getCachedAddr(dest.DestinationIP, dest.DestinationPort)
 				if err != nil {
 					logger.Error("Failed to resolve destination address: %v", err)
 					continue
@@ -571,7 +718,7 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD
 
 		// Forward to all peers
 		for _, dest := range proxyMapping.Destinations {
-			destAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", dest.DestinationIP, dest.DestinationPort))
+			destAddr, err := s.getCachedAddr(dest.DestinationIP, dest.DestinationPort)
 			if err != nil {
 				logger.Error("Failed to resolve destination address: %v", err)
 				continue
@@ -839,7 +986,7 @@ func (s *UDPProxyServer) clearSessionsForIP(ip string) {
 		s.wgSessions.Delete(key)
 	}
 
-	logger.Info("Cleared %d sessions for WG IP: %s", len(keysToDelete), ip)
+	logger.Debug("Cleared %d sessions for WG IP: %s", len(keysToDelete), ip)
 }
 
 // // clearProxyMappingsForWGIP removes all proxy mappings that have destinations pointing to a specific WireGuard IP
@@ -1030,6 +1177,30 @@ 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()