From d21c09c84f9e8c6a06732a2ff92ad912ff17129a Mon Sep 17 00:00:00 2001
From: Laurence <laurence.jones@live.co.uk>
Date: Fri, 13 Mar 2026 16:36:56 +0000
Subject: [PATCH] refactor(proxy): simplify tunnel tracking with mutex-only
 approach

Remove atomic counter in favor of simple int protected by mutex.
Eliminates race condition complexity and recheck logic.
---
 proxy/proxy.go | 18 +++++++-----------
 1 file changed, 7 insertions(+), 11 deletions(-)

diff --git a/proxy/proxy.go b/proxy/proxy.go
index 35759a0..4d814e9 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -14,7 +14,6 @@ import (
 	"strconv"
 	"strings"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/fosrl/gerbil/logger"
@@ -76,7 +75,7 @@ type SNIProxy struct {
 type activeTunnel struct {
 	ctx    context.Context
 	cancel context.CancelFunc
-	count  atomic.Int64
+	count  int // protected by activeTunnelsLock
 }
 
 // readOnlyConn is a wrapper for io.Reader that implements net.Conn
@@ -600,21 +599,18 @@ func (p *SNIProxy) handleConnection(clientConn net.Conn) {
 		tunnel = &activeTunnel{ctx: ctx, cancel: cancel}
 		p.activeTunnels[hostname] = tunnel
 	}
-	tunnel.count.Add(1)
+	tunnel.count++
 	tunnelCtx := tunnel.ctx
 	p.activeTunnelsLock.Unlock()
 
 	defer func() {
-		// Decrement count atomically; if we're the last connection, clean up
-		if tunnel.count.Add(-1) == 0 {
+		p.activeTunnelsLock.Lock()
+		tunnel.count--
+		if tunnel.count == 0 {
 			tunnel.cancel()
-			p.activeTunnelsLock.Lock()
-			// Only delete if the map still points to our tunnel
-			if p.activeTunnels[hostname] == tunnel {
-				delete(p.activeTunnels, hostname)
-			}
-			p.activeTunnelsLock.Unlock()
+			delete(p.activeTunnels, hostname)
 		}
+		p.activeTunnelsLock.Unlock()
 	}()
 
 	// Start bidirectional data transfer with tunnel context