Add pp to relay

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

.github/workflows/cicd.yml

@@ -1,52 +1,161 @@
 name: CI/CD Pipeline
 
+# CI/CD workflow for building, publishing, mirroring, signing container images and building release binaries.
+# Actions are pinned to specific SHAs to reduce supply-chain risk. This workflow triggers on tag push events.
+
+permissions:
+  contents: read
+  packages: write      # for GHCR push
+  id-token: write      # for Cosign Keyless (OIDC) Signing
+
+# Required secrets:
+# - DOCKER_HUB_USERNAME / DOCKER_HUB_ACCESS_TOKEN: push to Docker Hub
+# - GITHUB_TOKEN: used for GHCR login and OIDC keyless signing
+# - COSIGN_PRIVATE_KEY / COSIGN_PASSWORD / COSIGN_PUBLIC_KEY: for key-based signing
+
 on:
     push:
         tags:
-            - "*"
+            - "[0-9]+.[0-9]+.[0-9]+"
+            - "[0-9]+.[0-9]+.[0-9]+.rc.[0-9]+"
+
+concurrency:
+  group: ${{ github.ref }}
+  cancel-in-progress: true
 
 jobs:
   release:
     name: Build and Release
-        runs-on: ubuntu-latest
+    runs-on: amd64-runner
+    # Job-level timeout to avoid runaway or stuck runs
+    timeout-minutes: 120
+    env:
+      # Target images
+      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@v5
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
 
       - name: Set up QEMU
-              uses: docker/setup-qemu-action@v3
+        uses: docker/setup-qemu-action@c7c53464625b32c7a7e944ae62b3e17d2b600130 # v3.7.0
 
       - name: Set up Docker Buildx
-              uses: docker/setup-buildx-action@v3
+        uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
 
       - name: Log in to Docker Hub
-              uses: docker/login-action@v3
+        uses: docker/login-action@5e57cd118135c172c3672efd75eb46360885c0ef # v3.6.0
         with:
+          registry: docker.io
           username: ${{ secrets.DOCKER_HUB_USERNAME }}
           password: ${{ secrets.DOCKER_HUB_ACCESS_TOKEN }}
 
       - name: Extract tag name
         id: get-tag
         run: echo "TAG=${GITHUB_REF#refs/tags/}" >> $GITHUB_ENV
+        shell: bash
 
       - name: Install Go
-              uses: actions/setup-go@v6
+        uses: actions/setup-go@7a3fe6cf4cb3a834922a1244abfce67bcef6a0c5 # v6.2.0
         with:
           go-version: 1.25
 
-            - name: Build and push Docker images
+      - name: Update version in main.go
+        run: |
+          TAG=${{ env.TAG }}
+          if [ -f main.go ]; then
+            sed -i 's/version_replaceme/'"$TAG"'/' main.go
+            echo "Updated main.go with version $TAG"
+          else
+            echo "main.go not found"
+          fi
+        shell: bash
+
+      - name: Build and push Docker images (Docker Hub)
         run: |
           TAG=${{ env.TAG }}
           make docker-build-release tag=$TAG
