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