From 6daeb9a59c2f05ca70e694ffe36c55e8423bd9a2 Mon Sep 17 00:00:00 2001
From: Rowan Seymour <rowanseymour@gmail.com>
Date: Mon, 15 Jan 2024 16:05:47 -0500
Subject: [PATCH] Panic if trying to close or start and already closed socket

---
 httpx/websocket.go      |  8 ++++++++
 httpx/websocket_test.go | 11 +++++++++++
 2 files changed, 19 insertions(+)

diff --git a/httpx/websocket.go b/httpx/websocket.go
index 8b9d439..8fd5022 100644
--- a/httpx/websocket.go
+++ b/httpx/websocket.go
@@ -99,6 +99,10 @@ func (s *socket) OnMessage(fn func([]byte)) { s.onMessage = fn }
 func (s *socket) OnClose(fn func(int))      { s.onClose = fn }
 
 func (s *socket) Start() {
+	if s.closingWithCode != 0 {
+		panic("can't start socket which is closed or closing")
+	}
+
 	s.conn.SetReadDeadline(time.Now().Add(maxReadWait))
 	s.conn.SetPongHandler(s.pong)
 
@@ -116,6 +120,10 @@ func (s *socket) Send(msg []byte) {
 }
 
 func (s *socket) Close(code int) {
+	if s.closingWithCode != 0 {
+		panic("can't close socket which is already closed or closing")
+	}
+
 	s.closingWithCode = code
 
 	s.outbox <- message{type_: websocket.CloseMessage, data: websocket.FormatCloseMessage(code, "")}
diff --git a/httpx/websocket_test.go b/httpx/websocket_test.go
index 696fbf7..c9a3f1f 100644
--- a/httpx/websocket_test.go
+++ b/httpx/websocket_test.go
@@ -105,6 +105,17 @@ func TestSocketMessages(t *testing.T) {
 
 	assert.Equal(t, 1001, serverCloseCode)
 	assert.Equal(t, 1001, connCloseCode)
+
+	// check we can no longer send to the socket or close it again, or restart it
+	assert.Panics(t, func() {
+		sock.Send([]byte("x"))
+	})
+	assert.Panics(t, func() {
+		sock.Close(1000)
+	})
+	assert.Panics(t, func() {
+		sock.Start()
+	})
 }
 
 func TestSocketClientCloseWithMessage(t *testing.T) {