Add optional tc

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

@@ -36,13 +36,13 @@ jobs:
 
     steps:
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
 
       - name: Set up QEMU
         uses: docker/setup-qemu-action@c7c53464625b32c7a7e944ae62b3e17d2b600130 # v3.7.0
 
       - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 # v3.11.1
+        uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
 
       - name: Log in to Docker Hub
         uses: docker/login-action@5e57cd118135c172c3672efd75eb46360885c0ef # v3.6.0
@@ -57,7 +57,7 @@ jobs:
         shell: bash
 
       - name: Install Go
-        uses: actions/setup-go@4dc6199c7b1a012772edbd06daecab0f50c9053c # v6.1.0
+        uses: actions/setup-go@7a3fe6cf4cb3a834922a1244abfce67bcef6a0c5 # v6.2.0
         with:
           go-version: 1.25
 
@@ -155,7 +155,7 @@ jobs:
         shell: bash
 
       - name: Upload artifacts from /bin
-        uses: actions/upload-artifact@330a01c490aca151604b8cf639adc76d48f6c5d4 # v5.0.0
+        uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
         with:
           name: binaries
           path: bin/

.github/workflows/test.yml

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

README.md

@@ -42,13 +42,12 @@ In single node (self hosted) Pangolin deployments this can be bypassed by using
 
 ## CLI Args
 
+Important:
 - `reachableAt`: How should the remote server reach Gerbil's API?
 - `generateAndSaveKeyTo`: Where to save the generated WireGuard private key to persist across restarts.
