From 3feb22529bca4c63bb0171f88e6dfa1b62850ddf Mon Sep 17 00:00:00 2001 From: Connor Hindley Date: Thu, 15 May 2025 23:28:48 -0500 Subject: [PATCH] feat: Add support for DO websocket_auto_response. Noticed this was missing from the rust apis. Basic bindgen, miniflare test, and sandbox example. --- worker-sandbox/src/router.rs | 8 +++ worker-sandbox/src/test/auto_response.rs | 49 +++++++++++++++++++ worker-sandbox/src/test/mod.rs | 1 + worker-sandbox/tests/auto_response.spec.ts | 10 ++++ worker-sandbox/tests/mf.ts | 1 + worker-sandbox/wrangler.toml | 1 + worker-sys/src/types.rs | 2 + worker-sys/src/types/durable_object/state.rs | 13 ++++- .../types/websocket_request_response_pair.rs | 17 +++++++ worker/src/durable.rs | 8 +++ 10 files changed, 109 insertions(+), 1 deletion(-) create mode 100644 worker-sandbox/src/test/auto_response.rs create mode 100644 worker-sandbox/tests/auto_response.spec.ts create mode 100644 worker-sys/src/types/websocket_request_response_pair.rs diff --git a/worker-sandbox/src/router.rs b/worker-sandbox/src/router.rs index a93862e8..d66f7e19 100644 --- a/worker-sandbox/src/router.rs +++ b/worker-sandbox/src/router.rs @@ -224,6 +224,10 @@ pub fn make_router(data: SomeSharedData, env: Env) -> axum::Router { "/analytics-engine", get(handler!(analytics_engine::handle_analytics_event)), ) + .route( + "/durable/auto-response", + get(handler!(crate::test::auto_response::handle_auto_response)), + ) .fallback(get(handler!(catchall))) .layer(Extension(env)) .layer(Extension(data)) @@ -364,6 +368,10 @@ pub fn make_router<'a>(data: SomeSharedData) -> Router<'a, SomeSharedData> { .delete_async("/r2/delete", handler!(r2::delete)) .get_async("/socket/failed", handler!(socket::handle_socket_failed)) .get_async("/socket/read", handler!(socket::handle_socket_read)) + .get_async( + "/durable/auto-response", + handler!(crate::test::auto_response::handle_auto_response), + ) .or_else_any_method_async("/*catchall", handler!(catchall)) } diff --git a/worker-sandbox/src/test/auto_response.rs b/worker-sandbox/src/test/auto_response.rs new file mode 100644 index 00000000..fc3b0cf7 --- /dev/null +++ b/worker-sandbox/src/test/auto_response.rs @@ -0,0 +1,49 @@ +use worker::*; + +#[durable_object] +pub struct AutoResponseObject { + state: State, +} + +#[durable_object] +impl DurableObject for AutoResponseObject { + fn new(state: State, _env: Env) -> Self { + Self { state } + } + + async fn fetch(&mut self, req: Request) -> Result { + match req.path().as_str() { + "/set" => { + // Configure ping -> pong auto-response for all websockets bound to this DO. + let pair = worker_sys::WebSocketRequestResponsePair::new("ping", "pong")?; + self.state.set_websocket_auto_response(&pair); + Response::ok("ok") + } + "/get" => { + if let Some(pair) = self.state.get_websocket_auto_response() { + let req_str = pair.request(); + let res_str = pair.response(); + Response::ok(format!("{req_str}:{res_str}")) + } else { + Response::ok("none") + } + } + _ => Response::error("Not Found", 404), + } + } +} + +// Route handler to exercise the Durable Object from tests. +#[worker::send] +pub async fn handle_auto_response( + _req: Request, + env: Env, + _data: crate::SomeSharedData, +) -> Result { + let namespace = env.durable_object("AUTO")?; + let stub = namespace.id_from_name("singleton")?.get_stub()?; + // Ensure auto-response is configured + stub.fetch_with_str("https://fake-host/set").await?; + // Retrieve and return it for assertion + stub.fetch_with_str("https://fake-host/get").await +} diff --git a/worker-sandbox/src/test/mod.rs b/worker-sandbox/src/test/mod.rs index ec112528..a883dd43 100644 --- a/worker-sandbox/src/test/mod.rs +++ b/worker-sandbox/src/test/mod.rs @@ -1,3 +1,4 @@ +pub mod auto_response; pub mod durable; pub mod export_durable_object; pub mod put_raw; diff --git a/worker-sandbox/tests/auto_response.spec.ts b/worker-sandbox/tests/auto_response.spec.ts new file mode 100644 index 00000000..12a474a4 --- /dev/null +++ b/worker-sandbox/tests/auto_response.spec.ts @@ -0,0 +1,10 @@ +import { describe, test, expect } from "vitest"; +import { mf } from "./mf"; + +describe("durable object websocket auto-response", () => { + test("set and get auto-response pair", async () => { + const resp = await mf.dispatchFetch("http://fake.host/durable/auto-response"); + const text = await resp.text(); + expect(text).toBe("ping:pong"); + }); +}); diff --git a/worker-sandbox/tests/mf.ts b/worker-sandbox/tests/mf.ts index 7886c5b4..a8833bac 100644 --- a/worker-sandbox/tests/mf.ts +++ b/worker-sandbox/tests/mf.ts @@ -59,6 +59,7 @@ export const mf = new Miniflare({ durableObjects: { COUNTER: "Counter", PUT_RAW_TEST_OBJECT: "PutRawTestObject", + AUTO: "AutoResponseObject", }, kvNamespaces: ["SOME_NAMESPACE", "FILE_SIZES", "TEST"], serviceBindings: { diff --git a/worker-sandbox/wrangler.toml b/worker-sandbox/wrangler.toml index 31d82116..ac4719f1 100644 --- a/worker-sandbox/wrangler.toml +++ b/worker-sandbox/wrangler.toml @@ -29,6 +29,7 @@ bindings = [ { name = "COUNTER", class_name = "Counter" }, { name = "ALARM", class_name = "AlarmObject" }, { name = "PUT_RAW_TEST_OBJECT", class_name = "PutRawTestObject" }, + { name = "AUTO", class_name = "AutoResponseObject" }, ] [[analytics_engine_datasets]] diff --git a/worker-sys/src/types.rs b/worker-sys/src/types.rs index 76630bbb..947fc11c 100644 --- a/worker-sys/src/types.rs +++ b/worker-sys/src/types.rs @@ -19,6 +19,7 @@ mod socket; mod tls_client_auth; mod version; mod websocket_pair; +mod websocket_request_response_pair; pub use ai::*; pub use analytics_engine::*; @@ -41,3 +42,4 @@ pub use socket::*; pub use tls_client_auth::*; pub use version::*; pub use websocket_pair::*; +pub use websocket_request_response_pair::*; diff --git a/worker-sys/src/types/durable_object/state.rs b/worker-sys/src/types/durable_object/state.rs index 9ecb79da..0e5507b1 100644 --- a/worker-sys/src/types/durable_object/state.rs +++ b/worker-sys/src/types/durable_object/state.rs @@ -1,6 +1,6 @@ use wasm_bindgen::prelude::*; -use crate::types::{DurableObjectId, DurableObjectStorage}; +use crate::types::{DurableObjectId, DurableObjectStorage, WebSocketRequestResponsePair}; #[wasm_bindgen] extern "C" { @@ -43,4 +43,15 @@ extern "C" { this: &DurableObjectState, ws: &web_sys::WebSocket, ) -> Result, JsValue>; + + #[wasm_bindgen(method, catch, js_name=setWebSocketAutoResponse)] + pub fn set_websocket_auto_response( + this: &DurableObjectState, + pair: &WebSocketRequestResponsePair, + ) -> Result<(), JsValue>; + + #[wasm_bindgen(method, catch, js_name=getWebSocketAutoResponse)] + pub fn get_websocket_auto_response( + this: &DurableObjectState, + ) -> Result, JsValue>; } diff --git a/worker-sys/src/types/websocket_request_response_pair.rs b/worker-sys/src/types/websocket_request_response_pair.rs new file mode 100644 index 00000000..07bff451 --- /dev/null +++ b/worker-sys/src/types/websocket_request_response_pair.rs @@ -0,0 +1,17 @@ +use wasm_bindgen::prelude::*; + +#[wasm_bindgen] +extern "C" { + #[wasm_bindgen(extends=js_sys::Object)] + #[derive(Debug, Clone, PartialEq, Eq)] + pub type WebSocketRequestResponsePair; + + #[wasm_bindgen(constructor, catch)] + pub fn new(request: &str, response: &str) -> Result; + + #[wasm_bindgen(method, getter)] + pub fn request(this: &WebSocketRequestResponsePair) -> String; + + #[wasm_bindgen(method, getter)] + pub fn response(this: &WebSocketRequestResponsePair) -> String; +} diff --git a/worker/src/durable.rs b/worker/src/durable.rs index b6751e16..5b5a1ca2 100644 --- a/worker/src/durable.rs +++ b/worker/src/durable.rs @@ -264,6 +264,14 @@ impl State { pub fn get_tags(&self, websocket: &WebSocket) -> Vec { self.inner.get_tags(websocket.as_ref()).unwrap() } + + pub fn set_websocket_auto_response(&self, pair: &worker_sys::WebSocketRequestResponsePair) { + self.inner.set_websocket_auto_response(pair).unwrap(); + } + + pub fn get_websocket_auto_response(&self) -> Option { + self.inner.get_websocket_auto_response().unwrap() + } } impl From for State {