Add pp to relay

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

.github/DISCUSSION_TEMPLATE/feature-requests.yml

@@ -0,0 +1,47 @@
+body:
+    - type: textarea
+      attributes:
+          label: Summary
+          description: A clear and concise summary of the requested feature.
+      validations:
+          required: true
+
+    - type: textarea
+      attributes:
+          label: Motivation
+          description: |
+              Why is this feature important?
+              Explain the problem this feature would solve or what use case it would enable.
+      validations:
+          required: true
+
+    - type: textarea
+      attributes:
+          label: Proposed Solution
+          description: |
+              How would you like to see this feature implemented?
+              Provide as much detail as possible about the desired behavior, configuration, or changes.
+      validations:
+          required: true
+
+    - type: textarea
+      attributes:
+          label: Alternatives Considered
+          description: Describe any alternative solutions or workarounds you've thought about.
+      validations:
+          required: false
+
+    - type: textarea
+      attributes:
+          label: Additional Context
+          description: Add any other context, mockups, or screenshots about the feature request here.
+      validations:
+          required: false
+
+    - type: markdown
+      attributes:
+          value: |
+              Before submitting, please:
+              - Check if there is an existing issue for this feature.
+              - Clearly explain the benefit and use case.
+              - Be as specific as possible to help contributors evaluate and implement.

.github/ISSUE_TEMPLATE/1.bug_report.yml

@@ -0,0 +1,51 @@
+name: Bug Report
+description: Create a bug report
+labels: []
+body:
+    - type: textarea
+      attributes:
+          label: Describe the Bug
+          description: A clear and concise description of what the bug is.
+      validations:
+          required: true
+
+    - type: textarea
+      attributes:
+          label: Environment
+          description: Please fill out the relevant details below for your environment.
+          value: |
+              - OS Type & Version: (e.g., Ubuntu 22.04)
+              - Pangolin Version:
+              - Gerbil Version:
+              - Traefik Version:
+              - Newt Version:
+              - Olm Version: (if applicable)
+      validations:
+          required: true
+
+    - type: textarea
+      attributes:
+          label: To Reproduce
+          description: |
+              Steps to reproduce the behavior, please provide a clear description of how to reproduce the issue, based on the linked minimal reproduction. Screenshots can be provided in the issue body below.
+
+              If using code blocks, make sure syntax highlighting is correct and double-check that the rendered preview is not broken.
+      validations:
+          required: true
+
+    - type: textarea
+      attributes:
+          label: Expected Behavior
+          description: A clear and concise description of what you expected to happen.
+      validations:
+          required: true
+
+    - type: markdown
+      attributes:
+          value: |
+              Before posting the issue go through the steps you've written down to make sure the steps provided are detailed and clear.
+
+    - type: markdown
+      attributes:
+          value: |
+              Contributors should be able to follow the steps provided in order to reproduce the bug.

.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1,8 @@
+blank_issues_enabled: false
+contact_links:
+    - name: Need help or have questions?
+      url: https://github.com/orgs/fosrl/discussions
+      about: Ask questions, get help, and discuss with other community members
+    - name: Request a Feature
+      url: https://github.com/orgs/fosrl/discussions/new?category=feature-requests
+      about: Feature requests should be opened as discussions so others can upvote and comment

.github/dependabot.yml

@@ -33,3 +33,8 @@ updates:
       minor-updates:
         update-types:
           - "minor"
+
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"

.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:
+  release:
-        name: Build and 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:
+    steps:
-            - name: Checkout code
+      - name: Checkout code
-              uses: actions/checkout@v3
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
 
-            - name: Set up QEMU
+      - name: Set up QEMU
-              uses: docker/setup-qemu-action@v3
+        uses: docker/setup-qemu-action@c7c53464625b32c7a7e944ae62b3e17d2b600130 # v3.7.0
 
-            - name: Set up Docker Buildx
+      - name: Set up Docker Buildx
-              uses: docker/setup-buildx-action@v2
+        uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
 
-            - name: Log in to Docker Hub
+      - name: Log in to Docker Hub
-              uses: docker/login-action@v2
+        uses: docker/login-action@5e57cd118135c172c3672efd75eb46360885c0ef # v3.6.0
-              with:
+        with:
-                  username: ${{ secrets.DOCKER_HUB_USERNAME }}
+          registry: docker.io
-                  password: ${{ secrets.DOCKER_HUB_ACCESS_TOKEN }}
+          username: ${{ secrets.DOCKER_HUB_USERNAME }}
+          password: ${{ secrets.DOCKER_HUB_ACCESS_TOKEN }}
 
-            - name: Extract tag name
+      - name: Extract tag name
-              id: get-tag
+        id: get-tag
-              run: echo "TAG=${GITHUB_REF#refs/tags/}" >> $GITHUB_ENV
+        run: echo "TAG=${GITHUB_REF#refs/tags/}" >> $GITHUB_ENV
+        shell: bash
 
-            - name: Install Go
+      - name: Install Go
-              uses: actions/setup-go@v4
+        uses: actions/setup-go@7a3fe6cf4cb3a834922a1244abfce67bcef6a0c5 # v6.2.0
-              with:
+        with:
-                  go-version: 1.23.1
+          go-version: 1.25
 
-            - name: Build and push Docker images
+      - name: Update version in main.go
-              run: |
+        run: |
-                  TAG=${{ env.TAG }}
+          TAG=${{ env.TAG }}
-                  make docker-build-release tag=$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 binaries
+      - name: Build and push Docker images (Docker Hub)
-              run: |
+        run: |
-                  make go-build-release
+          TAG=${{ env.TAG }}
+          make docker-build-release tag=$TAG
+          echo "Built & pushed to: ${{ env.DOCKERHUB_IMAGE }}:${TAG}"
+        shell: bash
 
-            - name: Upload artifacts from /bin
+      - name: Login in to GHCR
-              uses: actions/upload-artifact@v4
+        uses: docker/login-action@5e57cd118135c172c3672efd75eb46360885c0ef # v3.6.0
-              with:
+        with:
-                  name: binaries
+          registry: ghcr.io
-                  path: bin/
+          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@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@v4
+      - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
 
       - name: Set up Go
-        uses: actions/setup-go@v4
+        uses: actions/setup-go@7a3fe6cf4cb3a834922a1244abfce67bcef6a0c5 # v6.2.0
         with:
-          go-version: '1.23'
+          go-version: 1.25
 
       - name: Build go
         run: go build

.go-version

@@ -0,0 +1 @@
+1.25

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

@@ -1,4 +1,4 @@
-FROM golang:1.24.3-alpine AS builder
+FROM golang:1.25-alpine AS builder
 
 # Set the working directory inside the container
 WORKDIR /app
@@ -16,18 +16,13 @@ COPY . .
 RUN CGO_ENABLED=0 GOOS=linux go build -o /gerbil
 
 # Start a new stage from scratch
-FROM ubuntu:24.04 AS runner
+FROM alpine:3.23 AS runner
 
