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#[async_trait]
14pub trait LiveComponent: Send + Sync + Default + 'static {
15 async fn mount(&mut self) {}
17
18 async fn handle_event(&mut self, _payload: Value) {}
21
22 fn render(&self) -> String;
26}
27
28const MAX_LIVE_MESSAGE_BYTES: usize = 64 * 1024;
34
35pub 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 component.mount().await;
53
54 while let Some(Ok(msg)) = rullst_ws.recv().await {
56 if let Ok(payload) = serde_json::from_str::<Value>(&msg) {
58 component.handle_event(payload).await;
60
61 let html = component.render();
63
64 if let Err(e) = rullst_ws.send_html(html).await {
66 eprintln!("Rullst Live WS Error: {}", e);
67 break; }
69 }
70 }
71 })
72}
73
74pub struct Live;
76
77impl Live {
78 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 let safe_path = ws_path
88 .replace('&', "&")
89 .replace('"', """)
90 .replace('\'', "'")
91 .replace('<', "<")
92 .replace('>', ">");
93
94 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&b=2\""));
218 assert!(html.contains("<h1>Live Demo</h1>"));
219 assert_ne!(html, "xyzzy");
220 assert_ne!(html, "");
221 }
222}