+          echo "Built & pushed to: ${{ env.DOCKERHUB_IMAGE }}:${TAG}"
+        shell: bash
+
+      - name: Login in to GHCR
+        uses: docker/login-action@5e57cd118135c172c3672efd75eb46360885c0ef # v3.6.0
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Install skopeo + jq
+        # skopeo: copy/inspect images between registries
+        # jq: JSON parsing tool used to extract digest values
+        run: |
+          sudo apt-get update -y
+          sudo apt-get install -y skopeo jq
+          skopeo --version
+        shell: bash
+
+      - name: Copy tag from Docker Hub to GHCR
+        # Mirror the already-built image (all architectures) to GHCR so we can sign it
+        run: |
+          set -euo pipefail
+          TAG=${{ env.TAG }}
+          echo "Copying ${{ env.DOCKERHUB_IMAGE }}:${TAG} -> ${{ env.GHCR_IMAGE }}:${TAG}"
+          skopeo copy --all --retry-times 3 \
+            docker://$DOCKERHUB_IMAGE:$TAG \
+            docker://$GHCR_IMAGE:$TAG
+        shell: bash
+
+      - name: Install cosign
+        # cosign is used to sign and verify container images (key and keyless)
+        uses: sigstore/cosign-installer@faadad0cce49287aee09b3a48701e75088a2c6ad # v4.0.0
+
+      - name: Dual-sign and verify (GHCR & Docker Hub)
+        # Sign each image by digest using keyless (OIDC) and key-based signing,
+        # then verify both the public key signature and the keyless OIDC signature.
+        env:
+          TAG: ${{ env.TAG }}
+          COSIGN_PRIVATE_KEY: ${{ secrets.COSIGN_PRIVATE_KEY }}
+          COSIGN_PASSWORD: ${{ secrets.COSIGN_PASSWORD }}
+          COSIGN_PUBLIC_KEY: ${{ secrets.COSIGN_PUBLIC_KEY }}
+          COSIGN_YES: "true"
+        run: |
+          set -euo pipefail
+
+          issuer="https://token.actions.githubusercontent.com"
+          id_regex="^https://github.com/${{ github.repository }}/.+"  # accept this repo (all workflows/refs)
+
+          for IMAGE in "${GHCR_IMAGE}" "${DOCKERHUB_IMAGE}"; do
+            echo "Processing ${IMAGE}:${TAG}"
+
+            DIGEST="$(skopeo inspect --retry-times 3 docker://${IMAGE}:${TAG} | jq -r '.Digest')"
+            REF="${IMAGE}@${DIGEST}"
+            echo "Resolved digest: ${REF}"
+
+            echo "==> cosign sign (keyless) --recursive ${REF}"
+            cosign sign --recursive "${REF}"
+
+            echo "==> cosign sign (key) --recursive ${REF}"
+            cosign sign --key env://COSIGN_PRIVATE_KEY --recursive "${REF}"
+
+            echo "==> cosign verify (public key) ${REF}"
+            cosign verify --key env://COSIGN_PUBLIC_KEY "${REF}" -o text
+
+            echo "==> cosign verify (keyless policy) ${REF}"
+            cosign verify \
+              --certificate-oidc-issuer "${issuer}" \
+              --certificate-identity-regexp "${id_regex}" \
+              "${REF}" -o text
+          done
+        shell: bash
 
       - name: Build binaries
         run: |
           make go-build-release
+        shell: bash
 
       - name: Upload artifacts from /bin
-              uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
         with:
           name: binaries
           path: bin/
-

.github/workflows/mirror.yaml

