diff --git a/pkg/gofr/websocket.go b/pkg/gofr/websocket.go index d418c42a94..c9e032f3d9 100644 --- a/pkg/gofr/websocket.go +++ b/pkg/gofr/websocket.go @@ -29,10 +29,16 @@ func (a *App) OverrideWebsocketUpgrader(wsUpgrader websocket.Upgrader) { // within the handler context. User can access the underlying WebSocket connection using `ctx.GetWebsocketConnection()`. func (a *App) WebSocket(route string, handler Handler) { a.GET(route, func(ctx *Context) (any, error) { - connID := ctx.Request.Context().Value(websocket.WSConnectionKey).(string) + connID, ok := ctx.Request.Context().Value(websocket.WSConnectionKey).(string) + if !ok { + // The request never went through the WebSocket upgrade middleware + // (e.g. a plain HTTP GET with no Upgrade headers), so no + // WSConnectionKey was ever set on the context. + return nil, websocket.ErrorNotWebSocketUpgrade{} + } conn := a.httpServer.ws.GetWebsocketConnection(connID) - if conn.Conn == nil { + if conn == nil || conn.Conn == nil { return nil, websocket.ErrorConnection } diff --git a/pkg/gofr/websocket/websocket.go b/pkg/gofr/websocket/websocket.go index 52e5bbf6ac..dbfaa6cf13 100644 --- a/pkg/gofr/websocket/websocket.go +++ b/pkg/gofr/websocket/websocket.go @@ -28,6 +28,24 @@ type Connection struct { // ErrorConnection is the connection error that occurs when websocket connection cannot be established. var ErrorConnection = errors.New("couldn't establish connection to web socket") +// ErrorNotWebSocketUpgrade is returned by a WebSocket route handler when the request never went +// through the WebSocket upgrade handshake (no WSConnectionKey on its context), as opposed to +// ErrorConnection, which means the upgrade happened but the resulting connection is missing or +// already closed. The former is a client error - the client simply never attempted an upgrade - +// while the latter is a server-side fault, so the two carry different HTTP status codes: this type +// implements StatusCode() to map to 400 (RFC 6455 4.2.1 calls for "an appropriate error code (such +// as 400 Bad Request)" on a failed handshake attempt), while ErrorConnection has no StatusCode() +// method and falls back to gofr's default 500. +type ErrorNotWebSocketUpgrade struct{} + +func (ErrorNotWebSocketUpgrade) Error() string { + return "request did not include a WebSocket upgrade" +} + +func (ErrorNotWebSocketUpgrade) StatusCode() int { + return http.StatusBadRequest +} + // The message types are defined in RFC 6455, section 11.8. const ( // TextMessage denotes a text data message. The text message payload is diff --git a/pkg/gofr/websocket_test.go b/pkg/gofr/websocket_test.go index e91db53a09..c8cef6d9e0 100644 --- a/pkg/gofr/websocket_test.go +++ b/pkg/gofr/websocket_test.go @@ -4,6 +4,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "net/http" "net/http/httptest" "reflect" @@ -70,6 +71,46 @@ func Test_WebSocket_Success(t *testing.T) { require.NoError(t, err) } +// Test_WebSocket_PlainHTTPRequestDoesNotPanic pins the fix for #3862: a plain +// HTTP request to a route registered via app.WebSocket (no Upgrade headers, +// so no WSConnectionKey is ever set on the request context) must get a clean +// 400 error response instead of panicking. Before the fix, the unsafe type +// assertion on WSConnectionKey panicked; the panic-recovery middleware then +// converted that into a generic 500 "Internal Server Error" response, which +// is indistinguishable by status code alone from a naive fix that also +// returned 500. The response is checked for both the 400 status (per RFC +// 6455 4.2.1 - a request that never attempted an upgrade is a client error, +// same as a request whose upgrade attempt failed, handled a few lines away +// in middleware/web_socket.go) and the specific message, since only the +// panic-recovery path produces the generic one. +func Test_WebSocket_PlainHTTPRequestDoesNotPanic(t *testing.T) { + testutil.NewServerConfigs(t) + + app := New() + + server := httptest.NewServer(app.httpServer.router) + defer server.Close() + + app.WebSocket("/ws", func(*Context) (any, error) { + return "unreachable: handler must not run without a WebSocket connection", nil + }) + + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, server.URL+"/ws", http.NoBody) + require.NoError(t, err) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + assert.Contains(t, string(body), "request did not include a WebSocket upgrade", + "expected the websocket.ErrorNotWebSocketUpgrade message, not a panic-recovery response") +} + func Test_AddWSService(t *testing.T) { port := testutil.GetFreePort(t) t.Setenv("HTTP_PORT", fmt.Sprint(port))