Compare commits

..

2 Commits

Author SHA1 Message Date
Laurence
14a3e7c531 Optimize lock usage in proxy connection handling
- Change activeTunnel.conns from slice to map for O(1) add/remove
- Improve lock scoping in UpdateLocalSNIs: use read lock for diff
  computation, minimize write lock hold time
- Move cache invalidation outside lock (go-cache is thread-safe)
2026-03-13 15:23:30 +00:00
Owen
b9261b8fea Add optional tc 2026-02-27 15:45:17 -08:00
8 changed files with 254 additions and 65 deletions

View File

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

View File

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

View File

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

View File

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

10
go.mod
View File

@@ -1,12 +1,12 @@
module github.com/fosrl/gerbil module github.com/fosrl/gerbil
go 1.26.0 go 1.25
require ( require (
github.com/patrickmn/go-cache v2.1.0+incompatible github.com/patrickmn/go-cache v2.1.0+incompatible
github.com/vishvananda/netlink v1.3.1 github.com/vishvananda/netlink v1.3.1
golang.org/x/crypto v0.49.0 golang.org/x/crypto v0.46.0
golang.org/x/sync v0.20.0 golang.org/x/sync v0.1.0
golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6 golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6
) )
@@ -17,7 +17,7 @@ require (
github.com/mdlayher/netlink v1.7.2 // indirect github.com/mdlayher/netlink v1.7.2 // indirect
github.com/mdlayher/socket v0.4.1 // indirect github.com/mdlayher/socket v0.4.1 // indirect
github.com/vishvananda/netns v0.0.5 // indirect github.com/vishvananda/netns v0.0.5 // indirect
golang.org/x/net v0.51.0 // indirect golang.org/x/net v0.47.0 // indirect
golang.org/x/sys v0.42.0 // indirect golang.org/x/sys v0.39.0 // indirect
golang.zx2c4.com/wireguard v0.0.0-20230325221338-052af4a8072b // indirect golang.zx2c4.com/wireguard v0.0.0-20230325221338-052af4a8072b // indirect
) )

16
go.sum
View File

@@ -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/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 h1:DfiHV+j8bA32MFM7bfEunvT8IAqQ/NzSJHtcmW5zdEY=
github.com/vishvananda/netns v0.0.5/go.mod h1:SpkAiCQRtJ6TvvxPnOSyH3BMl6unz3xZlaprSwhNNJM= github.com/vishvananda/netns v0.0.5/go.mod h1:SpkAiCQRtJ6TvvxPnOSyH3BMl6unz3xZlaprSwhNNJM=
golang.org/x/crypto v0.49.0 h1:+Ng2ULVvLHnJ/ZFEq4KdcDd/cfjrrjjNSXNzxg0Y4U4= golang.org/x/crypto v0.46.0 h1:cKRW/pmt1pKAfetfu+RCEvjvZkA9RimPbh7bhFjGVBU=
golang.org/x/crypto v0.49.0/go.mod h1:ErX4dUh2UM+CFYiXZRTcMpEcN8b/1gxEuv3nODoYtCA= golang.org/x/crypto v0.46.0/go.mod h1:Evb/oLKmMraqjZ2iQTwDwvCtJkczlDuTmdJXoZVzqU0=
golang.org/x/net v0.51.0 h1:94R/GTO7mt3/4wIKpcR5gkGmRLOuE/2hNGeWq/GBIFo= golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY=
golang.org/x/net v0.51.0/go.mod h1:aamm+2QF5ogm02fjy5Bb7CQ0WMt1/WVM7FtyaTLlA9Y= golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU=
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o=
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= 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.2.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.10.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.10.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo= golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk=
golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= 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 h1:J1CaxgLerRR5lgx3wnr6L04cJFbWoceSK9JWBdglINo=
golang.zx2c4.com/wireguard v0.0.0-20230325221338-052af4a8072b/go.mod h1:tqur9LnfstdR9ep2LaJT4lFUl0EjlHtge+gAjmsHUG4= 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= golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6 h1:CawjfCvYQH2OU3/TnxLx97WDSUDRABfT18pCOYwc2GE=

228
main.go
View File

@@ -33,15 +33,16 @@ import (
) )
var ( var (
interfaceName string interfaceName string
listenAddr string listenAddr string
mtuInt int mtuInt int
lastReadings = make(map[string]PeerReading) lastReadings = make(map[string]PeerReading)
mu sync.Mutex mu sync.Mutex
wgMu sync.Mutex // Protects WireGuard operations wgMu sync.Mutex // Protects WireGuard operations
notifyURL string notifyURL string
proxyRelay *relay.UDPProxyServer proxyRelay *relay.UDPProxyServer
proxySNI *proxy.SNIProxy proxySNI *proxy.SNIProxy
doTrafficShaping bool
) )
type WgConfig struct { type WgConfig struct {
@@ -151,6 +152,7 @@ func main() {
localOverridesStr = os.Getenv("LOCAL_OVERRIDES") localOverridesStr = os.Getenv("LOCAL_OVERRIDES")
trustedUpstreamsStr = os.Getenv("TRUSTED_UPSTREAMS") trustedUpstreamsStr = os.Getenv("TRUSTED_UPSTREAMS")
proxyProtocolStr := os.Getenv("PROXY_PROTOCOL") proxyProtocolStr := os.Getenv("PROXY_PROTOCOL")
doTrafficShapingStr := os.Getenv("DO_TRAFFIC_SHAPING")
if interfaceName == "" { if interfaceName == "" {
flag.StringVar(&interfaceName, "interface", "wg0", "Name of the WireGuard interface") flag.StringVar(&interfaceName, "interface", "wg0", "Name of the WireGuard interface")
@@ -222,6 +224,13 @@ func main() {
flag.BoolVar(&proxyProtocol, "proxy-protocol", true, "Enable PROXY protocol v1 for preserving client IP") 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() flag.Parse()
logger.Init() logger.Init()
@@ -886,17 +895,23 @@ func addPeerInternal(peer Peer) error {
return fmt.Errorf("failed to parse public key: %v", err) 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 // parse allowed IPs into array of net.IPNet
var allowedIPs []net.IPNet var allowedIPs []net.IPNet
var wgIPs []string var wgIPs []string
for _, ipStr := range peer.AllowedIPs { for _, ipStr := range peer.AllowedIPs {
logger.Debug("Parsing AllowedIP: %s", ipStr)
_, ipNet, err := net.ParseCIDR(ipStr) _, ipNet, err := net.ParseCIDR(ipStr)
if err != nil { 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) return fmt.Errorf("failed to parse allowed IP: %v", err)
} }
allowedIPs = append(allowedIPs, *ipNet) allowedIPs = append(allowedIPs, *ipNet)
// Extract the IP address from the CIDR for relay cleanup // 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{ peerConfig := wgtypes.PeerConfig{
@@ -912,6 +927,18 @@ func addPeerInternal(peer Peer) error {
return fmt.Errorf("failed to add peer: %v", err) 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 // Clear relay connections for the peer's WireGuard IPs
if proxyRelay != nil { if proxyRelay != nil {
for _, wgIP := range wgIPs { for _, wgIP := range wgIPs {
@@ -956,19 +983,17 @@ func removePeerInternal(publicKey string) error {
return fmt.Errorf("failed to parse public key: %v", err) 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 var wgIPs []string
if proxyRelay != nil { device, err := wgClient.Device(interfaceName)
device, err := wgClient.Device(interfaceName) if err == nil {
if err == nil { for _, peer := range device.Peers {
for _, peer := range device.Peers { if peer.PublicKey.String() == publicKey {
if peer.PublicKey.String() == publicKey { // Extract WireGuard IPs from this peer's allowed IPs
// Extract WireGuard IPs from this peer's allowed IPs for _, allowedIP := range peer.AllowedIPs {
for _, allowedIP := range peer.AllowedIPs { wgIPs = append(wgIPs, allowedIP.IP.String())
wgIPs = append(wgIPs, allowedIP.IP.String())
}
break
} }
break
} }
} }
} }
@@ -986,6 +1011,15 @@ func removePeerInternal(publicKey string) error {
return fmt.Errorf("failed to remove peer: %v", err) 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 // Clear relay connections for the peer's WireGuard IPs
if proxyRelay != nil { if proxyRelay != nil {
for _, wgIP := range wgIPs { for _, wgIP := range wgIPs {
@@ -1315,3 +1349,155 @@ func monitorMemory(limit uint64) {
time.Sleep(5 * time.Second) 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
}

View File

@@ -72,7 +72,7 @@ type SNIProxy struct {
} }
type activeTunnel struct { type activeTunnel struct {
conns []net.Conn conns map[net.Conn]struct{}
} }
// readOnlyConn is a wrapper for io.Reader that implements net.Conn // readOnlyConn is a wrapper for io.Reader that implements net.Conn
@@ -592,26 +592,19 @@ func (p *SNIProxy) handleConnection(clientConn net.Conn) {
p.activeTunnelsLock.Lock() p.activeTunnelsLock.Lock()
tunnel, ok := p.activeTunnels[hostname] tunnel, ok := p.activeTunnels[hostname]
if !ok { if !ok {
tunnel = &activeTunnel{} tunnel = &activeTunnel{conns: make(map[net.Conn]struct{})}
p.activeTunnels[hostname] = tunnel p.activeTunnels[hostname] = tunnel
} }
tunnel.conns = append(tunnel.conns, actualClientConn) tunnel.conns[actualClientConn] = struct{}{}
p.activeTunnelsLock.Unlock() p.activeTunnelsLock.Unlock()
defer func() { defer func() {
// Remove this conn from active tunnels // Remove this conn from active tunnels - O(1) with map
p.activeTunnelsLock.Lock() p.activeTunnelsLock.Lock()
if tunnel, ok := p.activeTunnels[hostname]; ok { if tunnel, ok := p.activeTunnels[hostname]; ok {
newConns := make([]net.Conn, 0, len(tunnel.conns)) delete(tunnel.conns, actualClientConn)
for _, c := range tunnel.conns { if len(tunnel.conns) == 0 {
if c != actualClientConn {
newConns = append(newConns, c)
}
}
if len(newConns) == 0 {
delete(p.activeTunnels, hostname) delete(p.activeTunnels, hostname)
} else {
tunnel.conns = newConns
} }
} }
p.activeTunnelsLock.Unlock() p.activeTunnelsLock.Unlock()
@@ -810,32 +803,42 @@ func (p *SNIProxy) ClearCache() {
// UpdateLocalSNIs updates the local SNIs and invalidates cache for changed domains // UpdateLocalSNIs updates the local SNIs and invalidates cache for changed domains
func (p *SNIProxy) UpdateLocalSNIs(fullDomains []string) { func (p *SNIProxy) UpdateLocalSNIs(fullDomains []string) {
newSNIs := make(map[string]struct{}) newSNIs := make(map[string]struct{}, len(fullDomains))
for _, domain := range fullDomains { for _, domain := range fullDomains {
newSNIs[domain] = struct{}{} newSNIs[domain] = struct{}{}
// Invalidate any cached route for this domain
p.cache.Delete(domain)
} }
// Update localSNIs // Get old SNIs with read lock to compute diff outside write lock
p.localSNIsLock.Lock() p.localSNIsLock.RLock()
oldSNIs := p.localSNIs
p.localSNIsLock.RUnlock()
// Compute removed SNIs outside the lock
removed := make([]string, 0) removed := make([]string, 0)
for sni := range p.localSNIs { for sni := range oldSNIs {
if _, stillLocal := newSNIs[sni]; !stillLocal { if _, stillLocal := newSNIs[sni]; !stillLocal {
removed = append(removed, sni) removed = append(removed, sni)
} }
} }
// Swap with minimal write lock hold time
p.localSNIsLock.Lock()
p.localSNIs = newSNIs p.localSNIs = newSNIs
p.localSNIsLock.Unlock() p.localSNIsLock.Unlock()
// Invalidate cache for new domains (cache is thread-safe)
for domain := range newSNIs {
p.cache.Delete(domain)
}
logger.Debug("Updated local SNIs, added %d, removed %d", len(newSNIs), len(removed)) logger.Debug("Updated local SNIs, added %d, removed %d", len(newSNIs), len(removed))
// Terminate tunnels for removed SNIs // Terminate tunnels for removed SNIs
if len(removed) > 0 { if len(removed) > 0 {
p.activeTunnelsLock.Lock() p.activeTunnelsLock.Lock()
for _, sni := range removed { for _, sni := range removed {
if tunnels, ok := p.activeTunnels[sni]; ok { if tunnel, ok := p.activeTunnels[sni]; ok {
for _, conn := range tunnels.conns { for conn := range tunnel.conns {
conn.Close() conn.Close()
} }
delete(p.activeTunnels, sni) delete(p.activeTunnels, sni)