Merge pull request #394 from Dstack-TEE/openapi

vmm: Add auto generated openapi docs

Author: kvinwang

.gitignore

@@ -8,3 +8,4 @@ node_modules/
 .venv
 /tmp
 .claude/settings.local.json
+__pycache__

Cargo.lock

@@ -112,6 +112,7 @@ hex = { version = "0.4.3", default-features = false }
 hex_fmt = "0.3.0"
 hex-literal = "1.0.0"
 prost = "0.13.5"
+prost-types = "0.13.5"
 scale = { version = "3.7.4", package = "parity-scale-codec", features = [
     "derive",
 ] }

Cargo.toml

@@ -1,3 +1,6 @@
 #![allow(async_fn_in_trait)]
 
+pub const FILE_DESCRIPTOR_SET: &[u8] =
+    include_bytes!(concat!(env!("OUT_DIR"), "/file_descriptor_set.bin"));
+
 include!(concat!(env!("OUT_DIR"), "/gateway.rs"));

gateway/rpc/src/generated.rs

@@ -1,3 +1,6 @@
 #![allow(async_fn_in_trait)]
 
+pub const FILE_DESCRIPTOR_SET: &[u8] =
+    include_bytes!(concat!(env!("OUT_DIR"), "/file_descriptor_set.bin"));
+
 include!(concat!(env!("OUT_DIR"), "/dstack_guest.rs"));

guest-agent/rpc/src/generated.rs

@@ -1,3 +1,6 @@
 #![allow(async_fn_in_trait)]
 
+pub const FILE_DESCRIPTOR_SET: &[u8] =
+    include_bytes!(concat!(env!("OUT_DIR"), "/file_descriptor_set.bin"));
+
 include!(concat!(env!("OUT_DIR"), "/guest_api.rs"));

guest-api/src/generated/mod.rs

@@ -1,3 +1,6 @@
 #![allow(async_fn_in_trait)]
 
+pub const FILE_DESCRIPTOR_SET: &[u8] =
+    include_bytes!(concat!(env!("OUT_DIR"), "/file_descriptor_set.bin"));
+
 include!(concat!(env!("OUT_DIR"), "/host_api.rs"));

host-api/src/generated/mod.rs

@@ -1,3 +1,6 @@
 #![allow(async_fn_in_trait)]
 
+pub const FILE_DESCRIPTOR_SET: &[u8] =
+    include_bytes!(concat!(env!("OUT_DIR"), "/file_descriptor_set.bin"));
+
 include!(concat!(env!("OUT_DIR"), "/kms.rs"));

kms/rpc/src/generated.rs

@@ -22,8 +22,10 @@ bon.workspace = true
 rocket-vsock-listener = { workspace = true, optional = true }
 serde.workspace = true
 x509-parser.workspace = true
+prost-types = { workspace = true, optional = true }
 
 [features]
 default = ["rocket", "client"]
 rocket = ["dep:rocket", "dep:rocket-vsock-listener"]
 client = ["reqwest"]
+openapi = ["dep:prost-types"]

ra-rpc/Cargo.toml

@@ -0,0 +1,73 @@
+# Generating OpenAPI docs for pRPC services
+
+This repository now ships a lightweight OpenAPI generator inside `ra-rpc`. It can
+derive a specification (plus a Swagger UI test page) directly from the protobuf
+descriptors that are already produced during the `prpc_build` step.
+
+## 1. Enable the feature
+
+Add the `openapi` feature when depending on `ra-rpc`:
+
+```toml
+[dependencies]
+ra-rpc = { path = "../ra-rpc", features = ["openapi", "rocket"] }
+```
+
+The `rocket` feature is optional if you only need the JSON document and plan to
+serve it through another framework.
+
+## 2. Export the descriptor from your `*-rpc` crate
+
+Every RPC crate already has access to `file_descriptor_set.bin`. Expose it so
+that application binaries can include it:
+
+```rust
+pub const FILE_DESCRIPTOR_SET: &[u8] =
+    include_bytes!(concat!(env!("OUT_DIR"), "/file_descriptor_set.bin"));
+```
+
+This repository now does that for all existing RPC crates.
+
+## 3. Build and mount the document
+
+```rust
+use ra_rpc::openapi::{
+    build_openapi_doc, DescriptorSource, DocumentInfo, ServiceConfig, SwaggerUiConfig,
+};
+
+fn openapi_doc() -> anyhow::Result<ra_rpc::openapi::OpenApiDoc> {
+    let descriptor = dstack_guest_agent_rpc::FILE_DESCRIPTOR_SET;
+    let sources = vec![DescriptorSource::new(
+        descriptor,
+        vec![
+            // Mounts /prpc/Worker.Version (prefix is optional when you don't trim).
+            ServiceConfig::new("Worker", "/prpc").with_method_prefix("Worker."),
+        ],
+    )];
+
+    let info = DocumentInfo::new("Guest Worker API", env!("CARGO_PKG_VERSION"))
+        .with_description("Auto generated from protobuf descriptors")
+        .add_server("https://example.com/prpc");
+
+    let ui = SwaggerUiConfig {
+        title: "Guest Worker RPC".into(),
+        ..Default::default()
+    };
+
+    build_openapi_doc(&sources, &info, ui)
+}
+```
+
+Serving it through Rocket is one line:
+
+```rust
+let openapi = openapi_doc()?;
+let rocket = ra_rpc::rocket_helper::mount_openapi_docs(rocket, openapi, "/rpc-docs");
+```
+
+* `GET /rpc-docs/openapi.json` returns the specification.
+* `GET /rpc-docs/docs` serves a Swagger UI page backed by the same spec.
+
+You can mount as many descriptor sources as you need (for example when the same
+binary exposes both admin and user RPC stacks). Just add more `DescriptorSource`
+entries that point to the relevant `FILE_DESCRIPTOR_SET` constants.