Support rocket_ws for WebSocket support

- New feature flag `rocket_sync_db_pools` for compatibility with [`rocket_sync_db_pools`](https://crates.io/crates/rocket_sync_db_pools). - New feature flag `rocket_ws` for compatibility with [`rocket_ws`](https://crates.io/crates/rocket_ws). - Added new example for WebSockets. - Added support for new [`Responder`](https://docs.rs/rocket/0.5.0/rocket/response/trait.Responder.html) types (implemented `OpenApiResponderInner`): - `rocket_ws::Channel<'o>` (when `rocket_ws` feature is enabled) - `rocket_ws::stream::MessageStream<'o, S>` (when `rocket_ws` feature is enabled) - Added support for new [`FromRequest`](https://docs.rs/rocket/0.5.0/rocket/request/trait.FromRequest.html) types (implemented `OpenApiFromRequest`): - `rocket_dyn_templates::Metadata<'r>` (when `rocket_dyn_templates` feature is enabled) - `rocket_sync_db_pools::example::ExampleDb` (when `rocket_sync_db_pools` feature is enabled) - `rocket_ws::WebSocket` (when `rocket_ws` feature is enabled)
GREsau · Dec 17, 2023 · fa77b69 · fa77b69
1 parent 22ea565
commit fa77b69
Show file tree

Hide file tree

Showing 11 changed files with 224 additions and 2 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -12,5 +12,6 @@ members = [
     "examples/dyn_templates",
     "examples/openapi_attributes",
     "examples/raw_identifiers",
+    "examples/websocket",
 ]
 resolver = "2"
diff --git a/README.md b/README.md
@@ -142,6 +142,8 @@ Rocket-Okapi:
 (when same Rocket feature flag is used.)
 - `rocket_dyn_templates`: Enable compatibility with [`rocket_dyn_templates`](https://crates.io/crates/rocket_dyn_templates).
 - `rocket_db_pools`: Enable compatibility with [`rocket_db_pools`](https://crates.io/crates/rocket_db_pools).
+- `rocket_sync_db_pools`: Enable compatibility with [`rocket_sync_db_pools`](https://crates.io/crates/rocket_sync_db_pools).
+- `rocket_ws`: Enable compatibility with [`rocket_ws`](https://crates.io/crates/rocket_ws).
 
 Note that not all feature flags from [`Schemars`][Schemars] are re-exported or enabled.
 So if you have objects for which the `JsonSchema` trait is not implemented,

diff --git a/examples/websocket/.gitignore b/examples/websocket/.gitignore
@@ -0,0 +1,4 @@
+/target
+**/*.rs.bk
+Cargo.lock
+/.idea
diff --git a/examples/websocket/Cargo.toml b/examples/websocket/Cargo.toml
@@ -0,0 +1,15 @@
+[package]
+name = "websocket"
+version = "0.1.0"
+authors = ["Ralph Bisschops <[email protected]>"]
+edition = "2021"
+
+[dependencies]
+rocket = { version = "=0.5.0", default-features = false, features = ["json"] }
+rocket_ws = "0.1.0"
+rocket_okapi = { path = "../../rocket-okapi", features = [
+    "swagger",
+    "rapidoc",
+    "rocket_ws",
+] }
+serde = "1.0"
diff --git a/examples/websocket/src/main.rs b/examples/websocket/src/main.rs
@@ -0,0 +1,105 @@
+use rocket::futures::{SinkExt, StreamExt};
+use rocket::get;
+use rocket::response::content::RawHtml;
+use rocket_okapi::settings::UrlObject;
+use rocket_okapi::{openapi, openapi_get_routes, rapidoc::*, swagger_ui::*};
+
+#[openapi]
+#[get("/")]
+fn test_websocket() -> RawHtml<&'static str> {
+    RawHtml(
+        r#"
+        <!DOCTYPE html>
+        <html>
+            <body>
+                Echo: <input type="text" id="echo_text" name="echo" size="10" />
+                <input type="button" value="Send" onclick="echo_send()" />
+                <br/>
+                <br/>
+                <p id="output"><p>
+                <script>
+                    // Create WebSocket connection.
+                    const hello_socket = new WebSocket("ws://localhost:8000/hello/bob");
+                    const echo_socket = new WebSocket("ws://localhost:8000/echo");
+                    const output = document.getElementById('output');
+                    
+                    // Listen for messages
+                    hello_socket.addEventListener("message", (event) => {
+                        console.log("Hello response: ", event.data);
+                        output.innerHTML += "Hello response: " + event.data + "<br/>";
+                    });
+                    echo_socket.addEventListener("message", (event) => {
+                        console.log("Echo response: ", event.data);
+                        output.innerHTML += "Echo response: " + event.data + "<br/>";
+                    });
+
+                    function echo_send(){
+                        echo_socket.send(document.getElementById('echo_text').value);
+                    }
+                </script>
+            </body>
+        </html>
+        "#,
+    )
+}
+
+#[openapi]
+#[get("/hello/<name>")]
+fn hello(ws: rocket_ws::WebSocket, name: &str) -> rocket_ws::Channel<'_> {
+    ws.channel(move |mut stream| {
+        Box::pin(async move {
+            let message = format!("Hello, {}!", name);
+            let _ = stream.send(message.into()).await;
+            Ok(())
+        })
+    })
+}
+
+#[openapi]
+#[get("/echo")]
+fn echo(ws: rocket_ws::WebSocket) -> rocket_ws::Channel<'static> {
+    ws.channel(move |mut stream| {
+        Box::pin(async move {
+            while let Some(message) = stream.next().await {
+                let _ = stream.send(message?).await;
+            }
+
+            Ok(())
+        })
+    })
+}
+
+#[rocket::main]
+async fn main() {
+    let launch_result = rocket::build()
+        .mount("/", openapi_get_routes![test_websocket, hello, echo,])
+        .mount(
+            "/swagger-ui/",
+            make_swagger_ui(&SwaggerUIConfig {
+                url: "../openapi.json".to_owned(),
+                ..Default::default()
+            }),
+        )
+        .mount(
+            "/rapidoc/",
+            make_rapidoc(&RapiDocConfig {
+                general: GeneralConfig {
+                    spec_urls: vec![UrlObject::new("General", "../openapi.json")],
+                    ..Default::default()
+                },
+                hide_show: HideShowConfig {
+                    allow_spec_url_load: false,
+                    allow_spec_file_load: false,
+                    allow_spec_file_download: true,
+                    ..Default::default()
+                },
+                ..Default::default()
+            }),
+        )
+        .launch()
+        .await;
+    match launch_result {
+        Ok(_) => println!("Rocket shut down gracefully."),
+        Err(err) => println!("Rocket had an error: {}", err),
+    };
+}
diff --git a/rocket-okapi-codegen/src/openapi_attr/mod.rs b/rocket-okapi-codegen/src/openapi_attr/mod.rs
@@ -336,6 +336,7 @@ fn create_route_operation_fn(
             let request_body = #request_body;
             // Add the security scheme that are quired for all the routes.
             let mut security_requirements = Vec::new();
+            let mut server_requirements = Vec::new();
 
             // Combine all parameters from all sources
             // Add all from `path_params` and `path_multi_param`
@@ -366,6 +367,15 @@ fn create_route_operation_fn(
                         // Add the security scheme that are quired for all the route.
                         security_requirements.push(requirement);
                     }
+                    // Add Server to this request.
+                    RequestHeaderInput::Server(url, description, variables) => {
+                        server_requirements.push(::rocket_okapi::okapi::openapi3::Server{
+                            url,
+                            description,
+                            variables,
+                            ..Default::default()
+                        });
+                    }
                     _ => {
                     }
                 }
@@ -377,6 +387,12 @@ fn create_route_operation_fn(
             } else {
                 Some(security_requirements)
             };
+            // Add `servers` section if list is not empty
+            let servers = if server_requirements.is_empty() {
+                None
+            } else {
+                Some(server_requirements)
+            };
             // Add route/endpoint to OpenApi object.
             gen.add_operation(::rocket_okapi::OperationInfo {
                 path: #path.to_owned(),
@@ -389,6 +405,7 @@ fn create_route_operation_fn(
                     summary: #title,
                     description: #desc,
                     security,
+                    servers,
                     tags: vec![#(#tags),*],
                     deprecated: #deprecated,
                     ..Default::default()

diff --git a/rocket-okapi/CHANGELOG.md b/rocket-okapi/CHANGELOG.md
@@ -25,6 +25,19 @@ This project follows the [Semantic Versioning standard](https://semver.org/).
 - Added support for new [`FromRequest`](https://docs.rs/rocket/0.5.0/rocket/request/trait.FromRequest.html)
   types (implemented `OpenApiFromRequest`):
   - `rocket::request::Outcome<T, T::Error>`
+- New feature flag `rocket_sync_db_pools` for compatibility with
+  [`rocket_sync_db_pools`](https://crates.io/crates/rocket_sync_db_pools).
+- New feature flag `rocket_ws` for compatibility with [`rocket_ws`](https://crates.io/crates/rocket_ws).
+- Added new example for WebSockets.
+- Added support for new [`Responder`](https://docs.rs/rocket/0.5.0/rocket/response/trait.Responder.html)
+  types (implemented `OpenApiResponderInner`):
+  - `rocket_ws::Channel<'o>` (when `rocket_ws` feature is enabled)
+  - `rocket_ws::stream::MessageStream<'o, S>` (when `rocket_ws` feature is enabled)
+- Added support for new [`FromRequest`](https://docs.rs/rocket/0.5.0/rocket/request/trait.FromRequest.html)
+  types (implemented `OpenApiFromRequest`):
+  - `rocket_dyn_templates::Metadata<'r>` (when `rocket_dyn_templates` feature is enabled)
+  - `rocket_sync_db_pools::example::ExampleDb` (when `rocket_sync_db_pools` feature is enabled)
+  - `rocket_ws::WebSocket` (when `rocket_ws` feature is enabled)
 
 ### Changed
 - `rocket-okapi` and `rocket-okapi-codegen` require `rocket v0.5.0`. (#132)

diff --git a/rocket-okapi/Cargo.toml b/rocket-okapi/Cargo.toml
@@ -23,6 +23,8 @@ log = "0.4"
 # time = { version = "0.2.27" }
 rocket_dyn_templates = { version = "=0.1.0", optional = true }
 rocket_db_pools = { version = "=0.1.0", optional = true }
+rocket_sync_db_pools = { version = "=0.1.0", optional = true }
+rocket_ws = { version = "=0.1.0", optional = true }
 
 [dev-dependencies]
 rocket_sync_db_pools = { version = "0.1.0", features = ["diesel_sqlite_pool"] }

diff --git a/rocket-okapi/src/request/from_request_impls.rs b/rocket-okapi/src/request/from_request_impls.rs
@@ -6,6 +6,7 @@ use std::result::Result as StdResult;
 // Implement `OpenApiFromRequest` for everything that implements `FromRequest`
 // Order is same as on:
 // https://docs.rs/rocket/0.5.0/rocket/request/trait.FromRequest.html#foreign-impls
+// https://api.rocket.rs/v0.5/rocket/request/trait.FromRequest.html#foreign-impls
 
 type Result = crate::Result<RequestHeaderInput>;
 
@@ -191,3 +192,37 @@ impl<'r, D: rocket_db_pools::Database> OpenApiFromRequest<'r> for rocket_db_pool
         Ok(RequestHeaderInput::None)
     }
 }
+
+#[cfg(feature = "rocket_dyn_templates")]
+impl<'r> OpenApiFromRequest<'r> for rocket_dyn_templates::Metadata<'r> {
+    fn from_request_input(_gen: &mut OpenApiGenerator, _name: String, _required: bool) -> Result {
+        Ok(RequestHeaderInput::None)
+    }
+}
+
+#[cfg(feature = "rocket_sync_db_pools")]
+impl<'r> OpenApiFromRequest<'r> for rocket_sync_db_pools::example::ExampleDb {
+    fn from_request_input(_gen: &mut OpenApiGenerator, _name: String, _required: bool) -> Result {
+        Ok(RequestHeaderInput::None)
+    }
+}
+
+#[cfg(feature = "rocket_ws")]
+impl<'r> OpenApiFromRequest<'r> for rocket_ws::WebSocket {
+    fn from_request_input(_gen: &mut OpenApiGenerator, _name: String, _required: bool) -> Result {
+        Ok(RequestHeaderInput::Server(
+            "ws://{server}/{base_path}".to_owned(),
+            Some("WebSocket connection".to_owned()),
+            okapi::map! {
+                "server".to_owned() => okapi::openapi3::ServerVariable {
+                    default: "127.0.0.1:8000".to_owned(),
+                    ..Default::default()
+                },
+                "base_path".to_owned() => okapi::openapi3::ServerVariable {
+                    default: "".to_owned(),
+                    ..Default::default()
+                },
+            },
+        ))
+    }
+}
diff --git a/rocket-okapi/src/request/mod.rs b/rocket-okapi/src/request/mod.rs
@@ -80,6 +80,18 @@ pub enum RequestHeaderInput {
     /// - [`SecurityScheme`] is global definition of the authentication (per OpenApi spec).
     /// - [`SecurityRequirement`] is the requirements for the route.
     Security(String, SecurityScheme, SecurityRequirement),
+    /// A server this resources is allocated on.
+    ///
+    /// Parameters:
+    /// - The url
+    /// - The description
+    /// - Variable mapping: A map between a variable name and its value.
+    /// The value is used for substitution in the server’s URL template.
+    Server(
+        String,
+        Option<String>,
+        okapi::Map<String, okapi::openapi3::ServerVariable>,
+    ),
 }
 
 // Re-export derive trait here for convenience.

diff --git a/rocket-okapi/src/response/responder_impls.rs b/rocket-okapi/src/response/responder_impls.rs
@@ -381,8 +381,6 @@ impl<T: Serialize + JsonSchema + Send> OpenApiResponderInner
     }
 }
 
-// From: https://docs.rs/rocket_dyn_templates/latest/rocket_dyn_templates/struct.Template.html#impl-Responder%3C'r,+'static%3E-for-Template
-
 /// Response is set to `String` (so `text/plain`) because the actual return type is unknown
 /// at compile time. The content type depends on the file extension, but this can change at runtime.
 #[cfg(feature = "rocket_dyn_templates")]
@@ -391,3 +389,21 @@ impl OpenApiResponderInner for rocket_dyn_templates::Template {
         String::responses(gen)
     }
 }
+
+/// A streaming channel, returned by [`rocket_ws::WebSocket::channel()`].
+#[cfg(feature = "rocket_ws")]
+impl<'r, 'o: 'r> OpenApiResponderInner for rocket_ws::Channel<'o> {
+    fn responses(gen: &mut OpenApiGenerator) -> Result {
+        // Response type is unknown at compile time.
+        <Vec<u8>>::responses(gen)
+    }
+}
+
+/// A `Stream` of `Messages``, returned by [`rocket_ws::WebSocket::stream()`], used via `Stream!`.
+#[cfg(feature = "rocket_ws")]
+impl<'r, 'o: 'r, S> OpenApiResponderInner for rocket_ws::stream::MessageStream<'o, S> {
+    fn responses(gen: &mut OpenApiGenerator) -> Result {
+        // Response type is unknown at compile time.
+        <Vec<u8>>::responses(gen)
+    }
+}