Implement websocket support in Dropshot as an Extractor and #[channel] macro

dropshot/Cargo.toml

@@ -27,6 +27,7 @@ rustls = "0.20.6"
 rustls-pemfile = "1.0.1"
 serde_json = "1.0.83"
 serde_urlencoded = "0.7.1"
+sha1 = "0.10.1"
 slog-async = "2.4.0"
 slog-bunyan = "2.4.0"
 slog-json = "2.6.1"
@@ -84,6 +85,10 @@ trybuild = "1.0.64"
 # Used by the https examples and tests
 pem = "1.1"
 rcgen = "0.9.3"
+# Used in a doc-test demonstrating the WebsocketUpgrade extractor.
+tungstenite = "0.17.3"
+tokio-tungstenite = "0.17.2"
+futures-util = "0.3.21"
 
 [dev-dependencies.rustls]
 version = "0.20.6"

dropshot/examples/websocket.rs

@@ -0,0 +1,95 @@
+// Copyright 2022 Oxide Computer Company
+/*!
+ * Example use of Dropshot with a websocket endpoint.
+ */
+
+use dropshot::channel;
+use dropshot::ApiDescription;
+use dropshot::ConfigDropshot;
+use dropshot::ConfigLogging;
+use dropshot::ConfigLoggingLevel;
+use dropshot::HttpServerStarter;
+use dropshot::Query;
+use dropshot::RequestContext;
+use dropshot::WebsocketConnection;
+use futures_util::SinkExt;
+use schemars::JsonSchema;
+use serde::Deserialize;
+use std::sync::Arc;
+use tungstenite::protocol::Role;
+use tungstenite::Message;
+
+#[tokio::main]
+async fn main() -> Result<(), String> {
+    /*
+     * We must specify a configuration with a bind address.  We'll use 127.0.0.1
+     * since it's available and won't expose this server outside the host.  We
+     * request port 0, which allows the operating system to pick any available
+     * port.
+     */
+    let config_dropshot: ConfigDropshot = Default::default();
+
+    /*
+     * For simplicity, we'll configure an "info"-level logger that writes to
+     * stderr assuming that it's a terminal.
+     */
+    let config_logging =
+        ConfigLogging::StderrTerminal { level: ConfigLoggingLevel::Info };
+    let log = config_logging
+        .to_logger("example-basic")
+        .map_err(|error| format!("failed to create logger: {}", error))?;
+
+    /*
+     * Build a description of the API.
+     */
+    let mut api = ApiDescription::new();
+    api.register(example_api_websocket_counter).unwrap();
+
+    /*
+     * Set up the server.
+     */
+    let server = HttpServerStarter::new(&config_dropshot, api, (), &log)
+        .map_err(|error| format!("failed to create server: {}", error))?
+        .start();
+
+    /*
+     * Wait for the server to stop.  Note that there's not any code to shut down
+     * this server, so we should never get past this point.
+     */
+    server.await
+}
+
+/*
+ * HTTP API interface
+ */
+
+#[derive(Deserialize, JsonSchema)]
+struct QueryParams {
+    start: Option<u8>,
+}
+
+/**
+ * An eternally-increasing sequence of bytes, wrapping on overflow, starting
+ * from the value given for the query parameter "start."
+ */
+#[channel {
+    protocol = WEBSOCKETS,
+    path = "/counter",
+}]
+async fn example_api_websocket_counter(
+    _rqctx: Arc<RequestContext<()>>,
+    upgraded: WebsocketConnection,
+    qp: Query<QueryParams>,
+) -> dropshot::WebsocketChannelResult {
+    let mut ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
+        upgraded.into_inner(),
+        Role::Server,
+        None,
+    )
+    .await;
+    let mut count = qp.into_inner().start.unwrap_or(0);
+    while ws.send(Message::Binary(vec![count])).await.is_ok() {
+        count = count.wrapping_add(1);
+    }
+    Ok(())
+}

dropshot/src/api_description.rs

@@ -45,7 +45,7 @@ pub struct ApiEndpoint<Context: ServerContext> {
     pub summary: Option<String>,
     pub description: Option<String>,
     pub tags: Vec<String>,
-    pub paginated: bool,
+    pub extension_mode: ExtensionMode,
     pub visible: bool,
 }
 
@@ -78,7 +78,7 @@ impl<'a, Context: ServerContext> ApiEndpoint<Context> {
             summary: None,
             description: None,
             tags: vec![],
-            paginated: func_parameters.paginated,
+            extension_mode: func_parameters.extension_mode,
             visible: true,
         }
     }
@@ -688,11 +688,20 @@ impl<Context: ServerContext> ApiDescription<Context> {
                 })
                 .next();
 
-            if endpoint.paginated {
-                operation.extensions.insert(
-                    crate::pagination::PAGINATION_EXTENSION.to_string(),
-                    serde_json::json! {true},
-                );
+            match endpoint.extension_mode {
+                ExtensionMode::None => {}
+                ExtensionMode::Paginated => {
+                    operation.extensions.insert(
+                        crate::pagination::PAGINATION_EXTENSION.to_string(),
+                        serde_json::json! {true},
+                    );
+                }
+                ExtensionMode::Websocket => {
+                    operation.extensions.insert(
+                        crate::websocket::WEBSOCKET_EXTENSION.to_string(),
+                        serde_json::json!({}),
+                    );
+                }
             }
 
             let response = if let Some(schema) = &endpoint.response.schema {
@@ -1579,6 +1588,22 @@ pub struct TagExternalDocs {
     pub url: String,
 }
 
+/**
+ * Dropshot/Progenitor features used by endpoints which are not a part of the base OpenAPI spec.
+ */
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub enum ExtensionMode {
+    None,
+    Paginated,
+    Websocket,
+}
+
+impl Default for ExtensionMode {
+    fn default() -> Self {
+        ExtensionMode::None
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::j2oas_schema;