From 277de285da316b1d2bca977ce919c4a5d0cab1ed Mon Sep 17 00:00:00 2001
From: lesismal <lesismal@gmail.com>
Date: Tue, 10 Oct 2023 20:58:17 +0800
Subject: [PATCH] websocket: serialize the order of the session set and get

---
 nbhttp/websocket/conn.go     | 22 +++++++++++++++++++++-
 nbhttp/websocket/upgrader.go |  1 +
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/nbhttp/websocket/conn.go b/nbhttp/websocket/conn.go
index c464fef6..558431e4 100644
--- a/nbhttp/websocket/conn.go
+++ b/nbhttp/websocket/conn.go
@@ -54,7 +54,8 @@ type Conn struct {
 
 	mux sync.Mutex
 
-	session interface{}
+	chSessionInited chan struct{}
+	session         interface{}
 
 	sendQueue     [][]byte
 	sendQueueSize int
@@ -96,6 +97,13 @@ func (c *Conn) IsAsyncWrite() bool {
 
 // Close .
 func (c *Conn) Close() error {
+	c.mux.Lock()
+	if c.chSessionInited != nil {
+		close(c.chSessionInited)
+		c.chSessionInited = nil
+	}
+	c.mux.Unlock()
+
 	if c.Conn == nil {
 		return nil
 	}
@@ -563,11 +571,23 @@ func (c *Conn) WriteMessage(messageType MessageType, data []byte) error {
 
 // Session returns user session.
 func (c *Conn) Session() interface{} {
+	c.mux.Lock()
+	ch := c.chSessionInited
+	c.mux.Unlock()
+	if ch != nil {
+		<-ch
+	}
 	return c.session
 }
 
 // SetSession sets user session.
 func (c *Conn) SetSession(session interface{}) {
+	c.mux.Lock()
+	if c.chSessionInited != nil {
+		close(c.chSessionInited)
+		c.chSessionInited = nil
+	}
+	c.mux.Unlock()
 	c.session = session
 }
 
diff --git a/nbhttp/websocket/upgrader.go b/nbhttp/websocket/upgrader.go
index a337ee31..635255c4 100644
--- a/nbhttp/websocket/upgrader.go
+++ b/nbhttp/websocket/upgrader.go
@@ -380,6 +380,7 @@ func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeade
 
 	if wsc.isBlockingMod {
 		if parser == nil {
+			wsc.chSessionInited = make(chan struct{})
 			go wsc.BlockingModReadLoop(u.BlockingModReadBufferSize)
 		}
 	}