@@ -0,0 +1,132 @@
+name: Mirror & Sign (Docker Hub to GHCR)
+
+on:
+  workflow_dispatch: {}
+
+permissions:
+  contents: read
+  packages: write
+  id-token: write   # for keyless OIDC
+
+env:
+  SOURCE_IMAGE: docker.io/fosrl/gerbil
+  DEST_IMAGE: ghcr.io/${{ github.repository_owner }}/${{ github.event.repository.name }}
+
+jobs:
+  mirror-and-dual-sign:
+    runs-on: amd64-runner
+    steps:
+      - name: Install skopeo + jq
+        run: |
+          sudo apt-get update -y
+          sudo apt-get install -y skopeo jq
+          skopeo --version
+
+      - name: Install cosign
+        uses: sigstore/cosign-installer@faadad0cce49287aee09b3a48701e75088a2c6ad # v4.0.0
+
+      - name: Input check
+        run: |
+          test -n "${SOURCE_IMAGE}" || (echo "SOURCE_IMAGE is empty" && exit 1)
+          echo "Source : ${SOURCE_IMAGE}"
+          echo "Target : ${DEST_IMAGE}"
+
+      # Auth for skopeo (containers-auth)
+      - name: Skopeo login to GHCR
+        run: |
+          skopeo login ghcr.io -u "${{ github.actor }}" -p "${{ secrets.GITHUB_TOKEN }}"
+
+      # Auth for cosign (docker-config)
+      - name: Docker login to GHCR (for cosign)
+        run: |
+          echo "${{ secrets.GITHUB_TOKEN }}" | docker login ghcr.io -u "${{ github.actor }}" --password-stdin
+
+      - name: List source tags
+        run: |
+          set -euo pipefail
+          skopeo list-tags --retry-times 3 docker://"${SOURCE_IMAGE}" \
+            | jq -r '.Tags[]' | sort -u > src-tags.txt
+          echo "Found source tags: $(wc -l < src-tags.txt)"
+          head -n 20 src-tags.txt || true
+
+      - name: List destination tags (skip existing)
+        run: |
+          set -euo pipefail
+          if skopeo list-tags --retry-times 3 docker://"${DEST_IMAGE}" >/tmp/dst.json 2>/dev/null; then
+            jq -r '.Tags[]' /tmp/dst.json | sort -u > dst-tags.txt
+          else
+            : > dst-tags.txt
+          fi
+          echo "Existing destination tags: $(wc -l < dst-tags.txt)"
+
+      - name: Mirror, dual-sign, and verify
+        env:
+          # keyless
+          COSIGN_YES: "true"
+          # key-based
+          COSIGN_PRIVATE_KEY: ${{ secrets.COSIGN_PRIVATE_KEY }}
+          COSIGN_PASSWORD:    ${{ secrets.COSIGN_PASSWORD }}
+          # verify
+          COSIGN_PUBLIC_KEY:  ${{ secrets.COSIGN_PUBLIC_KEY }}
+        run: |
+          set -euo pipefail
+          copied=0; skipped=0; v_ok=0; errs=0
+
+          issuer="https://token.actions.githubusercontent.com"
+          id_regex="^https://github.com/${{ github.repository }}/.+"
+
+          while read -r tag; do
+            [ -z "$tag" ] && continue
+
+            if grep -Fxq "$tag" dst-tags.txt; then
+              echo "::notice ::Skip (exists) ${DEST_IMAGE}:${tag}"
+              skipped=$((skipped+1))
+              continue
+            fi
+
+            echo "==> Copy ${SOURCE_IMAGE}:${tag} → ${DEST_IMAGE}:${tag}"
+            if ! skopeo copy --all --retry-times 3 \
+                 docker://"${SOURCE_IMAGE}:${tag}" docker://"${DEST_IMAGE}:${tag}"; then
+              echo "::warning title=Copy failed::${SOURCE_IMAGE}:${tag}"
+              errs=$((errs+1)); continue
+            fi
+            copied=$((copied+1))
+
+            digest="$(skopeo inspect --retry-times 3 docker://"${DEST_IMAGE}:${tag}" | jq -r '.Digest')"
+            ref="${DEST_IMAGE}@${digest}"
+
+            echo "==> cosign sign (keyless) --recursive ${ref}"
+            if ! cosign sign --recursive "${ref}"; then
+              echo "::warning title=Keyless sign failed::${ref}"
+              errs=$((errs+1))
+            fi
+
+            echo "==> cosign sign (key) --recursive ${ref}"
+            if ! cosign sign --key env://COSIGN_PRIVATE_KEY --recursive "${ref}"; then
+              echo "::warning title=Key sign failed::${ref}"
+              errs=$((errs+1))
+            fi
+
+            echo "==> cosign verify (public key) ${ref}"
+            if ! cosign verify --key env://COSIGN_PUBLIC_KEY "${ref}" -o text; then
+              echo "::warning title=Verify(pubkey) failed::${ref}"
+              errs=$((errs+1))
+            fi
+
+            echo "==> cosign verify (keyless policy) ${ref}"
+            if ! cosign verify \
+                 --certificate-oidc-issuer "${issuer}" \
+                 --certificate-identity-regexp "${id_regex}" \
+                 "${ref}" -o text; then
+              echo "::warning title=Verify(keyless) failed::${ref}"
+              errs=$((errs+1))
+            else
+              v_ok=$((v_ok+1))
+            fi
+          done < src-tags.txt
+
+          echo "---- Summary ----"
+          echo "Copied        : $copied"
+          echo "Skipped       : $skipped"
+          echo "Verified OK   : $v_ok"
+          echo "Errors        : $errs"

