From e2ff65fd7cca9533628369161fcea8dbf9237d90 Mon Sep 17 00:00:00 2001
From: David Pedersen <david.pdrsn@gmail.com>
Date: Mon, 11 Jul 2022 14:33:55 +0200
Subject: [PATCH] Return rejection if `WebSocket` request couldn't be upgraded
 (#1135)

* Return rejection if `WebSocket` request couldn't be upgraded

* make rejection about connection upgrades

* changelog links
---
 axum/CHANGELOG.md      |  7 ++++-
 axum/src/extract/ws.rs | 64 +++++++++++++++++++++++++++++++++++++-----
 2 files changed, 63 insertions(+), 8 deletions(-)

diff --git a/axum/CHANGELOG.md b/axum/CHANGELOG.md
index af1a3598..42350097 100644
--- a/axum/CHANGELOG.md
+++ b/axum/CHANGELOG.md
@@ -7,7 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 # Unreleased
 
-- None.
+- **fixed:** If `WebSocketUpgrade` cannot upgrade the connection it will return a
+  `WebSocketUpgradeRejection::ConnectionNotUpgradable` rejection ([#1135])
+- **changed:** `WebSocketUpgradeRejection` has a new variant `ConnectionNotUpgradable`
+  variant ([#1135])
+
+[#1135]: https://github.com/tokio-rs/axum/pull/1135
 
 # 0.5.12 (10. July, 2022)
 
diff --git a/axum/src/extract/ws.rs b/axum/src/extract/ws.rs
index 7ac75ab6..5d8f6005 100644
--- a/axum/src/extract/ws.rs
+++ b/axum/src/extract/ws.rs
@@ -267,14 +267,15 @@ where
             return Err(InvalidWebSocketVersionHeader.into());
         }
 
-        let sec_websocket_key =
-            if let Some(key) = req.headers_mut().remove(header::SEC_WEBSOCKET_KEY) {
-                key
-            } else {
-                return Err(WebSocketKeyHeaderMissing.into());
-            };
+        let sec_websocket_key = req
+            .headers_mut()
+            .remove(header::SEC_WEBSOCKET_KEY)
+            .ok_or(WebSocketKeyHeaderMissing)?;
 
-        let on_upgrade = req.extensions_mut().remove::<OnUpgrade>().unwrap();
+        let on_upgrade = req
+            .extensions_mut()
+            .remove::<OnUpgrade>()
+            .ok_or(ConnectionNotUpgradable)?;
 
         let sec_websocket_protocol = req.headers().get(header::SEC_WEBSOCKET_PROTOCOL).cloned();
 
@@ -564,6 +565,20 @@ pub mod rejection {
         pub struct WebSocketKeyHeaderMissing;
     }
 
+    define_rejection! {
+        #[status = UPGRADE_REQUIRED]
+        #[body = "WebSocket request couldn't be upgraded since no upgrade state was present"]
+        /// Rejection type for [`WebSocketUpgrade`](super::WebSocketUpgrade).
+        ///
+        /// This rejection is returned if the connection cannot be upgraded for example if the
+        /// request is HTTP/1.0.
+        ///
+        /// See [MDN] for more details about connection upgrades.
+        ///
+        /// [MDN]: https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Upgrade
+        pub struct ConnectionNotUpgradable;
+    }
+
     composite_rejection! {
         /// Rejection used for [`WebSocketUpgrade`](super::WebSocketUpgrade).
         ///
@@ -575,6 +590,7 @@ pub mod rejection {
             InvalidUpgradeHeader,
             InvalidWebSocketVersionHeader,
             WebSocketKeyHeaderMissing,
+            ConnectionNotUpgradable,
         }
     }
 }
@@ -640,3 +656,37 @@ pub mod close_code {
     /// action.
     pub const AGAIN: u16 = 1013;
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{body::Body, routing::get};
+    use http::{Request, Version};
+    use tower::ServiceExt;
+
+    #[tokio::test]
+    async fn rejects_http_1_0_requests() {
+        let svc = get(|ws: Result<WebSocketUpgrade, WebSocketUpgradeRejection>| {
+            let rejection = ws.unwrap_err();
+            assert!(matches!(
+                rejection,
+                WebSocketUpgradeRejection::ConnectionNotUpgradable(_)
+            ));
+            std::future::ready(())
+        });
+
+        let req = Request::builder()
+            .version(Version::HTTP_10)
+            .method(Method::GET)
+            .header("upgrade", "websocket")
+            .header("connection", "Upgrade")
+            .header("sec-websocket-key", "6D69KGBOr4Re+Nj6zx9aQA==")
+            .header("sec-websocket-version", "13")
+            .body(Body::empty())
+            .unwrap();
+
+        let res = svc.oneshot(req).await.unwrap();
+
+        assert_eq!(res.status(), StatusCode::OK);
+    }
+}