feat(http): support websocket server (#481)

Author: StellarisW

Cargo.lock

@@ -133,6 +133,7 @@ webpki-roots = "0.26"
 tokio-rustls = "0.25"
 native-tls = "0.2"
 tokio-native-tls = "0.3"
+tokio-tungstenite="0.23.1"
 
 [profile.release]
 opt-level = 3

Cargo.toml

@@ -59,9 +59,13 @@ tokio = { workspace = true, features = [
 tokio-util = { workspace = true, features = ["io"] }
 tracing.workspace = true
 
+# =====optional=====
 # server optional
 matchit = { workspace = true, optional = true }
 
+# protocol optional
+tokio-tungstenite = { workspace = true, optional = true }
+
 # tls optional
 tokio-rustls = { workspace = true, optional = true }
 tokio-native-tls = { workspace = true, optional = true }
@@ -86,11 +90,13 @@ default = []
 default_client = ["client", "json"]
 default_server = ["server", "query", "form", "json"]
 
-full = ["client", "server", "rustls", "cookie", "query", "form", "json", "tls"]
+full = ["client", "server", "rustls", "cookie", "query", "form", "json", "tls", "ws"]
 
 client = ["hyper/client", "hyper/http1"] # client core
 server = ["hyper/server", "hyper/http1", "dep:matchit"] # server core
 
+ws = ["dep:tokio-tungstenite"]
+
 tls = ["rustls"]
 __tls = []
 rustls = ["__tls", "dep:tokio-rustls", "volo/rustls"]

volo-http/Cargo.toml

@@ -76,7 +76,7 @@ impl Error for GenericRejectionError {}
 
 impl GenericRejectionError {
     /// Convert the [`GenericRejectionError`] to the corresponding [`StatusCode`]
-    pub fn to_status_code(self) -> StatusCode {
+    pub fn to_status_code(&self) -> StatusCode {
         match self {
             Self::BodyCollectionError => StatusCode::INTERNAL_SERVER_ERROR,
             Self::InvalidContentType => StatusCode::UNSUPPORTED_MEDIA_TYPE,
@@ -99,3 +99,69 @@ pub fn body_collection_error() -> ExtractBodyError {
 pub fn invalid_content_type() -> ExtractBodyError {
     ExtractBodyError::Generic(GenericRejectionError::InvalidContentType)
 }
+
+/// Rejection used for [`WebSocketUpgrade`](crate::server::utils::WebSocketUpgrade).
+#[derive(Debug)]
+#[non_exhaustive]
+pub enum WebSocketUpgradeRejectionError {
+    /// The request method must be `GET`
+    MethodNotGet,
+    /// The HTTP version is not supported
+    InvalidHttpVersion,
+    /// The `Connection` header is invalid
+    InvalidConnectionHeader,
+    /// The `Upgrade` header is invalid
+    InvalidUpgradeHeader,
+    /// The `Sec-WebSocket-Version` header is invalid
+    InvalidWebSocketVersionHeader,
+    /// The `Sec-WebSocket-Key` header is missing
+    WebSocketKeyHeaderMissing,
+    /// The connection is not upgradable
+    ConnectionNotUpgradable,
+}
+
+impl WebSocketUpgradeRejectionError {
+    /// Convert the [`WebSocketUpgradeRejectionError`] to the corresponding [`StatusCode`]
+    fn to_status_code(&self) -> StatusCode {
+        match self {
+            Self::MethodNotGet => StatusCode::METHOD_NOT_ALLOWED,
+            Self::InvalidHttpVersion => StatusCode::HTTP_VERSION_NOT_SUPPORTED,
+            Self::InvalidConnectionHeader => StatusCode::BAD_REQUEST,
+            Self::InvalidUpgradeHeader => StatusCode::BAD_REQUEST,
+            Self::InvalidWebSocketVersionHeader => StatusCode::BAD_REQUEST,
+            Self::WebSocketKeyHeaderMissing => StatusCode::BAD_REQUEST,
+            Self::ConnectionNotUpgradable => StatusCode::UPGRADE_REQUIRED,
+        }
+    }
+}
+
+impl Error for WebSocketUpgradeRejectionError {}
+
+impl fmt::Display for WebSocketUpgradeRejectionError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::MethodNotGet => write!(f, "Request method must be 'GET'"),
+            Self::InvalidHttpVersion => {
+                write!(f, "Http version not support, only support HTTP 1.1 for now")
+            }
+            Self::InvalidConnectionHeader => {
+                write!(f, "Connection header did not include 'upgrade'")
+            }
+            Self::InvalidUpgradeHeader => write!(f, "`Upgrade` header did not include 'websocket'"),
+            Self::InvalidWebSocketVersionHeader => {
+                write!(f, "`Sec-WebSocket-Version` header did not include '13'")
+            }
+            Self::WebSocketKeyHeaderMissing => write!(f, "`Sec-WebSocket-Key` header missing"),
+            Self::ConnectionNotUpgradable => write!(
+                f,
+                "WebSocket request couldn't be upgraded since no upgrade state was present"
+            ),
+        }
+    }
+}
+
+impl IntoResponse for WebSocketUpgradeRejectionError {
+    fn into_response(self) -> ServerResponse {
+        self.to_status_code().into_response()
+    }
+}

volo-http/src/error/server.rs

@@ -443,13 +443,15 @@ async fn serve_conn<S>(
     let notified = exit_notify.notified();
     tokio::pin!(notified);
 
-    let mut http_conn = server.serve_connection(TokioIo::new(conn), service);
+    let mut http_conn = server
+        .serve_connection(TokioIo::new(conn), service)
+        .with_upgrades();
 
     tokio::select! {
         _ = &mut notified => {
             tracing::trace!("[VOLO] closing a pending connection");
             // Graceful shutdown.
-            hyper::server::conn::http1::Connection::graceful_shutdown(
+            hyper::server::conn::http1::UpgradeableConnection::graceful_shutdown(
                 Pin::new(&mut http_conn)
             );
             // Continue to poll this connection until shutdown can finish.

volo-http/src/server/mod.rs

@@ -5,3 +5,8 @@ mod serve_dir;
 
 pub use file_response::FileResponse;
 pub use serve_dir::ServeDir;
+
+#[cfg(feature = "ws")]
+pub mod ws;
+#[cfg(feature = "ws")]
+pub use self::ws::{Config as WebSocketConfig, Message, WebSocket, WebSocketUpgrade};