1use std::sync::Arc;
2use std::time::Duration;
3
4use async_trait::async_trait;
5use axum::{
6 body::{
7 to_bytes,
8 Body,
9 Bytes,
10 },
11 extract::{
12 ws::{
13 Message,
14 WebSocket,
15 WebSocketUpgrade,
16 },
17 State,
18 },
19 http::{
20 header::{
21 AUTHORIZATION,
22 CONTENT_TYPE,
23 },
24 Request,
25 },
26 response::IntoResponse,
27 Router,
28};
29use base64::{
30 engine::general_purpose::STANDARD,
31 Engine,
32};
33use futures_util::{
34 stream::SplitSink,
35 SinkExt,
36 StreamExt,
37};
38use pipedash_core::application::ObservedNotification;
39use pipedash_core::domain::pairing::validate_device;
40use pipedash_core::domain::NotificationPreferences;
41use pipedash_core::event::{
42 CoreEvent,
43 EventBus,
44};
45use pipedash_core::CoreContext;
46use serde::{
47 Deserialize,
48 Serialize,
49};
50use serde_json::json;
51use tokio::sync::{
52 broadcast,
53 mpsc,
54 watch,
55 Semaphore,
56};
57use tokio::task::JoinSet;
58use tower::ServiceExt;
59
60use crate::state::AppState;
61
62const AUTH_TIMEOUT: Duration = Duration::from_secs(10);
63const PING_INTERVAL: Duration = Duration::from_secs(20);
64const WRITE_TIMEOUT: Duration = Duration::from_secs(30);
65const MAX_IN_FLIGHT: usize = 16;
66const MAX_ADMITTED: usize = 64;
67const OUTBOX_CAPACITY: usize = 64;
68const MAX_RESPONSE_BYTES: usize = 64 * 1024 * 1024;
69
70const AUTH_FAILED_MESSAGE: &str = "This device isn't paired with this Pipedash.";
71const RESPONSE_TOO_LARGE: &str = r#"{"code":"RESPONSE_TOO_LARGE","message":"The response is too large to send to this device."}"#;
72const INVALID_PATH: &str = r#"{"code":"BAD_REQUEST","message":"Invalid request path."}"#;
73const TOO_MANY_REQUESTS: &str = r#"{"code":"TOO_MANY_REQUESTS","message":"This device sent too many requests at once. Try again."}"#;
74
75#[derive(Deserialize)]
76#[serde(tag = "type", rename_all = "kebab-case")]
77enum ClientFrame {
78 Auth {
79 token: String,
80 },
81 Claim {
82 code: String,
83 name: String,
84 platform: String,
85 },
86 Request {
87 id: u64,
88 method: String,
89 path: String,
90 #[serde(default)]
91 body: Option<String>,
92 },
93 NotificationPreferences {
94 preferences: NotificationPreferences,
95 #[serde(default)]
96 after: Option<u64>,
97 },
98}
99
100#[derive(Serialize)]
101struct ResponseFrame<'a> {
102 #[serde(rename = "type")]
103 kind: &'static str,
104 id: u64,
105 status: u16,
106 #[serde(skip_serializing_if = "Option::is_none")]
107 body: Option<&'a str>,
108 #[serde(rename = "bodyBase64", skip_serializing_if = "Option::is_none")]
109 body_base64: Option<String>,
110}
111
112enum Identity {
113 Anonymous,
114 Operator(String),
115 Device { id: String, token: String },
116}
117
118pub struct WebSocketEventBus {
119 tx: broadcast::Sender<CoreEvent>,
120}
121
122impl WebSocketEventBus {
123 pub fn new() -> Self {
124 let (tx, _) = broadcast::channel(1024);
125 Self { tx }
126 }
127
128 pub fn from_sender(tx: broadcast::Sender<CoreEvent>) -> Self {
129 Self { tx }
130 }
131
132 pub fn subscribe(&self) -> broadcast::Receiver<CoreEvent> {
133 self.tx.subscribe()
134 }
135}
136
137impl Default for WebSocketEventBus {
138 fn default() -> Self {
139 Self::new()
140 }
141}
142
143#[async_trait]
144impl EventBus for WebSocketEventBus {
145 async fn emit(&self, event: CoreEvent) {
146 let _ = self.tx.send(event);
147 }
148
149 async fn emit_to(&self, _target: &str, event: CoreEvent) {
150 let _ = self.tx.send(event);
151 }
152}
153
154fn event_frame(event: &CoreEvent) -> Option<String> {
155 serde_json::to_string(&json!({
156 "type": event.event_name(),
157 "payload": event.to_json_payload(),
158 }))
159 .ok()
160}
161
162fn text(value: serde_json::Value) -> Message {
163 Message::Text(value.to_string().into())
164}
165
166fn response_frame(id: u64, status: u16, body: &[u8]) -> Message {
167 let (body, body_base64) = match std::str::from_utf8(body) {
168 Ok(text) => (Some(text), None),
169 Err(_) => (None, Some(STANDARD.encode(body))),
170 };
171 let frame = ResponseFrame {
172 kind: "response",
173 id,
174 status,
175 body,
176 body_base64,
177 };
178
179 Message::Text(serde_json::to_string(&frame).unwrap_or_default().into())
180}
181
182async fn dispatch(
183 router: Router, method: &str, path: &str, body: Option<String>, token: Option<&str>,
184) -> (u16, Bytes) {
185 let invalid = || (400, Bytes::from_static(INVALID_PATH.as_bytes()));
186 if !path.starts_with('/') || path.starts_with("//") {
187 return invalid();
188 }
189
190 let mut request = Request::builder()
191 .method(method)
192 .uri(format!("/api/v1{path}"));
193 if let Some(token) = token {
194 request = request.header(AUTHORIZATION, format!("Bearer {token}"));
195 }
196 if body.is_some() {
197 request = request.header(CONTENT_TYPE, "application/json");
198 }
199 let Ok(request) = request.body(body.map_or_else(Body::empty, Body::from)) else {
200 return invalid();
201 };
202
203 let Ok(response) = router.oneshot(request).await;
204 let (parts, body) = response.into_parts();
205 match to_bytes(body, MAX_RESPONSE_BYTES).await {
206 Ok(bytes) => (parts.status.as_u16(), bytes),
207 Err(_) => (500, Bytes::from_static(RESPONSE_TOO_LARGE.as_bytes())),
208 }
209}
210
211async fn write_loop(
212 mut sender: SplitSink<WebSocket, Message>, mut outbox: mpsc::Receiver<Message>,
213) {
214 let mut ping =
215 tokio::time::interval_at(tokio::time::Instant::now() + PING_INTERVAL, PING_INTERVAL);
216 loop {
217 let message = tokio::select! {
218 message = outbox.recv() => match message {
219 Some(message) => message,
220 None => break,
221 },
222 _ = ping.tick() => Message::Ping(Bytes::new()),
223 };
224 let closing = matches!(message, Message::Close(_));
225 if !matches!(
226 tokio::time::timeout(WRITE_TIMEOUT, sender.send(message)).await,
227 Ok(Ok(()))
228 ) || closing
229 {
230 break;
231 }
232 }
233}
234
235async fn next_event(
236 events: &mut Option<broadcast::Receiver<CoreEvent>>,
237) -> Result<CoreEvent, broadcast::error::RecvError> {
238 match events {
239 Some(events) => events.recv().await,
240 None => std::future::pending().await,
241 }
242}
243
244fn touch_device(state: &AppState, id: String) {
245 let state = state.clone();
246 tokio::spawn(async move {
247 let Ok(core) = state.core().await else {
248 return;
249 };
250 if let Err(error) = core.devices().touch(&id).await {
251 tracing::warn!(%error, "Failed to record device activity");
252 }
253 });
254}
255
256struct Connection {
257 state: AppState,
258 out: mpsc::Sender<Message>,
259 identity: Option<Identity>,
260 events: Option<broadcast::Receiver<CoreEvent>>,
261 router: Option<Router>,
262 permits: Arc<Semaphore>,
263 admitted: Arc<Semaphore>,
264 requests: JoinSet<()>,
265 preferences: NotificationPreferences,
266 last_notification: u64,
267}
268
269impl Connection {
270 fn new(state: AppState, out: mpsc::Sender<Message>) -> Self {
271 let identity = (!state.requires_auth()).then_some(Identity::Anonymous);
272 let events = identity.as_ref().map(|_| state.ws_event_bus.subscribe());
273
274 Self {
275 state,
276 out,
277 identity,
278 events,
279 router: None,
280 permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT)),
281 admitted: Arc::new(Semaphore::new(MAX_ADMITTED)),
282 requests: JoinSet::new(),
283 preferences: NotificationPreferences::default(),
284 last_notification: 0,
285 }
286 }
287
288 async fn reject(&self, code: &str, message: &str) {
289 let _ = self
290 .out
291 .send(text(json!({
292 "type": "auth-error",
293 "payload": { "code": code, "message": message },
294 })))
295 .await;
296 let _ = self.out.send(Message::Close(None)).await;
297 }
298
299 async fn authenticated(&mut self, identity: Identity, issued_token: Option<&str>) -> bool {
300 let device_id = match &identity {
301 Identity::Device { id, .. } => Some(id.as_str()),
302 _ => None,
303 };
304 let mut payload = json!({
305 "host": if self.state.desktop_host.is_some() { "desktop" } else { "server" },
306 "name": self.state.desktop_host.as_deref(),
307 "deviceId": device_id,
308 });
309 if let Some(token) = issued_token {
310 payload["token"] = token.into();
311 }
312
313 if self.events.is_none() {
314 self.events = Some(self.state.ws_event_bus.subscribe());
315 }
316 self.identity = Some(identity);
317 self.router = None;
318
319 self.out
320 .send(text(json!({ "type": "auth-ok", "payload": payload })))
321 .await
322 .is_ok()
323 }
324
325 async fn on_text(&mut self, raw: &str) -> bool {
326 let Ok(frame) = serde_json::from_str::<ClientFrame>(raw) else {
327 if self.identity.is_none() {
328 self.reject("AUTH_FAILED", AUTH_FAILED_MESSAGE).await;
329 return false;
330 }
331 tracing::trace!("Ignoring an unrecognized WebSocket frame");
332 return true;
333 };
334
335 match frame {
336 ClientFrame::Auth { token } => self.on_auth(&token).await,
337 ClientFrame::Claim {
338 code,
339 name,
340 platform,
341 } => self.on_claim(&code, &name, &platform).await,
342 ClientFrame::Request {
343 id,
344 method,
345 path,
346 body,
347 } => self.on_request(id, method, path, body).await,
348 ClientFrame::NotificationPreferences { preferences, after } => {
349 self.on_notification_preferences(preferences, after).await
350 }
351 }
352 }
353
354 async fn on_auth(&mut self, token: &str) -> bool {
355 if self.events.is_none() {
356 self.events = Some(self.state.ws_event_bus.subscribe());
357 }
358 let anonymous = matches!(self.identity, Some(Identity::Anonymous));
359
360 let identity = if self
361 .state
362 .api_token()
363 .is_some_and(|expected| crate::token_matches(token, &expected))
364 {
365 Identity::Operator(token.to_string())
366 } else {
367 match self.state.paired_device(token).await {
368 Ok(Some(device)) => {
369 touch_device(&self.state, device.id.clone());
370 Identity::Device {
371 id: device.id,
372 token: token.to_string(),
373 }
374 }
375 Ok(None) if anonymous => Identity::Anonymous,
376 Ok(None) => {
377 let config_error = self.state.inner.read().await.config_error.is_some();
378 let no_core = self.state.core().await.is_err();
379 if no_core && self.state.desktop_host.is_some() {
380 self.reject("HOST_LOCKED", crate::HOST_LOCKED_MESSAGE).await;
381 } else if no_core && config_error {
382 self.reject("SERVER_NOT_READY", crate::SERVER_NOT_READY_MESSAGE)
383 .await;
384 } else {
385 self.reject("AUTH_FAILED", AUTH_FAILED_MESSAGE).await;
386 }
387 return false;
388 }
389 Err(_) => {
390 self.reject("SERVER_NOT_READY", crate::SERVER_NOT_READY_MESSAGE)
391 .await;
392 return false;
393 }
394 }
395 };
396
397 self.authenticated(identity, None).await
398 }
399
400 async fn on_claim(&mut self, code: &str, name: &str, platform: &str) -> bool {
401 let Ok(core) = self.state.core().await else {
402 if self.state.desktop_host.is_some() {
403 self.reject("HOST_LOCKED", crate::HOST_LOCKED_MESSAGE).await;
404 } else {
405 self.reject("SERVER_NOT_READY", crate::SERVER_NOT_READY_MESSAGE)
406 .await;
407 }
408 return false;
409 };
410 let name = match validate_device(name, platform) {
411 Ok(_) if crate::demo::enabled() => crate::demo::device_name(platform),
412 Ok(name) => name,
413 Err((code, message)) => {
414 self.reject(code, message).await;
415 return false;
416 }
417 };
418 if !self.state.pairing_codes.take(code) {
419 self.reject(
420 "PAIRING_CODE_INVALID",
421 "This code expired or was already used. Create a new code on the server.",
422 )
423 .await;
424 return false;
425 }
426
427 match core.pair_device(&name, platform).await {
428 Ok((device, token)) => {
429 let identity = Identity::Device {
430 id: device.id,
431 token: token.clone(),
432 };
433 self.authenticated(identity, Some(&token)).await
434 }
435 Err(error) => {
436 tracing::warn!(%error, "Failed to pair a device");
437 self.reject(
438 "PAIRING_FAILED",
439 "Pipedash couldn't pair this device. Try again.",
440 )
441 .await;
442 false
443 }
444 }
445 }
446
447 async fn on_request(
448 &mut self, id: u64, method: String, path: String, body: Option<String>,
449 ) -> bool {
450 let token = match &self.identity {
451 None => {
452 self.reject("AUTH_FAILED", AUTH_FAILED_MESSAGE).await;
453 return false;
454 }
455 Some(Identity::Anonymous) => None,
456 Some(Identity::Operator(token)) | Some(Identity::Device { token, .. }) => {
457 Some(token.clone())
458 }
459 };
460
461 let state = &self.state;
462 let router = self
463 .router
464 .get_or_insert_with(|| crate::host_router(state.clone()))
465 .clone();
466 let Ok(admitted) = Arc::clone(&self.admitted).try_acquire_owned() else {
467 let _ = self
468 .out
469 .try_send(response_frame(id, 429, TOO_MANY_REQUESTS.as_bytes()));
470 return true;
471 };
472 let permits = Arc::clone(&self.permits);
473 let out = self.out.clone();
474 self.requests.spawn(async move {
475 let _admitted = admitted;
476 let Ok(_permit) = permits.acquire_owned().await else {
477 return;
478 };
479 let (status, body) = dispatch(router, &method, &path, body, token.as_deref()).await;
480 let _ = out.send(response_frame(id, status, &body)).await;
481 });
482
483 true
484 }
485
486 async fn on_notification_preferences(
487 &mut self, mut preferences: NotificationPreferences, after: Option<u64>,
488 ) -> bool {
489 match self.identity {
490 None => {
491 self.reject("AUTH_FAILED", AUTH_FAILED_MESSAGE).await;
492 return false;
493 }
494 Some(Identity::Device { .. }) => return true,
495 Some(_) => {}
496 }
497 preferences.keep_fired(&self.preferences);
498 self.preferences = preferences;
499
500 let Some(after) = after else {
501 return true;
502 };
503 let Ok(core) = self.state.core().await else {
504 return true;
505 };
506 for observed in core.notifications_after(after) {
507 if !self.show(&core, &observed).await {
508 return false;
509 }
510 }
511
512 true
513 }
514
515 async fn on_notification(&mut self, observed: &ObservedNotification) -> bool {
516 match self.state.core().await {
517 Ok(core) => self.show(&core, observed).await,
518 Err(_) => true,
519 }
520 }
521
522 async fn show(&mut self, core: &CoreContext, observed: &ObservedNotification) -> bool {
523 if observed.id <= self.last_notification {
524 return true;
525 }
526 let shown = match &self.identity {
527 None => return true,
528 Some(Identity::Device { id, .. }) => {
529 let mut preferences = match core.devices().notification_preferences().await {
530 Ok(mut stored) => stored.remove(id).unwrap_or_default(),
531 Err(error) => {
532 tracing::warn!(%error, "Couldn't load device notification preferences");
533 return true;
534 }
535 };
536 observed.admit(&mut preferences)
537 }
538 Some(_) => observed.admit(&mut self.preferences),
539 };
540 let Some(shown) = shown else {
541 return true;
542 };
543 self.last_notification = observed.id;
544
545 self.out
546 .send(text(json!({ "type": "notification", "payload": shown })))
547 .await
548 .is_ok()
549 }
550
551 async fn on_event(&mut self, event: Result<CoreEvent, broadcast::error::RecvError>) -> bool {
552 match event {
553 Ok(CoreEvent::NotificationObserved { notification }) => {
554 self.on_notification(¬ification).await
555 }
556 Ok(event) => {
557 if let (CoreEvent::DeviceRevoked { id }, Some(Identity::Device { id: own, .. })) =
558 (&event, &self.identity)
559 {
560 if id == own {
561 self.reject("DEVICE_REMOVED", "This device was removed.")
562 .await;
563 return false;
564 }
565 }
566
567 match event_frame(&event) {
568 Some(json) => self.out.send(Message::Text(json.into())).await.is_ok(),
569 None => true,
570 }
571 }
572 Err(broadcast::error::RecvError::Lagged(skipped)) => {
573 tracing::warn!(skipped, "WebSocket client fell behind, closing");
574 let _ = self.out.send(Message::Close(None)).await;
575 false
576 }
577 Err(broadcast::error::RecvError::Closed) => false,
578 }
579 }
580}
581
582pub async fn ws_handler(ws: WebSocketUpgrade, State(state): State<AppState>) -> impl IntoResponse {
583 ws.on_upgrade(move |socket| handle_socket(socket, state))
584}
585
586pub(crate) async fn closed(closing: &mut watch::Receiver<bool>) {
587 let _ = closing.wait_for(|closed| *closed).await;
588}
589
590async fn handle_socket(socket: WebSocket, state: AppState) {
591 let (sender, mut receiver) = socket.split();
592 let (out, outbox) = mpsc::channel(OUTBOX_CAPACITY);
593 let writer = tokio::spawn(write_loop(sender, outbox));
594 let mut conn = Connection::new(state, out);
595 let mut closing = conn.state.closing.subscribe();
596 let auth_deadline = tokio::time::sleep(AUTH_TIMEOUT);
597 tokio::pin!(auth_deadline);
598
599 loop {
600 tokio::select! {
601 _ = conn.out.closed() => break,
602 _ = closed(&mut closing) => {
603 let _ = conn.out.send(Message::Close(None)).await;
604 break;
605 }
606 frame = receiver.next() => match frame {
607 Some(Ok(Message::Text(raw))) => {
608 if !conn.on_text(&raw).await {
609 break;
610 }
611 }
612 Some(Ok(Message::Close(_))) | None => {
613 tracing::debug!("WebSocket client disconnected");
614 break;
615 }
616 Some(Err(error)) => {
617 tracing::warn!("WebSocket error: {}", error);
618 break;
619 }
620 Some(Ok(_)) => {}
621 },
622 event = next_event(&mut conn.events) => {
623 if !conn.on_event(event).await {
624 break;
625 }
626 }
627 Some(_) = conn.requests.join_next(), if !conn.requests.is_empty() => {}
628 _ = &mut auth_deadline, if conn.identity.is_none() => {
629 conn.reject("AUTH_TIMEOUT", "Authentication timed out.").await;
630 break;
631 }
632 }
633 }
634
635 drop(conn);
636 let _ = writer.await;
637}
638
639#[cfg(test)]
640mod tests {
641 use std::sync::Arc;
642 use std::time::Duration;
643
644 use axum::routing::get;
645 use axum::Router;
646 use serde_json::{
647 json,
648 Value,
649 };
650 use tokio::io::{
651 AsyncReadExt,
652 AsyncWriteExt,
653 };
654 use tokio::net::{
655 TcpListener,
656 TcpStream,
657 };
658
659 use super::*;
660 use crate::test_support::hosted;
661
662 async fn read_frame(stream: &mut TcpStream) -> Option<(u8, Vec<u8>)> {
663 let mut head = [0u8; 2];
664 stream.read_exact(&mut head).await.ok()?;
665 let len = match head[1] & 0x7f {
666 126 => stream.read_u16().await.ok()? as usize,
667 127 => stream.read_u64().await.ok()? as usize,
668 len => len as usize,
669 };
670 let mut payload = vec![0u8; len];
671 stream.read_exact(&mut payload).await.ok()?;
672 Some((head[0] & 0x0f, payload))
673 }
674
675 async fn upgrade(addr: std::net::SocketAddr, path: &str) -> TcpStream {
676 let mut stream = TcpStream::connect(addr).await.unwrap();
677 let handshake = format!(
678 "GET {path} HTTP/1.1\r\nHost: {addr}\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\
679 Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n\r\n"
680 );
681 stream.write_all(handshake.as_bytes()).await.unwrap();
682 let mut response = Vec::new();
683 while !response.ends_with(b"\r\n\r\n") {
684 response.push(stream.read_u8().await.unwrap());
685 }
686 assert!(response.starts_with(b"HTTP/1.1 101"));
687 tokio::time::sleep(Duration::from_millis(100)).await;
688
689 stream
690 }
691
692 async fn connect_state(state: AppState) -> TcpStream {
693 let app = Router::new()
694 .route("/ws", get(ws_handler))
695 .with_state(state);
696 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
697 let addr = listener.local_addr().unwrap();
698 tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
699
700 upgrade(addr, "/ws").await
701 }
702
703 async fn connect(bus: &Arc<WebSocketEventBus>) -> TcpStream {
704 connect_state(AppState::setup_mode(bus.clone())).await
705 }
706
707 async fn send_json(stream: &mut TcpStream, value: Value) {
708 let payload = value.to_string().into_bytes();
709 let mask = [7u8, 11, 13, 17];
710 let mut frame = vec![0x81];
711 match payload.len() {
712 len if len < 126 => frame.push(0x80 | len as u8),
713 len => {
714 frame.push(0x80 | 126);
715 frame.extend((len as u16).to_be_bytes());
716 }
717 }
718 frame.extend(mask);
719 frame.extend(payload.iter().enumerate().map(|(i, b)| b ^ mask[i % 4]));
720 stream.write_all(&frame).await.unwrap();
721 }
722
723 async fn next_message(stream: &mut TcpStream) -> Value {
724 let (opcode, payload) = tokio::time::timeout(Duration::from_secs(5), read_frame(stream))
725 .await
726 .expect("no message arrived")
727 .expect("the socket closed");
728 assert_eq!(opcode, 0x1);
729 serde_json::from_slice(&payload).unwrap()
730 }
731
732 async fn assert_closed(stream: &mut TcpStream) {
733 let closed = tokio::time::timeout(Duration::from_secs(5), async {
734 loop {
735 match read_frame(stream).await {
736 None | Some((0x8, _)) => break,
737 Some(_) => {}
738 }
739 }
740 })
741 .await;
742 assert!(closed.is_ok(), "the socket stayed open");
743 }
744
745 async fn claim(stream: &mut TcpStream, code: &str) -> Value {
746 send_json(
747 stream,
748 json!({ "type": "claim", "code": code, "name": "Pixel", "platform": "android" }),
749 )
750 .await;
751 next_message(stream).await
752 }
753
754 #[tokio::test]
755 async fn lagged_client_gets_closed_so_it_reconnects() {
756 let bus = Arc::new(WebSocketEventBus::new());
757 let mut stream = connect(&bus).await;
758
759 for _ in 0..1100 {
760 bus.emit(CoreEvent::ProvidersChanged).await;
761 }
762
763 let closed = tokio::time::timeout(Duration::from_secs(5), async {
764 loop {
765 match read_frame(&mut stream).await {
766 None => break,
767 Some((0x8, _)) => break,
768 Some(_) => {}
769 }
770 }
771 })
772 .await;
773 assert!(closed.is_ok(), "socket stayed open after the client lagged");
774 }
775
776 #[tokio::test]
777 async fn live_run_events_reach_clients_with_their_payloads() {
778 let bus = Arc::new(WebSocketEventBus::new());
779 let mut stream = connect(&bus).await;
780 let run: pipedash_core::domain::PipelineRun = serde_json::from_value(serde_json::json!({
781 "id": "p__3", "pipeline_id": "p", "run_number": 3, "status": "running",
782 "started_at": "2026-01-01T00:00:00Z", "duration_seconds": 12, "logs_url": ""
783 }))
784 .unwrap();
785
786 bus.emit(CoreEvent::RunLogTreeChanged {
787 pipeline_id: "p".into(),
788 run_number: 3,
789 })
790 .await;
791 bus.emit(CoreEvent::RunUpdated {
792 pipeline_id: "p".into(),
793 run_number: 3,
794 run: run.clone(),
795 })
796 .await;
797
798 assert_eq!(
799 next_message(&mut stream).await,
800 serde_json::json!({
801 "type": "run-log-tree-changed",
802 "payload": { "pipelineId": "p", "runNumber": 3 },
803 })
804 );
805 assert_eq!(
806 next_message(&mut stream).await,
807 serde_json::json!({
808 "type": "run-updated",
809 "payload": {
810 "pipelineId": "p",
811 "runNumber": 3,
812 "run": serde_json::to_value(&run).unwrap(),
813 },
814 })
815 );
816 }
817
818 #[tokio::test]
819 async fn a_browser_auth_frame_gets_auth_ok_when_no_token_is_required() {
820 let bus = Arc::new(WebSocketEventBus::new());
821 let mut stream = connect(&bus).await;
822
823 send_json(&mut stream, json!({ "type": "auth", "token": "anything" })).await;
824
825 assert_eq!(
826 next_message(&mut stream).await,
827 json!({
828 "type": "auth-ok",
829 "payload": { "host": "server", "name": null, "deviceId": null },
830 })
831 );
832 }
833
834 #[tokio::test]
835 async fn the_operator_token_authenticates_and_requests_carry_it() {
836 let state = AppState::setup_mode(Arc::new(WebSocketEventBus::new()));
837 state.set_api_token("vault-password");
838
839 let mut stream = connect_state(state.clone()).await;
840 send_json(&mut stream, json!({ "type": "auth", "token": "wrong" })).await;
841 let rejected = next_message(&mut stream).await;
842 assert_eq!(rejected["type"], "auth-error");
843 assert_eq!(rejected["payload"]["code"], "AUTH_FAILED");
844 assert_closed(&mut stream).await;
845
846 let mut stream = connect_state(state).await;
847 send_json(
848 &mut stream,
849 json!({ "type": "auth", "token": "vault-password" }),
850 )
851 .await;
852 assert_eq!(next_message(&mut stream).await["type"], "auth-ok");
853 send_json(
854 &mut stream,
855 json!({ "type": "request", "id": 5, "method": "GET", "path": "/providers" }),
856 )
857 .await;
858 let response = next_message(&mut stream).await;
859
860 assert_eq!(response["type"], "response");
861 assert_eq!(response["id"], 5);
862 assert_eq!(response["status"], 503);
863 }
864
865 #[tokio::test]
866 async fn a_request_before_authenticating_is_refused() {
867 let state = AppState::setup_mode(Arc::new(WebSocketEventBus::new()));
868 state.set_api_token("vault-password");
869 let mut stream = connect_state(state).await;
870
871 send_json(
872 &mut stream,
873 json!({ "type": "request", "id": 1, "method": "GET", "path": "/providers" }),
874 )
875 .await;
876
877 assert_eq!(
878 next_message(&mut stream).await["payload"]["code"],
879 "AUTH_FAILED"
880 );
881 assert_closed(&mut stream).await;
882 }
883
884 #[tokio::test]
885 async fn a_claim_pairs_the_device_and_serves_requests_then_the_token_reconnects() {
886 let host = hosted("ws-claim").await;
887 let (code, _) = host.state.pairing_codes.issue().unwrap();
888 let mut stream = connect_state(host.state.clone()).await;
889
890 let ok = claim(&mut stream, &code).await;
891
892 assert_eq!(ok["type"], "auth-ok");
893 assert_eq!(ok["payload"]["host"], "desktop");
894 assert_eq!(ok["payload"]["name"], "test-mac");
895 let device_id = ok["payload"]["deviceId"].as_str().unwrap().to_string();
896 let token = ok["payload"]["token"].as_str().unwrap().to_string();
897 assert_eq!(token.len(), 64);
898
899 send_json(
900 &mut stream,
901 json!({ "type": "request", "id": 1, "method": "GET", "path": "/devices" }),
902 )
903 .await;
904 let response = next_message(&mut stream).await;
905 assert_eq!(response["id"], 1);
906 assert_eq!(response["status"], 200);
907 let body: Value = serde_json::from_str(response["body"].as_str().unwrap()).unwrap();
908 assert_eq!(body["devices"][0]["id"], device_id);
909 assert_eq!(body["devices"][0]["current"], true);
910
911 let mut again = connect_state(host.state.clone()).await;
912 send_json(&mut again, json!({ "type": "auth", "token": token })).await;
913 let ok = next_message(&mut again).await;
914 assert_eq!(ok["payload"]["deviceId"], device_id);
915 assert!(ok["payload"].get("token").is_none());
916
917 let mut stranger = connect_state(host.state.clone()).await;
918 send_json(&mut stranger, json!({ "type": "auth", "token": "unknown" })).await;
919 assert_eq!(
920 next_message(&mut stranger).await["payload"]["code"],
921 "AUTH_FAILED"
922 );
923 }
924
925 #[tokio::test]
926 async fn stopping_the_host_closes_connected_sockets() {
927 let state = AppState::hosted(Arc::new(WebSocketEventBus::new()), None, "test-mac");
928 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
929 let addr = listener.local_addr().unwrap();
930 let (stop, stopped) = tokio::sync::oneshot::channel::<()>();
931 let server = tokio::spawn(crate::serve_host(listener, state, async move {
932 let _ = stopped.await;
933 }));
934 let mut stream = upgrade(addr, "/api/v1/ws").await;
935
936 stop.send(()).unwrap();
937
938 assert_closed(&mut stream).await;
939 tokio::time::timeout(Duration::from_secs(5), server)
940 .await
941 .expect("serve_host kept running")
942 .unwrap()
943 .unwrap();
944 }
945
946 #[tokio::test]
947 async fn a_busy_socket_rejects_excess_requests_and_still_closes() {
948 let host = crate::test_support::server("ws-busy").await;
949 let mut stream = connect_state(host.state.clone()).await;
950 let blocked = host.state.inner.write().await;
951
952 for id in 1..=MAX_ADMITTED as u64 + 1 {
953 send_json(
954 &mut stream,
955 json!({ "type": "request", "id": id, "method": "GET", "path": "/providers" }),
956 )
957 .await;
958 }
959
960 let rejected = next_message(&mut stream).await;
961 assert_eq!(rejected["id"], MAX_ADMITTED as u64 + 1);
962 assert_eq!(rejected["status"], 429);
963
964 host.state.closing.send_replace(true);
965 assert_closed(&mut stream).await;
966 drop(blocked);
967 }
968
969 #[tokio::test]
970 async fn a_shared_host_refuses_socket_upgrades_from_other_origins() {
971 let state = AppState::hosted(Arc::new(WebSocketEventBus::new()), None, "test-mac");
972 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
973 let addr = listener.local_addr().unwrap();
974 tokio::spawn(crate::serve_host(listener, state, std::future::pending()));
975
976 let mut stream = TcpStream::connect(addr).await.unwrap();
977 let handshake = format!(
978 "GET /api/v1/ws HTTP/1.1\r\nHost: {addr}\r\nOrigin: https://attacker.example\r\n\
979 Upgrade: websocket\r\nConnection: Upgrade\r\n\
980 Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n\r\n"
981 );
982 stream.write_all(handshake.as_bytes()).await.unwrap();
983 let mut status = [0u8; 12];
984 stream.read_exact(&mut status).await.unwrap();
985 assert_eq!(&status, b"HTTP/1.1 403");
986
987 upgrade(addr, "/api/v1/ws").await;
988 }
989
990 #[tokio::test]
991 async fn a_device_token_on_a_server_without_an_api_token_becomes_a_device() {
992 let host = crate::test_support::server("ws-open-device").await;
993 let core = host.state.core().await.unwrap();
994 let (device, token) = core.pair_device("Pixel", "android").await.unwrap();
995 let mut stream = connect_state(host.state.clone()).await;
996
997 send_json(&mut stream, json!({ "type": "auth", "token": token })).await;
998 let ok = next_message(&mut stream).await;
999 assert_eq!(ok["payload"]["deviceId"], device.id);
1000
1001 send_json(
1002 &mut stream,
1003 json!({ "type": "request", "id": 1, "method": "GET", "path": "/devices" }),
1004 )
1005 .await;
1006 let response = next_message(&mut stream).await;
1007 let body: Value = serde_json::from_str(response["body"].as_str().unwrap()).unwrap();
1008 assert_eq!(body["devices"][0]["current"], true);
1009
1010 core.revoke_device(&device.id).await.unwrap();
1011 let removed = next_message(&mut stream).await;
1012 assert_eq!(removed["payload"]["code"], "DEVICE_REMOVED");
1013 assert_closed(&mut stream).await;
1014 }
1015
1016 #[tokio::test]
1017 async fn a_failed_device_lookup_is_not_reported_as_a_revoked_device() {
1018 let host = hosted("ws-lookup-failure").await;
1019 let core = host.state.core().await.unwrap();
1020 let (_, token) = core.pair_device("Pixel", "android").await.unwrap();
1021 let pipedash_core::infrastructure::database::DatabasePool::Sqlite(pool) =
1022 core.provider_service.repository().cache_pool().clone()
1023 else {
1024 panic!("expected a SQLite pool");
1025 };
1026 pool.close().await;
1027 let mut stream = connect_state(host.state.clone()).await;
1028
1029 send_json(&mut stream, json!({ "type": "auth", "token": token })).await;
1030
1031 let rejected = next_message(&mut stream).await;
1032 assert_eq!(rejected["payload"]["code"], "SERVER_NOT_READY");
1033 assert_closed(&mut stream).await;
1034 }
1035
1036 #[tokio::test]
1037 async fn a_locked_host_answers_host_locked_instead_of_forgetting_the_device() {
1038 let state = AppState::hosted(Arc::new(WebSocketEventBus::new()), None, "test-mac");
1039 let mut stream = connect_state(state).await;
1040
1041 send_json(&mut stream, json!({ "type": "auth", "token": "any" })).await;
1042
1043 let rejected = next_message(&mut stream).await;
1044 assert_eq!(rejected["payload"]["code"], "HOST_LOCKED");
1045 assert_eq!(rejected["payload"]["message"], crate::HOST_LOCKED_MESSAGE);
1046 assert_closed(&mut stream).await;
1047 }
1048
1049 #[tokio::test]
1050 async fn a_server_with_a_broken_config_does_not_forget_its_devices() {
1051 let state = AppState::config_error(Arc::new(WebSocketEventBus::new()), "broken".into());
1052 state.set_api_token("vault-password");
1053 let mut stream = connect_state(state).await;
1054
1055 send_json(
1056 &mut stream,
1057 json!({ "type": "auth", "token": "device-token" }),
1058 )
1059 .await;
1060
1061 let rejected = next_message(&mut stream).await;
1062 assert_eq!(rejected["payload"]["code"], "SERVER_NOT_READY");
1063 assert_closed(&mut stream).await;
1064 }
1065
1066 #[tokio::test]
1067 async fn a_pairing_code_can_be_claimed_only_once() {
1068 let host = hosted("ws-claim-once").await;
1069 let (code, _) = host.state.pairing_codes.issue().unwrap();
1070
1071 let mut first = connect_state(host.state.clone()).await;
1072 assert_eq!(claim(&mut first, &code).await["type"], "auth-ok");
1073
1074 let mut second = connect_state(host.state.clone()).await;
1075 let rejected = claim(&mut second, &code).await;
1076 assert_eq!(rejected["type"], "auth-error");
1077 assert_eq!(rejected["payload"]["code"], "PAIRING_CODE_INVALID");
1078 assert_closed(&mut second).await;
1079 }
1080
1081 #[tokio::test]
1082 async fn a_claim_with_invalid_device_details_is_rejected_without_burning_the_code() {
1083 let host = hosted("ws-claim-invalid").await;
1084 let (code, _) = host.state.pairing_codes.issue().unwrap();
1085 let mut stream = connect_state(host.state.clone()).await;
1086
1087 send_json(
1088 &mut stream,
1089 json!({ "type": "claim", "code": code, "name": "Pixel", "platform": "palm" }),
1090 )
1091 .await;
1092
1093 assert_eq!(
1094 next_message(&mut stream).await["payload"]["code"],
1095 "INVALID_PLATFORM"
1096 );
1097 assert!(host.state.pairing_codes.take(&code));
1098 }
1099
1100 #[tokio::test]
1101 async fn revoking_a_device_closes_its_socket_with_device_removed() {
1102 let host = hosted("ws-revoke").await;
1103 let (code, _) = host.state.pairing_codes.issue().unwrap();
1104 let mut stream = connect_state(host.state.clone()).await;
1105 let ok = claim(&mut stream, &code).await;
1106 let device_id = ok["payload"]["deviceId"].as_str().unwrap().to_string();
1107
1108 host.state
1109 .core()
1110 .await
1111 .unwrap()
1112 .revoke_device(&device_id)
1113 .await
1114 .unwrap();
1115
1116 let removed = next_message(&mut stream).await;
1117 assert_eq!(removed["type"], "auth-error");
1118 assert_eq!(removed["payload"]["code"], "DEVICE_REMOVED");
1119 assert_closed(&mut stream).await;
1120 }
1121
1122 #[tokio::test]
1123 async fn a_connection_that_never_authenticates_is_closed_after_ten_seconds() {
1124 let host = hosted("ws-timeout").await;
1125 let mut stream = connect_state(host.state.clone()).await;
1126
1127 tokio::time::pause();
1128 tokio::time::advance(AUTH_TIMEOUT + Duration::from_secs(1)).await;
1129
1130 let timed_out = next_message(&mut stream).await;
1131 assert_eq!(timed_out["type"], "auth-error");
1132 assert_eq!(timed_out["payload"]["code"], "AUTH_TIMEOUT");
1133 assert_closed(&mut stream).await;
1134 }
1135
1136 fn failed_pipeline(name: &str) -> pipedash_core::Pipeline {
1137 serde_json::from_value(json!({
1138 "id": format!("jenkins__1__{name}"),
1139 "provider_id": 1,
1140 "provider_type": "jenkins",
1141 "name": name,
1142 "status": "failed",
1143 "last_run": null,
1144 "last_updated": chrono::Utc::now(),
1145 "repository": "",
1146 "branch": null,
1147 "workflow_file": null,
1148 }))
1149 .unwrap()
1150 }
1151
1152 async fn started(name: &str) -> (crate::test_support::Hosted, CoreContext) {
1153 let host = crate::test_support::server(name).await;
1154 let core = host.state.core().await.unwrap();
1155 core.start_background_tasks().await;
1156
1157 (host, core)
1158 }
1159
1160 async fn fail(core: &CoreContext, names: &[&str]) {
1161 core.event_bus
1162 .emit(CoreEvent::PipelinesFailed {
1163 pipelines: names.iter().copied().map(failed_pipeline).collect(),
1164 })
1165 .await;
1166 }
1167
1168 async fn recorded(core: &CoreContext, count: usize) -> Vec<ObservedNotification> {
1169 tokio::time::timeout(Duration::from_secs(5), async {
1170 loop {
1171 let seen = core.notifications_after(0);
1172 if seen.len() >= count {
1173 return seen;
1174 }
1175 tokio::time::sleep(Duration::from_millis(10)).await;
1176 }
1177 })
1178 .await
1179 .expect("the recorder never observed the failures")
1180 }
1181
1182 async fn next_notification(stream: &mut TcpStream) -> Value {
1183 loop {
1184 let message = next_message(stream).await;
1185 assert_ne!(message["type"], "notification-observed");
1186 if message["type"] == "notification" {
1187 return message["payload"].clone();
1188 }
1189 }
1190 }
1191
1192 async fn synced(stream: &mut TcpStream) {
1193 send_json(
1194 stream,
1195 json!({ "type": "request", "id": 99, "method": "GET", "path": "/providers" }),
1196 )
1197 .await;
1198 loop {
1199 let message = next_message(stream).await;
1200 if message["type"] == "response" && message["id"] == 99 {
1201 return;
1202 }
1203 }
1204 }
1205
1206 #[tokio::test]
1207 async fn a_browser_that_muted_a_pipeline_gets_no_notification_for_it() {
1208 let (host, core) = started("ws-notify-muted").await;
1209 host.state.set_api_token("vault-password");
1210 let mut stream = connect_state(host.state.clone()).await;
1211 send_json(
1212 &mut stream,
1213 json!({ "type": "auth", "token": "vault-password" }),
1214 )
1215 .await;
1216 assert_eq!(next_message(&mut stream).await["type"], "auth-ok");
1217 send_json(
1218 &mut stream,
1219 json!({
1220 "type": "notification-preferences",
1221 "preferences": { "muted_pipelines": ["jenkins__1__muted"] },
1222 }),
1223 )
1224 .await;
1225 synced(&mut stream).await;
1226
1227 fail(&core, &["muted", "other"]).await;
1228
1229 let shown = next_notification(&mut stream).await;
1230 assert_eq!(shown["pipeline_id"], "jenkins__1__other");
1231 assert_eq!(shown["kind"], "failed");
1232 assert_eq!(shown["tone"], "danger");
1233 assert_eq!(shown["fired"], Value::Null);
1234 }
1235
1236 #[tokio::test]
1237 async fn a_reconnecting_browser_gets_the_missed_notifications_once() {
1238 let (host, core) = started("ws-notify-replay").await;
1239 fail(&core, &["a", "b"]).await;
1240 let seen = recorded(&core, 2).await;
1241 let mut stream = connect_state(host.state.clone()).await;
1242
1243 send_json(
1244 &mut stream,
1245 json!({ "type": "notification-preferences", "preferences": {}, "after": seen[0].id }),
1246 )
1247 .await;
1248 let replayed = next_notification(&mut stream).await;
1249 assert_eq!(replayed["id"], seen[1].id);
1250 assert_eq!(replayed["pipeline_id"], "jenkins__1__b");
1251
1252 fail(&core, &["c"]).await;
1253 assert_eq!(
1254 next_notification(&mut stream).await["pipeline_id"],
1255 "jenkins__1__c"
1256 );
1257
1258 send_json(
1259 &mut stream,
1260 json!({ "type": "notification-preferences", "preferences": {}, "after": 0 }),
1261 )
1262 .await;
1263 synced(&mut stream).await;
1264 fail(&core, &["d"]).await;
1265 assert_eq!(
1266 next_notification(&mut stream).await["pipeline_id"],
1267 "jenkins__1__d"
1268 );
1269 }
1270
1271 #[tokio::test]
1272 async fn a_replay_goes_through_the_preferences_the_browser_sent() {
1273 let (host, core) = started("ws-notify-replay-muted").await;
1274 fail(&core, &["muted", "other"]).await;
1275 recorded(&core, 2).await;
1276 let mut stream = connect_state(host.state.clone()).await;
1277
1278 send_json(
1279 &mut stream,
1280 json!({
1281 "type": "notification-preferences",
1282 "preferences": { "muted_pipelines": ["jenkins__1__muted"] },
1283 "after": 0,
1284 }),
1285 )
1286 .await;
1287
1288 assert_eq!(
1289 next_notification(&mut stream).await["pipeline_id"],
1290 "jenkins__1__other"
1291 );
1292 }
1293
1294 #[tokio::test]
1295 async fn a_device_connection_uses_the_preferences_stored_for_it() {
1296 let (host, core) = started("ws-notify-device").await;
1297 let (device, token) = core.pair_device("Pixel", "android").await.unwrap();
1298 core.devices()
1299 .set_notification_preferences(
1300 &device.id,
1301 &NotificationPreferences {
1302 muted_pipelines: vec!["jenkins__1__muted".into()],
1303 ..NotificationPreferences::default()
1304 },
1305 )
1306 .await
1307 .unwrap();
1308 let mut stream = connect_state(host.state.clone()).await;
1309 send_json(&mut stream, json!({ "type": "auth", "token": token })).await;
1310 assert_eq!(next_message(&mut stream).await["type"], "auth-ok");
1311 send_json(
1312 &mut stream,
1313 json!({
1314 "type": "notification-preferences",
1315 "preferences": { "muted_pipelines": ["jenkins__1__other"] },
1316 }),
1317 )
1318 .await;
1319 synced(&mut stream).await;
1320
1321 fail(&core, &["muted", "other"]).await;
1322
1323 assert_eq!(
1324 next_notification(&mut stream).await["pipeline_id"],
1325 "jenkins__1__other"
1326 );
1327 let stored = core.devices().notification_preferences().await.unwrap();
1328 assert_eq!(stored[&device.id].muted_pipelines, ["jenkins__1__muted"]);
1329 }
1330
1331 #[tokio::test]
1332 async fn the_raw_notification_event_is_never_sent_and_needs_a_core() {
1333 let bus = Arc::new(WebSocketEventBus::new());
1334 let mut stream = connect(&bus).await;
1335 let notice = pipedash_core::application::LocalNotification {
1336 key: "pipeline:x".into(),
1337 title: "x failed".into(),
1338 body: String::new(),
1339 pipeline_id: Some("x".into()),
1340 provider_id: Some(1),
1341 provider_ids: vec![1],
1342 kind: pipedash_core::domain::AlertKind::Failed,
1343 finish: None,
1344 level: None,
1345 link: None,
1346 };
1347
1348 bus.emit(CoreEvent::NotificationObserved {
1349 notification: ObservedNotification {
1350 id: 1,
1351 at: chrono::Utc::now(),
1352 notice,
1353 },
1354 })
1355 .await;
1356 bus.emit(CoreEvent::ProvidersChanged).await;
1357
1358 assert_eq!(next_message(&mut stream).await["type"], "providers-changed");
1359 }
1360}