From d8e316299e3b029efd4c6cd8e34d97478a8fc8ef Mon Sep 17 00:00:00 2001
From: Abhishek Kondur <abhi98.kondur@gmail.com>
Date: Mon, 31 Aug 2026 19:28:04 +0530
Subject: [PATCH] Fix: Capture the done channel in a local variable and close
 that instead. Added regression tests in

---
 uplink/websocket.go      |  7 +++-
 uplink/websocket_test.go | 81 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 86 insertions(+), 2 deletions(-)
 create mode 100644 uplink/websocket_test.go

diff --git a/uplink/websocket.go b/uplink/websocket.go
index 601f2ab..7c7cb2e 100644
--- a/uplink/websocket.go
+++ b/uplink/websocket.go
@@ -105,11 +105,14 @@ func (c *WebSocketConn) StartPingLoop(ctx context.Context) {
 	}
 	pctx, cancel := context.WithCancel(ctx)
 	c.pingCancel = cancel
-	c.pingDone = make(chan struct{})
+	done := make(chan struct{})
+	c.pingDone = done
 	c.pingMu.Unlock()
 
 	go func() {
-		defer close(c.pingDone)
+		// Close the local done channel, not c.pingDone: stopPing nils the field
+		// before waiting, so defer close(c.pingDone) would panic on a nil channel.
+		defer close(done)
 		t := time.NewTicker(c.pingInterval)
 		defer t.Stop()
 		for {
diff --git a/uplink/websocket_test.go b/uplink/websocket_test.go
new file mode 100644
index 0000000..48e5f05
--- /dev/null
+++ b/uplink/websocket_test.go
@@ -0,0 +1,81 @@
+package uplink
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+func TestWebSocketConn_StartPingLoopCloseNoPanic(t *testing.T) {
+	upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ws, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		defer ws.Close()
+		conn := NewWebSocketConn(ws, DefaultMaxMessageSize)
+		conn.pingInterval = 10 * time.Millisecond
+		conn.StartPingLoop(r.Context())
+		// Keep the connection open briefly so the ping loop is running,
+		// then Close (stopPing) races with the goroutine exit.
+		time.Sleep(30 * time.Millisecond)
+		_ = conn.Close()
+	}))
+	defer srv.Close()
+
+	wsURL := "ws" + srv.URL[len("http"):]
+	cli, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cli.Close()
+
+	// Drain until server closes; panic in the ping loop would fail the test process.
+	deadline := time.Now().Add(2 * time.Second)
+	for time.Now().Before(deadline) {
+		_ = cli.SetReadDeadline(time.Now().Add(50 * time.Millisecond))
+		if _, _, err := cli.ReadMessage(); err != nil {
+			return
+		}
+	}
+}
+
+func TestWebSocketConn_stopPingIdempotent(t *testing.T) {
+	upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
+	done := make(chan struct{})
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		defer close(done)
+		ws, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		conn := NewWebSocketConn(ws, DefaultMaxMessageSize)
+		conn.pingInterval = time.Hour
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+		conn.StartPingLoop(ctx)
+		conn.stopPing()
+		conn.stopPing() // second stop must not hang or panic
+		_ = conn.ws.Close()
+	}))
+	defer srv.Close()
+
+	wsURL := "ws" + srv.URL[len("http"):]
+	cli, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	cli.Close()
+
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("server handler did not finish")
+	}
+}