Skip to main content

rullst_core/
live.rs

1use crate::ws::WebSocket;
2use async_trait::async_trait;
3use axum::extract::ws::WebSocketUpgrade;
4use axum::response::IntoResponse;
5use serde_json::Value;
6
7#[cfg(test)]
8mod macro_dispatch_tests;
9
10/// Rullst Live Component (Server-Driven UI)
11/// Inspired by Phoenix LiveView and Laravel Livewire, allowing you to write
12/// interactive components entirely in Rust, updated in real-time via WebSockets.
13#[async_trait]
14pub trait LiveComponent: Send + Sync + Default + 'static {
15    /// Called on the first render (both on the initial HTTP load and when the WebSocket connection opens).
16    async fn mount(&mut self) {}
17
18    /// Processes JSON events originating from the frontend via WebSocket.
19    /// HTMX will by default send a JSON payload containing headers and submitted values (hx-vals, forms).
20    async fn handle_event(&mut self, _payload: Value) {}
21
22    /// Renders the current state of the component as an HTML String.
23    /// REQUIRED: The root of the rendered string MUST have a unique `id` attribute
24    /// so that HTMX knows exactly which DOM node to update.
25    fn render(&self) -> String;
26}
27
28/// Largest WebSocket message or frame accepted by [`live_ws_handler`].
29///
30/// HTMX event payloads are small JSON documents. Without an explicit bound the
31/// socket inherited tungstenite's 64 MiB message default, and every message is
32/// parsed into a `serde_json::Value`.
33const MAX_LIVE_MESSAGE_BYTES: usize = 64 * 1024;
34
35/// Generic Axum handler for the WebSocket route of a Rullst Live component.
36/// It will instantiate the component, call `mount`, and enter the event-listening loop.
37///
38/// Incoming frames and messages are limited to 64 KiB. A larger message ends
39/// the session before it is buffered in full or parsed; `serde_json` also keeps
40/// its default nesting limit of 128. Origin checks, authentication, connection
41/// caps and idle timeouts remain route-level application policy. The
42/// `live::recovery` API applies stricter bounds and admission control.
43pub async fn live_ws_handler<C: LiveComponent>(ws: WebSocketUpgrade) -> impl IntoResponse {
44    let ws = ws
45        .max_message_size(MAX_LIVE_MESSAGE_BYTES)
46        .max_frame_size(MAX_LIVE_MESSAGE_BYTES);
47    ws.on_upgrade(|socket| async move {
48        let mut rullst_ws = WebSocket::new(socket);
49        let mut component = C::default();
50
51        // Mount the initial state in the WebSocket session
52        component.mount().await;
53
54        // Continuous loop receiving events from the frontend (HTMX ws-ext)
55        while let Some(Ok(msg)) = rullst_ws.recv().await {
56            // HTMX sends messages in JSON format with headers and input values
57            if let Ok(payload) = serde_json::from_str::<Value>(&msg) {
58                // Forward the event to the component lifecycle
59                component.handle_event(payload).await;
60
61                // Re-render the HTML after the possible state mutation
62                let html = component.render();
63
64                // Push the new HTML via WebSocket. HTMX will hot-swap it automatically using the root ID.
65                if let Err(e) = rullst_ws.send_html(html).await {
66                    eprintln!("Rullst Live WS Error: {}", e);
67                    break; // Client disconnected or network failure
68                }
69            }
70        }
71    })
72}
73
74/// Utility to facilitate mounting a Live component in a normal HTTP page.
75pub struct Live;
76
77impl Live {
78    /// Generates the wrapper `<div>` tag that activates the `hx-ext="ws"` HTMX extension.
79    /// It pre-renders (`mount` + `render`) on the first load. Search-engine
80    /// behavior still depends on the rendered document and deployment.
81    pub async fn mount<C: LiveComponent>(ws_path: &str) -> String {
82        let mut comp = C::default();
83        comp.mount().await;
84        let html = comp.render();
85
86        // HTML escape ws_path to prevent path/attribute injection
87        let safe_path = ws_path
88            .replace('&', "&amp;")
89            .replace('"', "&quot;")
90            .replace('\'', "&#x27;")
91            .replace('<', "&lt;")
92            .replace('>', "&gt;");
93
94        // Wrap the component in an invisible div that instructs HTMX to open the WebSocket
95        format!(
96            "<div hx-ext=\"ws\" ws-connect=\"{}\">\n{}\n</div>",
97            safe_path, html
98        )
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105
106    #[derive(Default)]
107    struct DummyComponent;
108
109    #[async_trait]
110    impl LiveComponent for DummyComponent {
111        fn render(&self) -> String {
112            "<h1>Live Demo</h1>".to_string()
113        }
114    }
115
116    #[derive(Default)]
117    struct CountingComponent {
118        events: usize,
119    }
120
121    #[async_trait]
122    impl LiveComponent for CountingComponent {
123        async fn handle_event(&mut self, _payload: Value) {
124            self.events += 1;
125        }
126
127        fn render(&self) -> String {
128            format!("<p id=\"events\">{}</p>", self.events)
129        }
130    }
131
132    #[tokio::test]
133    #[allow(clippy::unwrap_used, clippy::expect_used)]
134    async fn legacy_handler_closes_oversized_messages_without_parsing_them() {
135        use futures_util::{SinkExt, StreamExt};
136        use tokio_tungstenite::tungstenite::Message;
137
138        let app = axum::Router::new().route(
139            "/ws",
140            axum::routing::get(live_ws_handler::<CountingComponent>),
141        );
142        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
143        let address = listener.local_addr().unwrap();
144        let server = tokio::spawn(async move { axum::serve(listener, app).await });
145        let (mut socket, _) = tokio_tungstenite::connect_async(format!("ws://{address}/ws"))
146            .await
147            .unwrap();
148        let wait = std::time::Duration::from_secs(5);
149
150        let bounded = format!("[{}0]", "0,".repeat(MAX_LIVE_MESSAGE_BYTES / 2 - 2));
151        assert!(bounded.len() <= MAX_LIVE_MESSAGE_BYTES);
152        socket.send(Message::Text(bounded.into())).await.unwrap();
153        let reply = tokio::time::timeout(wait, socket.next()).await.unwrap();
154        assert!(
155            matches!(&reply, Some(Ok(Message::Text(html))) if html.as_str() == "<p id=\"events\">1</p>")
156        );
157
158        let oversized = format!("[{}0]", "0,".repeat(MAX_LIVE_MESSAGE_BYTES / 2));
159        assert!(oversized.len() > MAX_LIVE_MESSAGE_BYTES);
160        let _ = socket.send(Message::Text(oversized.into())).await;
161        let reply = tokio::time::timeout(wait, socket.next()).await.unwrap();
162        assert!(
163            !matches!(reply, Some(Ok(Message::Text(_)))),
164            "an oversized message must close the session instead of being rendered"
165        );
166        server.abort();
167    }
168
169    #[tokio::test]
170    #[allow(clippy::unwrap_used, clippy::expect_used)]
171    async fn legacy_handler_keeps_the_session_across_control_frames() {
172        use futures_util::{SinkExt, StreamExt};
173        use tokio_tungstenite::tungstenite::Message;
174
175        let app = axum::Router::new().route(
176            "/ws",
177            axum::routing::get(live_ws_handler::<CountingComponent>),
178        );
179        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
180        let address = listener.local_addr().unwrap();
181        let server = tokio::spawn(async move { axum::serve(listener, app).await });
182        let (mut socket, _) = tokio_tungstenite::connect_async(format!("ws://{address}/ws"))
183            .await
184            .unwrap();
185
186        socket
187            .send(Message::Ping(b"keepalive".to_vec().into()))
188            .await
189            .unwrap();
190        socket
191            .send(Message::Pong(b"unsolicited".to_vec().into()))
192            .await
193            .unwrap();
194        socket.send(Message::Text("{}".into())).await.unwrap();
195
196        let rendered = tokio::time::timeout(std::time::Duration::from_secs(5), async {
197            loop {
198                match socket.next().await {
199                    Some(Ok(Message::Pong(_))) => continue,
200                    other => break other,
201                }
202            }
203        })
204        .await
205        .unwrap();
206        assert!(
207            matches!(&rendered, Some(Ok(Message::Text(html))) if html.as_str() == "<p id=\"events\">1</p>"),
208            "control frames must not end the session: {rendered:?}"
209        );
210        server.abort();
211    }
212
213    #[tokio::test]
214    async fn test_live_mount() {
215        let html = Live::mount::<DummyComponent>("/ws/demo?a=1&b=2").await;
216        assert!(html.contains("hx-ext=\"ws\""));
217        assert!(html.contains("ws-connect=\"/ws/demo?a=1&amp;b=2\""));
218        assert!(html.contains("<h1>Live Demo</h1>"));
219        assert_ne!(html, "xyzzy");
220        assert_ne!(html, "");
221    }
222}