Merge pull request #68 from fosrl/go

Bump go version

.github/workflows/cicd.yml

@@ -36,7 +36,7 @@ jobs:
 
     steps:
       - name: Checkout code
-        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
+        uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
 
       - name: Set up QEMU
         uses: docker/setup-qemu-action@c7c53464625b32c7a7e944ae62b3e17d2b600130 # v3.7.0
@@ -45,7 +45,7 @@ jobs:
         uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
 
       - name: Log in to Docker Hub
-        uses: docker/login-action@5e57cd118135c172c3672efd75eb46360885c0ef # v3.6.0
+        uses: docker/login-action@b45d80f862d83dbcd57f89517bcf500b2ab88fb2 # v4.0.0
         with:
           registry: docker.io
           username: ${{ secrets.DOCKER_HUB_USERNAME }}
@@ -57,7 +57,7 @@ jobs:
         shell: bash
 
       - name: Install Go
-        uses: actions/setup-go@7a3fe6cf4cb3a834922a1244abfce67bcef6a0c5 # v6.2.0
+        uses: actions/setup-go@4b73464bb391d4059bd26b0524d20df3927bd417 # v6.3.0
         with:
           go-version: 1.25
 
@@ -80,7 +80,7 @@ jobs:
         shell: bash
 
       - name: Login in to GHCR
-        uses: docker/login-action@5e57cd118135c172c3672efd75eb46360885c0ef # v3.6.0
+        uses: docker/login-action@b45d80f862d83dbcd57f89517bcf500b2ab88fb2 # v4.0.0
         with:
           registry: ghcr.io
           username: ${{ github.actor }}
@@ -108,7 +108,7 @@ jobs:
 
       - name: Install cosign
         # cosign is used to sign and verify container images (key and keyless)
-        uses: sigstore/cosign-installer@faadad0cce49287aee09b3a48701e75088a2c6ad # v4.0.0
+        uses: sigstore/cosign-installer@ba7bc0a3fef59531c69a25acd34668d6d3fe6f22 # v4.1.0
 
       - name: Dual-sign and verify (GHCR & Docker Hub)
         # Sign each image by digest using keyless (OIDC) and key-based signing,
@@ -155,7 +155,7 @@ jobs:
         shell: bash
 
       - name: Upload artifacts from /bin
-        uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
+        uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0
         with:
           name: binaries
           path: bin/

.github/workflows/mirror.yaml

@@ -23,7 +23,7 @@ jobs:
           skopeo --version
 
       - name: Install cosign
-        uses: sigstore/cosign-installer@faadad0cce49287aee09b3a48701e75088a2c6ad # v4.0.0
+        uses: sigstore/cosign-installer@ba7bc0a3fef59531c69a25acd34668d6d3fe6f22 # v4.1.0
 
       - name: Input check
         run: |

.github/workflows/test.yml

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

Dockerfile

@@ -1,4 +1,4 @@
-FROM golang:1.25-alpine AS builder
+FROM golang:1.26-alpine AS builder
 
 # Set the working directory inside the container
 WORKDIR /app

go.mod

@@ -1,12 +1,12 @@
 module github.com/fosrl/gerbil
 
-go 1.25
+go 1.26.0
 
 require (
 	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/vishvananda/netlink v1.3.1
-	golang.org/x/crypto v0.46.0
+	golang.org/x/crypto v0.49.0
-	golang.org/x/sync v0.1.0
+	golang.org/x/sync v0.20.0
 	golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6
 )
 
@@ -17,7 +17,7 @@ require (
 	github.com/mdlayher/netlink v1.7.2 // indirect
 	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/net v0.51.0 // indirect
-	golang.org/x/sys v0.39.0 // indirect
+	golang.org/x/sys v0.42.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.46.0 h1:cKRW/pmt1pKAfetfu+RCEvjvZkA9RimPbh7bhFjGVBU=
+golang.org/x/crypto v0.49.0 h1:+Ng2ULVvLHnJ/ZFEq4KdcDd/cfjrrjjNSXNzxg0Y4U4=
-golang.org/x/crypto v0.46.0/go.mod h1:Evb/oLKmMraqjZ2iQTwDwvCtJkczlDuTmdJXoZVzqU0=
+golang.org/x/crypto v0.49.0/go.mod h1:ErX4dUh2UM+CFYiXZRTcMpEcN8b/1gxEuv3nODoYtCA=
-golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY=
+golang.org/x/net v0.51.0 h1:94R/GTO7mt3/4wIKpcR5gkGmRLOuE/2hNGeWq/GBIFo=
-golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU=
+golang.org/x/net v0.51.0/go.mod h1:aamm+2QF5ogm02fjy5Bb7CQ0WMt1/WVM7FtyaTLlA9Y=
-golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o=
+golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
-golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
 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.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk=
+golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo=
-golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
+golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
 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

@@ -42,7 +42,6 @@ var (
 	notifyURL     string
 	proxyRelay    *relay.UDPProxyServer
 	proxySNI      *proxy.SNIProxy
-	doTrafficShaping bool
 )
 
 type WgConfig struct {
@@ -152,7 +151,6 @@ 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")
@@ -224,13 +222,6 @@ 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()
@@ -355,8 +346,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 +379,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)
@@ -897,23 +886,17 @@ 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
-		extractedIP := ipNet.IP.String()
+		wgIPs = append(wgIPs, ipNet.IP.String())
-		wgIPs = append(wgIPs, extractedIP)
-		logger.Debug("Extracted IP %s from AllowedIP %s", extractedIP, ipStr)
 	}
 
 	peerConfig := wgtypes.PeerConfig{
@@ -929,18 +912,6 @@ 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 {
@@ -985,8 +956,9 @@ 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 and bandwidth limits
+	// Get current peer info before removing to clear relay connections
 	var wgIPs []string
+	if proxyRelay != nil {
 		device, err := wgClient.Device(interfaceName)
 		if err == nil {
 			for _, peer := range device.Peers {
@@ -999,6 +971,7 @@ func removePeerInternal(publicKey string) error {
 				}
 			}
 		}
+	}
 
 	peerConfig := wgtypes.PeerConfig{
 		PublicKey: pubKey,
@@ -1013,15 +986,6 @@ 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 {
@@ -1351,155 +1315,3 @@ func monitorMemory(limit uint64) {
 		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,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
@@ -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))
 

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,
 		serverURL:   serverURL,
 		privateKey:  privateKey,
-		packetChan:       make(chan Packet, 50000), // Increased from 1000 to handle high throughput
+		packetChan:  make(chan Packet, 1000),
 		ReachableAt: reachableAt,
 		ctx:         ctx,
 		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()