websocket: serialize the order of the session set and get

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
 }
 

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)
 		}
 	}