From 7a56b23cd8cec931cb13ceb54df493ace2e6b99e Mon Sep 17 00:00:00 2001
From: lesismal <lesismal@gmail.com>
Date: Thu, 12 Oct 2023 00:52:17 +0800
Subject: [PATCH] websocket: add SessionWithContext, typo

---
 nbhttp/websocket/conn.go | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/nbhttp/websocket/conn.go b/nbhttp/websocket/conn.go
index eed59125..4a9ba72f 100644
--- a/nbhttp/websocket/conn.go
+++ b/nbhttp/websocket/conn.go
@@ -6,6 +6,7 @@ package websocket
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"fmt"
 	"io"
@@ -564,6 +565,11 @@ func (c *Conn) WriteMessage(messageType MessageType, data []byte) error {
 
 // Session returns user session.
 func (c *Conn) Session() interface{} {
+	return c.session
+}
+
+// SessionWithLock returns user session with lock, returns as soon as the session has been seted.
+func (c *Conn) SessionWithLock() interface{} {
 	c.mux.Lock()
 	ch := c.chSessionInited
 	c.mux.Unlock()
@@ -573,8 +579,19 @@ func (c *Conn) Session() interface{} {
 	return c.session
 }
 
-// SessionWithLock returns user session with lock.
-func (c *Conn) SessionWithLock() interface{} {
+// SessionWithContext returns user session, returns as soon as the session has been seted or
+// waits until the context is done.
+func (c *Conn) SessionWithContext(ctx context.Context) interface{} {
+	c.mux.Lock()
+	ch := c.chSessionInited
+	c.mux.Unlock()
+	if ch != nil {
+		select {
+		case <-ch:
+		case <-ctx.Done():
+		}
+
+	}
 	return c.session
 }