-- `remoteConfig` (optional): Remote config location to HTTP get the JSON based config from. See `example_config.json`
-- `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: `:3004`
@@ -66,7 +65,6 @@ Note: You must use either `config` or `remoteConfig` to configure WireGuard.
 All CLI arguments can also be provided via environment variables:
 
 - `INTERFACE`: Name of the WireGuard interface
-- `CONFIG`: Path to local configuration file
 - `REMOTE_CONFIG`: URL of the remote config server
 - `LISTEN`: Address to listen on for HTTP server
 - `GENERATE_AND_SAVE_KEY_TO`: Path to save generated private key

config_example.json

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

go.mod

@@ -5,7 +5,8 @@ go 1.25
 require (
 	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/vishvananda/netlink v1.3.1
-	golang.org/x/crypto v0.45.0
+	golang.org/x/crypto v0.46.0
+	golang.org/x/sync v0.1.0
 	golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6
 )
 
@@ -17,7 +18,6 @@ require (
 	github.com/mdlayher/socket v0.4.1 // indirect
 	github.com/vishvananda/netns v0.0.5 // indirect
 	golang.org/x/net v0.47.0 // indirect
-	golang.org/x/sync v0.1.0 // indirect
-	golang.org/x/sys v0.38.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.45.0 h1:jMBrvKuj23MTlT0bQEOBcAE0mjg8mK9RXFhRH6nyF3Q=
-golang.org/x/crypto v0.45.0/go.mod h1:XTGrrkGJve7CYK7J8PEww4aY7gM3qMCElcJQ8n8JdX4=
+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.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc=
-golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
+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

@@ -8,11 +8,15 @@ import (
 	"flag"
 	"fmt"
 	"io"
+	"log"
 	"net"
 	"net/http"
+	_ "net/http/pprof"
 	"os"
 	"os/exec"
 	"os/signal"
+	"runtime"
+	"runtime/pprof"
 	"strconv"
 	"strings"
 	"sync"
@@ -29,20 +33,22 @@ import (
 )
 
 var (
-	interfaceName string
-	listenAddr    string
-	mtuInt        int
-	lastReadings  = make(map[string]PeerReading)
-	mu            sync.Mutex
-	wgMu          sync.Mutex // Protects WireGuard operations
-	notifyURL     string
-	proxyRelay    *relay.UDPProxyServer
-	proxySNI      *proxy.SNIProxy
+	interfaceName    string
+	listenAddr       string
+	mtuInt           int
+	lastReadings     = make(map[string]PeerReading)
+	mu               sync.Mutex
+	wgMu             sync.Mutex // Protects WireGuard operations
+	notifyURL        string
+	proxyRelay       *relay.UDPProxyServer
+	proxySNI         *proxy.SNIProxy
+	doTrafficShaping bool
 )
 
 type WgConfig struct {
 	PrivateKey string `json:"privateKey"`
 	ListenPort int    `json:"listenPort"`
+	RelayPort  int    `json:"relayPort"`
 	IpAddress  string `json:"ipAddress"`
 	Peers      []Peer `json:"peers"`
 }
@@ -111,6 +117,8 @@ func parseLogLevel(level string) logger.LogLevel {
 }
 
 func main() {
+	go monitorMemory(1024 * 1024 * 512) // trigger if memory usage exceeds 512MB
+
 	var (
 		err                  error
 		wgconfig             WgConfig
@@ -144,6 +152,7 @@ func main() {
 	localOverridesStr = os.Getenv("LOCAL_OVERRIDES")
 	trustedUpstreamsStr = os.Getenv("TRUSTED_UPSTREAMS")
 	proxyProtocolStr := os.Getenv("PROXY_PROTOCOL")
+	doTrafficShapingStr := os.Getenv("DO_TRAFFIC_SHAPING")
 
 	if interfaceName == "" {
 		flag.StringVar(&interfaceName, "interface", "wg0", "Name of the WireGuard interface")
@@ -215,6 +224,13 @@ func main() {
 		flag.BoolVar(&proxyProtocol, "proxy-protocol", true, "Enable PROXY protocol v1 for preserving client IP")
 	}
 
+	if doTrafficShapingStr != "" {
+		doTrafficShaping = strings.ToLower(doTrafficShapingStr) == "true"
+	}
+	if doTrafficShapingStr == "" {
+		flag.BoolVar(&doTrafficShaping, "do-traffic-shaping", false, "Whether to set up traffic shaping rules for peers (requires tc command and root privileges)")
+	}
+
 	flag.Parse()
 
 	logger.Init()
@@ -340,7 +356,11 @@ func main() {
 	})
 
 	// Start the UDP proxy server
-	proxyRelay = relay.NewUDPProxyServer(groupCtx, ":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)
@@ -382,6 +402,7 @@ func main() {
 	http.HandleFunc("/update-proxy-mapping", handleUpdateProxyMapping)
 	http.HandleFunc("/update-destinations", handleUpdateDestinations)
 	http.HandleFunc("/update-local-snis", handleUpdateLocalSNIs)
+	http.HandleFunc("/healthz", handleHealthz)
 	logger.Info("Starting HTTP server on %s", listenAddr)
 
 	// HTTP server with graceful shutdown on context cancel
@@ -543,6 +564,10 @@ func ensureWireguardInterface(wgconfig WgConfig) error {
 		logger.Warn("Failed to ensure MSS clamping: %v", err)
 	}
 
+	if err := ensureWireguardFirewall(); err != nil {
+		logger.Warn("Failed to ensure WireGuard firewall rules: %v", err)
+	}
+
 	logger.Info("WireGuard interface %s created and configured", interfaceName)
 
 	return nil
@@ -711,6 +736,113 @@ func ensureMSSClamping() error {
 	return nil
 }
 
+func ensureWireguardFirewall() error {
+	// Rules to enforce:
+	// 1. Allow established/related connections (responses to our outbound traffic)
+	// 2. Allow ICMP ping packets
+	// 3. Drop all other inbound traffic from peers
+
+	// Define the rules we want to ensure exist
+	rules := [][]string{
+		// Allow established and related connections (responses to outbound traffic)
+		{
+			"-A", "INPUT",
+			"-i", interfaceName,
+			"-m", "conntrack",
+			"--ctstate", "ESTABLISHED,RELATED",
+			"-j", "ACCEPT",
+		},
+		// Allow ICMP ping requests
+		{
+			"-A", "INPUT",
+			"-i", interfaceName,
+			"-p", "icmp",
+			"--icmp-type", "8",
+			"-j", "ACCEPT",
+		},
+		// Drop all other inbound traffic from WireGuard interface
+		{
+			"-A", "INPUT",
+			"-i", interfaceName,
+			"-j", "DROP",
+		},
+	}
+
+	// First, try to delete any existing rules for this interface
+	for _, rule := range rules {
+		deleteArgs := make([]string, len(rule))
+		copy(deleteArgs, rule)
+		// Change -A to -D for deletion
+		for i, arg := range deleteArgs {
+			if arg == "-A" {
+				deleteArgs[i] = "-D"
+				break
+			}
+		}
+
+		deleteCmd := exec.Command("/usr/sbin/iptables", deleteArgs...)
+		logger.Debug("Attempting to delete existing firewall rule: %v", deleteArgs)
+
+		// Try deletion multiple times to handle multiple existing rules
+		for i := 0; i < 5; i++ {
+			out, err := deleteCmd.CombinedOutput()
+			if err != nil {
+				if exitErr, ok := err.(*exec.ExitError); ok {
+					logger.Debug("Deletion stopped: %v (output: %s)", exitErr.String(), string(out))
+				}
+				break // No more rules to delete
+			}
+			logger.Info("Deleted existing firewall rule (attempt %d)", i+1)
+		}
+	}
+
+	// Now add the rules
+	var errors []error
+	for i, rule := range rules {
+		addCmd := exec.Command("/usr/sbin/iptables", rule...)
+		logger.Info("Adding WireGuard firewall rule %d: %v", i+1, rule)
+
+		if out, err := addCmd.CombinedOutput(); err != nil {
+			errMsg := fmt.Sprintf("Failed to add firewall rule %d: %v (output: %s)", i+1, err, string(out))
+			logger.Error("%s", errMsg)
+			errors = append(errors, fmt.Errorf("%s", errMsg))
+			continue
+		}
+
+		// Verify the rule was added by checking
+		checkArgs := make([]string, len(rule))
+		copy(checkArgs, rule)
+		// Change -A to -C for check
+		for j, arg := range checkArgs {
+			if arg == "-A" {
+				checkArgs[j] = "-C"
+				break
+			}
+		}
+
+		checkCmd := exec.Command("/usr/sbin/iptables", checkArgs...)
+		if out, err := checkCmd.CombinedOutput(); err != nil {
+			errMsg := fmt.Sprintf("Rule verification failed for rule %d: %v (output: %s)", i+1, err, string(out))
+			logger.Error("%s", errMsg)
+			errors = append(errors, fmt.Errorf("%s", errMsg))
+			continue
+		}
+
+		logger.Info("Successfully added and verified WireGuard firewall rule %d", i+1)
+	}
+
+	if len(errors) > 0 {
+		var errMsgs []string
+		for _, err := range errors {
+			errMsgs = append(errMsgs, err.Error())
+		}
+		return fmt.Errorf("WireGuard firewall setup encountered errors:\n%s", strings.Join(errMsgs, "\n"))
+	}
+
+	logger.Info("WireGuard firewall rules successfully configured for interface %s", interfaceName)
+	return nil
+}
+
 func handlePeer(w http.ResponseWriter, r *http.Request) {
 	switch r.Method {
 	case http.MethodPost:
@@ -722,6 +854,15 @@ func handlePeer(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+func handleHealthz(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodGet {
+		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
+		return
+	}
+	w.WriteHeader(http.StatusOK)
+	_, _ = w.Write([]byte("ok"))
+}
+
 func handleAddPeer(w http.ResponseWriter, r *http.Request) {
 	var peer Peer
 	if err := json.NewDecoder(r.Body).Decode(&peer); err != nil {
@@ -754,17 +895,23 @@ func addPeerInternal(peer Peer) error {
 		return fmt.Errorf("failed to parse public key: %v", err)
 	}
 
+	logger.Debug("Adding peer %s with AllowedIPs: %v", peer.PublicKey, peer.AllowedIPs)
+
 	// parse allowed IPs into array of net.IPNet
 	var allowedIPs []net.IPNet
 	var wgIPs []string
 	for _, ipStr := range peer.AllowedIPs {
+		logger.Debug("Parsing AllowedIP: %s", ipStr)
 		_, ipNet, err := net.ParseCIDR(ipStr)
 		if err != nil {
+			logger.Warn("Failed to parse allowed IP '%s' for peer %s: %v", ipStr, peer.PublicKey, err)
 			return fmt.Errorf("failed to parse allowed IP: %v", err)
 		}
 		allowedIPs = append(allowedIPs, *ipNet)
 		// Extract the IP address from the CIDR for relay cleanup
-		wgIPs = append(wgIPs, ipNet.IP.String())
+		extractedIP := ipNet.IP.String()
+		wgIPs = append(wgIPs, extractedIP)
+		logger.Debug("Extracted IP %s from AllowedIP %s", extractedIP, ipStr)
 	}
 
 	peerConfig := wgtypes.PeerConfig{
@@ -780,6 +927,18 @@ func addPeerInternal(peer Peer) error {
 		return fmt.Errorf("failed to add peer: %v", err)
 	}
 
+	// Setup bandwidth limiting for each peer IP
+	if doTrafficShaping {
+		logger.Debug("doTrafficShaping is true, setting up bandwidth limits for %d IPs", len(wgIPs))
+		for _, wgIP := range wgIPs {
+			if err := setupPeerBandwidthLimit(wgIP); err != nil {
+				logger.Warn("Failed to setup bandwidth limit for peer IP %s: %v", wgIP, err)
+			}
+		}
+	} else {
+		logger.Debug("doTrafficShaping is false, skipping bandwidth limit setup")
+	}
+
 	// Clear relay connections for the peer's WireGuard IPs
 	if proxyRelay != nil {
 		for _, wgIP := range wgIPs {
@@ -824,19 +983,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 proxyRelay != nil {
-		device, err := wgClient.Device(interfaceName)
-		if err == nil {
-			for _, peer := range device.Peers {
-				if peer.PublicKey.String() == publicKey {
-					// Extract WireGuard IPs from this peer's allowed IPs
-					for _, allowedIP := range peer.AllowedIPs {
-						wgIPs = append(wgIPs, allowedIP.IP.String())
-					}
-					break
+	device, err := wgClient.Device(interfaceName)
+	if err == nil {
+		for _, peer := range device.Peers {
+			if peer.PublicKey.String() == publicKey {
+				// Extract WireGuard IPs from this peer's allowed IPs
+				for _, allowedIP := range peer.AllowedIPs {
+					wgIPs = append(wgIPs, allowedIP.IP.String())
 				}
+				break
 			}
 		}
 	}
@@ -854,6 +1011,15 @@ func removePeerInternal(publicKey string) error {
 		return fmt.Errorf("failed to remove peer: %v", err)
 	}
 
+	// Remove bandwidth limits for each peer IP
+	if doTrafficShaping {
+		for _, wgIP := range wgIPs {
+			if err := removePeerBandwidthLimit(wgIP); err != nil {
+				logger.Warn("Failed to remove bandwidth limit for peer IP %s: %v", wgIP, err)
+			}
+		}
+	}
+
 	// Clear relay connections for the peer's WireGuard IPs
 	if proxyRelay != nil {
 		for _, wgIP := range wgIPs {
@@ -1161,3 +1327,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

@@ -758,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)
@@ -778,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)

relay/relay.go

@@ -464,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)
 			}
 		}
 
@@ -839,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