initial websocket implmentation

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]

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<RequestContext<()>>,
+) -> Result<HttpResponseUpgradedWebSocket, HttpError> {
+    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)
+}

dropshot/src/error.rs

@@ -310,6 +310,12 @@ impl Error for HttpError {
     }
 }
 
+impl From<tokio_tungstenite::tungstenite::error::Error> for HttpError {
+    fn from(error: tokio_tungstenite::tungstenite::error::Error) -> Self {
+        HttpError::for_internal_error(format!("{}", error))
+    }
+}
+
 #[cfg(test)]
 mod test {
     use crate::HttpErrorResponseBody;