Merge branch 'main' of github.com:fosrl/gerbil

Bumps the prod-minor-updates group with 1 update: [golang.org/x/crypto](https://github.com/golang/crypto).


Updates `golang.org/x/crypto` from 0.45.0 to 0.46.0
- [Commits](https://github.com/golang/crypto/compare/v0.45.0...v0.46.0)

---
updated-dependencies:
- dependency-name: golang.org/x/crypto
  dependency-version: 0.46.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: prod-minor-updates
...

Signed-off-by: dependabot[bot] <support@github.com>

.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/workflows/cicd.yml

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

CONTRIBUTING.md

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

Dockerfile

@@ -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,7 +20,7 @@ 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
 
@@ -42,16 +42,15 @@ In single node (self hosted) Pangolin deployments this can be bypassed by using
 
 ## CLI Args
 
+Important:
 - `reachableAt`: How should the remote server reach Gerbil's API?
 - `generateAndSaveKeyTo`: Where to save the generated WireGuard private key to persist across restarts.
-- `remoteConfig` (optional): Remote config location to HTTP get the JSON based config from. See `example_config.json`
-- `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.
+- `remoteConfig`: Remote config location to HTTP get the JSON based config from.
 
+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
@@ -66,7 +65,6 @@ Note: You must use either `config` or `remoteConfig` to configure WireGuard.
 All CLI arguments can also be provided via environment variables:
 
 - `INTERFACE`: Name of the WireGuard interface
-- `CONFIG`: Path to local configuration file
 - `REMOTE_CONFIG`: URL of the remote config server
 - `LISTEN`: Address to listen on for HTTP server
 - `GENERATE_AND_SAVE_KEY_TO`: Path to save generated private key
@@ -84,7 +82,7 @@ Example:
 
 ```bash
 ./gerbil \
---reachableAt=http://gerbil:3003 \
+--reachableAt=http://gerbil:3004 \
 --generateAndSaveKeyTo=/var/config/key \
 --remoteConfig=http://pangolin:3001/api/v1/
 ```
@@ -96,7 +94,7 @@ services:
     container_name: gerbil
     restart: unless-stopped
     command:
-      - --reachableAt=http://gerbil:3003
+      - --reachableAt=http://gerbil:3004
       - --generateAndSaveKeyTo=/var/config/key
       - --remoteConfig=http://pangolin:3001/api/v1/
     volumes:

SECURITY.md

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

config_example.json

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

go.mod

@@ -3,8 +3,10 @@ module github.com/fosrl/gerbil
 go 1.25
 
 require (
+	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/vishvananda/netlink v1.3.1
-	golang.org/x/crypto v0.36.0
+	golang.org/x/crypto v0.46.0
+	golang.org/x/sync v0.1.0
 	golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6
 )
 
@@ -14,10 +16,8 @@ require (
 	github.com/mdlayher/genetlink v1.3.2 // indirect
 	github.com/mdlayher/netlink v1.7.2 // indirect
 	github.com/mdlayher/socket v0.4.1 // indirect
-	github.com/patrickmn/go-cache v2.1.0+incompatible // indirect
 	github.com/vishvananda/netns v0.0.5 // indirect
-	golang.org/x/net v0.38.0 // indirect
-	golang.org/x/sync v0.1.0 // indirect
-	golang.org/x/sys v0.31.0 // indirect
+	golang.org/x/net v0.47.0 // indirect
+	golang.org/x/sys v0.39.0 // indirect
 	golang.zx2c4.com/wireguard v0.0.0-20230325221338-052af4a8072b // indirect
 )

go.sum

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

main.go

@@ -2,15 +2,21 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
+	"errors"
 	"flag"
 	"fmt"
 	"io"
+	"log"
 	"net"
 	"net/http"
+	_ "net/http/pprof"
 	"os"
 	"os/exec"
 	"os/signal"
+	"runtime"
+	"runtime/pprof"
 	"strconv"
 	"strings"
 	"sync"
@@ -21,6 +27,7 @@ import (
 	"github.com/fosrl/gerbil/proxy"
 	"github.com/fosrl/gerbil/relay"
 	"github.com/vishvananda/netlink"
+	"golang.org/x/sync/errgroup"
 	"golang.zx2c4.com/wireguard/wgctrl"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 )
@@ -40,6 +47,7 @@ var (
 type WgConfig struct {
 	PrivateKey string `json:"privateKey"`
 	ListenPort int    `json:"listenPort"`
+	RelayPort  int    `json:"relayPort"`
 	IpAddress  string `json:"ipAddress"`
 	Peers      []Peer `json:"peers"`
 }
@@ -108,6 +116,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
@@ -121,6 +131,7 @@ func main() {
 		localProxyAddr       string
 		localProxyPort       int
 		localOverridesStr    string
+		trustedUpstreamsStr  string
 		proxyProtocol        bool
 	)
 
@@ -138,6 +149,7 @@ func main() {
 	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")
 
 	if interfaceName == "" {
@@ -150,7 +162,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
@@ -160,9 +172,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)")
 	}
@@ -197,6 +211,9 @@ func main() {
 	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"
@@ -210,6 +227,26 @@ func main() {
 	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)
@@ -301,10 +338,20 @@ 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)
+
+	// Periodic bandwidth reporting
+	group.Go(func() error {
+		return periodicBandwidthCheck(groupCtx, remoteConfigURL+"/gerbil/receive-bandwidth")
+	})
 
 	// Start the UDP proxy server
-	proxyRelay = relay.NewUDPProxyServer(":21820", remoteConfigURL, key, reachableAt)
+	relayPort := wgconfig.RelayPort
+	if relayPort == 0 {
+		relayPort = 21820 // in case there is no relay port set, use 21820
+	}
+	proxyRelay = relay.NewUDPProxyServer(groupCtx, fmt.Sprintf(":%d", relayPort), remoteConfigURL, key, reachableAt)
 	err = proxyRelay.Start()
 	if err != nil {
 		logger.Fatal("Failed to start UDP proxy server: %v", err)
@@ -323,7 +370,16 @@ func main() {
 		logger.Info("Local overrides configured: %v", localOverrides)
 	}
 
-	proxySNI, err = proxy.NewSNIProxy(sniProxyPort, remoteConfigURL, key.PublicKey().String(), localProxyAddr, localProxyPort, localOverrides, proxyProtocol)
+	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)
+	}
+
+	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)
 	}
@@ -337,20 +393,42 @@ func main() {
 	http.HandleFunc("/update-proxy-mapping", handleUpdateProxyMapping)
 	http.HandleFunc("/update-destinations", handleUpdateDestinations)
 	http.HandleFunc("/update-local-snis", handleUpdateLocalSNIs)
+	http.HandleFunc("/healthz", handleHealthz)
 	logger.Info("Starting HTTP server on %s", listenAddr)
 
-	// Run HTTP server in a goroutine
-	go func() {
-		if err := http.ListenAndServe(listenAddr, nil); err != nil {
-			logger.Error("HTTP server failed: %v", err)
+	// HTTP server with graceful shutdown on context cancel
+	server := &http.Server{
+		Addr:    listenAddr,
+		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
-	sigCh := make(chan os.Signal, 1)
-	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
-	<-sigCh
-	logger.Info("Shutting down servers...")
+	// Wait for all goroutines to finish
+	if err := group.Wait(); err != nil && !errors.Is(err, context.Canceled) {
+		logger.Error("Service exited with error: %v", err)
+	} else if errors.Is(err, context.Canceled) {
+		logger.Info("Context cancelled, shutting down")
+	}
 }
 
 func loadRemoteConfig(url string, key wgtypes.Key, reachableAt string) (WgConfig, error) {
@@ -477,6 +555,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
@@ -607,7 +689,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
 		}
@@ -624,7 +706,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
 		}
@@ -645,6 +727,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:
@@ -656,6 +845,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 {
@@ -945,13 +1143,18 @@ func handleUpdateLocalSNIs(w http.ResponseWriter, r *http.Request) {
 	})
 }
 
-func periodicBandwidthCheck(endpoint string) {
+func periodicBandwidthCheck(ctx context.Context, endpoint string) error {
 	ticker := time.NewTicker(10 * time.Second)
 	defer ticker.Stop()
 
-	for range ticker.C {
-		if err := reportPeerBandwidth(endpoint); err != nil {
-			logger.Info("Failed to report peer bandwidth: %v", err)
+	for {
+		select {
+		case <-ticker.C:
+			if err := reportPeerBandwidth(endpoint); err != nil {
+				logger.Info("Failed to report peer bandwidth: %v", err)
+			}
+		case <-ctx.Done():
+			return ctx.Err()
 		}
 	}
 }
@@ -971,8 +1174,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,
@@ -1029,14 +1237,7 @@ func calculatePeerBandwidth() ([]PeerBandwidth, error) {
 
 	// Clean up old peers
 	for publicKey := range lastReadings {
-		found := false
-		for _, peer := range device.Peers {
-			if peer.PublicKey.String() == publicKey {
-				found = true
-				break
-			}
-		}
-		if !found {
+		if _, exists := currentPeerKeys[publicKey]; !exists {
 			delete(lastReadings, publicKey)
 		}
 	}
@@ -1092,3 +1293,25 @@ 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)
+	}
+}

proxy/proxy.go

@@ -11,6 +11,7 @@ import (
 	"log"
 	"net"
 	"net/http"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -31,6 +32,16 @@ type RouteAPIResponse struct {
 	Endpoints []string `json:"endpoints"`
 }
 
+// ProxyProtocolInfo holds information parsed from incoming PROXY protocol header
+type ProxyProtocolInfo struct {
+	Protocol     string // TCP4 or TCP6
+	SrcIP        string
+	DestIP       string
+	SrcPort      int
+	DestPort     int
+	OriginalConn net.Conn // The original connection after PROXY protocol parsing
+}
+
 // SNIProxy represents the main proxy server
 type SNIProxy struct {
 	port            int
@@ -55,6 +66,9 @@ type SNIProxy 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 {
@@ -75,6 +89,194 @@ func (conn readOnlyConn) SetDeadline(t time.Time) error      { return nil }
 func (conn readOnlyConn) SetReadDeadline(t time.Time) error  { return nil }
 func (conn readOnlyConn) SetWriteDeadline(t time.Time) error { return nil }
 
+// parseProxyProtocolHeader parses a PROXY protocol v1 header from the connection
+func (p *SNIProxy) parseProxyProtocolHeader(conn net.Conn) (*ProxyProtocolInfo, net.Conn, error) {
+	// Check if the connection comes from a trusted upstream
+	remoteHost, _, err := net.SplitHostPort(conn.RemoteAddr().String())
+	if err != nil {
+		return nil, conn, fmt.Errorf("failed to parse remote address: %w", err)
+	}
+
+	// Resolve the remote IP to hostname to check if it's trusted
+	// For simplicity, we'll check the IP directly in trusted upstreams
+	// In production, you might want to do reverse DNS lookup
+	if _, isTrusted := p.trustedUpstreams[remoteHost]; !isTrusted {
+		// Not from trusted upstream, return original connection
+		return nil, conn, nil
+	}
+
+	// Set read timeout for PROXY protocol parsing
+	if err := conn.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil {
+		return nil, conn, fmt.Errorf("failed to set read deadline: %w", err)
+	}
+
+	// Read the first line (PROXY protocol header)
+	buffer := make([]byte, 512) // PROXY protocol header should be much smaller
+	n, err := conn.Read(buffer)
+	if err != nil {
+		// If we can't read from trusted upstream, treat as regular connection
+		logger.Debug("Could not read from trusted upstream %s, treating as regular connection: %v", remoteHost, err)
+		// Clear read timeout before returning
+		if clearErr := conn.SetReadDeadline(time.Time{}); clearErr != nil {
+			logger.Debug("Failed to clear read deadline: %v", clearErr)
+		}
+		return nil, conn, nil
+	}
+
+	// Find the end of the first line (CRLF)
+	headerEnd := bytes.Index(buffer[:n], []byte("\r\n"))
+	if headerEnd == -1 {
+		// No PROXY protocol header found, treat as regular TLS connection
+		// Return the connection with the buffered data prepended
+		logger.Debug("No PROXY protocol header from trusted upstream %s, treating as regular TLS connection", remoteHost)
+
+		// Clear read timeout
+		if err := conn.SetReadDeadline(time.Time{}); err != nil {
+			logger.Debug("Failed to clear read deadline: %v", err)
+		}
+
+		// Create a reader that includes the buffered data + original connection
+		newReader := io.MultiReader(bytes.NewReader(buffer[:n]), conn)
+		wrappedConn := &proxyProtocolConn{
+			Conn:   conn,
+			reader: newReader,
+		}
+		return nil, wrappedConn, nil
+	}
+
+	headerLine := string(buffer[:headerEnd])
+	remainingData := buffer[headerEnd+2 : n]
+
+	// Parse PROXY protocol line: "PROXY TCP4/TCP6 srcIP destIP srcPort destPort"
+	parts := strings.Fields(headerLine)
+	if len(parts) != 6 || parts[0] != "PROXY" {
+		// Check for PROXY UNKNOWN
+		if len(parts) == 2 && parts[0] == "PROXY" && parts[1] == "UNKNOWN" {
+			// PROXY UNKNOWN - use original connection info
+			return nil, conn, nil
+		}
+		// Invalid PROXY protocol, but might be regular TLS - treat as such
+		logger.Debug("Invalid PROXY protocol from trusted upstream %s, treating as regular TLS connection: %s", remoteHost, headerLine)
+
+		// Clear read timeout
+		if err := conn.SetReadDeadline(time.Time{}); err != nil {
+			logger.Debug("Failed to clear read deadline: %v", err)
+		}
+
+		// Return the connection with all buffered data prepended
+		newReader := io.MultiReader(bytes.NewReader(buffer[:n]), conn)
+		wrappedConn := &proxyProtocolConn{
+			Conn:   conn,
+			reader: newReader,
+		}
+		return nil, wrappedConn, nil
+	}
+
+	protocol := parts[1]
+	srcIP := parts[2]
+	destIP := parts[3]
+	srcPort, err := strconv.Atoi(parts[4])
+	if err != nil {
+		return nil, conn, fmt.Errorf("invalid source port in PROXY header: %s", parts[4])
+	}
+	destPort, err := strconv.Atoi(parts[5])
+	if err != nil {
+		return nil, conn, fmt.Errorf("invalid destination port in PROXY header: %s", parts[5])
+	}
+
+	// Create a new reader that includes remaining data + original connection
+	var newReader io.Reader
+	if len(remainingData) > 0 {
+		newReader = io.MultiReader(bytes.NewReader(remainingData), conn)
+	} else {
+		newReader = conn
+	}
+
+	// Create a wrapper connection that reads from the combined reader
+	wrappedConn := &proxyProtocolConn{
+		Conn:   conn,
+		reader: newReader,
+	}
+
+	proxyInfo := &ProxyProtocolInfo{
+		Protocol:     protocol,
+		SrcIP:        srcIP,
+		DestIP:       destIP,
+		SrcPort:      srcPort,
+		DestPort:     destPort,
+		OriginalConn: wrappedConn,
+	}
+
+	// Clear read timeout
+	if err := conn.SetReadDeadline(time.Time{}); err != nil {
+		return nil, conn, fmt.Errorf("failed to clear read deadline: %w", err)
+	}
+
+	return proxyInfo, wrappedConn, nil
+}
+
+// proxyProtocolConn wraps a connection to read from a custom reader
+type proxyProtocolConn struct {
+	net.Conn
+	reader io.Reader
+}
+
+func (c *proxyProtocolConn) Read(b []byte) (int, error) {
+	return c.reader.Read(b)
+}
+
+// buildProxyProtocolHeaderFromInfo creates a PROXY protocol v1 header using ProxyProtocolInfo
+func (p *SNIProxy) buildProxyProtocolHeaderFromInfo(proxyInfo *ProxyProtocolInfo, targetAddr net.Addr) string {
+	targetTCP, ok := targetAddr.(*net.TCPAddr)
+	if !ok {
+		// Fallback for unknown address types
+		return "PROXY UNKNOWN\r\n"
+	}
+
+	// Use the original client information from the PROXY protocol
+	var targetIP string
+	var protocol string
+
+	// Parse source IP to determine protocol family
+	srcIP := net.ParseIP(proxyInfo.SrcIP)
+	if srcIP == nil {
+		return "PROXY UNKNOWN\r\n"
+	}
+
+	if srcIP.To4() != nil {
+		// Source is IPv4, use TCP4 protocol
+		protocol = "TCP4"
+		if targetTCP.IP.To4() != nil {
+			// Target is also IPv4, use as-is
+			targetIP = targetTCP.IP.String()
+		} else {
+			// Target is IPv6, but we need IPv4 for consistent protocol family
+			if targetTCP.IP.IsLoopback() {
+				targetIP = "127.0.0.1"
+			} else {
+				targetIP = "127.0.0.1" // Safe fallback
+			}
+		}
+	} else {
+		// Source is IPv6, use TCP6 protocol
+		protocol = "TCP6"
+		if targetTCP.IP.To4() != nil {
+			// Target is IPv4, convert to IPv6 representation
+			targetIP = "::ffff:" + targetTCP.IP.String()
+		} else {
+			// Target is also IPv6, use as-is
+			targetIP = targetTCP.IP.String()
+		}
+	}
+
+	return fmt.Sprintf("PROXY %s %s %s %d %d\r\n",
+		protocol,
+		proxyInfo.SrcIP,
+		targetIP,
+		proxyInfo.SrcPort,
+		targetTCP.Port)
+}
+
 // buildProxyProtocolHeader creates a PROXY protocol v1 header
 func buildProxyProtocolHeader(clientAddr, targetAddr net.Addr) string {
 	clientTCP, ok := clientAddr.(*net.TCPAddr)
@@ -131,7 +333,7 @@ func buildProxyProtocolHeader(clientAddr, targetAddr net.Addr) string {
 }
 
 // NewSNIProxy creates a new SNI proxy instance
-func NewSNIProxy(port int, remoteConfigURL, publicKey, localProxyAddr string, localProxyPort int, localOverrides []string, proxyProtocol bool) (*SNIProxy, error) {
+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
@@ -142,19 +344,36 @@ func NewSNIProxy(port int, remoteConfigURL, publicKey, localProxyAddr string, lo
 		}
 	}
 
+	// 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),
+		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
@@ -270,14 +489,35 @@ func (p *SNIProxy) handleConnection(clientConn net.Conn) {
 
 	logger.Debug("Accepted connection from %s", clientConn.RemoteAddr())
 
+	// Check for PROXY protocol from trusted upstream
+	var proxyInfo *ProxyProtocolInfo
+	var actualClientConn net.Conn = clientConn
+
+	if len(p.trustedUpstreams) > 0 {
+		var err error
+		proxyInfo, actualClientConn, err = p.parseProxyProtocolHeader(clientConn)
+		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 := clientConn.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil {
+	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(clientConn)
+	hostname, clientReader, err := p.extractSNI(actualClientConn)
 	if err != nil {
 		logger.Debug("SNI extraction failed: %v", err)
 		return
@@ -291,13 +531,20 @@ func (p *SNIProxy) handleConnection(clientConn net.Conn) {
 	logger.Debug("SNI hostname detected: %s", hostname)
 
 	// Remove read timeout for normal operation
-	if err := clientConn.SetReadDeadline(time.Time{}); err != nil {
+	if err := actualClientConn.SetReadDeadline(time.Time{}); err != nil {
 		logger.Debug("Failed to clear read deadline: %v", err)
 		return
 	}
 
-	// Get routing information
-	route, err := p.getRoute(hostname, clientConn.RemoteAddr().String())
+	// 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
@@ -325,7 +572,14 @@ func (p *SNIProxy) handleConnection(clientConn net.Conn) {
 
 	// Send PROXY protocol header if enabled
 	if p.proxyProtocol {
-		proxyHeader := buildProxyProtocolHeader(clientConn.RemoteAddr(), targetConn.LocalAddr())
+		var proxyHeader string
+		if proxyInfo != nil {
+			// Use original client info from PROXY protocol
+			proxyHeader = p.buildProxyProtocolHeaderFromInfo(proxyInfo, targetConn.LocalAddr())
+		} else {
+			// Use direct client connection info
+			proxyHeader = buildProxyProtocolHeader(clientConn.RemoteAddr(), targetConn.LocalAddr())
+		}
 		logger.Debug("Sending PROXY protocol header: %s", strings.TrimSpace(proxyHeader))
 
 		if _, err := targetConn.Write([]byte(proxyHeader)); err != nil {
@@ -341,7 +595,7 @@ func (p *SNIProxy) handleConnection(clientConn net.Conn) {
 		tunnel = &activeTunnel{}
 		p.activeTunnels[hostname] = tunnel
 	}
-	tunnel.conns = append(tunnel.conns, clientConn)
+	tunnel.conns = append(tunnel.conns, actualClientConn)
 	p.activeTunnelsLock.Unlock()
 
 	defer func() {
@@ -350,7 +604,7 @@ func (p *SNIProxy) handleConnection(clientConn net.Conn) {
 		if tunnel, ok := p.activeTunnels[hostname]; ok {
 			newConns := make([]net.Conn, 0, len(tunnel.conns))
 			for _, c := range tunnel.conns {
-				if c != clientConn {
+				if c != actualClientConn {
 					newConns = append(newConns, c)
 				}
 			}
@@ -364,7 +618,7 @@ func (p *SNIProxy) handleConnection(clientConn net.Conn) {
 	}()
 
 	// Start bidirectional data transfer
-	p.pipe(clientConn, targetConn, clientReader)
+	p.pipe(actualClientConn, targetConn, clientReader)
 }
 
 // getRoute retrieves routing information for a hostname
@@ -504,14 +758,20 @@ func (p *SNIProxy) pipe(clientConn, targetConn net.Conn, clientReader io.Reader)
 	var wg sync.WaitGroup
 	wg.Add(2)
 
+	// closeOnce ensures we only close connections once
+	var closeOnce sync.Once
+	closeConns := func() {
+		closeOnce.Do(func() {
+			// Close both connections to unblock any pending reads
+			clientConn.Close()
+			targetConn.Close()
+		})
+	}
+
 	// Copy data from client to target (using the buffered reader)
 	go func() {
 		defer wg.Done()
-		defer func() {
-			if tcpConn, ok := targetConn.(*net.TCPConn); ok {
-				tcpConn.CloseWrite()
-			}
-		}()
+		defer closeConns()
 
 		// Use a large buffer for better performance
 		buf := make([]byte, 32*1024)
@@ -524,11 +784,7 @@ func (p *SNIProxy) pipe(clientConn, targetConn net.Conn, clientReader io.Reader)
 	// Copy data from target to client
 	go func() {
 		defer wg.Done()
-		defer func() {
-			if tcpConn, ok := clientConn.(*net.TCPConn); ok {
-				tcpConn.CloseWrite()
-			}
-		}()
+		defer closeConns()
 
 		// Use a large buffer for better performance
 		buf := make([]byte, 32*1024)

proxy/proxy_test.go

@@ -76,3 +76,44 @@ func TestBuildProxyProtocolHeaderUnknownType(t *testing.T) {
 		t.Errorf("Expected %q, got %q", expected, result)
 	}
 }
+
+func TestBuildProxyProtocolHeaderFromInfo(t *testing.T) {
+	proxy, err := NewSNIProxy(8443, "", "", "127.0.0.1", 443, nil, true, nil)
+	if err != nil {
+		t.Fatalf("Failed to create SNI proxy: %v", err)
+	}
+
+	// Test IPv4 case
+	proxyInfo := &ProxyProtocolInfo{
+		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 := proxy.buildProxyProtocolHeaderFromInfo(proxyInfo, 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
+	proxyInfo = &ProxyProtocolInfo{
+		Protocol: "TCP6",
+		SrcIP:    "2001:db8::1",
+		DestIP:   "2001:db8::2",
+		SrcPort:  12345,
+		DestPort: 443,
+	}
+
+	targetAddr, _ = net.ResolveTCPAddr("tcp6", "[::1]:8080")
+	header = proxy.buildProxyProtocolHeaderFromInfo(proxyInfo, targetAddr)
+
+	expected = "PROXY TCP6 2001:db8::1 ::1 12345 8080\r\n"
+	if header != expected {
+		t.Errorf("Expected header '%s', got '%s'", expected, header)
+	}
+}

relay/relay.go

@@ -2,6 +2,7 @@ package relay
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"encoding/json"
 	"fmt"
@@ -24,20 +25,22 @@ type EncryptedHolePunchMessage struct {
 }
 
 type HolePunchMessage struct {
-	OlmID  string `json:"olmId"`
-	NewtID string `json:"newtId"`
-	Token  string `json:"token"`
+	OlmID     string `json:"olmId"`
+	NewtID    string `json:"newtId"`
+	Token     string `json:"token"`
+	PublicKey string `json:"publicKey"`
 }
 
 type ClientEndpoint struct {
-	OlmID       string `json:"olmId"`
-	NewtID      string `json:"newtId"`
-	Token       string `json:"token"`
-	IP          string `json:"ip"`
-	Port        int    `json:"port"`
-	Timestamp   int64  `json:"timestamp"`
-	ReachableAt string `json:"reachableAt"`
-	PublicKey   string `json:"publicKey"`
+	OlmID             string `json:"olmId"`
+	NewtID            string `json:"newtId"`
+	Token             string `json:"token"`
+	IP                string `json:"ip"`
+	Port              int    `json:"port"`
+	Timestamp         int64  `json:"timestamp"`
+	ReachableAt       string `json:"reachableAt"`
+	ExitNodePublicKey string `json:"exitNodePublicKey"`
+	ClientPublicKey   string `json:"publicKey"`
 }
 
 // Updated to support multiple destination peers
@@ -58,12 +61,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"
 }
@@ -101,22 +144,30 @@ 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
 	// ReachableAt is the URL where this server can be reached
 	ReachableAt string
 }
 
-// NewUDPProxyServer initializes the server with a buffered packet channel.
-func NewUDPProxyServer(addr, serverURL string, privateKey wgtypes.Key, reachableAt string) *UDPProxyServer {
+// NewUDPProxyServer initializes the server with a buffered packet channel and derived context.
+func NewUDPProxyServer(parentCtx context.Context, addr, serverURL string, privateKey wgtypes.Key, reachableAt string) *UDPProxyServer {
+	ctx, cancel := context.WithCancel(parentCtx)
 	return &UDPProxyServer{
 		addr:        addr,
 		serverURL:   serverURL,
 		privateKey:  privateKey,
 		packetChan:  make(chan Packet, 1000),
 		ReachableAt: reachableAt,
+		ctx:         ctx,
+		cancel:      cancel,
 	}
 }
 
@@ -156,21 +207,58 @@ func (s *UDPProxyServer) Start() error {
 	// Start the proxy mapping cleanup routine
 	go s.cleanupIdleProxyMappings()
 
+	// Start the communication pattern cleanup routine
+	go s.cleanupIdleCommunicationPatterns()
+
 	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)
-			continue
+			// If we're shutting down, exit
+			select {
+			case <-s.ctx.Done():
+				bufferPool.Put(buf[:1500])
+				return
+			default:
+				logger.Error("Error reading UDP packet: %v", err)
+				bufferPool.Put(buf[:1500])
+				continue
+			}
 		}
 		s.packetChan <- Packet{data: buf[:n], remoteAddr: remoteAddr, n: n}
 	}
@@ -219,14 +307,15 @@ func (s *UDPProxyServer) packetWorker() {
 			}
 
 			endpoint := ClientEndpoint{
-				NewtID:      msg.NewtID,
-				OlmID:       msg.OlmID,
-				Token:       msg.Token,
-				IP:          packet.remoteAddr.IP.String(),
-				Port:        packet.remoteAddr.Port,
-				Timestamp:   time.Now().Unix(),
-				ReachableAt: s.ReachableAt,
-				PublicKey:   s.privateKey.PublicKey().String(),
+				NewtID:            msg.NewtID,
+				OlmID:             msg.OlmID,
+				Token:             msg.Token,
+				IP:                packet.remoteAddr.IP.String(),
+				Port:              packet.remoteAddr.Port,
+				Timestamp:         time.Now().Unix(),
+				ReachableAt:       s.ReachableAt,
+				ExitNodePublicKey: s.privateKey.PublicKey().String(),
+				ClientPublicKey:   msg.PublicKey,
 			}
 			logger.Debug("Created endpoint from packet remoteAddr %s: IP=%s, Port=%d", packet.remoteAddr.String(), endpoint.IP, endpoint.Port)
 			s.notifyServer(endpoint)
@@ -375,7 +464,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)
 			}
 		}
 
@@ -425,13 +514,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 {
-				// Found matching session
-				destAddr = session.DestAddr
-
-				// Update last seen time
-				session.LastSeen = time.Now()
-				s.wgSessions.Store(k, session)
+			// Check if session matches (read lock for check)
+			if session.GetSenderIndex() == receiverIndex {
+				// Found matching session - get dest addr and update last seen
+				destAddr = session.GetDestAddr()
+				session.UpdateLastSeen()
 				return false // stop iteration
 			}
 			return true // continue iteration
@@ -445,6 +532,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)
@@ -465,6 +555,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)
@@ -548,6 +641,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)
 			}
 		}
 
@@ -562,49 +658,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 {
-		now := time.Now()
-		s.connections.Range(func(key, value interface{}) bool {
-			destConn := value.(*DestinationConn)
-			if now.Sub(destConn.lastUsed) > 10*time.Minute {
-				destConn.conn.Close()
-				s.connections.Delete(key)
-			}
-			return true
-		})
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			now := time.Now()
+			s.connections.Range(func(key, value interface{}) bool {
+				destConn := value.(*DestinationConn)
+				if now.Sub(destConn.lastUsed) > 10*time.Minute {
+					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()
-		s.wgSessions.Range(func(key, value interface{}) bool {
-			session := value.(*WireGuardSession)
-			if now.Sub(session.LastSeen) > 15*time.Minute {
-				s.wgSessions.Delete(key)
-				logger.Debug("Removed idle session: %s", key)
-			}
-			return true
-		})
+
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			now := time.Now()
+			s.wgSessions.Range(func(key, value interface{}) bool {
+				session := value.(*WireGuardSession)
+				// 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 {
-		now := time.Now()
-		s.proxyMappings.Range(func(key, value interface{}) bool {
-			mapping := value.(ProxyMapping)
-			// Remove mappings that haven't been used in 30 minutes
-			if now.Sub(mapping.LastUsed) > 30*time.Minute {
-				s.proxyMappings.Delete(key)
-				logger.Debug("Removed idle proxy mapping: %s", key)
-			}
-			return true
-		})
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			now := time.Now()
+			s.proxyMappings.Range(func(key, value interface{}) bool {
+				mapping := value.(ProxyMapping)
+				// Remove mappings that haven't been used in 30 minutes
+				if now.Sub(mapping.LastUsed) > 30*time.Minute {
+					s.proxyMappings.Delete(key)
+					logger.Debug("Removed idle proxy mapping: %s", key)
+				}
+				return true
+			})
+		case <-s.ctx.Done():
+			return
+		}
 	}
 }
 
@@ -709,8 +825,9 @@ func (s *UDPProxyServer) clearSessionsForIP(ip string) {
 		keyStr := key.(string)
 		session := value.(*WireGuardSession)
 
-		// Check if the session's destination address contains the WG IP
-		if session.DestAddr != nil && session.DestAddr.IP.String() == ip {
+		// Check if the session's destination address contains the WG IP (thread-safe)
+		destAddr := session.GetDestAddr()
+		if destAddr != nil && destAddr.IP.String() == ip {
 			keysToDelete = append(keysToDelete, keyStr)
 			logger.Debug("Marking session for deletion for WG IP %s: %s", ip, keyStr)
 		}
@@ -722,7 +839,7 @@ func (s *UDPProxyServer) clearSessionsForIP(ip string) {
 		s.wgSessions.Delete(key)
 	}
 
-	logger.Info("Cleared %d sessions for WG IP: %s", len(keysToDelete), ip)
+	logger.Debug("Cleared %d sessions for WG IP: %s", len(keysToDelete), ip)
 }
 
 // // clearProxyMappingsForWGIP removes all proxy mappings that have destinations pointing to a specific WireGuard IP
@@ -823,3 +940,121 @@ 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
+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
+		}
+	}
+}