From dc96c9588265f4785cf507ba2532e004566c069d Mon Sep 17 00:00:00 2001
From: lololozhkin <alexand3r.lozhkin@gmail.com>
Date: Wed, 15 Feb 2023 20:52:03 +0500
Subject: [PATCH] socketmode: awaiting of message receiver goroutine in run
 method to avoid panics

---
 socketmode/socket_mode_managed_conn.go | 39 ++++++++++++++++++++++++--
 1 file changed, 36 insertions(+), 3 deletions(-)

diff --git a/socketmode/socket_mode_managed_conn.go b/socketmode/socket_mode_managed_conn.go
index b259fd624..216ee307e 100644
--- a/socketmode/socket_mode_managed_conn.go
+++ b/socketmode/socket_mode_managed_conn.go
@@ -123,9 +123,10 @@ func (smc *Client) run(ctx context.Context, connectionCount int) error {
 		}
 	}()
 
-	// We don't wait on runMessageReceiver because it doesn't block on a select with the context,
-	// so we'd have to wait for the ReadJSON to time out, which can take a while.
+	// Need to wait for runMessageReceiver to avoid panic described in https://github.com/slack-go/slack/issues/1125
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		defer cancel()
 
 		// The receiver reads WebSocket messages, and enqueues parsed Socket Mode requests to be handled by
@@ -439,8 +440,40 @@ func (smc *Client) receiveMessagesInto(ctx context.Context, conn *websocket.Conn
 	smc.Debugf("Starting to receive message")
 	defer smc.Debugf("Finished to receive message")
 
+	var err error
+
 	event := json.RawMessage{}
-	err := conn.ReadJSON(&event)
+
+	readJsonErr := make(chan error)
+	go func() {
+		select {
+		case readJsonErr <- conn.ReadJSON(&event):
+			// if conn.ReadJSON method returns after ctx.Done(), no one will read the unbuffered channel
+			// so, to avoid goroutines leak we need to check for ctx.Done() here too.
+			// need to say here, that conn.ReadJSON will really return after ctx.Done(), because in that case
+			// conn is closed
+			break
+		case <-ctx.Done():
+			// just need to listen ctx.Done, nothing has to be done here
+			break
+		}
+	}()
+
+	select {
+	case err = <-readJsonErr:
+		// we have awaited response from conn.ReadJSON, so, we can handle error now
+		break
+	case <-ctx.Done():
+		// context cancellation signal got, closing connection and returning
+		cerr := conn.Close()
+		if cerr != nil {
+			smc.Debugf("Failed to close connection: %v", cerr)
+		}
+
+		return ctx.Err()
+	}
+
+	// after select above, the err will be set to result from conn.ReadJSON, handling
 
 	// check if the connection was closed.
 	if websocket.IsUnexpectedCloseError(err) {