From dbec7b71a8b52a358e2d17fc9132bd6da1890e7b Mon Sep 17 00:00:00 2001 From: Kirill Sysoev Date: Sat, 13 Apr 2024 20:36:07 +0800 Subject: [PATCH] Fix race condition in setting up ws connection --- channel/channel.go | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/channel/channel.go b/channel/channel.go index 3ac426c..aeda62d 100644 --- a/channel/channel.go +++ b/channel/channel.go @@ -42,17 +42,17 @@ func (c *DefaultChannel) Path() string { // Handler returns http.Handler for channel func (c *DefaultChannel) Handler() http.Handler { - var ctx context.Context + return c.setContext(c.wrapMiddleware(c.wsConnectionHandler())) +} - saveCtx := func(next http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - ctx = r.Context() - next.ServeHTTP(w, r) - }) - } +// wsConnectionHandler handles the WebSocket connection and sets up the necessary components for communication. +func (c *DefaultChannel) wsConnectionHandler() http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ctx := r.Context() - wsHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - ws, err := websocket.Accept(w, r, nil) + ws, err := websocket.Accept(w, r, &websocket.AcceptOptions{ + OriginPatterns: []string{"*"}, + }) if err != nil { return @@ -61,8 +61,6 @@ func (c *DefaultChannel) Handler() http.Handler { conn := c.connRegistry.AddConnection(ctx, ws, c.disptacher.Dispatch) conn.HandleRequests() }) - - return c.setContext(c.wrapMiddleware(saveCtx(wsHandler))) } // SetContext sets context for channel