Port tokio-rs/axum#3141

finnbear · Jan 15, 2025 · 68681f3 · 68681f3
1 parent ca4acfd
commit 68681f3
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/src/upgrade.rs b/src/upgrade.rs
@@ -184,15 +184,15 @@ impl<F> WebSocketUpgrade<F> {
             callback(socket).await;
         });
 
-        if let Some(sec_websocket_key) = &self.sec_websocket_key {
+        let mut response = if let Some(sec_websocket_key) = &self.sec_websocket_key {
             // If `sec_websocket_key` was `Some`, we are using HTTP/1.1.
 
             #[allow(clippy::declare_interior_mutable_const)]
             const UPGRADE: HeaderValue = HeaderValue::from_static("upgrade");
             #[allow(clippy::declare_interior_mutable_const)]
             const WEBSOCKET: HeaderValue = HeaderValue::from_static("websocket");
 
-            let mut builder = Response::builder()
+            let builder = Response::builder()
                 .status(StatusCode::SWITCHING_PROTOCOLS)
                 .header(header::CONNECTION, UPGRADE)
                 .header(header::UPGRADE, WEBSOCKET)
@@ -201,17 +201,21 @@ impl<F> WebSocketUpgrade<F> {
                     sign(sec_websocket_key.as_bytes()),
                 );
 
-            if let Some(protocol) = self.protocol {
-                builder = builder.header(header::SEC_WEBSOCKET_PROTOCOL, protocol);
-            }
-
             builder.body(Body::empty()).unwrap()
         } else {
             // Otherwise, we are HTTP/2+. As established in RFC 9113 section 8.5, we just respond
             // with a 2XX with an empty body:
             // <https://datatracker.ietf.org/doc/html/rfc9113#name-the-connect-method>.
             Response::new(Body::empty())
+        };
+
+        if let Some(protocol) = self.protocol {
+            response
+                .headers_mut()
+                .insert(header::SEC_WEBSOCKET_PROTOCOL, protocol);
         }
+
+        response
     }
 }