.github/workflows/test.yml

@@ -1,5 +1,8 @@
 name: Run Tests
 
+permissions:
+  contents: read
+
 on:
   pull_request:
     branches:
@@ -8,15 +11,15 @@ on:
 
 jobs:
   test:
-    runs-on: ubuntu-latest
+    runs-on: amd64-runner
 
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
 
       - name: Set up Go
-        uses: actions/setup-go@v6
+        uses: actions/setup-go@7a3fe6cf4cb3a834922a1244abfce67bcef6a0c5 # v6.2.0
         with:
-          go-version: '1.25'
+          go-version: 1.25
 
       - name: Build go
         run: go build

CONTRIBUTING.md

@@ -4,11 +4,7 @@ Contributions are welcome!
 
 Please see the contribution and local development guide on the docs page before getting started:
 
-https://docs.fossorial.io/development
+https://docs.pangolin.net/development/contributing
-
-For ideas about what features to work on and our future plans, please see the roadmap:
-
-https://docs.fossorial.io/roadmap
 
 ### Licensing Considerations
 

Dockerfile

@@ -16,7 +16,7 @@ COPY . .
 RUN CGO_ENABLED=0 GOOS=linux go build -o /gerbil
 
 # Start a new stage from scratch
-FROM alpine:3.22 AS runner
+FROM alpine:3.23 AS runner
 
 RUN apk add --no-cache iptables iproute2
 

README.md

