[v1.x] Support TransportSecuritySettings in the WebSocket server transport#2992
Conversation
…sport Brings websocket_server() to parity with the SSE and Streamable HTTP transports: it now accepts the same optional security_settings argument and validates Host/Origin headers via TransportSecurityMiddleware before accepting the handshake. As with the other transports, validation is opt-in (None leaves it disabled). TransportSecurityMiddleware.validate_request() now accepts any starlette.requests.HTTPConnection so a WebSocket can be passed directly; Request is a subclass so existing call sites are unchanged.
There was a problem hiding this comment.
I didn't find any issues — the implementation mirrors the SSE transport's validation flow and is well covered by the new tests — but since this adds a new public parameter and touches the transport security path (DNS rebinding protection), it's worth a maintainer's look.
Extended reasoning...
Overview
This PR adds an optional security_settings: TransportSecuritySettings parameter to the deprecated websocket_server() transport, running TransportSecurityMiddleware.validate_request() before the WebSocket handshake is accepted. It also widens the validate_request() type hint from Request to HTTPConnection (a superclass, so existing call sites are unaffected), and adds a new test module mirroring the existing SSE security tests.
Security risks
The change is itself a security feature (Host/Origin allow-list parity for WebSocket endpoints). The default of None preserves existing behaviour, so there is no regression risk for current users. The rejection path (pre-accept close() + ValueError) matches the SSE transport's connect_sse() behaviour. Risk would only arise if the validation were silently bypassable, but the in-process test asserts websocket.close is sent before any websocket.accept, and the end-to-end tests verify a 403 on rejected handshakes.
Level of scrutiny
Moderate. The diff is small and follows an established pattern, and the type-hint change in transport_security.py is behaviour-neutral. However, it introduces a new public API parameter and modifies a security-relevant module, which under repository norms warrants a maintainer review rather than automatic approval. There is also a minor design question a maintainer may want to weigh in on: adding a new feature to a transport that is already deprecated and slated for removal in 2.0.
Other factors
No bugs were found by the automated review. Test coverage is good: default-off, explicit-disable, allowed/blocked Origin, blocked Host, and a unit-level pre-accept rejection check. There are no prior reviewer comments on the PR.
The SSE and Streamable HTTP server transports both accept an optional
security_settings: TransportSecuritySettingsand runTransportSecurityMiddleware.validate_request()before establishing a session. The (deprecated) WebSocket transport never grew the same hook, so there was no SDK-level way to apply the same Host/Origin allow-list to a WebSocket endpoint.Motivation and Context
Brings
websocket_server()to parity with the other two HTTP-based server transports so thatTransportSecuritySettingscan be applied uniformly.How Has This Been Tested?
New
tests/server/test_websocket_security.pymirrorstests/server/test_sse_security.py: default (off) accepts any Origin, configured allow-list rejects non-matching Origin/Host with HTTP 403, matching Origin is accepted, and explicit disable accepts everything. An in-process test asserts the rejection happens viawebsocket.closebeforewebsocket.accept.Breaking Changes
None. The new
security_settingsparameter defaults toNone, which preserves the existing behaviour.TransportSecurityMiddleware.validate_request()now acceptsstarlette.requests.HTTPConnectioninstead ofRequest;Requestis a subclass, so existing call sites are unchanged.Types of changes
Checklist
Additional context
When validation fails the handshake is rejected with a pre-accept
close()(which the ASGI server renders as HTTP 403) andValueError("Request validation failed")is raised, matchingSseServerTransport.connect_sse(). A pre-acceptclose()is used rather thanWebSocket.send_denial_response()because the latter is unavailable on the starlette floor for this branch.AI Disclaimer