diff --git a/Cargo.lock b/Cargo.lock index dcf779510..93d759bbc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -72,7 +72,16 @@ dependencies = [ "block-padding", "byte-tools", "byteorder", - "generic-array", + "generic-array 0.12.4", +] + +[[package]] +name = "block-buffer" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf7fe51849ea569fd452f37822f606a5cabb684dc918707a0193fd4664ff324" +dependencies = [ + "generic-array 0.14.5", ] [[package]] @@ -165,6 +174,15 @@ version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5827cebf4670468b8772dd191856768aedcb1b0278a04f989f7766351917b9dc" +[[package]] +name = "cpufeatures" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59a6001667ab124aebae2a495118e11d30984c3a653e99d86d58971708cf5e4b" +dependencies = [ + "libc", +] + [[package]] name = "crossbeam-channel" version = "0.5.1" @@ -185,13 +203,33 @@ dependencies = [ "lazy_static", ] +[[package]] +name = "crypto-common" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57952ca27b5e3606ff4dd79b0020231aaf9d6aa76dc05fd30137538c50bd3ce8" +dependencies = [ + "generic-array 0.14.5", + "typenum", +] + [[package]] name = "digest" version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f3d0c8c8752312f9713efd397ff63acb9f85585afbf179282e720e7704954dd5" dependencies = [ - "generic-array", + "generic-array 0.12.4", +] + +[[package]] +name = "digest" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2fb860ca6fafa5552fb6d0e816a69c8e49f0908bf524e30a90d97c85892d506" +dependencies = [ + "block-buffer 0.10.2", + "crypto-common", ] [[package]] @@ -237,6 +275,7 @@ dependencies = [ "dropshot_endpoint", "expectorate", "futures", + "headers", "hostname", "http", "hyper", @@ -267,6 +306,7 @@ dependencies = [ "tempfile", "tokio", "tokio-rustls", + "tokio-tungstenite", "toml", "trybuild", "usdt", @@ -449,6 +489,16 @@ dependencies = [ "typenum", ] +[[package]] +name = "generic-array" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd48d33ec7f05fbfa152300fdad764757cbded343c1aa1cff2fbaf4134851803" +dependencies = [ + "typenum", + "version_check", +] + [[package]] name = "getrandom" version = "0.2.3" @@ -491,6 +541,31 @@ version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db0d4cf898abf0081f964436dc980e96670a0f36863e4b83aaacdb65c9d7ccc3" +[[package]] +name = "headers" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cff78e5788be1e0ab65b04d306b2ed5092c815ec97ec70f4ebd5aee158aa55d" +dependencies = [ + "base64", + "bitflags", + "bytes", + "headers-core", + "http", + "httpdate", + "mime", + "sha-1 0.10.0", +] + +[[package]] +name = "headers-core" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7f66481bfee273957b1f20485a4ff3362987f85b2c236580d81b4eb7a326429" +dependencies = [ + "http", +] + [[package]] name = "hermit-abi" version = "0.1.19" @@ -915,7 +990,7 @@ checksum = "54be6e404f5317079812fc8f9f5279de376d8856929e21c184ecf6bbd692a11d" dependencies = [ "maplit", "pest", - "sha-1", + "sha-1 0.8.2", ] [[package]] @@ -1288,12 +1363,23 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f7d94d0bede923b3cea61f3f1ff57ff8cdfd77b400fb8f9998949e0cf04163df" dependencies = [ - "block-buffer", - "digest", + "block-buffer 0.7.3", + "digest 0.8.1", "fake-simd", "opaque-debug", ] +[[package]] +name = "sha-1" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "028f48d513f9678cda28f6e4064755b3fbb2af6acd672f2c209b62323f7aea0f" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest 0.10.3", +] + [[package]] name = "signal-hook-registry" version = "1.4.0" @@ -1599,6 +1685,18 @@ dependencies = [ "webpki", ] +[[package]] +name = "tokio-tungstenite" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06cda1232a49558c46f8a504d5b93101d42c0bf7f911f12a105ba48168f821ae" +dependencies = [ + "futures-util", + "log", + "tokio", + "tungstenite", +] + [[package]] name = "tokio-util" version = "0.6.8" @@ -1668,6 +1766,25 @@ dependencies = [ "toml", ] +[[package]] +name = "tungstenite" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d96a2dea40e7570482f28eb57afbe42d97551905da6a9400acc5c328d24004f5" +dependencies = [ + "base64", + "byteorder", + "bytes", + "http", + "httparse", + "log", + "rand", + "sha-1 0.10.0", + "thiserror", + "url", + "utf-8", +] + [[package]] name = "typenum" version = "1.14.0" @@ -1802,6 +1919,12 @@ dependencies = [ "usdt-impl", ] +[[package]] +name = "utf-8" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" + [[package]] name = "uuid" version = "1.1.2" diff --git a/dropshot/Cargo.toml b/dropshot/Cargo.toml index 00aee1bf3..0d2da99f6 100644 --- a/dropshot/Cargo.toml +++ b/dropshot/Cargo.toml @@ -16,6 +16,7 @@ async-trait = "0.1.56" base64 = "0.13.0" bytes = "1" futures = "0.3.21" +headers = "0.3.7" hostname = "0.3.0" http = "0.2.8" indexmap = "1.9.1" @@ -32,6 +33,7 @@ slog-bunyan = "2.4.0" slog-json = "2.6.1" slog-term = "2.9.0" tokio-rustls = "0.23.4" +tokio-tungstenite = "0.17.1" toml = "0.5.9" [dependencies.chrono] diff --git a/dropshot/examples/websocket.rs b/dropshot/examples/websocket.rs new file mode 100644 index 000000000..e08c2fb21 --- /dev/null +++ b/dropshot/examples/websocket.rs @@ -0,0 +1,45 @@ +use dropshot::{ + endpoint, ApiDescription, HttpError, HttpResponseUpgradedWebSocket, + RequestContext, +}; +use futures::FutureExt; +use futures::StreamExt; +use std::sync::Arc; + +fn main() -> Result<(), String> { + /* + * Build a description of the API. + */ + let mut api = ApiDescription::new(); + api.register(websocket).unwrap(); + + api.openapi("WebSocket Echo", "") + .write(&mut std::io::stdout()) + .map_err(|e| e.to_string())?; + + Ok(()) +} + +#[allow(unused_variables)] +#[endpoint { + method = GET, + path = "/echo", +}] +/// Echo a message back to the client. +async fn websocket( + rqctx: Arc>, +) -> Result { + let (response, _) = rqctx + .upgrade_websocket(None, |ws| { + // Just echo all messages back... + let (tx, rx) = ws.split(); + rx.forward(tx).map(|result| { + if let Err(e) = result { + eprintln!("websocket error: {:?}", e); + } + }) + }) + .await?; + + Ok(response) +} diff --git a/dropshot/src/error.rs b/dropshot/src/error.rs index 5fb33bfb2..420cc6439 100644 --- a/dropshot/src/error.rs +++ b/dropshot/src/error.rs @@ -310,6 +310,12 @@ impl Error for HttpError { } } +impl From for HttpError { + fn from(error: tokio_tungstenite::tungstenite::error::Error) -> Self { + HttpError::for_internal_error(format!("{}", error)) + } +} + #[cfg(test)] mod test { use crate::HttpErrorResponseBody; diff --git a/dropshot/src/handler.rs b/dropshot/src/handler.rs index 88669bbe2..d525bad58 100644 --- a/dropshot/src/handler.rs +++ b/dropshot/src/handler.rs @@ -50,10 +50,14 @@ use crate::pagination::PaginationParams; use crate::pagination::PAGINATION_PARAM_SENTINEL; use crate::router::VariableSet; use crate::to_map::to_map; +use crate::WebSocketConfig; use async_trait::async_trait; use bytes::Bytes; use futures::lock::Mutex; +use futures::FutureExt; +use futures::TryFutureExt; +use headers::HeaderMapExt; use http::HeaderMap; use http::StatusCode; use hyper::Body; @@ -74,6 +78,7 @@ use std::future::Future; use std::marker::PhantomData; use std::num::NonZeroU32; use std::sync::Arc; +use tokio::task::JoinHandle; /** * Type alias for the result returned by HTTP handler functions. @@ -147,6 +152,72 @@ impl RequestContext { */ .unwrap_or(server_config.page_default_nitems)) } + + /// Upgrade a connection to a websocket connection. This automatically handles making + /// sure the request has the correct headers and creating the response with the correct + /// headers. + pub async fn upgrade_websocket( + &self, + config: Option, + func: F, + ) -> Result< + (HttpResponseUpgradedWebSocket, JoinHandle>), + HttpError, + > + where + F: FnOnce(crate::websocket::WebSocket) -> U + Send + 'static, + U: Future + Send + 'static, + { + // We want to parse the headers from the request context. + let mut request = self.request.lock().await; + + // We want to make sure we have the websocket upgrade header. + let h: Option = request.headers().typed_get(); + if let Some(h) = h { + if !h.contains(&http::header::UPGRADE) { + return Err(HttpError::for_bad_request( + None, + "Connection header did not include 'upgrade'".to_string(), + )); + } + } else { + return Err(HttpError::for_bad_request( + None, + "Connection header not sent".to_string(), + )); + } + + // Now get the secure websocket key. + let h: Option = request.headers().typed_get(); + let websocket_key = if let Some(h) = h { + h + } else { + return Err(HttpError::for_bad_request( + None, + "Websocket key not sent".to_string(), + )); + }; + + // Spawn a task to handle the websocket connection. + // Use the hyper feature of upgrading a connection. + let fut = hyper::upgrade::on(&mut *request) + .and_then(move |upgraded| { + crate::websocket::WebSocket::from_raw_socket( + // Pass the upgraded object as the base layer stream of the Websocket. + upgraded, + tokio_tungstenite::tungstenite::protocol::Role::Server, + config, + ) + .map(Ok) + }) + .and_then(move |websocket| { + // Run the function that was passed to us. + func(websocket).map(Ok) + }); + let h = tokio::task::spawn(fut); + + Ok((HttpResponseUpgradedWebSocket(websocket_key), h)) + } } /** @@ -1414,6 +1485,35 @@ impl From for HttpHandlerResult { } } +/** + * `HttpResponseUpgradedWebSocket` wraps the value passed into the request header as the Secure WebSocket + * Key. + * It denotes an HTTP 101 "Switching Protocols". This is used for the response from + * connections upgraded to a WebSocket. It automatically handles setting the correct response + * headers. + */ +pub struct HttpResponseUpgradedWebSocket(pub headers::SecWebsocketKey); +impl HttpCodedResponse for HttpResponseUpgradedWebSocket { + type Body = Empty; + const STATUS_CODE: StatusCode = StatusCode::SWITCHING_PROTOCOLS; + const DESCRIPTION: &'static str = "upgraded to WebSocket"; +} +impl From for HttpHandlerResult { + fn from(r: HttpResponseUpgradedWebSocket) -> HttpHandlerResult { + /* Compute the body. */ + let mut result: http::Response = + HttpResponseUpgradedWebSocket::for_object(Empty)?; + /* Add in the headers we need for an upgraded connection. */ + result.headers_mut().typed_insert(headers::Connection::upgrade()); + result.headers_mut().typed_insert(headers::Upgrade::websocket()); + result + .headers_mut() + .typed_insert(headers::SecWebsocketAccept::from(r.0)); + + Ok(result) + } +} + #[derive(Serialize, JsonSchema)] pub struct NoHeaders {} diff --git a/dropshot/src/lib.rs b/dropshot/src/lib.rs index d758c5988..3f0b8005a 100644 --- a/dropshot/src/lib.rs +++ b/dropshot/src/lib.rs @@ -611,6 +611,7 @@ mod router; mod server; mod to_map; mod type_util; +mod websocket; pub mod test_util; @@ -643,6 +644,7 @@ pub use handler::HttpResponseDeleted; pub use handler::HttpResponseHeaders; pub use handler::HttpResponseOk; pub use handler::HttpResponseUpdatedNoContent; +pub use handler::HttpResponseUpgradedWebSocket; pub use handler::NoHeaders; pub use handler::Path; pub use handler::Query; @@ -664,6 +666,7 @@ pub use pagination::ResultsPage; pub use pagination::WhichPage; pub use server::ServerContext; pub use server::{HttpServer, HttpServerStarter}; +pub use websocket::{WebSocket, WebSocketConfig, WebSocketMessage}; /* * Users of the `endpoint` macro need the following macros: diff --git a/dropshot/src/websocket.rs b/dropshot/src/websocket.rs new file mode 100644 index 000000000..3de669176 --- /dev/null +++ b/dropshot/src/websocket.rs @@ -0,0 +1,287 @@ +//! Websocket specific functions. + +use std::{ + borrow::Cow, + pin::Pin, + task::{Context, Poll}, +}; + +use futures::{future, ready, FutureExt, Sink, Stream}; +use tokio_tungstenite::{tungstenite::protocol, WebSocketStream}; + +use crate::HttpError; + +/// The configuration for the websocket. +pub struct WebSocketConfig { + /// The size of the send queue. You can use it to turn on/off the backpressure features. `None` means here that the size of the queue is unlimited. The default value is the unlimited queue. + pub max_send_queue: Option, + /// The maximum size of a message. `None` means no size limit. The default value is 64 MiB which should be reasonably big for all normal use-cases but small enough to prevent memory eating by a malicious user. + pub max_message_size: Option, + /// The maximum size of a single message frame. `None` means no size limit. The limit is for frame payload NOT including the frame header. The default value is 16 MiB which should be reasonably big for all normal use-cases but small enough to prevent memory eating by a malicious user. + pub max_frame_size: Option, + /// When set to `true`, the server will accept and handle unmasked frames from the client. According to the RFC 6455, the server must close the connection to the client in such cases, however it seems like there are some popular libraries that are sending unmasked frames, ignoring the RFC. By default this option is set to `false`, i.e. according to RFC 6455. + pub accept_unmasked_frames: bool, +} + +/// A websocket `Stream` and `Sink`. +/// +/// Ping messages sent from the client will be handled internally by replying with a Pong message. +/// Close messages need to be handled explicitly: usually by closing the `Sink` end of the +/// `WebSocket`. +/// +/// **Note!** +/// Due to rust futures nature, pings won't be handled until read part of `WebSocket` is polled +pub struct WebSocket { + inner: WebSocketStream, +} + +impl WebSocket { + pub async fn from_raw_socket( + upgraded: hyper::upgrade::Upgraded, + role: protocol::Role, + config: Option, + ) -> Self { + WebSocketStream::from_raw_socket( + upgraded, + role, + config.map(|c| protocol::WebSocketConfig { + max_send_queue: c.max_send_queue, + max_message_size: c.max_message_size, + max_frame_size: c.max_frame_size, + accept_unmasked_frames: c.accept_unmasked_frames, + }), + ) + .map(|inner| WebSocket { inner }) + .await + } + + /// Gracefully close this websocket. + pub async fn close(mut self) -> Result<(), HttpError> { + future::poll_fn(|cx| Pin::new(&mut self).poll_close(cx)).await + } +} + +impl Stream for WebSocket { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + match ready!(Pin::new(&mut self.inner).poll_next(cx)) { + Some(Ok(item)) => { + Poll::Ready(Some(Ok(WebSocketMessage { inner: item }))) + } + Some(Err(e)) => Poll::Ready(Some(Err(e.into()))), + None => Poll::Ready(None), + } + } +} + +impl Sink for WebSocket { + type Error = HttpError; + + fn poll_ready( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + match ready!(Pin::new(&mut self.inner).poll_ready(cx)) { + Ok(()) => Poll::Ready(Ok(())), + Err(e) => Poll::Ready(Err(e.into())), + } + } + + fn start_send( + mut self: Pin<&mut Self>, + item: WebSocketMessage, + ) -> Result<(), HttpError> { + match Pin::new(&mut self.inner).start_send(item.inner) { + Ok(()) => Ok(()), + Err(e) => Err(e.into()), + } + } + + fn poll_flush( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + match ready!(Pin::new(&mut self.inner).poll_flush(cx)) { + Ok(()) => Poll::Ready(Ok(())), + Err(e) => Poll::Ready(Err(e.into())), + } + } + + fn poll_close( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + match ready!(Pin::new(&mut self.inner).poll_close(cx)) { + Ok(()) => Poll::Ready(Ok(())), + Err(err) => Poll::Ready(Err(err.into())), + } + } +} + +impl std::fmt::Debug for WebSocket { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WebSocket").finish() + } +} + +/// A WebSocket message. +#[derive(Eq, PartialEq, Clone)] +pub struct WebSocketMessage { + inner: protocol::Message, +} + +impl WebSocketMessage { + /// Construct a new Text `Message`. + pub fn text>(s: S) -> Self { + WebSocketMessage { inner: protocol::Message::text(s) } + } + + /// Construct a new Json `Message` + pub fn json(j: T) -> Result + where + T: serde::Serialize, + { + Ok(Self { + inner: protocol::Message::text(serde_json::to_string(&j).map_err( + |e| { + HttpError::for_bad_request( + None, + format!( + "unable to serialize string as JSON body: {}", + e + ), + ) + }, + )?), + }) + } + + /// Construct a new Binary `Message`. + pub fn binary>>(v: V) -> Self { + WebSocketMessage { inner: protocol::Message::binary(v) } + } + + /// Construct a new Ping `Message`. + pub fn ping>>(v: V) -> Self { + WebSocketMessage { inner: protocol::Message::Ping(v.into()) } + } + + /// Construct a new Pong `Message`. + /// + /// Note that one rarely needs to manually construct a Pong message because the underlying tungstenite socket + /// automatically responds to the Ping messages it receives. Manual construction might still be useful in some cases + /// like in tests or to send unidirectional heartbeats. + pub fn pong>>(v: V) -> Self { + WebSocketMessage { inner: protocol::Message::Pong(v.into()) } + } + + /// Construct the default Close `Message`. + pub fn close() -> Self { + WebSocketMessage { inner: protocol::Message::Close(None) } + } + + /// Construct a Close `Message` with a code and reason. + pub fn close_with( + code: impl Into, + reason: impl Into>, + ) -> Self { + WebSocketMessage { + inner: protocol::Message::Close(Some( + protocol::frame::CloseFrame { + code: protocol::frame::coding::CloseCode::from(code.into()), + reason: reason.into(), + }, + )), + } + } + + /// Returns true if this message is a Text message. + pub fn is_text(&self) -> bool { + self.inner.is_text() + } + + /// Returns true if this message is a Binary message. + pub fn is_binary(&self) -> bool { + self.inner.is_binary() + } + + /// Returns true if this message a is a Close message. + pub fn is_close(&self) -> bool { + self.inner.is_close() + } + + /// Returns true if this message is a Ping message. + pub fn is_ping(&self) -> bool { + self.inner.is_ping() + } + + /// Returns true if this message is a Pong message. + pub fn is_pong(&self) -> bool { + self.inner.is_pong() + } + + /// Try to get the close frame (close code and reason) + pub fn close_frame(&self) -> Option<(u16, &str)> { + if let protocol::Message::Close(Some(ref close_frame)) = self.inner { + Some((close_frame.code.into(), close_frame.reason.as_ref())) + } else { + None + } + } + + /// Try to get a reference to the string text, if this is a Text message. + pub fn to_str(&self) -> Result<&str, HttpError> { + match self.inner { + protocol::Message::Text(ref s) => Ok(s), + _ => Err(HttpError::for_internal_error( + "websocket message is not a text message".to_string(), + )), + } + } + + /// Return the bytes of this message, if the message can contain data. + pub fn as_bytes(&self) -> &[u8] { + match self.inner { + protocol::Message::Text(ref s) => s.as_bytes(), + protocol::Message::Binary(ref v) => v, + protocol::Message::Ping(ref v) => v, + protocol::Message::Pong(ref v) => v, + protocol::Message::Frame(ref frame) => frame.payload().as_slice(), + protocol::Message::Close(_) => &[], + } + } + + /// Return the type as decoded json. + pub fn as_json(&self) -> Result + where + T: serde::de::DeserializeOwned, + { + serde_json::from_slice(self.as_bytes()).map_err(|e| { + HttpError::for_internal_error(format!( + "deserializing websocket message as JSON failed: {}", + e + )) + }) + } + + /// Destructure this message into binary data. + pub fn into_bytes(self) -> Vec { + self.inner.into_data() + } +} + +impl std::fmt::Debug for WebSocketMessage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + std::fmt::Debug::fmt(&self.inner, f) + } +} + +impl From for Vec { + fn from(m: WebSocketMessage) -> Self { + m.into_bytes() + } +} diff --git a/dropshot/tests/test_websocket.rs b/dropshot/tests/test_websocket.rs new file mode 100644 index 000000000..5511cd02e --- /dev/null +++ b/dropshot/tests/test_websocket.rs @@ -0,0 +1,139 @@ +// Copyright 2021 Oxide Computer Company + +//! Test cases for websockets. + +use dropshot::{ + endpoint, ApiDescription, HttpError, HttpResponseUpgradedWebSocket, + RequestContext, +}; +use futures::{FutureExt, SinkExt, StreamExt}; +use http::{Method, StatusCode}; +use std::sync::Arc; + +extern crate slog; + +pub mod common; + +fn api() -> ApiDescription { + let mut api = ApiDescription::new(); + api.register(websocket).unwrap(); + api +} + +#[allow(unused_variables)] +#[endpoint { + method = GET, + path = "/echo", +}] +/// Echo a message back to the client. +async fn websocket( + rqctx: Arc>, +) -> Result { + let (response, _) = rqctx + .upgrade_websocket(None, |ws| { + // Just echo all messages back... + let (tx, rx) = ws.split(); + rx.forward(tx).map(|result| { + if let Err(e) = result { + eprintln!("websocket error: {:?}", e); + } + }) + }) + .await?; + + Ok(response) +} + +#[tokio::test] +async fn test_websocket_server_websocket_client() { + let api = api(); + let testctx = common::test_setup("websocket_server_websocket_client", api); + let client = &testctx.client_testctx; + + let url = client.url("/echo").to_string().replace("http://", "ws://"); + + // Actually make a websocket connection. + let (mut ws_stream, _) = + tokio_tungstenite::connect_async(url).await.unwrap(); + + // Write a message. + ws_stream + .send(tokio_tungstenite::tungstenite::Message::Text( + "Hello, world!".to_string(), + )) + .await + .unwrap(); + + // Get the first message. + let msg = ws_stream.next().await.unwrap().unwrap(); + if let tokio_tungstenite::tungstenite::Message::Text(text) = msg { + assert_eq!(text, "Hello, world!"); + } else { + unreachable!(); + } + + testctx.teardown().await; +} + +#[tokio::test] +async fn test_websocket_server_no_upgrade_header() { + let api = api(); + let testctx = common::test_setup("websocket_server_no_upgrade_header", api); + let client = &testctx.client_testctx; + + if let Err(e) = client + .make_request_no_body(Method::GET, "/echo", StatusCode::BAD_REQUEST) + .await + { + assert_eq!(e.message, "Connection header not sent"); + } else { + unreachable!(); + } +} + +#[tokio::test] +async fn test_websocket_server_upgrade_header_incorrect() { + let api = api(); + let testctx = + common::test_setup("websocket_server_upgrade_header_incorrect", api); + let client = &testctx.client_testctx; + + let uri = client.url("/echo"); + let request = hyper::Request::builder() + .method(Method::GET) + .header("Connection", "BAD") + .uri(uri) + .body(hyper::Body::empty()) + .expect("attempted to construct invalid request"); + + if let Err(e) = + client.make_request_with_request(request, StatusCode::BAD_REQUEST).await + { + assert_eq!(e.message, "Connection header did not include 'upgrade'"); + } else { + unreachable!(); + } +} + +#[tokio::test] +async fn test_websocket_server_no_websocket_key() { + let api = api(); + let testctx = common::test_setup("websocket_server_no_websocket_key", api); + let client = &testctx.client_testctx; + + let uri = client.url("/echo"); + let request = hyper::Request::builder() + .method(Method::GET) + .header("Connection", "Upgrade") + .uri(uri) + .body(hyper::Body::empty()) + .expect("attempted to construct invalid request"); + + if let Err(e) = + client.make_request_with_request(request, StatusCode::BAD_REQUEST).await + { + assert_eq!(e.message, "Websocket key not sent"); + } else { + unreachable!(); + } +}