-RUN apt-get update && apt-get install -y iptables iproute2 && rm -rf /var/lib/apt/lists/*
+RUN apk add --no-cache iptables iproute2
 
-# Copy the pre-built binary file from the previous stage and the entrypoint script
 COPY --from=builder /gerbil /usr/local/bin/
 COPY entrypoint.sh /
 
 RUN chmod +x /entrypoint.sh
-
-# Copy the entrypoint script
 ENTRYPOINT ["/entrypoint.sh"]
-
-# Command to run the executable
 CMD ["gerbil"]

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
 
@@ -20,34 +20,51 @@ Gerbil will create the peers defined in the config on the WireGuard interface. T
 
 ### Report Bandwidth
 
-Bytes transmitted in and out of each peer are collected every 10 seconds, and incremental usage is reported via the "reportBandwidthTo" endpoint. This can be used to track data usage of each peer on the remote server.
+Bytes transmitted in and out of each peer are collected every 10 seconds, and incremental usage is reported via the api endpoint. This can be used to track data usage of each peer on the remote server.
 
 ### Handle client relaying
 
 Gerbil listens on port 21820 for incoming UDP hole punch packets to orchestrate NAT hole punching between olm and newt clients. Additionally, it handles relaying data through the gerbil server down to the newt. This is accomplished by scanning each packet for headers and handling them appropriately.
 
+### SNI Proxy
+
+Gerbil includes an SNI (Server Name Indication) proxy that enables intelligent routing of HTTPS traffic between Pangolin nodes. When a TLS connection comes in, the proxy extracts the hostname from the SNI extension and queries Pangolin to determine the correct routing destination. This allows seamless routing of web traffic through the WireGuard mesh network:
+
+- If the hostname is configured for local handling (via local overrides or local SNIs), traffic is routed to the local proxy
+- Otherwise, the proxy queries Pangolin's routing API to determine which node should handle the traffic
+- Supports caching of routing decisions to improve performance
+- Handles connection pooling and graceful shutdown
+- Optional PROXY protocol v1 support to preserve original client IP addresses when forwarding to downstream proxies (HAProxy, Nginx, etc.)
+
+The PROXY protocol allows downstream proxies to know the real client IP address instead of seeing the SNI proxy's IP. When enabled with `--proxy-protocol`, the SNI proxy will prepend a PROXY protocol header to each connection containing the original client's IP and port information.
+
+In single node (self hosted) Pangolin deployments this can be bypassed by using port 443:443 to route to Traefik instead of the SNI proxy at 8443.
+
 ## 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: `:3003`
+- `listen` (optional): Port to listen on for HTTP server. Default: `:3004`
 - `log-level` (optional): The log level to use (DEBUG, INFO, WARN, ERROR, FATAL). Default: `INFO`
 - `mtu` (optional): MTU of the WireGuard interface. Default: `1280`
 - `notify` (optional): URL to notify on peer changes
+- `sni-port` (optional): Port for the SNI proxy to listen on. Default: `8443`
+- `local-proxy` (optional): Address for local proxy when routing local traffic. Default: `localhost`
+- `local-proxy-port` (optional): Port for local proxy when routing local traffic. Default: `443`
+- `local-overrides` (optional): Comma-separated list of domain names that should always be routed to the local proxy
+- `proxy-protocol` (optional): Enable PROXY protocol v1 for preserving client IP addresses when forwarding to downstream proxies. Default: `false`
 
 ## Environment Variables
 
 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
@@ -55,15 +72,19 @@ All CLI arguments can also be provided via environment variables:
 - `LOG_LEVEL`: Log level (DEBUG, INFO, WARN, ERROR, FATAL)
 - `MTU`: MTU of the WireGuard interface
 - `NOTIFY_URL`: URL to notify on peer changes
+- `SNI_PORT`: Port for the SNI proxy to listen on
+- `LOCAL_PROXY`: Address for local proxy when routing local traffic
+- `LOCAL_PROXY_PORT`: Port for local proxy when routing local traffic
+- `LOCAL_OVERRIDES`: Comma-separated list of domain names that should always be routed to the local proxy
+- `PROXY_PROTOCOL`: Enable PROXY protocol v1 for preserving client IP addresses (true/false)
 
 Example:
 
 ```bash
 ./gerbil \
---reachableAt=http://gerbil:3003 \
+--reachableAt=http://gerbil:3004 \
 --generateAndSaveKeyTo=/var/config/key \
---remoteConfig=http://pangolin:3001/api/v1/gerbil/get-config \
+--remoteConfig=http://pangolin:3001/api/v1/
---reportBandwidthTo=http://pangolin:3001/api/v1/gerbil/receive-bandwidth
 ```
 
 ```yaml
@@ -73,10 +94,9 @@ 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/gerbil/get-config
+      - --remoteConfig=http://pangolin:3001/api/v1/
-      - --reportBandwidthTo=http://pangolin:3001/api/v1/gerbil/receive-bandwidth
     volumes:
       - ./config/:/var/config
     cap_add:
@@ -85,6 +105,7 @@ services:
     ports:
       - 51820:51820/udp
       - 21820:21820/udp
+      - 443:8443/tcp  # SNI proxy port
 ```
 
 ## Build

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

@@ -1,11 +1,12 @@
 module github.com/fosrl/gerbil
 
-go 1.23.1
+go 1.25
-
-toolchain go1.23.2
 
 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/sync v0.1.0
 	golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6
 )
 
@@ -16,9 +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/crypto v0.36.0 // indirect
+	golang.org/x/net v0.47.0 // indirect
-	golang.org/x/net v0.38.0 // indirect
+	golang.org/x/sys v0.39.0 // indirect
-	golang.org/x/sync v0.1.0 // indirect
-	golang.org/x/sys v0.31.0 // indirect
 	golang.zx2c4.com/wireguard v0.0.0-20230325221338-052af4a8072b // indirect
 )

go.sum

@@ -10,20 +10,22 @@ github.com/mdlayher/socket v0.4.1 h1:eM9y2/jlbs1M615oshPQOHZzj6R6wMT7bX5NPiQvn2U
 github.com/mdlayher/socket v0.4.1/go.mod h1:cAqeGjoufqdxWkD7DkpyS+wcefOtmu5OQ8KuoJGIReA=
 github.com/mikioh/ipaddr v0.0.0-20190404000644-d465c8ab6721 h1:RlZweED6sbSArvlE924+mUcZuXKLBHA35U7LN621Bws=
 github.com/mikioh/ipaddr v0.0.0-20190404000644-d465c8ab6721/go.mod h1:Ickgr2WtCLZ2MDGd4Gr0geeCH5HybhRJbonOgQpvSxc=
+github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
+github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/vishvananda/netlink v1.3.1 h1:3AEMt62VKqz90r0tmNhog0r/PpWKmrEShJU0wJW6bV0=
 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.36.0 h1:AnAEvhDddvBdpY+uR+MyHmuZzzNqXSe/GvuDeob5L34=
+golang.org/x/crypto v0.46.0 h1:cKRW/pmt1pKAfetfu+RCEvjvZkA9RimPbh7bhFjGVBU=
-golang.org/x/crypto v0.36.0/go.mod h1:Y4J0ReaxCR1IMaabaSMugxJES1EpwhBHhv2bDHklZvc=
+golang.org/x/crypto v0.46.0/go.mod h1:Evb/oLKmMraqjZ2iQTwDwvCtJkczlDuTmdJXoZVzqU0=
-golang.org/x/net v0.38.0 h1:vRMAPTMaeGqVhG5QyLJHqNDwecKTomGeqbnfZyKlBI8=
+golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY=
-golang.org/x/net v0.38.0/go.mod h1:ivrbrMbzFq5J41QOQh0siUuly180yBYtLp+CKbEaFx8=
+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.31.0 h1:ioabZlmFYtWhL+TRYpcnNlLwhyxaM9kWTDEmfnprqik=
+golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk=
-golang.org/x/sys v0.31.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
+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"
@@ -18,26 +24,31 @@ import (
 	"time"
 
 	"github.com/fosrl/gerbil/logger"
+	"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"
 )
 
 var (
-	interfaceName string
+	interfaceName    string
-	listenAddr    string
+	listenAddr       string
-	mtuInt        int
+	mtuInt           int
-	lastReadings  = make(map[string]PeerReading)
+	lastReadings     = make(map[string]PeerReading)
-	mu            sync.Mutex
+	mu               sync.Mutex
-	wgMu          sync.Mutex // Protects WireGuard operations
+	wgMu             sync.Mutex // Protects WireGuard operations
-	notifyURL     string
+	notifyURL        string
-	proxyServer   *relay.UDPProxyServer
+	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"`
 }
@@ -106,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
@@ -115,6 +128,12 @@ func main() {
 		reachableAt          string
 		logLevel             string
 		mtu                  string
+		sniProxyPort         int
+		localProxyAddr       string
+		localProxyPort       int
+		localOverridesStr    string
+		trustedUpstreamsStr  string
+		proxyProtocol        bool
 	)
 
 	interfaceName = os.Getenv("INTERFACE")
@@ -127,6 +146,14 @@ func main() {
 	mtu = os.Getenv("MTU")
 	notifyURL = os.Getenv("NOTIFY_URL")
 
+	sniProxyPortStr := os.Getenv("SNI_PORT")
+	localProxyAddr = os.Getenv("LOCAL_PROXY")
+	localProxyPortStr := os.Getenv("LOCAL_PROXY_PORT")
+	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")
 	}
@@ -137,7 +164,7 @@ func main() {
 		flag.StringVar(&remoteConfigURL, "remoteConfig", "", "URL of the Pangolin server")
 	}
 	if listenAddr == "" {
-		flag.StringVar(&listenAddr, "listen", ":3003", "Address to listen on")
+		flag.StringVar(&listenAddr, "listen", "", "DEPRECATED (overridden by reachableAt): Address to listen on")
 	}
 	// DEPRECATED AND UNSED: reportBandwidthTo
 	// allow reportBandwidthTo to be passed but dont do anything with it just thow it away
@@ -147,9 +174,11 @@ func main() {
 	if generateAndSaveKeyTo == "" {
 		flag.StringVar(&generateAndSaveKeyTo, "generateAndSaveKeyTo", "", "Path to save generated private key")
 	}
+
 	if reachableAt == "" {
 		flag.StringVar(&reachableAt, "reachableAt", "", "Endpoint of the http server to tell remote config about")
 	}
+
 	if logLevel == "" {
 		flag.StringVar(&logLevel, "log-level", "INFO", "Log level (DEBUG, INFO, WARN, ERROR, FATAL)")
 	}
@@ -159,11 +188,74 @@ func main() {
 	if notifyURL == "" {
 		flag.StringVar(&notifyURL, "notify", "", "URL to notify on peer changes")
 	}
+
+	if sniProxyPortStr != "" {
+		if port, err := strconv.Atoi(sniProxyPortStr); err == nil {
+			sniProxyPort = port
+		}
+	}
+	if sniProxyPortStr == "" {
+		flag.IntVar(&sniProxyPort, "sni-port", 8443, "Port to listen on")
+	}
+
+	if localProxyAddr == "" {
+		flag.StringVar(&localProxyAddr, "local-proxy", "localhost", "Local proxy address")
+	}
+
+	if localProxyPortStr != "" {
+		if port, err := strconv.Atoi(localProxyPortStr); err == nil {
+			localProxyPort = port
+		}
+	}
+	if localProxyPortStr == "" {
+		flag.IntVar(&localProxyPort, "local-proxy-port", 443, "Local proxy port")
+	}
+	if localOverridesStr != "" {
+		flag.StringVar(&localOverridesStr, "local-overrides", "", "Comma-separated list of local overrides for SNI proxy")
+	}
+	if trustedUpstreamsStr == "" {
+		flag.StringVar(&trustedUpstreamsStr, "trusted-upstreams", "", "Comma-separated list of trusted upstream proxy domain names/IPs that can send PROXY protocol")
+	}
+
+	if proxyProtocolStr != "" {
+		proxyProtocol = strings.ToLower(proxyProtocolStr) == "true"
+	}
+	if proxyProtocolStr == "" {
+		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://") {
+			parts := strings.Split(reachableAt, ":")
+			if len(parts) == 3 {
+				port := parts[2]
+				if strings.Contains(port, "/") {
+					port = strings.Split(port, "/")[0]
+				}
+				listenAddr = ":" + port
+			}
+		}
+	} else if listenAddr == "" {
+		listenAddr = ":3003"
+	}
+
 	mtuInt, err = strconv.Atoi(mtu)
 	if err != nil {
 		logger.Fatal("Failed to parse MTU: %v", err)
@@ -255,34 +347,99 @@ 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
-	proxyServer = relay.NewUDPProxyServer(":21820", remoteConfigURL, key, reachableAt)
+	group.Go(func() error {
-	err = proxyServer.Start()
+		return periodicBandwidthCheck(groupCtx, remoteConfigURL+"/gerbil/receive-bandwidth")
+	})
+
+	// Parse local overrides and trusted upstreams early so that both the relay
+	// and the SNI proxy share the same configuration values.
+	var localOverrides []string
+	if localOverridesStr != "" {
+		localOverrides = strings.Split(localOverridesStr, ",")
+		for i, domain := range localOverrides {
+			localOverrides[i] = strings.TrimSpace(domain)
+		}
+		logger.Info("Local overrides configured: %v", localOverrides)
+	}
+
+	var trustedUpstreams []string
+	if trustedUpstreamsStr != "" {
+		trustedUpstreams = strings.Split(trustedUpstreamsStr, ",")
+		for i, upstream := range trustedUpstreams {
+			trustedUpstreams[i] = strings.TrimSpace(upstream)
+		}
+		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 proxyServer.Stop()
+	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)
+	}
+
+	if err := proxySNI.Start(); err != nil {
+		logger.Fatal("Failed to start proxy: %v", err)
+	}
 
 	// Set up HTTP server
 	http.HandleFunc("/peer", handlePeer)
 	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) {
@@ -409,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
@@ -539,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
 		}
@@ -556,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
 		}
@@ -577,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:
@@ -588,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 {
@@ -620,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{
@@ -646,10 +929,22 @@ func addPeerInternal(peer Peer) error {
 		return fmt.Errorf("failed to add peer: %v", err)
 	}
 
-	// Clear relay connections for the peer's WireGuard IPs
+	// Setup bandwidth limiting for each peer IP
-	if proxyServer != nil {
+	if doTrafficShaping {
+		logger.Debug("doTrafficShaping is true, setting up bandwidth limits for %d IPs", len(wgIPs))
 		for _, wgIP := range wgIPs {
-			proxyServer.OnPeerAdded(wgIP)
+			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 {
+			proxyRelay.OnPeerAdded(wgIP)
 		}
 	}
 
@@ -690,19 +985,17 @@ 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 proxyServer != nil {
+	device, err := wgClient.Device(interfaceName)
-		device, err := wgClient.Device(interfaceName)
+	if err == nil {
-		if err == nil {
+		for _, peer := range device.Peers {
-			for _, peer := range device.Peers {
+			if peer.PublicKey.String() == publicKey {
-				if peer.PublicKey.String() == publicKey {
+				// Extract WireGuard IPs from this peer's allowed IPs
-					// Extract WireGuard IPs from this peer's allowed IPs
+				for _, allowedIP := range peer.AllowedIPs {
-					for _, allowedIP := range peer.AllowedIPs {
+					wgIPs = append(wgIPs, allowedIP.IP.String())
-						wgIPs = append(wgIPs, allowedIP.IP.String())
-					}
-					break
 				}
+				break
 			}
 		}
 	}
@@ -720,10 +1013,19 @@ func removePeerInternal(publicKey string) error {
 		return fmt.Errorf("failed to remove peer: %v", err)
 	}
 
-	// Clear relay connections for the peer's WireGuard IPs
+	// Remove bandwidth limits for each peer IP
-	if proxyServer != nil {
+	if doTrafficShaping {
 		for _, wgIP := range wgIPs {
-			proxyServer.OnPeerRemoved(wgIP)
+			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 {
+			proxyRelay.OnPeerRemoved(wgIP)
 		}
 	}
 
@@ -760,13 +1062,13 @@ func handleUpdateProxyMapping(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Update the proxy mappings in the relay server
-	if proxyServer == nil {
+	if proxyRelay == nil {
 		logger.Error("Proxy server is not available")
 		http.Error(w, "Proxy server is not available", http.StatusInternalServerError)
 		return
 	}
 
-	updatedCount := proxyServer.UpdateDestinationInMappings(update.OldDestination, update.NewDestination)
+	updatedCount := proxyRelay.UpdateDestinationInMappings(update.OldDestination, update.NewDestination)
 
 	logger.Info("Updated %d proxy mappings: %s:%d -> %s:%d",
 		updatedCount,
@@ -830,13 +1132,13 @@ func handleUpdateDestinations(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Update the proxy mappings in the relay server
-	if proxyServer == nil {
+	if proxyRelay == nil {
 		logger.Error("Proxy server is not available")
 		http.Error(w, "Proxy server is not available", http.StatusInternalServerError)
 		return
 	}
 
-	proxyServer.UpdateProxyMapping(request.SourceIP, request.SourcePort, request.Destinations)
+	proxyRelay.UpdateProxyMapping(request.SourceIP, request.SourcePort, request.Destinations)
 
 	logger.Info("Updated proxy mapping for %s:%d with %d destinations",
 		request.SourceIP, request.SourcePort, len(request.Destinations))
@@ -851,13 +1153,44 @@ func handleUpdateDestinations(w http.ResponseWriter, r *http.Request) {
 	})
 }
 
-func periodicBandwidthCheck(endpoint string) {
+// UpdateLocalSNIsRequest represents the JSON payload for updating local SNIs
+type UpdateLocalSNIsRequest struct {
+	FullDomains []string `json:"fullDomains"`
+}
+
+func handleUpdateLocalSNIs(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodPost {
+		logger.Error("Invalid method: %s", r.Method)
+		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
+		return
+	}
+
+	var req UpdateLocalSNIsRequest
+	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+		http.Error(w, "Invalid JSON payload", http.StatusBadRequest)
+		return
+	}
+
+	proxySNI.UpdateLocalSNIs(req.FullDomains)
+
+	w.WriteHeader(http.StatusOK)
+	json.NewEncoder(w).Encode(map[string]interface{}{
+		"status": "Local SNIs updated successfully",
+	})
+}
+
+func periodicBandwidthCheck(ctx context.Context, endpoint string) error {
 	ticker := time.NewTicker(10 * time.Second)
 	defer ticker.Stop()
 
-	for range ticker.C {
+	for {
-		if err := reportPeerBandwidth(endpoint); err != nil {
+		select {
-			logger.Info("Failed to report peer bandwidth: %v", err)
+		case <-ticker.C:
+			if err := reportPeerBandwidth(endpoint); err != nil {
+				logger.Info("Failed to report peer bandwidth: %v", err)
+			}
+		case <-ctx.Done():
+			return ctx.Err()
 		}
 	}
 }
@@ -877,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,
@@ -935,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)
 		}
 	}
@@ -998,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

@@ -0,0 +1,594 @@
+package proxy
+
+import (
+	"bytes"
+	"context"
+	"crypto/tls"
+	"encoding/json"
+	"fmt"
+	"hash/fnv"
+	"io"
+	"log"
+	"net"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/fosrl/gerbil/logger"
+	"github.com/fosrl/gerbil/proxyproto"
+	"github.com/patrickmn/go-cache"
+)
+
+// RouteRecord represents a routing configuration
+type RouteRecord struct {
+	Hostname   string
+	TargetHost string
+	TargetPort int
+}
+
+// RouteAPIResponse represents the response from the route API
+type RouteAPIResponse struct {
+	Endpoints []string `json:"endpoints"`
+}
+
+// SNIProxy represents the main proxy server
+type SNIProxy struct {
+	port            int
+	cache           *cache.Cache
+	listener        net.Listener
+	ctx             context.Context
+	cancel          context.CancelFunc
+	wg              sync.WaitGroup
+	localProxyAddr  string
+	localProxyPort  int
+	remoteConfigURL string
+	publicKey       string
+	proxyProtocol   bool // Enable PROXY protocol v1
+
+	// New fields for fast local SNI lookup
+	localSNIs     map[string]struct{}
+	localSNIsLock sync.RWMutex
+
+	// Local overrides for domains that should always use local proxy
+	localOverrides map[string]struct{}
+
+	// Track active tunnels by SNI
+	activeTunnels     map[string]*activeTunnel
+	activeTunnelsLock sync.Mutex
+
+	// Trusted upstream proxies that can send PROXY protocol
+	trustedUpstreams map[string]struct{}
+}
+
+type activeTunnel struct {
+	conns []net.Conn
+}
+
+// readOnlyConn is a wrapper for io.Reader that implements net.Conn
+type readOnlyConn struct {
+	reader io.Reader
+}
+
+func (conn readOnlyConn) Read(p []byte) (int, error)         { return conn.reader.Read(p) }
+func (conn readOnlyConn) Write(p []byte) (int, error)        { return 0, io.ErrClosedPipe }
+func (conn readOnlyConn) Close() error                       { return nil }
+func (conn readOnlyConn) LocalAddr() net.Addr                { return nil }
+func (conn readOnlyConn) RemoteAddr() net.Addr               { return nil }
+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 }
+
+// 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())
+
+	// Create local overrides map
+	overridesMap := make(map[string]struct{})
+	for _, domain := range localOverrides {
+		if domain != "" {
+			overridesMap[domain] = struct{}{}
+		}
+	}
+
+	// Create trusted upstreams map
+	trustedMap := make(map[string]struct{})
+	for _, upstream := range trustedUpstreams {
+		if upstream != "" {
+			// Add both the domain and potentially resolved IPs
+			trustedMap[upstream] = struct{}{}
+
+			// Try to resolve the domain to IPs and add them too
+			if ips, err := net.LookupIP(upstream); err == nil {
+				for _, ip := range ips {
+					trustedMap[ip.String()] = struct{}{}
+				}
+			}
+		}
+	}
+
+	proxy := &SNIProxy{
+		port:             port,
+		cache:            cache.New(3*time.Second, 10*time.Minute),
+		ctx:              ctx,
+		cancel:           cancel,
+		localProxyAddr:   localProxyAddr,
+		localProxyPort:   localProxyPort,
+		remoteConfigURL:  remoteConfigURL,
+		publicKey:        publicKey,
+		proxyProtocol:    proxyProtocol,
+		localSNIs:        make(map[string]struct{}),
+		localOverrides:   overridesMap,
+		activeTunnels:    make(map[string]*activeTunnel),
+		trustedUpstreams: trustedMap,
+	}
+
+	return proxy, nil
+}
+
+// Start begins listening for connections
+func (p *SNIProxy) Start() error {
+	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", p.port))
+	if err != nil {
+		return fmt.Errorf("failed to listen on port %d: %w", p.port, err)
+	}
+
+	p.listener = listener
+	logger.Debug("SNI Proxy listening on port %d", p.port)
+
+	// Accept connections in a goroutine
+	go p.acceptConnections()
+
+	return nil
+}
+
+// Stop gracefully shuts down the proxy
+func (p *SNIProxy) Stop() error {
+	log.Println("Stopping SNI Proxy...")
+
+	p.cancel()
+
+	if p.listener != nil {
+		p.listener.Close()
+	}
+
+	// Wait for all goroutines to finish with timeout
+	done := make(chan struct{})
+	go func() {
+		p.wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		log.Println("All connections closed gracefully")
+	case <-time.After(30 * time.Second):
+		log.Println("Timeout waiting for connections to close")
+	}
+
+	log.Println("SNI Proxy stopped")
+	return nil
+}
+
+// acceptConnections handles incoming connections
+func (p *SNIProxy) acceptConnections() {
+	for {
+		conn, err := p.listener.Accept()
+		if err != nil {
+			select {
+			case <-p.ctx.Done():
+				return
+			default:
+				logger.Debug("Accept error: %v", err)
+				continue
+			}
+		}
+
+		p.wg.Add(1)
+		go p.handleConnection(conn)
+	}
+}
+
+// readClientHello reads and parses the TLS ClientHello message
+func (p *SNIProxy) readClientHello(reader io.Reader) (*tls.ClientHelloInfo, error) {
+	var hello *tls.ClientHelloInfo
+	err := tls.Server(readOnlyConn{reader: reader}, &tls.Config{
+		GetConfigForClient: func(argHello *tls.ClientHelloInfo) (*tls.Config, error) {
+			hello = new(tls.ClientHelloInfo)
+			*hello = *argHello
+			return nil, nil
+		},
+	}).Handshake()
+	if hello == nil {
+		return nil, err
+	}
+	return hello, nil
+}
+
+// peekClientHello reads the ClientHello while preserving the data for forwarding
+func (p *SNIProxy) peekClientHello(reader io.Reader) (*tls.ClientHelloInfo, io.Reader, error) {
+	peekedBytes := new(bytes.Buffer)
+	hello, err := p.readClientHello(io.TeeReader(reader, peekedBytes))
+	if err != nil {
+		return nil, nil, err
+	}
+	return hello, io.MultiReader(peekedBytes, reader), nil
+}
+
+// extractSNI extracts the SNI hostname from the TLS ClientHello
+func (p *SNIProxy) extractSNI(conn net.Conn) (string, io.Reader, error) {
+	clientHello, clientReader, err := p.peekClientHello(conn)
+	if err != nil {
+		return "", nil, fmt.Errorf("failed to peek ClientHello: %w", err)
+	}
+
+	if clientHello.ServerName == "" {
+		return "", clientReader, fmt.Errorf("no SNI hostname found in ClientHello")
+	}
+
+	return clientHello.ServerName, clientReader, nil
+}
+
+// handleConnection processes a single client connection
+func (p *SNIProxy) handleConnection(clientConn net.Conn) {
+	defer p.wg.Done()
+	defer clientConn.Close()
+
+	logger.Debug("Accepted connection from %s", clientConn.RemoteAddr())
+
+	// Check for PROXY protocol from trusted upstream
+	var proxyInfo *proxyproto.Info
+	var actualClientConn net.Conn = clientConn
+
+	if len(p.trustedUpstreams) > 0 {
+		var err error
+		proxyInfo, actualClientConn, err = proxyproto.ParseV1Header(clientConn, p.trustedUpstreams)
+		if err != nil {
+			logger.Debug("Failed to parse PROXY protocol: %v", err)
+			return
+		}
+		if proxyInfo != nil {
+			logger.Debug("Received PROXY protocol from trusted upstream: %s:%d -> %s:%d",
+				proxyInfo.SrcIP, proxyInfo.SrcPort, proxyInfo.DestIP, proxyInfo.DestPort)
+		} else {
+			// No PROXY protocol detected, but connection is from trusted upstream
+			// This is fine - treat as regular connection
+			logger.Debug("No PROXY protocol detected from trusted upstream, treating as regular connection")
+		}
+	}
+
+	// Set read timeout for SNI extraction
+	if err := actualClientConn.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil {
+		logger.Debug("Failed to set read deadline: %v", err)
+		return
+	}
+
+	// Extract SNI hostname
+	hostname, clientReader, err := p.extractSNI(actualClientConn)
+	if err != nil {
+		logger.Debug("SNI extraction failed: %v", err)
+		return
+	}
+
+	if hostname == "" {
+		log.Println("No SNI hostname found")
+		return
+	}
+
+	logger.Debug("SNI hostname detected: %s", hostname)
+
+	// Remove read timeout for normal operation
+	if err := actualClientConn.SetReadDeadline(time.Time{}); err != nil {
+		logger.Debug("Failed to clear read deadline: %v", err)
+		return
+	}
+
+	// Get routing information - use original client address if available from PROXY protocol
+	var clientAddrStr string
+	if proxyInfo != nil {
+		clientAddrStr = fmt.Sprintf("%s:%d", proxyInfo.SrcIP, proxyInfo.SrcPort)
+	} else {
+		clientAddrStr = clientConn.RemoteAddr().String()
+	}
+
+	route, err := p.getRoute(hostname, clientAddrStr)
+	if err != nil {
+		logger.Debug("Failed to get route for %s: %v", hostname, err)
+		return
+	}
+
+	if route == nil {
+		logger.Debug("No route found for hostname: %s", hostname)
+		return
+	}
+
+	logger.Debug("Routing %s to %s:%d", hostname, route.TargetHost, route.TargetPort)
+
+	// Connect to target server
+	targetConn, err := net.DialTimeout("tcp",
+		fmt.Sprintf("%s:%d", route.TargetHost, route.TargetPort),
+		10*time.Second)
+	if err != nil {
+		logger.Debug("Failed to connect to target %s:%d: %v",
+			route.TargetHost, route.TargetPort, err)
+		return
+	}
+	defer targetConn.Close()
+
+	logger.Debug("Connected to target: %s:%d", route.TargetHost, route.TargetPort)
+
+	// Send PROXY protocol header if enabled
+	if p.proxyProtocol {
+		var proxyHeader string
+		if proxyInfo != nil {
+			// Use original client info from PROXY protocol
+			proxyHeader = proxyproto.BuildV1HeaderFromInfo(proxyInfo, targetConn.LocalAddr())
+		} else {
+			// Use direct client connection info
+			proxyHeader = proxyproto.BuildV1Header(clientConn.RemoteAddr(), targetConn.LocalAddr())
+		}
+		logger.Debug("Sending PROXY protocol header: %s", strings.TrimSpace(proxyHeader))
+
+		if _, err := targetConn.Write([]byte(proxyHeader)); err != nil {
+			logger.Debug("Failed to send PROXY protocol header: %v", err)
+			return
+		}
+	}
+
+	// Track this tunnel by SNI
+	p.activeTunnelsLock.Lock()
+	tunnel, ok := p.activeTunnels[hostname]
+	if !ok {
+		tunnel = &activeTunnel{}
+		p.activeTunnels[hostname] = tunnel
+	}
+	tunnel.conns = append(tunnel.conns, actualClientConn)
+	p.activeTunnelsLock.Unlock()
+
+	defer func() {
+		// Remove this conn from active tunnels
+		p.activeTunnelsLock.Lock()
+		if tunnel, ok := p.activeTunnels[hostname]; ok {
+			newConns := make([]net.Conn, 0, len(tunnel.conns))
+			for _, c := range tunnel.conns {
+				if c != actualClientConn {
+					newConns = append(newConns, c)
+				}
+			}
+			if len(newConns) == 0 {
+				delete(p.activeTunnels, hostname)
+			} else {
+				tunnel.conns = newConns
+			}
+		}
+		p.activeTunnelsLock.Unlock()
+	}()
+
+	// Start bidirectional data transfer
+	p.pipe(actualClientConn, targetConn, clientReader)
+}
+
+// getRoute retrieves routing information for a hostname
+func (p *SNIProxy) getRoute(hostname, clientAddr string) (*RouteRecord, error) {
+	// Check local overrides first
+	if _, isOverride := p.localOverrides[hostname]; isOverride {
+		logger.Debug("Local override matched for hostname: %s", hostname)
+		return &RouteRecord{
+			Hostname:   hostname,
+			TargetHost: p.localProxyAddr,
+			TargetPort: p.localProxyPort,
+		}, nil
+	}
+
+	// Fast path: check if hostname is in localSNIs
+	p.localSNIsLock.RLock()
+	_, isLocal := p.localSNIs[hostname]
+	p.localSNIsLock.RUnlock()
+	if isLocal {
+		return &RouteRecord{
+			Hostname:   hostname,
+			TargetHost: p.localProxyAddr,
+			TargetPort: p.localProxyPort,
+		}, nil
+	}
+
+	// Check cache first
+	if cached, found := p.cache.Get(hostname); found {
+		if cached == nil {
+			return nil, nil // Cached negative result
+		}
+		logger.Debug("Cache hit for hostname: %s", hostname)
+		return cached.(*RouteRecord), nil
+	}
+
+	logger.Debug("Cache miss for hostname: %s, querying API", hostname)
+
+	// Query API with timeout
+	ctx, cancel := context.WithTimeout(p.ctx, 5*time.Second)
+	defer cancel()
+
+	// Construct API URL (without hostname in path)
+	apiURL := fmt.Sprintf("%s/gerbil/get-resolved-hostname", p.remoteConfigURL)
+
+	// Create request body with hostname and public key
+	requestBody := map[string]string{
+		"hostname":  hostname,
+		"publicKey": p.publicKey,
+	}
+
+	jsonBody, err := json.Marshal(requestBody)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal request body: %w", err)
+	}
+
+	// Create HTTP request
+	req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewBuffer(jsonBody))
+	if err != nil {
+		return nil, fmt.Errorf("failed to create request: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+
+	// Make HTTP request
+	client := &http.Client{Timeout: 5 * time.Second}
+	resp, err := client.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("API request failed: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode == http.StatusNotFound {
+		// Cache negative result for shorter time (1 minute)
+		p.cache.Set(hostname, nil, 1*time.Minute)
+		return nil, nil
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("API returned status %d", resp.StatusCode)
+	}
+
+	// Parse response
+	var apiResponse RouteAPIResponse
+	if err := json.NewDecoder(resp.Body).Decode(&apiResponse); err != nil {
+		return nil, fmt.Errorf("failed to decode API response: %w", err)
+	}
+
+	endpoints := apiResponse.Endpoints
+
+	// Default target configuration
+	targetHost := p.localProxyAddr
+	targetPort := p.localProxyPort
+
+	// If no endpoints returned, use local node
+	if len(endpoints) == 0 {
+		logger.Debug("No endpoints returned for hostname: %s, using local node", hostname)
+	} else {
+		// Select endpoint using consistent hashing for stickiness
+		selectedEndpoint := p.selectStickyEndpoint(clientAddr, endpoints)
+		targetHost = selectedEndpoint
+		targetPort = 443 // Default HTTPS port
+		logger.Debug("Selected endpoint %s for hostname %s from client %s", selectedEndpoint, hostname, clientAddr)
+	}
+
+	route := &RouteRecord{
+		Hostname:   hostname,
+		TargetHost: targetHost,
+		TargetPort: targetPort,
+	}
+
+	// Cache the result
+	p.cache.Set(hostname, route, cache.DefaultExpiration)
+	logger.Debug("Cached route for hostname: %s", hostname)
+
+	return route, nil
+}
+
+// selectStickyEndpoint selects an endpoint using consistent hashing to ensure
+// the same client always routes to the same endpoint for load balancing
+func (p *SNIProxy) selectStickyEndpoint(clientAddr string, endpoints []string) string {
+	if len(endpoints) == 0 {
+		return p.localProxyAddr
+	}
+	if len(endpoints) == 1 {
+		return endpoints[0]
+	}
+
+	// Use FNV hash for consistent selection based on client address
+	hash := fnv.New32a()
+	hash.Write([]byte(clientAddr))
+	index := hash.Sum32() % uint32(len(endpoints))
+
+	return endpoints[index]
+}
+
+// pipe handles bidirectional data transfer between connections
+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 closeConns()
+
+		// Use a large buffer for better performance
+		buf := make([]byte, 32*1024)
+		_, err := io.CopyBuffer(targetConn, clientReader, buf)
+		if err != nil && err != io.EOF {
+			logger.Debug("Copy client->target error: %v", err)
+		}
+	}()
+
+	// Copy data from target to client
+	go func() {
+		defer wg.Done()
+		defer closeConns()
+
+		// Use a large buffer for better performance
+		buf := make([]byte, 32*1024)
+		_, err := io.CopyBuffer(clientConn, targetConn, buf)
+		if err != nil && err != io.EOF {
+			logger.Debug("Copy target->client error: %v", err)
+		}
+	}()
+
+	wg.Wait()
+}
+
+// GetCacheStats returns cache statistics
+func (p *SNIProxy) GetCacheStats() (int, int) {
+	return p.cache.ItemCount(), len(p.cache.Items())
+}
+
+// ClearCache clears all cached entries
+func (p *SNIProxy) ClearCache() {
+	p.cache.Flush()
+	log.Println("Cache cleared")
+}
+
+// UpdateLocalSNIs updates the local SNIs and invalidates cache for changed domains
+func (p *SNIProxy) UpdateLocalSNIs(fullDomains []string) {
+	newSNIs := make(map[string]struct{})
+	for _, domain := range fullDomains {
+		newSNIs[domain] = struct{}{}
+		// Invalidate any cached route for this domain
+		p.cache.Delete(domain)
+	}
+
+	// Update localSNIs
+	p.localSNIsLock.Lock()
+	removed := make([]string, 0)
+	for sni := range p.localSNIs {
+		if _, stillLocal := newSNIs[sni]; !stillLocal {
+			removed = append(removed, sni)
+		}
+	}
+	p.localSNIs = newSNIs
+	p.localSNIsLock.Unlock()
+
+	logger.Debug("Updated local SNIs, added %d, removed %d", len(newSNIs), len(removed))
+
+	// Terminate tunnels for removed SNIs
+	if len(removed) > 0 {
+		p.activeTunnelsLock.Lock()
+		for _, sni := range removed {
+			if tunnels, ok := p.activeTunnels[sni]; ok {
+				for _, conn := range tunnels.conns {
+					conn.Close()
+				}
+				delete(p.activeTunnels, sni)
+				logger.Debug("Closed tunnels for SNI target change: %s", sni)
+			}
+		}
+		p.activeTunnelsLock.Unlock()
+	}
+}

proxy/proxy_test.go

@@ -0,0 +1,205 @@
+package proxy
+
+import (
+	"net"
+	"testing"
+
+	"github.com/fosrl/gerbil/proxyproto"
+)
+
+func TestBuildProxyProtocolHeader(t *testing.T) {
+	tests := []struct {
+		name       string
+		clientAddr string
+		targetAddr string
+		expected   string
+	}{
+		{
+			name:       "IPv4 client and target",
+			clientAddr: "192.168.1.100:12345",
+			targetAddr: "10.0.0.1:443",
+			expected:   "PROXY TCP4 192.168.1.100 10.0.0.1 12345 443\r\n",
+		},
+		{
+			name:       "IPv6 client and target",
+			clientAddr: "[2001:db8::1]:12345",
+			targetAddr: "[2001:db8::2]:443",
+			expected:   "PROXY TCP6 2001:db8::1 2001:db8::2 12345 443\r\n",
+		},
+		{
+			name:       "IPv4 client with IPv6 loopback target",
+			clientAddr: "192.168.1.100:12345",
+			targetAddr: "[::1]:443",
+			expected:   "PROXY TCP4 192.168.1.100 127.0.0.1 12345 443\r\n",
+		},
+		{
+			name:       "IPv4 client with IPv6 target",
+			clientAddr: "192.168.1.100:12345",
+			targetAddr: "[2001:db8::2]:443",
+			expected:   "PROXY TCP4 192.168.1.100 127.0.0.1 12345 443\r\n",
+		},
+		{
+			name:       "IPv6 client with IPv4 target",
+			clientAddr: "[2001:db8::1]:12345",
+			targetAddr: "10.0.0.1:443",
+			expected:   "PROXY TCP6 2001:db8::1 ::ffff:10.0.0.1 12345 443\r\n",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			clientTCP, err := net.ResolveTCPAddr("tcp", tt.clientAddr)
+			if err != nil {
+				t.Fatalf("Failed to resolve client address: %v", err)
+			}
+
+			targetTCP, err := net.ResolveTCPAddr("tcp", tt.targetAddr)
+			if err != nil {
+				t.Fatalf("Failed to resolve target address: %v", err)
+			}
+
+			result := proxyproto.BuildV1Header(clientTCP, targetTCP)
+			if result != tt.expected {
+				t.Errorf("Expected %q, got %q", tt.expected, result)
+			}
+		})
+	}
+}
+
+func TestBuildProxyProtocolHeaderUnknownType(t *testing.T) {
+	// Test with non-TCP address type
+	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)
+	expected := "PROXY UNKNOWN\r\n"
+
+	if result != expected {
+		t.Errorf("Expected %q, got %q", expected, result)
+	}
+}
+
+func TestBuildProxyProtocolHeaderFromInfo(t *testing.T) {
+	// Test IPv4 case
+	info := &proxyproto.Info{
+		Protocol: "TCP4",
+		SrcIP:    "10.0.0.1",
+		DestIP:   "192.168.1.100",
+		SrcPort:  12345,
+		DestPort: 443,
+	}
+
+	targetAddr, _ := net.ResolveTCPAddr("tcp", "127.0.0.1:8080")
+	header := proxyproto.BuildV1HeaderFromInfo(info, targetAddr)
+
+	expected := "PROXY TCP4 10.0.0.1 127.0.0.1 12345 8080\r\n"
+	if header != expected {
+		t.Errorf("Expected header '%s', got '%s'", expected, header)
+	}
+
+	// Test IPv6 case
+	info = &proxyproto.Info{
+		Protocol: "TCP6",
+		SrcIP:    "2001:db8::1",
+		DestIP:   "2001:db8::2",
+		SrcPort:  12345,
+		DestPort: 443,
+	}
+
+	targetAddr, _ = net.ResolveTCPAddr("tcp6", "[::1]:8080")
+	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"
@@ -24,19 +28,22 @@ type EncryptedHolePunchMessage struct {
 }
 
 type HolePunchMessage struct {
-	OlmID  string `json:"olmId"`
+	OlmID     string `json:"olmId"`
-	NewtID string `json:"newtId"`
+	NewtID    string `json:"newtId"`
-	Token  string `json:"token"`
+	Token     string `json:"token"`
+	PublicKey string `json:"publicKey"`
 }
 
 type ClientEndpoint struct {
-	OlmID       string `json:"olmId"`
+	OlmID             string `json:"olmId"`
-	NewtID      string `json:"newtId"`
+	NewtID            string `json:"newtId"`
-	Token       string `json:"token"`
+	Token             string `json:"token"`
-	IP          string `json:"ip"`
+	IP                string `json:"ip"`
-	Port        int    `json:"port"`
+	Port              int    `json:"port"`
-	Timestamp   int64  `json:"timestamp"`
+	Timestamp         int64  `json:"timestamp"`
-	ReachableAt string `json:"reachableAt"`
+	ReachableAt       string `json:"reachableAt"`
+	ExitNodePublicKey string `json:"exitNodePublicKey"`
+	ClientPublicKey   string `json:"publicKey"`
 }
 
 // Updated to support multiple destination peers
@@ -57,12 +64,52 @@ 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
+	ToDestination  *net.UDPAddr // The destination address
+	ClientIndex    uint32       // The receiver index seen from client
+	DestIndex      uint32       // The receiver index seen from destination
+	LastFromClient time.Time    // Last packet from client to destination
+	LastFromDest   time.Time    // Last packet from destination to client
+	PacketCount    int          // Number of packets observed
+}
+
 type InitialMappings struct {
 	Mappings map[string]ProxyMapping `json:"mappings"` // key is "ip:port"
 }
@@ -74,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
@@ -100,22 +154,66 @@ 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"
 	wgSessions sync.Map
+	// 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,
+		addr:             addr,
-		serverURL:   serverURL,
+		serverURL:        serverURL,
-		privateKey:  privateKey,
+		privateKey:       privateKey,
-		packetChan:  make(chan Packet, 1000),
+		packetChan:       make(chan Packet, 50000), // Increased from 1000 to handle high throughput
-		ReachableAt: reachableAt,
+		ReachableAt:      reachableAt,
+		ctx:              ctx,
+		cancel:           cancel,
+		proxyProtocol:    proxyProtocol,
+		trustedUpstreams: trustedMap,
 	}
 }
 
@@ -137,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()
 	}
@@ -155,21 +258,61 @@ func (s *UDPProxyServer) Start() error {
 	// Start the proxy mapping cleanup routine
 	go s.cleanupIdleProxyMappings()
 
+	// 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 {
-			logger.Error("Error reading UDP packet: %v", err)
+			// If we're shutting down, exit
-			continue
+			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}
 	}
@@ -178,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])
@@ -202,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
@@ -218,16 +417,19 @@ func (s *UDPProxyServer) packetWorker() {
 			}
 
 			endpoint := ClientEndpoint{
-				NewtID:      msg.NewtID,
+				NewtID:            msg.NewtID,
-				OlmID:       msg.OlmID,
+				OlmID:             msg.OlmID,
-				Token:       msg.Token,
+				Token:             msg.Token,
-				IP:          packet.remoteAddr.IP.String(),
+				IP:                effectiveAddr.IP.String(),
-				Port:        packet.remoteAddr.Port,
+				Port:              effectiveAddr.Port,
-				Timestamp:   time.Now().Unix(),
+				Timestamp:         time.Now().Unix(),
-				ReachableAt: s.ReachableAt,
+				ReachableAt:       s.ReachableAt,
+				ExitNodePublicKey: s.privateKey.PublicKey().String(),
+				ClientPublicKey:   msg.PublicKey,
 			}
-			logger.Debug("Created endpoint from packet remoteAddr %s: IP=%s, Port=%d", packet.remoteAddr.String(), endpoint.IP, endpoint.Port)
+			logger.Debug("Created endpoint from packet remoteAddr %s: IP=%s, Port=%d", effectiveAddr.String(), endpoint.IP, endpoint.Port)
 			s.notifyServer(endpoint)
+			s.clearSessionsForIP(endpoint.IP) // Clear sessions for this IP to allow re-establishment
 		}
 		// Return the buffer to the pool for reuse.
 		bufferPool.Put(packet.data[:1500])
@@ -324,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 {
@@ -355,10 +594,10 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD
 	switch messageType {
 	case WireGuardMessageTypeHandshakeInitiation:
 		// Initial handshake: forward to all peers
-		logger.Debug("Forwarding handshake initiation from %s (sender index: %d)", remoteAddr, senderIndex)
+		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
@@ -372,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)
 			}
 		}
 
@@ -394,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
@@ -414,7 +653,7 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD
 
 	case WireGuardMessageTypeTransportData:
 		// Data packet: forward only to the established session peer
-		logger.Debug("Received transport data with receiver index %d from %s", receiverIndex, remoteAddr)
+		// logger.Debug("Received transport data with receiver index %d from %s", receiverIndex, remoteAddr)
 
 		// Look up the session based on the receiver index
 		var destAddr *net.UDPAddr
@@ -422,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
@@ -442,6 +679,9 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD
 				return
 			}
 
+			// Track communication pattern for session rebuilding
+			s.trackCommunicationPattern(remoteAddr, destAddr, receiverIndex, true)
+
 			_, err = conn.Write(packet)
 			if err != nil {
 				logger.Debug("Failed to forward transport data: %v", err)
@@ -450,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
@@ -462,6 +702,9 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD
 					continue
 				}
 
+				// Track communication pattern for session rebuilding
+				s.trackCommunicationPattern(remoteAddr, destAddr, receiverIndex, true)
+
 				_, err = conn.Write(packet)
 				if err != nil {
 					logger.Debug("Failed to forward transport data: %v", err)
@@ -475,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
@@ -545,6 +788,9 @@ func (s *UDPProxyServer) handleResponses(conn *net.UDPConn, destAddr *net.UDPAdd
 					LastSeen:      time.Now(),
 				})
 				logger.Debug("Stored session mapping: %s -> %s", sessionKey, destAddr.String())
+			} else if ok && buffer[0] == WireGuardMessageTypeTransportData {
+				// Track communication pattern for session rebuilding (reverse direction)
+				s.trackCommunicationPattern(destAddr, remoteAddr, receiverIndex, false)
 			}
 		}
 
@@ -559,49 +805,69 @@ 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()
-		now := time.Now()
+	for {
-		s.connections.Range(func(key, value interface{}) bool {
+		select {
-			destConn := value.(*DestinationConn)
+		case <-ticker.C:
-			if now.Sub(destConn.lastUsed) > 10*time.Minute {
+			now := time.Now()
-				destConn.conn.Close()
+			s.connections.Range(func(key, value interface{}) bool {
-				s.connections.Delete(key)
+				destConn := value.(*DestinationConn)
-			}
+				if now.Sub(destConn.lastUsed) > 10*time.Minute {
-			return true
+					destConn.conn.Close()
-		})
+					s.connections.Delete(key)
+				}
+				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 {
+
-		now := time.Now()
+	defer ticker.Stop()
-		s.wgSessions.Range(func(key, value interface{}) bool {
+	for {
-			session := value.(*WireGuardSession)
+		select {
-			if now.Sub(session.LastSeen) > 15*time.Minute {
+		case <-ticker.C:
-				s.wgSessions.Delete(key)
+			now := time.Now()
-				logger.Debug("Removed idle session: %s", key)
+			s.wgSessions.Range(func(key, value interface{}) bool {
-			}
+				session := value.(*WireGuardSession)
-			return true
+				// 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()
-		now := time.Now()
+	for {
-		s.proxyMappings.Range(func(key, value interface{}) bool {
+		select {
-			mapping := value.(ProxyMapping)
+		case <-ticker.C:
-			// Remove mappings that haven't been used in 30 minutes
+			now := time.Now()
-			if now.Sub(mapping.LastUsed) > 30*time.Minute {
+			s.proxyMappings.Range(func(key, value interface{}) bool {
-				s.proxyMappings.Delete(key)
+				mapping := value.(ProxyMapping)
-				logger.Debug("Removed idle proxy mapping: %s", key)
+				// Remove mappings that haven't been used in 30 minutes
-			}
+				if now.Sub(mapping.LastUsed) > 30*time.Minute {
-			return true
+					s.proxyMappings.Delete(key)
-		})
+					logger.Debug("Removed idle proxy mapping: %s", key)
+				}
+				return true
+			})
+		case <-s.ctx.Done():
+			return
+		}
 	}
 }
 
@@ -660,7 +926,7 @@ func (s *UDPProxyServer) UpdateProxyMapping(sourceIP string, sourcePort int, des
 func (s *UDPProxyServer) OnPeerAdded(wgIP string) {
 	logger.Info("Clearing connections for added peer with WG IP: %s", wgIP)
 	s.clearConnectionsForWGIP(wgIP)
-	s.clearSessionsForWGIP(wgIP)
+	// s.clearSessionsForWGIP(wgIP) THE DEST ADDR IS NOT THE WG IP, SO THIS IS NOT NEEDED
 	// s.clearProxyMappingsForWGIP(wgIP)
 }
 
@@ -668,7 +934,7 @@ func (s *UDPProxyServer) OnPeerAdded(wgIP string) {
 func (s *UDPProxyServer) OnPeerRemoved(wgIP string) {
 	logger.Info("Clearing connections for removed peer with WG IP: %s", wgIP)
 	s.clearConnectionsForWGIP(wgIP)
-	s.clearSessionsForWGIP(wgIP)
+	// s.clearSessionsForWGIP(wgIP) THE DEST ADDR IS NOT THE WG IP, SO THIS IS NOT NEEDED
 	// s.clearProxyMappingsForWGIP(wgIP)
 }
 
@@ -699,17 +965,18 @@ func (s *UDPProxyServer) clearConnectionsForWGIP(wgIP string) {
 }
 
 // clearSessionsForWGIP removes all WireGuard sessions associated with a specific WireGuard IP
-func (s *UDPProxyServer) clearSessionsForWGIP(wgIP string) {
+func (s *UDPProxyServer) clearSessionsForIP(ip string) {
 	var keysToDelete []string
 
 	s.wgSessions.Range(func(key, value interface{}) bool {
 		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() == wgIP {
+		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", wgIP, keyStr)
+			logger.Debug("Marking session for deletion for WG IP %s: %s", ip, keyStr)
 		}
 		return true
 	})
@@ -719,7 +986,7 @@ func (s *UDPProxyServer) clearSessionsForWGIP(wgIP string) {
 		s.wgSessions.Delete(key)
 	}
 
-	logger.Info("Cleared %d sessions for WG IP: %s", len(keysToDelete), wgIP)
+	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
@@ -820,3 +1087,145 @@ func (s *UDPProxyServer) UpdateDestinationInMappings(oldDest, newDest PeerDestin
 
 	return updatedCount
 }
+
+// trackCommunicationPattern tracks bidirectional communication patterns to rebuild sessions
+func (s *UDPProxyServer) trackCommunicationPattern(fromAddr, toAddr *net.UDPAddr, receiverIndex uint32, fromClient bool) {
+	var clientAddr, destAddr *net.UDPAddr
+	var clientIndex, destIndex uint32
+
+	if fromClient {
+		clientAddr = fromAddr
+		destAddr = toAddr
+		clientIndex = receiverIndex
+		destIndex = 0 // We don't know the destination index yet
+	} else {
+		clientAddr = toAddr
+		destAddr = fromAddr
+		clientIndex = 0 // We don't know the client index yet
+		destIndex = receiverIndex
+	}
+
+	patternKey := fmt.Sprintf("%s-%s", clientAddr.String(), destAddr.String())
+	now := time.Now()
+
+	if existingPattern, ok := s.commPatterns.Load(patternKey); ok {
+		pattern := existingPattern.(*CommunicationPattern)
+
+		// Update the pattern
+		if fromClient {
+			pattern.LastFromClient = now
+			if pattern.ClientIndex == 0 {
+				pattern.ClientIndex = clientIndex
+			}
+		} else {
+			pattern.LastFromDest = now
+			if pattern.DestIndex == 0 {
+				pattern.DestIndex = destIndex
+			}
+		}
+
+		pattern.PacketCount++
+		s.commPatterns.Store(patternKey, pattern)
+
+		// Check if we have bidirectional communication and can rebuild a session
+		s.tryRebuildSession(pattern)
+	} else {
+		// Create new pattern
+		pattern := &CommunicationPattern{
+			FromClient:    clientAddr,
+			ToDestination: destAddr,
+			ClientIndex:   clientIndex,
+			DestIndex:     destIndex,
+			PacketCount:   1,
+		}
+
+		if fromClient {
+			pattern.LastFromClient = now
+		} else {
+			pattern.LastFromDest = now
+		}
+
+		s.commPatterns.Store(patternKey, pattern)
+	}
+}
+
+// tryRebuildSession attempts to rebuild a WireGuard session from communication patterns
+func (s *UDPProxyServer) tryRebuildSession(pattern *CommunicationPattern) {
+	// Check if we have bidirectional communication within a reasonable time window
+	timeDiff := pattern.LastFromClient.Sub(pattern.LastFromDest)
+	if timeDiff < 0 {
+		timeDiff = -timeDiff
+	}
+
+	// Only rebuild if we have recent bidirectional communication and both indices
+	if timeDiff < 30*time.Second && pattern.ClientIndex != 0 && pattern.DestIndex != 0 && pattern.PacketCount >= 4 {
+		// Create session mapping: client's index maps to destination
+		sessionKey := fmt.Sprintf("%d:%d", pattern.DestIndex, pattern.ClientIndex)
+
+		// Check if we already have this session
+		if _, exists := s.wgSessions.Load(sessionKey); !exists {
+			s.wgSessions.Store(sessionKey, &WireGuardSession{
+				ReceiverIndex: pattern.DestIndex,
+				SenderIndex:   pattern.ClientIndex,
+				DestAddr:      pattern.ToDestination,
+				LastSeen:      time.Now(),
+			})
+			logger.Info("Rebuilt WireGuard session from communication pattern: %s -> %s (packets: %d)",
+				sessionKey, pattern.ToDestination.String(), pattern.PacketCount)
+		}
+	}
+}
+
+// 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()
+	for {
+		select {
+		case <-ticker.C:
+			now := time.Now()
+			s.commPatterns.Range(func(key, value interface{}) bool {
+				pattern := value.(*CommunicationPattern)
+
+				// Get the most recent activity
+				lastActivity := pattern.LastFromClient
+				if pattern.LastFromDest.After(lastActivity) {
+					lastActivity = pattern.LastFromDest
+				}
+
+				// Remove patterns that haven't had activity in 20 minutes
+				if now.Sub(lastActivity) > 20*time.Minute {
+					s.commPatterns.Delete(key)
+					logger.Debug("Removed idle communication pattern: %s", key)
+				}
+				return true
+			})
+		case <-s.ctx.Done():
+			return
+		}
+	}
+}