Add websockets

Cargo.toml

@@ -32,11 +32,12 @@ unstable = []
 __internal__bench = []
 
 [dependencies]
-async-h1 = { version = "2.0.1", optional = true }
+async-h1 = { version = "2.1.2", optional = true }
 async-sse = "4.0.0"
 async-std = { version = "1.6.0", features = ["unstable"] }
+async-tungstenite = "0.8"
 femme = "2.0.1"
-http-types = "2.2.1"
+http-types = "2.4.0"
 kv-log-macro = "1.0.4"
 serde = "1.0.102"
 serde_json = "1.0.41"
@@ -45,6 +46,7 @@ logtest = "2.0.0"
 async-trait = "0.1.36"
 futures-util = "0.3.5"
 pin-project-lite = "0.1.7"
+websocket_handshake = "0.1.0"
 
 [dev-dependencies]
 async-std = { version = "1.6.0", features = ["unstable", "attributes"] }
@@ -65,3 +67,7 @@ required-features = ["unstable"]
 name = "router"
 harness = false
 
+[patch.crates-io]
+http-types = { git = "https://github.com/Yarn/http-types.git", rev = "4ba7c547d37f7e172e69a6bd619336e067f70103" }
+async-h1 = { git = "https://github.com/Yarn/async-h1.git", rev = "5bca9bea599e3f500e608fcfea32399985c83f31" }
+

src/lib.rs

@@ -213,6 +213,7 @@ pub mod log;
 pub mod prelude;
 pub mod security;
 pub mod sse;
+pub mod websocket;
 pub mod utils;
 
 pub use endpoint::Endpoint;

src/websocket/mod.rs

@@ -0,0 +1,2 @@
+mod upgrade;
+pub use upgrade::upgrade;

src/websocket/upgrade.rs

@@ -0,0 +1,58 @@
+use crate::http::{Body, StatusCode};
+use crate::{Request, Response, Result, Error};
+
+use std::fmt;
+
+use async_std::future::Future;
+
+use async_tungstenite::WebSocketStream;
+use async_tungstenite::tungstenite::protocol::Role;
+
+pub struct Handle {
+    ws: WebSocketStream<Box<dyn http_types::ReadWrite>>,
+}
+
+impl Handle {
+    pub fn into_inner(self) -> WebSocketStream<Box<dyn http_types::ReadWrite>> {
+        self.ws
+    }
+}
+
+impl fmt::Debug for Handle {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Handle")
+            .field("ws", &"...")
+            .finish()
+    }
+}
+
+/// Upgrade an existing HTTP connection to a websocket connection.
+pub fn upgrade<F, Fut, State>(req: Request<State>, handler: F) -> Result<Response>
+where
+    State: Clone + Send + Sync + 'static,
+    F: FnOnce(Request<State>, Handle) -> Fut + Send + Sync + 'static,
+    Fut: Future<Output = Result<()>> + Send + Sync + 'static,
+{
+    if req.version() != Some(http_types::Version::Http1_1) {
+        return Err(Error::from_str(
+            StatusCode::NotAcceptable,
+            "Websockets are not supported on http version != 1.1.",
+        ))
+    }
+
+    let mut res: Response = websocket_handshake::http1_1::check_request_headers(req.as_ref())
+        .map_err(|err| Error::new(StatusCode::BadRequest, err))?
+        .make_response().into();
+
+    let body = Body::io(|io| async move {
+        let ws = WebSocketStream::from_raw_socket(io, Role::Server, None).await;
+
+        let handle = Handle {
+            ws: ws
+        };
+        (handler)(req, handle).await.unwrap();
+    });
+    res.set_body(body);
+
+    Ok(res)
+}