@@ -6,7 +6,7 @@ Gerbil is a simple [WireGuard](https://www.wireguard.com/) interface management
 
 Gerbil works with Pangolin, Newt, and Olm as part of the larger system. See documentation below:
 
--   [Full Documentation](https://docs.fossorial.io)
+-   [Full Documentation](https://docs.pangolin.net)
 
 ## Key Functions
 
@@ -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
@@ -96,7 +94,7 @@ services:
     container_name: gerbil
     restart: unless-stopped
     command:
-      - --reachableAt=http://gerbil:3003
+      - --reachableAt=http://gerbil:3004
       - --generateAndSaveKeyTo=/var/config/key
       - --remoteConfig=http://pangolin:3001/api/v1/
     volumes:

SECURITY.md

@@ -3,7 +3,7 @@
 If you discover a security vulnerability, please follow the steps below to responsibly disclose it to us:
 
 1. **Do not create a public GitHub issue or discussion post.** This could put the security of other users at risk.
-2. Send a detailed report to [security@fossorial.io](mailto:security@fossorial.io) or send a **private** message to a maintainer on [Discord](https://discord.gg/HCJR8Xhme4). Include:
+2. Send a detailed report to [security@pangolin.net](mailto:security@pangolin.net) or send a **private** message to a maintainer on [Discord](https://discord.gg/HCJR8Xhme4). Include:
 
 - Description and location of the vulnerability.
 - Potential impact of the vulnerability.

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.43.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
 )
 
@@ -16,8 +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.45.0 // 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.37.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.43.0 h1:dduJYIi3A3KOfdGOHX8AVZ/jGiyPa3IbBozJ5kNuE04=
+golang.org/x/crypto v0.46.0 h1:cKRW/pmt1pKAfetfu+RCEvjvZkA9RimPbh7bhFjGVBU=
-golang.org/x/crypto v0.43.0/go.mod h1:BFbav4mRNlXJL4wNeejLpWxB7wMbc79PdRGhWKncxR0=
+golang.org/x/crypto v0.46.0/go.mod h1:Evb/oLKmMraqjZ2iQTwDwvCtJkczlDuTmdJXoZVzqU0=
-golang.org/x/net v0.45.0 h1:RLBg5JKixCy82FtLJpeNlVM0nrSqpCRYzVU1n8kj0tM=
+golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY=
-golang.org/x/net v0.45.0/go.mod h1:ECOoLqd5U3Lhyeyo/QDCEVQ4sNgYsqvCZ722XogGieY=
+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.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ=
+golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk=
-golang.org/x/sys v0.37.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

@@ -2,15 +2,21 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
+	"errors"
 	"flag"
 	"fmt"
 	"io"
+	"log"
 	"net"
 	"net/http"
+	_ "net/http/pprof"
 	"os"
 	"os/exec"
 	"os/signal"
+	"runtime"
+	"runtime/pprof"
 	"strconv"
 	"strings"
 	"sync"
@@ -21,6 +27,7 @@ import (
 	"github.com/fosrl/gerbil/proxy"
 	"github.com/fosrl/gerbil/relay"
 	"github.com/vishvananda/netlink"
+	"golang.org/x/sync/errgroup"
 	"golang.zx2c4.com/wireguard/wgctrl"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 )
@@ -35,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"`
 }
@@ -108,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
@@ -141,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")
@@ -212,11 +224,22 @@ 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()
 	logger.GetLogger().SetLevel(parseLogLevel(logLevel))
 
+	// Base context for the application; cancel on SIGINT/SIGTERM
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stop()
+
 	// try to parse as http://host:port and set the listenAddr to the :port from this reachableAt.
 	if reachableAt != "" && listenAddr == "" {
 		if strings.HasPrefix(reachableAt, "http://") || strings.HasPrefix(reachableAt, "https://") {
@@ -324,19 +347,16 @@ func main() {
 	// Ensure the WireGuard peers exist
 	ensureWireguardPeers(wgconfig.Peers)
 
-	go periodicBandwidthCheck(remoteConfigURL + "/gerbil/receive-bandwidth")
+	// Child error group derived from base context
+	group, groupCtx := errgroup.WithContext(ctx)
 
-	// Start the UDP proxy server
+	// Periodic bandwidth reporting
-	proxyRelay = relay.NewUDPProxyServer(":21820", remoteConfigURL, key, reachableAt)
+	group.Go(func() error {
-	err = proxyRelay.Start()
+		return periodicBandwidthCheck(groupCtx, remoteConfigURL+"/gerbil/receive-bandwidth")
-	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
+	// Parse local overrides and trusted upstreams early so that both the relay
-	// 		 SO YOU DON'T NEED TO SET THIS SEPARATELY
+	// and the SNI proxy share the same configuration values.
-	// Parse local overrides
 	var localOverrides []string
 	if localOverridesStr != "" {
 		localOverrides = strings.Split(localOverridesStr, ",")
@@ -355,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)
@@ -369,20 +404,42 @@ 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)
 
-	// Run HTTP server in a goroutine
+	// HTTP server with graceful shutdown on context cancel
-	go func() {
+	server := &http.Server{
-		if err := http.ListenAndServe(listenAddr, nil); err != nil {
+		Addr:    listenAddr,
-			logger.Error("HTTP server failed: %v", err)
+		Handler: nil,
 	}
-	}()
+	group.Go(func() error {
+		// http.ErrServerClosed is returned on graceful shutdown; not an error for us
+		if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+			return err
+		}
+		return nil
+	})
+	group.Go(func() error {
+		<-groupCtx.Done()
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		_ = server.Shutdown(shutdownCtx)
+		// Stop background components as the context is canceled
+		if proxySNI != nil {
+			_ = proxySNI.Stop()
+		}
+		if proxyRelay != nil {
+			proxyRelay.Stop()
+		}
+		return nil
+	})
 
-	// Keep the main goroutine running
+	// Wait for all goroutines to finish
-	sigCh := make(chan os.Signal, 1)
+	if err := group.Wait(); err != nil && !errors.Is(err, context.Canceled) {
-	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
+		logger.Error("Service exited with error: %v", err)
-	<-sigCh
+	} else if errors.Is(err, context.Canceled) {
-	logger.Info("Shutting down servers...")
+		logger.Info("Context cancelled, shutting down")
+	}
 }
 
 func loadRemoteConfig(url string, key wgtypes.Key, reachableAt string) (WgConfig, error) {
@@ -509,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
@@ -639,7 +700,7 @@ func ensureMSSClamping() error {
 		if out, err := addCmd.CombinedOutput(); err != nil {
 			errMsg := fmt.Sprintf("Failed to add MSS clamping rule for chain %s: %v (output: %s)",
 				chain, err, string(out))
-			logger.Error(errMsg)
+			logger.Error("%s", errMsg)
 			errors = append(errors, fmt.Errorf("%s", errMsg))
 			continue
 		}
@@ -656,7 +717,7 @@ func ensureMSSClamping() error {
 		if out, err := checkCmd.CombinedOutput(); err != nil {
 			errMsg := fmt.Sprintf("Rule verification failed for chain %s: %v (output: %s)",
 				chain, err, string(out))
-			logger.Error(errMsg)
+			logger.Error("%s", errMsg)
 			errors = append(errors, fmt.Errorf("%s", errMsg))
 			continue
 		}
@@ -677,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:
@@ -688,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 {
@@ -720,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{
@@ -746,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 {
@@ -790,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 {
@@ -805,7 +999,6 @@ func removePeerInternal(publicKey string) error {
 			}
 		}
 	}
-	}
 
 	peerConfig := wgtypes.PeerConfig{
 		PublicKey: pubKey,
@@ -820,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 {
@@ -977,14 +1179,19 @@ func handleUpdateLocalSNIs(w http.ResponseWriter, r *http.Request) {
 	})
 }
 
-func periodicBandwidthCheck(endpoint string) {
+func periodicBandwidthCheck(ctx context.Context, endpoint string) error {
 	ticker := time.NewTicker(10 * time.Second)
 	defer ticker.Stop()
 
-	for range ticker.C {
+	for {
+		select {
+		case <-ticker.C:
 			if err := reportPeerBandwidth(endpoint); err != nil {
 				logger.Info("Failed to report peer bandwidth: %v", err)
 			}
+		case <-ctx.Done():
+			return ctx.Err()
+		}
 	}
 }
 
@@ -1003,8 +1210,13 @@ func calculatePeerBandwidth() ([]PeerBandwidth, error) {
 	mu.Lock()
 	defer mu.Unlock()
 
+	// Track the set of peers currently present on the device to prune stale readings efficiently
+	currentPeerKeys := make(map[string]struct{}, len(device.Peers))
+
 	for _, peer := range device.Peers {
 		publicKey := peer.PublicKey.String()
+		currentPeerKeys[publicKey] = struct{}{}
+
 		currentReading := PeerReading{
 			BytesReceived:    peer.ReceiveBytes,
 			BytesTransmitted: peer.TransmitBytes,
@@ -1061,14 +1273,7 @@ func calculatePeerBandwidth() ([]PeerBandwidth, error) {
 
 	// Clean up old peers
 	for publicKey := range lastReadings {
-		found := false
+		if _, exists := currentPeerKeys[publicKey]; !exists {
-		for _, peer := range device.Peers {
-			if peer.PublicKey.String() == publicKey {
-				found = true
-				break
-			}
-		}
-		if !found {
 			delete(lastReadings, publicKey)
 		}
 	}
@@ -1124,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

@@ -2,16 +2,20 @@ package relay
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"encoding/json"
 	"fmt"
 	"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"
@@ -27,6 +31,7 @@ type HolePunchMessage struct {
 	OlmID     string `json:"olmId"`
 	NewtID    string `json:"newtId"`
 	Token     string `json:"token"`
+	PublicKey string `json:"publicKey"`
 }
 
 type ClientEndpoint struct {
@@ -37,7 +42,8 @@ type ClientEndpoint struct {
 	Port              int    `json:"port"`
 	Timestamp         int64  `json:"timestamp"`
 	ReachableAt       string `json:"reachableAt"`
-	PublicKey   string `json:"publicKey"`
+	ExitNodePublicKey string `json:"exitNodePublicKey"`
+	ClientPublicKey   string `json:"publicKey"`
 }
 
 // Updated to support multiple destination peers
@@ -58,12 +64,41 @@ type DestinationConn struct {
 
 // Type for storing WireGuard handshake information
 type WireGuardSession struct {
+	mu            sync.RWMutex
 	ReceiverIndex uint32
 	SenderIndex   uint32
 	DestAddr      *net.UDPAddr
 	LastSeen      time.Time
 }
 
+// GetSenderIndex returns the SenderIndex in a thread-safe manner
+func (s *WireGuardSession) GetSenderIndex() uint32 {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.SenderIndex
+}
+
+// GetDestAddr returns the DestAddr in a thread-safe manner
+func (s *WireGuardSession) GetDestAddr() *net.UDPAddr {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.DestAddr
+}
+
+// GetLastSeen returns the LastSeen timestamp in a thread-safe manner
+func (s *WireGuardSession) GetLastSeen() time.Time {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.LastSeen
+}
+
+// UpdateLastSeen updates the LastSeen timestamp in a thread-safe manner
+func (s *WireGuardSession) UpdateLastSeen() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.LastSeen = time.Now()
+}
+
 // Type for tracking bidirectional communication patterns to rebuild sessions
 type CommunicationPattern struct {
 	FromClient     *net.UDPAddr // The client address
@@ -86,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
@@ -112,6 +154,8 @@ type UDPProxyServer struct {
 	connections   sync.Map // map[string]*DestinationConn where key is destination "ip:port"
 	privateKey    wgtypes.Key
 	packetChan    chan Packet
+	ctx           context.Context
+	cancel        context.CancelFunc
 
 	// Session tracking for WireGuard peers
 	// Key format: "senderIndex:receiverIndex"
@@ -119,18 +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.
+// NewUDPProxyServer initializes the server with a buffered packet channel and derived context.
-func NewUDPProxyServer(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,
 	}
 }
 
@@ -152,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()
 	}
@@ -173,22 +261,59 @@ 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
 }
 
 func (s *UDPProxyServer) Stop() {
-	s.conn.Close()
+	// Signal all background goroutines to stop
+	if s.cancel != nil {
+		s.cancel()
+	}
+	// Close listener to unblock reads
+	if s.conn != nil {
+		_ = s.conn.Close()
+	}
+	// Close all downstream UDP connections
+	s.connections.Range(func(key, value interface{}) bool {
+		if dc, ok := value.(*DestinationConn); ok && dc.conn != nil {
+			_ = dc.conn.Close()
+		}
+		return true
+	})
+	// Close packet channel to stop workers
+	select {
+	case <-s.ctx.Done():
+	default:
+	}
+	close(s.packetChan)
 }
 
 // readPackets continuously reads from the UDP socket and pushes packets into the channel.
 func (s *UDPProxyServer) readPackets() {
 	for {
+		// Exit promptly if context is canceled
+		select {
+		case <-s.ctx.Done():
+			return
+		default:
+		}
 		buf := bufferPool.Get().([]byte)
 		n, remoteAddr, err := s.conn.ReadFromUDP(buf)
 		if err != nil {
+			// If we're shutting down, exit
+			select {
+			case <-s.ctx.Done():
+				bufferPool.Put(buf[:1500])
+				return
+			default:
 				logger.Error("Error reading UDP packet: %v", err)
+				bufferPool.Put(buf[:1500])
 				continue
 			}
+		}
 		s.packetChan <- Packet{data: buf[:n], remoteAddr: remoteAddr, n: n}
 	}
 }
@@ -196,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])
@@ -220,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
@@ -239,13 +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,
-				PublicKey:   s.privateKey.PublicKey().String(),
+				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
 		}
@@ -344,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 {
@@ -378,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
@@ -392,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)
 			}
 		}
 
@@ -414,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
@@ -442,13 +661,11 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD
 		// First check for existing sessions to see if we know where to send this packet
 		s.wgSessions.Range(func(k, v interface{}) bool {
 			session := v.(*WireGuardSession)
-			if session.SenderIndex == receiverIndex {
+			// Check if session matches (read lock for check)
-				// Found matching session
+			if session.GetSenderIndex() == receiverIndex {
-				destAddr = session.DestAddr
+				// Found matching session - get dest addr and update last seen
-
+				destAddr = session.GetDestAddr()
-				// Update last seen time
+				session.UpdateLastSeen()
-				session.LastSeen = time.Now()
-				s.wgSessions.Store(k, session)
 				return false // stop iteration
 			}
 			return true // continue iteration
@@ -473,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
@@ -501,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
@@ -588,7 +805,10 @@ func (s *UDPProxyServer) handleResponses(conn *net.UDPConn, destAddr *net.UDPAdd
 // Add a cleanup method to periodically remove idle connections
 func (s *UDPProxyServer) cleanupIdleConnections() {
 	ticker := time.NewTicker(5 * time.Minute)
-	for range ticker.C {
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
 			now := time.Now()
 			s.connections.Range(func(key, value interface{}) bool {
 				destConn := value.(*DestinationConn)
@@ -598,29 +818,43 @@ func (s *UDPProxyServer) cleanupIdleConnections() {
 				}
 				return true
 			})
+		case <-s.ctx.Done():
+			return
+		}
 	}
 }
 
 // New method to periodically remove idle sessions
 func (s *UDPProxyServer) cleanupIdleSessions() {
 	ticker := time.NewTicker(5 * time.Minute)
-	for range ticker.C {
+
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
 			now := time.Now()
 			s.wgSessions.Range(func(key, value interface{}) bool {
 				session := value.(*WireGuardSession)
-			if now.Sub(session.LastSeen) > 15*time.Minute {
+				// Use thread-safe method to read LastSeen
+				if now.Sub(session.GetLastSeen()) > 15*time.Minute {
 					s.wgSessions.Delete(key)
 					logger.Debug("Removed idle session: %s", key)
 				}
 				return true
 			})
+		case <-s.ctx.Done():
+			return
+		}
 	}
 }
 
 // New method to periodically remove idle proxy mappings
 func (s *UDPProxyServer) cleanupIdleProxyMappings() {
 	ticker := time.NewTicker(10 * time.Minute)
-	for range ticker.C {
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
 			now := time.Now()
 			s.proxyMappings.Range(func(key, value interface{}) bool {
 				mapping := value.(ProxyMapping)
@@ -631,6 +865,9 @@ func (s *UDPProxyServer) cleanupIdleProxyMappings() {
 				}
 				return true
 			})
+		case <-s.ctx.Done():
+			return
+		}
 	}
 }
 
@@ -735,8 +972,9 @@ func (s *UDPProxyServer) clearSessionsForIP(ip string) {
 		keyStr := key.(string)
 		session := value.(*WireGuardSession)
 
-		// Check if the session's destination address contains the WG IP
+		// Check if the session's destination address contains the WG IP (thread-safe)
-		if session.DestAddr != nil && session.DestAddr.IP.String() == ip {
+		destAddr := session.GetDestAddr()
+		if destAddr != nil && destAddr.IP.String() == ip {
 			keysToDelete = append(keysToDelete, keyStr)
 			logger.Debug("Marking session for deletion for WG IP %s: %s", ip, keyStr)
 		}
@@ -748,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
@@ -926,14 +1164,12 @@ func (s *UDPProxyServer) tryRebuildSession(pattern *CommunicationPattern) {
 
 		// Check if we already have this session
 		if _, exists := s.wgSessions.Load(sessionKey); !exists {
-			session := &WireGuardSession{
+			s.wgSessions.Store(sessionKey, &WireGuardSession{
 				ReceiverIndex: pattern.DestIndex,
 				SenderIndex:   pattern.ClientIndex,
 				DestAddr:      pattern.ToDestination,
 				LastSeen:      time.Now(),
-			}
+			})
-
-			s.wgSessions.Store(sessionKey, session)
 			logger.Info("Rebuilt WireGuard session from communication pattern: %s -> %s (packets: %d)",
 				sessionKey, pattern.ToDestination.String(), pattern.PacketCount)
 		}
@@ -941,9 +1177,36 @@ 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)
-	for range ticker.C {
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
 			now := time.Now()
 			s.commPatterns.Range(func(key, value interface{}) bool {
 				pattern := value.(*CommunicationPattern)
@@ -961,5 +1224,8 @@ func (s *UDPProxyServer) cleanupIdleCommunicationPatterns() {
 				}
 				return true
 			})
+		case <-s.ctx.Done():
+			return
+		}
 	}
 }