From 55f56c9bcc16b0d07f3ad80a0bf170ae8569cdd9 Mon Sep 17 00:00:00 2001
From: "lilong.129" <lilong.129@bytedance.com>
Date: Tue, 20 Aug 2024 20:48:19 +0800
Subject: [PATCH] fix: wsConnMap mutex lock

---
 hrp/step_websocket.go | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/hrp/step_websocket.go b/hrp/step_websocket.go
index aaf4abed..57c0c2e2 100644
--- a/hrp/step_websocket.go
+++ b/hrp/step_websocket.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"fmt"
 	"net/http"
+	"sync"
 	"testing"
 	"time"
 	"unsafe"
@@ -17,11 +18,18 @@ import (
 )
 
 var (
+	wsMutex           sync.Mutex
 	wsConnMap         map[string]*websocket.Conn // save all websocket connections
 	pongResponseChan  chan string                // channel used to receive pong response message
 	closeResponseChan chan *wsCloseRespObject    // channel used to receive close response message
 )
 
+func init() {
+	wsConnMap = make(map[string]*websocket.Conn)
+	pongResponseChan = make(chan string, 1)
+	closeResponseChan = make(chan *wsCloseRespObject, 1)
+}
+
 const (
 	wsOpen         ActionType = "open"
 	wsPing         ActionType = "ping"
@@ -426,12 +434,6 @@ func runStepWebSocket(r *SessionRunner, step *TStep) (stepResult *StepResult, er
 }
 
 func getWsClient(url string) *websocket.Conn {
-	if wsConnMap == nil {
-		wsConnMap = make(map[string]*websocket.Conn)
-		pongResponseChan = make(chan string, 1)
-		closeResponseChan = make(chan *wsCloseRespObject, 1)
-	}
-
 	if client, ok := wsConnMap[url]; ok {
 		return client
 	}
@@ -500,7 +502,10 @@ func openWithTimeout(urlStr string, requestHeader http.Header, r *SessionRunner,
 
 			return nil
 		})
+		wsMutex.Lock()
 		wsConnMap[urlStr] = conn
+		wsMutex.Unlock()
+
 		openResponseChan <- resp
 	}()