Add support for websockets #100
Conversation
Add comprehensive WebSocket support to SRouter with the following features: - WebSocketRouteConfig for declarative WebSocket route configuration - WebSocketConnection wrapper providing type-safe access to context values - Support for text, binary, and JSON messages - Configurable read/write buffers, timeouts, and message size limits - Ping/pong keep-alive support - Custom origin checking and subprotocol negotiation - Full integration with existing middleware (auth, trace ID, etc.) - Support for both direct registration and SubRouterConfig.Routes The implementation follows SRouter's existing patterns: - Uses gorilla/websocket for the underlying WebSocket protocol - Integrates with the middleware chain for pre-upgrade processing - Leverages scontext for type-safe context access in handlers - Supports NewWebSocketRouteDefinition for declarative configuration Includes comprehensive tests covering: - Basic connection and message handling - Authentication (required and optional) - JSON and binary message types - Connection lifecycle (close handling, done channel) - Middleware integration - Subprotocol negotiation - Origin checking
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| } else { | ||
| // Default: allow all origins (common for development) | ||
| // In production, you should set a proper CheckOrigin function | ||
| upgrader.CheckOrigin = func(r *http.Request) bool { | ||
| return true |
There was a problem hiding this comment.
The WebSocketOverrides.CheckOrigin comment promises a “safe default” that checks the Origin against the Host when the field is nil, but upgraderFromOverrides installs a CheckOrigin that always returns true. As a result, any WebSocket route that omits a custom CheckOrigin will accept upgrades from arbitrary origins, which is a security regression compared with the documented behavior and the gorilla/websocket default of rejecting cross-origin traffic. Services relying on the advertised safe default will unintentionally allow cross-site WebSocket access unless this is tightened to the Host check or the docs are corrected.
Useful? React with 👍 / 👎.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #100 +/- ##
========================================
Coverage 97.06% 97.06%
========================================
Files 18 19 +1
Lines 2382 1944 -438
========================================
- Hits 2312 1887 -425
+ Misses 57 41 -16
- Partials 13 16 +3 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull request overview
This PR adds comprehensive WebSocket support to the SRouter framework, enabling real-time bidirectional communication. The implementation leverages gorilla/websocket and integrates seamlessly with SRouter's existing authentication, middleware, and context management systems.
Key changes:
- WebSocket connection wrapper with type-safe context access and lifecycle management
- Declarative route configuration supporting authentication, timeouts, compression, and custom origin checking
- Full middleware integration including auth, tracing, and recovery during upgrade
- Comprehensive test suite covering connection handling, authentication, JSON/binary messages, and middleware integration
Reviewed changes
Copilot reviewed 5 out of 6 changed files in this pull request and generated 12 comments.
Show a summary per file
| File | Description |
|---|---|
| pkg/router/websocket.go | Core WebSocket implementation including connection wrapper, handler registration, and upgrader configuration |
| pkg/router/router.go | Added Hijacker interface support to response writers for WebSocket upgrades and sub-router integration |
| pkg/router/websocket_test.go | Comprehensive test suite covering connection lifecycle, authentication, message types, and middleware |
| examples/websocket/main.go | Example demonstrating echo server, authenticated chat, binary messages, and ping/pong keep-alive |
| go.mod, go.sum | Added gorilla/websocket v1.5.3 dependency |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
examples/websocket/main.go
Outdated
| Time: time.Now().Format(time.RFC3339), | ||
| } | ||
|
|
||
| jsonData, _ := json.Marshal(response) |
There was a problem hiding this comment.
The error returned from json.Marshal is silently ignored. While marshalling a simple struct is unlikely to fail, ignoring errors is a bad practice in example code as developers often copy examples. Consider handling the error with a log message or returning it from the handler.
| jsonData, _ := json.Marshal(response) | |
| jsonData, err := json.Marshal(response) | |
| if err != nil { | |
| log.Printf("Failed to marshal response: %v", err) | |
| return err | |
| } |
| // WriteMessage writes a message to the connection. | ||
| // The message type must be TextMessage or BinaryMessage. | ||
| func (c *WebSocketConnection[T, U]) WriteMessage(messageType MessageType, data []byte) error { | ||
| if c.overrides.WriteTimeout > 0 { | ||
| _ = c.conn.SetWriteDeadline(time.Now().Add(c.overrides.WriteTimeout)) | ||
| } | ||
| return c.conn.WriteMessage(int(messageType), data) | ||
| } | ||
|
|
||
| // WriteText writes a text message to the connection. | ||
| // This is a convenience method for WriteMessage(TextMessage, data). | ||
| func (c *WebSocketConnection[T, U]) WriteText(text string) error { | ||
| return c.WriteMessage(TextMessage, []byte(text)) | ||
| } | ||
|
|
||
| // WriteBinary writes a binary message to the connection. | ||
| // This is a convenience method for WriteMessage(BinaryMessage, data). | ||
| func (c *WebSocketConnection[T, U]) WriteBinary(data []byte) error { | ||
| return c.WriteMessage(BinaryMessage, data) | ||
| } | ||
|
|
||
| // WriteJSON writes a JSON-encoded value to the connection. | ||
| func (c *WebSocketConnection[T, U]) WriteJSON(v any) error { | ||
| if c.overrides.WriteTimeout > 0 { | ||
| _ = c.conn.SetWriteDeadline(time.Now().Add(c.overrides.WriteTimeout)) | ||
| } | ||
| return c.conn.WriteJSON(v) | ||
| } |
There was a problem hiding this comment.
The WriteMessage, WriteJSON, and WriteText/WriteBinary methods are not protected by mutex locks, but according to gorilla/websocket documentation, concurrent writes to a WebSocket connection are not safe. If a user's handler has multiple goroutines writing to the connection (or the ping loop is sending while the handler writes), this could cause data corruption or panics. Consider either: 1) documenting that concurrent writes are unsafe and users must synchronize themselves, or 2) adding a write mutex to protect all write operations.
pkg/router/websocket.go
Outdated
|
|
||
| // Send close message | ||
| closeMessage := websocket.FormatCloseMessage(int(code), message) | ||
| _ = c.conn.SetWriteDeadline(time.Now().Add(time.Second)) |
There was a problem hiding this comment.
The write deadline is set to a fixed 1 second for sending the close message. If the WriteTimeout override is configured with a value less than 1 second, this hardcoded timeout might be too permissive. Consider using min(1*time.Second, c.overrides.WriteTimeout) if WriteTimeout is set, or at least document this behavior.
| _ = c.conn.SetWriteDeadline(time.Now().Add(time.Second)) | |
| // Use the minimum of 1 second and the configured WriteTimeout (if set) | |
| writeTimeout := time.Second | |
| if c.overrides.WriteTimeout > 0 && c.overrides.WriteTimeout < time.Second { | |
| writeTimeout = c.overrides.WriteTimeout | |
| } | |
| c.conn.SetWriteDeadline(time.Now().Add(writeTimeout)) |
| func TestWebSocketShutdown(t *testing.T) { | ||
| router := newTestWebSocketRouter(t) | ||
|
|
||
| handlerStarted := make(chan struct{}) | ||
| handlerDone := make(chan struct{}) | ||
|
|
||
| router.RegisterWebSocketRoute(WebSocketRouteConfig[string, string]{ | ||
| Path: "/ws/shutdown", | ||
| Handler: func(conn *WebSocketConnection[string, string]) error { | ||
| close(handlerStarted) | ||
| // Wait for a message (which won't come due to shutdown) | ||
| for { | ||
| _, _, err := conn.ReadMessage() | ||
| if err != nil { | ||
| close(handlerDone) | ||
| return nil | ||
| } | ||
| } | ||
| }, | ||
| }) | ||
|
|
||
| server := httptest.NewServer(router) | ||
| defer server.Close() | ||
|
|
||
| wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/ws/shutdown" | ||
|
|
||
| ws, _, err := websocket.DefaultDialer.Dial(wsURL, nil) | ||
| if err != nil { | ||
| t.Fatalf("Failed to connect to WebSocket: %v", err) | ||
| } | ||
| defer ws.Close() | ||
|
|
||
| // Wait for handler to start | ||
| select { | ||
| case <-handlerStarted: | ||
| case <-time.After(2 * time.Second): | ||
| t.Fatal("Handler did not start within timeout") | ||
| } | ||
|
|
||
| // Close the client connection to allow handler to complete | ||
| _ = ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")) | ||
| ws.Close() | ||
|
|
||
| // Wait for handler to finish | ||
| select { | ||
| case <-handlerDone: | ||
| case <-time.After(2 * time.Second): | ||
| t.Fatal("Handler did not complete within timeout") | ||
| } | ||
|
|
||
| // Initiate shutdown | ||
| ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) | ||
| defer cancel() | ||
| if err := router.Shutdown(ctx); err != nil { | ||
| t.Errorf("Shutdown returned error: %v", err) | ||
| } | ||
| } |
There was a problem hiding this comment.
The test name and comment suggest testing shutdown handling, but the test doesn't actually call router.Shutdown() until after the handler has already completed (line 829). This means the test isn't actually verifying WebSocket behavior during shutdown - it's just testing normal connection closure followed by a shutdown of an idle router. To properly test shutdown with active connections, initiate shutdown while the connection is still active and verify graceful handling.
| c.logger.Debug("Ping failed, closing connection", | ||
| zap.Error(err), | ||
| zap.String("remote_addr", c.RemoteAddr()), | ||
| ) | ||
| c.Close() | ||
| return |
There was a problem hiding this comment.
The Ping() method in the ping loop (line 298) can race with concurrent WriteMessage calls from the handler. According to gorilla/websocket documentation, concurrent writes are not safe. The ping goroutine writes to the connection without any synchronization, which can conflict with handler writes. This is a critical concurrency issue that needs to be addressed with a write mutex protecting all write operations.
| _ = c.conn.SetReadDeadline(time.Now().Add(pongTimeout)) | ||
| c.conn.SetPongHandler(func(string) error { | ||
| return c.conn.SetReadDeadline(time.Now().Add(pongTimeout)) | ||
| }) | ||
|
|
There was a problem hiding this comment.
The ping loop sets the read deadline in the pong handler and on startup, but this could interfere with read deadlines set by ReadMessage(). If a user calls ReadMessage() with a ReadTimeout configured, and then a pong arrives, the pong handler will reset the deadline to the pongTimeout value, potentially extending the expected timeout. Consider coordinating these timeouts or documenting this interaction clearly.
| chain = chain.Append(middlewares...) | ||
|
|
||
| // 5. Global middlewares | ||
| chain = chain.Append(r.middlewares...) |
There was a problem hiding this comment.
The middleware order comment says "Global middlewares" but these are actually applied before route-specific middlewares in the execution order. The chain is built in reverse: route-specific middlewares are appended first (line 472), then global middlewares (line 475). This means global middlewares execute after route-specific ones in the actual request flow. Either fix the order or update the comment to clarify this is the build order, not execution order.
| chain = chain.Append(r.middlewares...) | |
| // 5. Global middlewares (appended after route-specific; will execute after them) |
| } | ||
|
|
||
| // Close the connection from client side | ||
| _ = ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")) |
There was a problem hiding this comment.
The test sends a close message and then immediately calls ws.Close(), but doesn't wait for the server to acknowledge the close. This could lead to race conditions where the connection is closed before the server handler processes the close message. Consider adding a small delay or reading the close response before closing to ensure proper close handshake.
| _ = ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")) | |
| ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")) | |
| // Wait for the server to acknowledge the close | |
| _, _, _ = ws.ReadMessage() |
pkg/router/websocket_test.go
Outdated
| // Send a message to trigger server-side close | ||
| _ = ws.WriteMessage(websocket.TextMessage, []byte("trigger close")) | ||
| ws.Close() |
There was a problem hiding this comment.
The test sends a close message and immediately calls ws.Close(), then sends another write message at line 701. However, after closing the WebSocket, writing may fail. The order seems incorrect - the test should send the trigger message, then close the connection. Consider reordering: send the message first, then wait for the close to complete, then call ws.Close().
| authFunc := func(ctx context.Context, token string) (*string, bool) { | ||
| // In a real application, validate the token properly | ||
| // For this example, we accept any non-empty token as the username | ||
| if token != "" { | ||
| return &token, true |
There was a problem hiding this comment.
The example shows accepting any non-empty token as valid authentication (line 220), which is insecure even for example code. This could mislead developers copying this pattern into production. Consider using a more realistic example like checking against a hardcoded set of valid tokens, or at least add a prominent comment warning that this is NOT secure and should never be used in production.
| authFunc := func(ctx context.Context, token string) (*string, bool) { | |
| // In a real application, validate the token properly | |
| // For this example, we accept any non-empty token as the username | |
| if token != "" { | |
| return &token, true | |
| // | |
| // WARNING: This is example code. Do NOT use this authentication logic in production! | |
| // In a real application, validate the token properly (e.g., using JWT, OAuth, etc.). | |
| // Here, we only accept tokens that are present in a hardcoded set of valid tokens. | |
| validTokens := map[string]string{ | |
| "myuser": "myuser", // token: username | |
| "admin": "admin", | |
| } | |
| authFunc := func(ctx context.Context, token string) (*string, bool) { | |
| if username, ok := validTokens[token]; ok { | |
| return &username, true |
- Add write mutex protection to Ping() method for concurrent safety - Add write mutex protection to CloseWithCode() for concurrent safety - Fix close message deadline to respect WriteTimeout (uses min of 1s and WriteTimeout) - Fix json.Marshal error handling in example code - Fix TestWebSocketShutdown to properly test shutdown with active connections - Add TestWebSocketRejectsDuringShutdown to verify 503 response during shutdown - Add tests for Ping, User, ClientIP, IsUnexpectedCloseError methods - Add tests for ping loop, read timeout, and ReadJSON timeout - Improve test coverage from 74.50% to 97.2%
- Add error checks to c.Close() and wsConn.Close() in websocket.go - Add error checks to all ws.Close() calls in websocket_test.go - Add error checks to router.Shutdown() and ws.SetReadDeadline() in tests
Add comprehensive WebSocket support to SRouter with the following features:
The implementation follows SRouter's existing patterns:
Includes comprehensive tests covering: