extension/websocket: fix Accept blocked when listener closed

lesismal · Jan 1, 2024 · be90e91 · be90e91
1 parent f4ea71f
commit be90e91
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/extension/protocol/websocket/websocket.go b/extension/protocol/websocket/websocket.go
@@ -46,7 +46,6 @@ func (ln *Listener) Handler(w http.ResponseWriter, r *http.Request) {
 	case <-ln.chClose:
 		c.Close()
 	}
-
 }
 
 // Close .
@@ -64,9 +63,10 @@ func (ln *Listener) Addr() net.Addr {
 
 // Accept .
 func (ln *Listener) Accept() (net.Conn, error) {
-	c := <-ln.chAccept
-	if c != nil {
+	select {
+	case c := <-ln.chAccept:
 		return c, nil
+	case <-ln.chClose:
 	}
 	return nil, ErrClosed
 }