From 72d737ef4d154a7ba9d63d6e1dadea78c8e91501 Mon Sep 17 00:00:00 2001
From: Aditya kumar singh <143548997+Adityakk9031@users.noreply.github.com>
Date: Tue, 21 Jul 2026 17:47:03 +0530
Subject: [PATCH] fix: prevent panic on double Stop() of UDPProxyServer (#51)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The UDPProxyServer.Stop() method was being called twice on every clean
shutdown — once via defer proxyRelay.Stop() in main() and once from
the shutdown errgroup goroutine — causing a 'close of closed channel'
panic.

The old guard using a bare select on s.ctx.Done() was ineffective: by
the second call the context was already cancelled, so both code paths
fell through unconditionally to close(s.packetChan), panicking.

Fix: wrap the entire Stop() body in a sync.Once so that all cleanup
(cancel, conn close, packetChan close) runs at most once regardless of
how many times or from how many goroutines Stop() is called.

Fixes #51
---
 relay/relay.go | 44 +++++++++++++++++++++++++-------------------
 1 file changed, 25 insertions(+), 19 deletions(-)

diff --git a/relay/relay.go b/relay/relay.go
index 4d0d4df..df48136 100644
--- a/relay/relay.go
+++ b/relay/relay.go
@@ -205,6 +205,10 @@ type UDPProxyServer struct {
 	packetChan      chan Packet
 	ctx             context.Context
 	cancel          context.CancelFunc
+	// stopOnce ensures Stop() is idempotent: packetChan is closed at most once
+	// even when Stop() is called concurrently or multiple times (e.g. via both
+	// defer and the shutdown goroutine in main).
+	stopOnce sync.Once
 
 	// Session tracking for WireGuard peers
 	// Key format: "senderIndex:receiverIndex"
@@ -305,28 +309,30 @@ func (s *UDPProxyServer) Start() error {
 	return nil
 }
 
+// Stop shuts down the UDPProxyServer. It is safe to call multiple times
+// and from concurrent goroutines; the internal cleanup runs at most once.
 func (s *UDPProxyServer) Stop() {
-	// 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()
+	s.stopOnce.Do(func() {
+		// Signal all background goroutines to stop
+		if s.cancel != nil {
+			s.cancel()
 		}
-		return true
+		// 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 unblock and stop all packetWorker goroutines.
+		// This must happen after cancel() so that readPackets() exits cleanly
+		// and no further sends to packetChan occur before we close it.
+		close(s.packetChan)
 	})
-	// 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.