Add customizable WebSocket subprotocol selection

[Lethe10137]

Motivation

Token-encoded subprotocol is a common practice for authentication in websocket, and though not a standard usage, it is even used by kubernetes,

However, current implementation of axum does not support read such ws subprotocol.

Solution

This adds a getter for the Sec-WebSocket-Protocol requested by the
client and a setter for the chosen subprotocol.

This enables more flexible selection logic (e.g. Token-encoded subprotocols or runtime selection) and avoids unnecessary allocations present in the existing protocols() helper.

And, in the use case of k8s, where the client init the ws connection as:

var ws = new WebSocket(
  "wss://<server>/api/v1/namespaces/myns/pods/mypod/logs?follow=true",
  [
    "base64url.bearer.authorization.k8s.io.bXl0b2tlbg",
    "base64.binary.k8s.io"
  ]
);

example usage could be:

async fn ws_handler(
    mut ws: WebSocketUpgrade,
    ConnectInfo(addr): ConnectInfo<SocketAddr>,
) -> impl IntoResponse {
    let token = ws
        .requested_protocols()
        .find_map(|protocol| protocol.strip_prefix("base64url.bearer.authorization.k8s.io."));

    if token.is_some_and(|token| token == "bXl0b2tl3bg") {
        ws.set_selected_protocol(HeaderValue::from_static("base64.binary.k8s.io"));
        ws.on_upgrade(move |socket| handle_socket(socket, addr))
    } else {
        // Reject the connection
        (axum::http::StatusCode::UNAUTHORIZED, "Unauthorized").into_response()
    }
}

[jplatte]

Why is this here? I think we generally shouldn't make any promises about allocations.

[Lethe10137]

axum/axum/src/extract/ws.rs

Lines 251 to 258 in 07f8312

	// FIXME: This will often allocate a new `String` and so is less efficient than it
	// could be. But that can't be fixed without breaking changes to the public API.
	.map(Into::into)
	.find(|protocol| {
	req_protocols
	.split(',')
	.any(|req_protocol| req_protocol.trim() == protocol)
	})

Actually, I tried to respond to the FIXME in protocols(). Yet indeed we can just remove these two lines.

[jplatte]

Extra reviews appreciated, because I don't know much about ws, especially this subprotocol stuff. cc @SabrinaJewson who has reviewed a lot of WS + SSE stuff before. Also @Turbo87 maybe?

[SabrinaJewson]

Hmmm. I can’t help but wonder if we should just offer an explicit iterator over the request’s protocols together with a setter for the chosen protocol header, in case for example the token authentication logic is async or fallible? e.g.

pub fn requested_protocols(&self) -> impl Iterator<Item = &str> {
    self.sec_websocket_protocol
        .as_ref()
        .and_then(|p| p.to_str().ok())
        .into_iter()
        .flat_map(|s| s.split(','))
        .map(|s| s.trim())
}
pub fn set_protocol(&mut self, protocol: HeaderValue) {
    self.protocol = Some(protocol);
}

[Lethe10137]

Hmmm. I can’t help but wonder if we should just offer an explicit iterator over the request’s protocols together with a setter for the chosen protocol header, in case for example the token authentication logic is async or fallible? e.g.

pub fn requested_protocols(&self) -> impl Iterator<Item = &str> {
    self.sec_websocket_protocol
        .as_ref()
        .and_then(|p| p.to_str().ok())
        .into_iter()
        .flat_map(|s| s.split(','))
        .map(|s| s.trim())
}
pub fn set_protocol(&mut self, protocol: HeaderValue) {
    self.protocol = Some(protocol);
}

Thanks for your time.

Indeed this would be much easier to understand.
I would do the following in the next commit:

• add these two interfaces
• remove the select_protocol() I added

Should I close this PR and open a new one, or just force-push on this PR?

[jplatte]

Force-push is fine.

[Lethe10137]

Force-push is fine.

I have force-pushed the code, with updated doc comments. And the use example in this PR has been correspondingly updated. Could you please have a review of that?

[Lethe10137]

@SabrinaJewson Could you please have a quick review of this? Following your insightful opinion, I added a getter and setter (as you wrote in the comment) in the current PR. These two methods are really useful for a project of mine,
where neither cookie nor url param is not feasible places to place the auth token for websocket.

[SabrinaJewson]

It looks good to me. Minor comment: I know I suggested the name originally, but since the getter is called selected_protocol, maybe set_selected_protocol is a more appropriate name for the setter?

[Lethe10137]

It looks good to me. Minor comment: I know I suggested the name originally, but since the getter is called selected_protocol, maybe set_selected_protocol is a more appropriate name for the setter?

Happy new year! I have changed the name for the setter in the latest commit.

[Lethe10137]

@SabrinaJewson Could you please tell me what further modifications are needed for this PR to be merged?

[SabrinaJewson]

Sorry – to be clear I’m not a maintainer, so you’ll have to page @jplatte.

[jplatte]

I'm not a big fan of just returning an empty iterator when the header value is non-utf8, but I guess it can be fixed later (seems extremely unlikely to happen in practice unless the client is actively testing edge case behavior).