pipedash_web/

ws.rs

use std::sync::Arc;
use std::time::Duration;

use async_trait::async_trait;
use axum::{
    body::{
        to_bytes,
        Body,
        Bytes,
    },
    extract::{
        ws::{
            Message,
            WebSocket,
            WebSocketUpgrade,
        },
        State,
    },
    http::{
        header::{
            AUTHORIZATION,
            CONTENT_TYPE,
        },
        Request,
    },
    response::IntoResponse,
    Router,
};
use base64::{
    engine::general_purpose::STANDARD,
    Engine,
};
use futures_util::{
    stream::SplitSink,
    SinkExt,
    StreamExt,
};
use pipedash_core::domain::pairing::validate_device;
use pipedash_core::event::{
    CoreEvent,
    EventBus,
};
use serde::{
    Deserialize,
    Serialize,
};
use serde_json::json;
use tokio::sync::{
    broadcast,
    mpsc,
    watch,
    Semaphore,
};
use tokio::task::JoinSet;
use tower::ServiceExt;

use crate::state::AppState;

const AUTH_TIMEOUT: Duration = Duration::from_secs(10);
const PING_INTERVAL: Duration = Duration::from_secs(20);
const WRITE_TIMEOUT: Duration = Duration::from_secs(30);
const MAX_IN_FLIGHT: usize = 16;
const MAX_ADMITTED: usize = 64;
const OUTBOX_CAPACITY: usize = 64;
const MAX_RESPONSE_BYTES: usize = 64 * 1024 * 1024;

const AUTH_FAILED_MESSAGE: &str = "This device isn't paired with this Pipedash.";
const RESPONSE_TOO_LARGE: &str = r#"{"code":"RESPONSE_TOO_LARGE","message":"The response is too large to send to this device."}"#;
const INVALID_PATH: &str = r#"{"code":"BAD_REQUEST","message":"Invalid request path."}"#;
const TOO_MANY_REQUESTS: &str = r#"{"code":"TOO_MANY_REQUESTS","message":"This device sent too many requests at once. Try again."}"#;

#[derive(Deserialize)]
#[serde(tag = "type", rename_all = "kebab-case")]
enum ClientFrame {
    Auth {
        token: String,
    },
    Claim {
        code: String,
        name: String,
        platform: String,
    },
    Request {
        id: u64,
        method: String,
        path: String,
        #[serde(default)]
        body: Option<String>,
    },
}

#[derive(Serialize)]
struct ResponseFrame<'a> {
    #[serde(rename = "type")]
    kind: &'static str,
    id: u64,
    status: u16,
    #[serde(skip_serializing_if = "Option::is_none")]
    body: Option<&'a str>,
    #[serde(rename = "bodyBase64", skip_serializing_if = "Option::is_none")]
    body_base64: Option<String>,
}

enum Identity {
    Anonymous,
    Operator(String),
    Device { id: String, token: String },
}

pub struct WebSocketEventBus {
    tx: broadcast::Sender<CoreEvent>,
}

impl WebSocketEventBus {
    pub fn new() -> Self {
        let (tx, _) = broadcast::channel(1024);
        Self { tx }
    }

    pub fn from_sender(tx: broadcast::Sender<CoreEvent>) -> Self {
        Self { tx }
    }

    pub fn subscribe(&self) -> broadcast::Receiver<CoreEvent> {
        self.tx.subscribe()
    }
}

impl Default for WebSocketEventBus {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl EventBus for WebSocketEventBus {
    async fn emit(&self, event: CoreEvent) {
        let _ = self.tx.send(event);
    }

    async fn emit_to(&self, _target: &str, event: CoreEvent) {
        let _ = self.tx.send(event);
    }
}

fn event_frame(event: &CoreEvent) -> Option<String> {
    serde_json::to_string(&json!({
        "type": event.event_name(),
        "payload": event.to_json_payload(),
    }))
    .ok()
}

fn text(value: serde_json::Value) -> Message {
    Message::Text(value.to_string().into())
}

fn response_frame(id: u64, status: u16, body: &[u8]) -> Message {
    let (body, body_base64) = match std::str::from_utf8(body) {
        Ok(text) => (Some(text), None),
        Err(_) => (None, Some(STANDARD.encode(body))),
    };
    let frame = ResponseFrame {
        kind: "response",
        id,
        status,
        body,
        body_base64,
    };

    Message::Text(serde_json::to_string(&frame).unwrap_or_default().into())
}

async fn dispatch(
    router: Router, method: &str, path: &str, body: Option<String>, token: Option<&str>,
) -> (u16, Bytes) {
    let invalid = || (400, Bytes::from_static(INVALID_PATH.as_bytes()));
    if !path.starts_with('/') || path.starts_with("//") {
        return invalid();
    }

    let mut request = Request::builder()
        .method(method)
        .uri(format!("/api/v1{path}"));
    if let Some(token) = token {
        request = request.header(AUTHORIZATION, format!("Bearer {token}"));
    }
    if body.is_some() {
        request = request.header(CONTENT_TYPE, "application/json");
    }
    let Ok(request) = request.body(body.map_or_else(Body::empty, Body::from)) else {
        return invalid();
    };

    let Ok(response) = router.oneshot(request).await;
    let (parts, body) = response.into_parts();
    match to_bytes(body, MAX_RESPONSE_BYTES).await {
        Ok(bytes) => (parts.status.as_u16(), bytes),
        Err(_) => (500, Bytes::from_static(RESPONSE_TOO_LARGE.as_bytes())),
    }
}

async fn write_loop(
    mut sender: SplitSink<WebSocket, Message>, mut outbox: mpsc::Receiver<Message>,
) {
    let mut ping =
        tokio::time::interval_at(tokio::time::Instant::now() + PING_INTERVAL, PING_INTERVAL);
    loop {
        let message = tokio::select! {
            message = outbox.recv() => match message {
                Some(message) => message,
                None => break,
            },
            _ = ping.tick() => Message::Ping(Bytes::new()),
        };
        let closing = matches!(message, Message::Close(_));
        if !matches!(
            tokio::time::timeout(WRITE_TIMEOUT, sender.send(message)).await,
            Ok(Ok(()))
        ) || closing
        {
            break;
        }
    }
}

async fn next_event(
    events: &mut Option<broadcast::Receiver<CoreEvent>>,
) -> Result<CoreEvent, broadcast::error::RecvError> {
    match events {
        Some(events) => events.recv().await,
        None => std::future::pending().await,
    }
}

fn touch_device(state: &AppState, id: String) {
    let state = state.clone();
    tokio::spawn(async move {
        let Ok(core) = state.core().await else {
            return;
        };
        if let Err(error) = core.devices().touch(&id).await {
            tracing::warn!(%error, "Failed to record device activity");
        }
    });
}

struct Connection {
    state: AppState,
    out: mpsc::Sender<Message>,
    identity: Option<Identity>,
    events: Option<broadcast::Receiver<CoreEvent>>,
    router: Option<Router>,
    permits: Arc<Semaphore>,
    admitted: Arc<Semaphore>,
    requests: JoinSet<()>,
}

impl Connection {
    fn new(state: AppState, out: mpsc::Sender<Message>) -> Self {
        let identity = (!state.requires_auth()).then_some(Identity::Anonymous);
        let events = identity.as_ref().map(|_| state.ws_event_bus.subscribe());

        Self {
            state,
            out,
            identity,
            events,
            router: None,
            permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT)),
            admitted: Arc::new(Semaphore::new(MAX_ADMITTED)),
            requests: JoinSet::new(),
        }
    }

    async fn reject(&self, code: &str, message: &str) {
        let _ = self
            .out
            .send(text(json!({
                "type": "auth-error",
                "payload": { "code": code, "message": message },
            })))
            .await;
        let _ = self.out.send(Message::Close(None)).await;
    }

    async fn authenticated(&mut self, identity: Identity, issued_token: Option<&str>) -> bool {
        let device_id = match &identity {
            Identity::Device { id, .. } => Some(id.as_str()),
            _ => None,
        };
        let mut payload = json!({
            "host": if self.state.desktop_host.is_some() { "desktop" } else { "server" },
            "name": self.state.desktop_host.as_deref(),
            "deviceId": device_id,
        });
        if let Some(token) = issued_token {
            payload["token"] = token.into();
        }

        if self.events.is_none() {
            self.events = Some(self.state.ws_event_bus.subscribe());
        }
        self.identity = Some(identity);
        self.router = None;

        self.out
            .send(text(json!({ "type": "auth-ok", "payload": payload })))
            .await
            .is_ok()
    }

    async fn on_text(&mut self, raw: &str) -> bool {
        let Ok(frame) = serde_json::from_str::<ClientFrame>(raw) else {
            if self.identity.is_none() {
                self.reject("AUTH_FAILED", AUTH_FAILED_MESSAGE).await;
                return false;
            }
            tracing::trace!("Ignoring an unrecognized WebSocket frame");
            return true;
        };

        match frame {
            ClientFrame::Auth { token } => self.on_auth(&token).await,
            ClientFrame::Claim {
                code,
                name,
                platform,
            } => self.on_claim(&code, &name, &platform).await,
            ClientFrame::Request {
                id,
                method,
                path,
                body,
            } => self.on_request(id, method, path, body).await,
        }
    }

    async fn on_auth(&mut self, token: &str) -> bool {
        if self.events.is_none() {
            self.events = Some(self.state.ws_event_bus.subscribe());
        }
        let anonymous = matches!(self.identity, Some(Identity::Anonymous));

        let identity = if self
            .state
            .api_token()
            .is_some_and(|expected| crate::token_matches(token, &expected))
        {
            Identity::Operator(token.to_string())
        } else {
            match self.state.paired_device(token).await {
                Ok(Some(device)) => {
                    touch_device(&self.state, device.id.clone());
                    Identity::Device {
                        id: device.id,
                        token: token.to_string(),
                    }
                }
                Ok(None) if anonymous => Identity::Anonymous,
                Ok(None) => {
                    let config_error = self.state.inner.read().await.config_error.is_some();
                    let no_core = self.state.core().await.is_err();
                    if no_core && self.state.desktop_host.is_some() {
                        self.reject("HOST_LOCKED", crate::HOST_LOCKED_MESSAGE).await;
                    } else if no_core && config_error {
                        self.reject("SERVER_NOT_READY", crate::SERVER_NOT_READY_MESSAGE)
                            .await;
                    } else {
                        self.reject("AUTH_FAILED", AUTH_FAILED_MESSAGE).await;
                    }
                    return false;
                }
                Err(_) => {
                    self.reject("SERVER_NOT_READY", crate::SERVER_NOT_READY_MESSAGE)
                        .await;
                    return false;
                }
            }
        };

        self.authenticated(identity, None).await
    }

    async fn on_claim(&mut self, code: &str, name: &str, platform: &str) -> bool {
        let Ok(core) = self.state.core().await else {
            if self.state.desktop_host.is_some() {
                self.reject("HOST_LOCKED", crate::HOST_LOCKED_MESSAGE).await;
            } else {
                self.reject("SERVER_NOT_READY", crate::SERVER_NOT_READY_MESSAGE)
                    .await;
            }
            return false;
        };
        let name = match validate_device(name, platform) {
            Ok(_) if crate::demo::enabled() => crate::demo::device_name(platform),
            Ok(name) => name,
            Err((code, message)) => {
                self.reject(code, message).await;
                return false;
            }
        };
        if !self.state.pairing_codes.take(code) {
            self.reject(
                "PAIRING_CODE_INVALID",
                "This code expired or was already used. Create a new code on the server.",
            )
            .await;
            return false;
        }

        match core.pair_device(&name, platform).await {
            Ok((device, token)) => {
                let identity = Identity::Device {
                    id: device.id,
                    token: token.clone(),
                };
                self.authenticated(identity, Some(&token)).await
            }
            Err(error) => {
                tracing::warn!(%error, "Failed to pair a device");
                self.reject(
                    "PAIRING_FAILED",
                    "Pipedash couldn't pair this device. Try again.",
                )
                .await;
                false
            }
        }
    }

    async fn on_request(
        &mut self, id: u64, method: String, path: String, body: Option<String>,
    ) -> bool {
        let token = match &self.identity {
            None => {
                self.reject("AUTH_FAILED", AUTH_FAILED_MESSAGE).await;
                return false;
            }
            Some(Identity::Anonymous) => None,
            Some(Identity::Operator(token)) | Some(Identity::Device { token, .. }) => {
                Some(token.clone())
            }
        };

        let state = &self.state;
        let router = self
            .router
            .get_or_insert_with(|| crate::host_router(state.clone()))
            .clone();
        let Ok(admitted) = Arc::clone(&self.admitted).try_acquire_owned() else {
            let _ = self
                .out
                .try_send(response_frame(id, 429, TOO_MANY_REQUESTS.as_bytes()));
            return true;
        };
        let permits = Arc::clone(&self.permits);
        let out = self.out.clone();
        self.requests.spawn(async move {
            let _admitted = admitted;
            let Ok(_permit) = permits.acquire_owned().await else {
                return;
            };
            let (status, body) = dispatch(router, &method, &path, body, token.as_deref()).await;
            let _ = out.send(response_frame(id, status, &body)).await;
        });

        true
    }

    async fn on_event(&mut self, event: Result<CoreEvent, broadcast::error::RecvError>) -> bool {
        match event {
            Ok(event) => {
                if let (CoreEvent::DeviceRevoked { id }, Some(Identity::Device { id: own, .. })) =
                    (&event, &self.identity)
                {
                    if id == own {
                        self.reject("DEVICE_REMOVED", "This device was removed.")
                            .await;
                        return false;
                    }
                }

                match event_frame(&event) {
                    Some(json) => self.out.send(Message::Text(json.into())).await.is_ok(),
                    None => true,
                }
            }
            Err(broadcast::error::RecvError::Lagged(skipped)) => {
                tracing::warn!(skipped, "WebSocket client fell behind, closing");
                let _ = self.out.send(Message::Close(None)).await;
                false
            }
            Err(broadcast::error::RecvError::Closed) => false,
        }
    }
}

pub async fn ws_handler(ws: WebSocketUpgrade, State(state): State<AppState>) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_socket(socket, state))
}

pub(crate) async fn closed(closing: &mut watch::Receiver<bool>) {
    let _ = closing.wait_for(|closed| *closed).await;
}

async fn handle_socket(socket: WebSocket, state: AppState) {
    let (sender, mut receiver) = socket.split();
    let (out, outbox) = mpsc::channel(OUTBOX_CAPACITY);
    let writer = tokio::spawn(write_loop(sender, outbox));
    let mut conn = Connection::new(state, out);
    let mut closing = conn.state.closing.subscribe();
    let auth_deadline = tokio::time::sleep(AUTH_TIMEOUT);
    tokio::pin!(auth_deadline);

    loop {
        tokio::select! {
            _ = conn.out.closed() => break,
            _ = closed(&mut closing) => {
                let _ = conn.out.send(Message::Close(None)).await;
                break;
            }
            frame = receiver.next() => match frame {
                Some(Ok(Message::Text(raw))) => {
                    if !conn.on_text(&raw).await {
                        break;
                    }
                }
                Some(Ok(Message::Close(_))) | None => {
                    tracing::debug!("WebSocket client disconnected");
                    break;
                }
                Some(Err(error)) => {
                    tracing::warn!("WebSocket error: {}", error);
                    break;
                }
                Some(Ok(_)) => {}
            },
            event = next_event(&mut conn.events) => {
                if !conn.on_event(event).await {
                    break;
                }
            }
            Some(_) = conn.requests.join_next(), if !conn.requests.is_empty() => {}
            _ = &mut auth_deadline, if conn.identity.is_none() => {
                conn.reject("AUTH_TIMEOUT", "Authentication timed out.").await;
                break;
            }
        }
    }

    drop(conn);
    let _ = writer.await;
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::time::Duration;

    use axum::routing::get;
    use axum::Router;
    use serde_json::{
        json,
        Value,
    };
    use tokio::io::{
        AsyncReadExt,
        AsyncWriteExt,
    };
    use tokio::net::{
        TcpListener,
        TcpStream,
    };

    use super::*;
    use crate::test_support::hosted;

    async fn read_frame(stream: &mut TcpStream) -> Option<(u8, Vec<u8>)> {
        let mut head = [0u8; 2];
        stream.read_exact(&mut head).await.ok()?;
        let len = match head[1] & 0x7f {
            126 => stream.read_u16().await.ok()? as usize,
            127 => stream.read_u64().await.ok()? as usize,
            len => len as usize,
        };
        let mut payload = vec![0u8; len];
        stream.read_exact(&mut payload).await.ok()?;
        Some((head[0] & 0x0f, payload))
    }

    async fn upgrade(addr: std::net::SocketAddr, path: &str) -> TcpStream {
        let mut stream = TcpStream::connect(addr).await.unwrap();
        let handshake = format!(
            "GET {path} HTTP/1.1\r\nHost: {addr}\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\
             Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n\r\n"
        );
        stream.write_all(handshake.as_bytes()).await.unwrap();
        let mut response = Vec::new();
        while !response.ends_with(b"\r\n\r\n") {
            response.push(stream.read_u8().await.unwrap());
        }
        assert!(response.starts_with(b"HTTP/1.1 101"));
        tokio::time::sleep(Duration::from_millis(100)).await;

        stream
    }

    async fn connect_state(state: AppState) -> TcpStream {
        let app = Router::new()
            .route("/ws", get(ws_handler))
            .with_state(state);
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });

        upgrade(addr, "/ws").await
    }

    async fn connect(bus: &Arc<WebSocketEventBus>) -> TcpStream {
        connect_state(AppState::setup_mode(bus.clone())).await
    }

    async fn send_json(stream: &mut TcpStream, value: Value) {
        let payload = value.to_string().into_bytes();
        let mask = [7u8, 11, 13, 17];
        let mut frame = vec![0x81];
        match payload.len() {
            len if len < 126 => frame.push(0x80 | len as u8),
            len => {
                frame.push(0x80 | 126);
                frame.extend((len as u16).to_be_bytes());
            }
        }
        frame.extend(mask);
        frame.extend(payload.iter().enumerate().map(|(i, b)| b ^ mask[i % 4]));
        stream.write_all(&frame).await.unwrap();
    }

    async fn next_message(stream: &mut TcpStream) -> Value {
        let (opcode, payload) = tokio::time::timeout(Duration::from_secs(5), read_frame(stream))
            .await
            .expect("no message arrived")
            .expect("the socket closed");
        assert_eq!(opcode, 0x1);
        serde_json::from_slice(&payload).unwrap()
    }

    async fn assert_closed(stream: &mut TcpStream) {
        let closed = tokio::time::timeout(Duration::from_secs(5), async {
            loop {
                match read_frame(stream).await {
                    None | Some((0x8, _)) => break,
                    Some(_) => {}
                }
            }
        })
        .await;
        assert!(closed.is_ok(), "the socket stayed open");
    }

    async fn claim(stream: &mut TcpStream, code: &str) -> Value {
        send_json(
            stream,
            json!({ "type": "claim", "code": code, "name": "Pixel", "platform": "android" }),
        )
        .await;
        next_message(stream).await
    }

    #[tokio::test]
    async fn lagged_client_gets_closed_so_it_reconnects() {
        let bus = Arc::new(WebSocketEventBus::new());
        let mut stream = connect(&bus).await;

        for _ in 0..1100 {
            bus.emit(CoreEvent::ProvidersChanged).await;
        }

        let closed = tokio::time::timeout(Duration::from_secs(5), async {
            loop {
                match read_frame(&mut stream).await {
                    None => break,
                    Some((0x8, _)) => break,
                    Some(_) => {}
                }
            }
        })
        .await;
        assert!(closed.is_ok(), "socket stayed open after the client lagged");
    }

    #[tokio::test]
    async fn live_run_events_reach_clients_with_their_payloads() {
        let bus = Arc::new(WebSocketEventBus::new());
        let mut stream = connect(&bus).await;
        let run: pipedash_core::domain::PipelineRun = serde_json::from_value(serde_json::json!({
            "id": "p__3", "pipeline_id": "p", "run_number": 3, "status": "running",
            "started_at": "2026-01-01T00:00:00Z", "duration_seconds": 12, "logs_url": ""
        }))
        .unwrap();

        bus.emit(CoreEvent::RunLogTreeChanged {
            pipeline_id: "p".into(),
            run_number: 3,
        })
        .await;
        bus.emit(CoreEvent::RunUpdated {
            pipeline_id: "p".into(),
            run_number: 3,
            run: run.clone(),
        })
        .await;

        assert_eq!(
            next_message(&mut stream).await,
            serde_json::json!({
                "type": "run-log-tree-changed",
                "payload": { "pipelineId": "p", "runNumber": 3 },
            })
        );
        assert_eq!(
            next_message(&mut stream).await,
            serde_json::json!({
                "type": "run-updated",
                "payload": {
                    "pipelineId": "p",
                    "runNumber": 3,
                    "run": serde_json::to_value(&run).unwrap(),
                },
            })
        );
    }

    #[tokio::test]
    async fn a_browser_auth_frame_gets_auth_ok_when_no_token_is_required() {
        let bus = Arc::new(WebSocketEventBus::new());
        let mut stream = connect(&bus).await;

        send_json(&mut stream, json!({ "type": "auth", "token": "anything" })).await;

        assert_eq!(
            next_message(&mut stream).await,
            json!({
                "type": "auth-ok",
                "payload": { "host": "server", "name": null, "deviceId": null },
            })
        );
    }

    #[tokio::test]
    async fn the_operator_token_authenticates_and_requests_carry_it() {
        let state = AppState::setup_mode(Arc::new(WebSocketEventBus::new()));
        state.set_api_token("vault-password");

        let mut stream = connect_state(state.clone()).await;
        send_json(&mut stream, json!({ "type": "auth", "token": "wrong" })).await;
        let rejected = next_message(&mut stream).await;
        assert_eq!(rejected["type"], "auth-error");
        assert_eq!(rejected["payload"]["code"], "AUTH_FAILED");
        assert_closed(&mut stream).await;

        let mut stream = connect_state(state).await;
        send_json(
            &mut stream,
            json!({ "type": "auth", "token": "vault-password" }),
        )
        .await;
        assert_eq!(next_message(&mut stream).await["type"], "auth-ok");
        send_json(
            &mut stream,
            json!({ "type": "request", "id": 5, "method": "GET", "path": "/providers" }),
        )
        .await;
        let response = next_message(&mut stream).await;

        assert_eq!(response["type"], "response");
        assert_eq!(response["id"], 5);
        assert_eq!(response["status"], 503);
    }

    #[tokio::test]
    async fn a_request_before_authenticating_is_refused() {
        let state = AppState::setup_mode(Arc::new(WebSocketEventBus::new()));
        state.set_api_token("vault-password");
        let mut stream = connect_state(state).await;

        send_json(
            &mut stream,
            json!({ "type": "request", "id": 1, "method": "GET", "path": "/providers" }),
        )
        .await;

        assert_eq!(
            next_message(&mut stream).await["payload"]["code"],
            "AUTH_FAILED"
        );
        assert_closed(&mut stream).await;
    }

    #[tokio::test]
    async fn a_claim_pairs_the_device_and_serves_requests_then_the_token_reconnects() {
        let host = hosted("ws-claim").await;
        let (code, _) = host.state.pairing_codes.issue().unwrap();
        let mut stream = connect_state(host.state.clone()).await;

        let ok = claim(&mut stream, &code).await;

        assert_eq!(ok["type"], "auth-ok");
        assert_eq!(ok["payload"]["host"], "desktop");
        assert_eq!(ok["payload"]["name"], "test-mac");
        let device_id = ok["payload"]["deviceId"].as_str().unwrap().to_string();
        let token = ok["payload"]["token"].as_str().unwrap().to_string();
        assert_eq!(token.len(), 64);

        send_json(
            &mut stream,
            json!({ "type": "request", "id": 1, "method": "GET", "path": "/devices" }),
        )
        .await;
        let response = next_message(&mut stream).await;
        assert_eq!(response["id"], 1);
        assert_eq!(response["status"], 200);
        let body: Value = serde_json::from_str(response["body"].as_str().unwrap()).unwrap();
        assert_eq!(body["devices"][0]["id"], device_id);
        assert_eq!(body["devices"][0]["current"], true);

        let mut again = connect_state(host.state.clone()).await;
        send_json(&mut again, json!({ "type": "auth", "token": token })).await;
        let ok = next_message(&mut again).await;
        assert_eq!(ok["payload"]["deviceId"], device_id);
        assert!(ok["payload"].get("token").is_none());

        let mut stranger = connect_state(host.state.clone()).await;
        send_json(&mut stranger, json!({ "type": "auth", "token": "unknown" })).await;
        assert_eq!(
            next_message(&mut stranger).await["payload"]["code"],
            "AUTH_FAILED"
        );
    }

    #[tokio::test]
    async fn stopping_the_host_closes_connected_sockets() {
        let state = AppState::hosted(Arc::new(WebSocketEventBus::new()), None, "test-mac");
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let (stop, stopped) = tokio::sync::oneshot::channel::<()>();
        let server = tokio::spawn(crate::serve_host(listener, state, async move {
            let _ = stopped.await;
        }));
        let mut stream = upgrade(addr, "/api/v1/ws").await;

        stop.send(()).unwrap();

        assert_closed(&mut stream).await;
        tokio::time::timeout(Duration::from_secs(5), server)
            .await
            .expect("serve_host kept running")
            .unwrap()
            .unwrap();
    }

    #[tokio::test]
    async fn a_busy_socket_rejects_excess_requests_and_still_closes() {
        let host = crate::test_support::server("ws-busy").await;
        let mut stream = connect_state(host.state.clone()).await;
        let blocked = host.state.inner.write().await;

        for id in 1..=MAX_ADMITTED as u64 + 1 {
            send_json(
                &mut stream,
                json!({ "type": "request", "id": id, "method": "GET", "path": "/providers" }),
            )
            .await;
        }

        let rejected = next_message(&mut stream).await;
        assert_eq!(rejected["id"], MAX_ADMITTED as u64 + 1);
        assert_eq!(rejected["status"], 429);

        host.state.closing.send_replace(true);
        assert_closed(&mut stream).await;
        drop(blocked);
    }

    #[tokio::test]
    async fn a_shared_host_refuses_socket_upgrades_from_other_origins() {
        let state = AppState::hosted(Arc::new(WebSocketEventBus::new()), None, "test-mac");
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        tokio::spawn(crate::serve_host(listener, state, std::future::pending()));

        let mut stream = TcpStream::connect(addr).await.unwrap();
        let handshake = format!(
            "GET /api/v1/ws HTTP/1.1\r\nHost: {addr}\r\nOrigin: https://attacker.example\r\n\
             Upgrade: websocket\r\nConnection: Upgrade\r\n\
             Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n\r\n"
        );
        stream.write_all(handshake.as_bytes()).await.unwrap();
        let mut status = [0u8; 12];
        stream.read_exact(&mut status).await.unwrap();
        assert_eq!(&status, b"HTTP/1.1 403");

        upgrade(addr, "/api/v1/ws").await;
    }

    #[tokio::test]
    async fn a_device_token_on_a_server_without_an_api_token_becomes_a_device() {
        let host = crate::test_support::server("ws-open-device").await;
        let core = host.state.core().await.unwrap();
        let (device, token) = core.pair_device("Pixel", "android").await.unwrap();
        let mut stream = connect_state(host.state.clone()).await;

        send_json(&mut stream, json!({ "type": "auth", "token": token })).await;
        let ok = next_message(&mut stream).await;
        assert_eq!(ok["payload"]["deviceId"], device.id);

        send_json(
            &mut stream,
            json!({ "type": "request", "id": 1, "method": "GET", "path": "/devices" }),
        )
        .await;
        let response = next_message(&mut stream).await;
        let body: Value = serde_json::from_str(response["body"].as_str().unwrap()).unwrap();
        assert_eq!(body["devices"][0]["current"], true);

        core.revoke_device(&device.id).await.unwrap();
        let removed = next_message(&mut stream).await;
        assert_eq!(removed["payload"]["code"], "DEVICE_REMOVED");
        assert_closed(&mut stream).await;
    }

    #[tokio::test]
    async fn a_failed_device_lookup_is_not_reported_as_a_revoked_device() {
        let host = hosted("ws-lookup-failure").await;
        let core = host.state.core().await.unwrap();
        let (_, token) = core.pair_device("Pixel", "android").await.unwrap();
        let pipedash_core::infrastructure::database::DatabasePool::Sqlite(pool) =
            core.provider_service.repository().cache_pool().clone()
        else {
            panic!("expected a SQLite pool");
        };
        pool.close().await;
        let mut stream = connect_state(host.state.clone()).await;

        send_json(&mut stream, json!({ "type": "auth", "token": token })).await;

        let rejected = next_message(&mut stream).await;
        assert_eq!(rejected["payload"]["code"], "SERVER_NOT_READY");
        assert_closed(&mut stream).await;
    }

    #[tokio::test]
    async fn a_locked_host_answers_host_locked_instead_of_forgetting_the_device() {
        let state = AppState::hosted(Arc::new(WebSocketEventBus::new()), None, "test-mac");
        let mut stream = connect_state(state).await;

        send_json(&mut stream, json!({ "type": "auth", "token": "any" })).await;

        let rejected = next_message(&mut stream).await;
        assert_eq!(rejected["payload"]["code"], "HOST_LOCKED");
        assert_eq!(rejected["payload"]["message"], crate::HOST_LOCKED_MESSAGE);
        assert_closed(&mut stream).await;
    }

    #[tokio::test]
    async fn a_server_with_a_broken_config_does_not_forget_its_devices() {
        let state = AppState::config_error(Arc::new(WebSocketEventBus::new()), "broken".into());
        state.set_api_token("vault-password");
        let mut stream = connect_state(state).await;

        send_json(
            &mut stream,
            json!({ "type": "auth", "token": "device-token" }),
        )
        .await;

        let rejected = next_message(&mut stream).await;
        assert_eq!(rejected["payload"]["code"], "SERVER_NOT_READY");
        assert_closed(&mut stream).await;
    }

    #[tokio::test]
    async fn a_pairing_code_can_be_claimed_only_once() {
        let host = hosted("ws-claim-once").await;
        let (code, _) = host.state.pairing_codes.issue().unwrap();

        let mut first = connect_state(host.state.clone()).await;
        assert_eq!(claim(&mut first, &code).await["type"], "auth-ok");

        let mut second = connect_state(host.state.clone()).await;
        let rejected = claim(&mut second, &code).await;
        assert_eq!(rejected["type"], "auth-error");
        assert_eq!(rejected["payload"]["code"], "PAIRING_CODE_INVALID");
        assert_closed(&mut second).await;
    }

    #[tokio::test]
    async fn a_claim_with_invalid_device_details_is_rejected_without_burning_the_code() {
        let host = hosted("ws-claim-invalid").await;
        let (code, _) = host.state.pairing_codes.issue().unwrap();
        let mut stream = connect_state(host.state.clone()).await;

        send_json(
            &mut stream,
            json!({ "type": "claim", "code": code, "name": "Pixel", "platform": "palm" }),
        )
        .await;

        assert_eq!(
            next_message(&mut stream).await["payload"]["code"],
            "INVALID_PLATFORM"
        );
        assert!(host.state.pairing_codes.take(&code));
    }

    #[tokio::test]
    async fn revoking_a_device_closes_its_socket_with_device_removed() {
        let host = hosted("ws-revoke").await;
        let (code, _) = host.state.pairing_codes.issue().unwrap();
        let mut stream = connect_state(host.state.clone()).await;
        let ok = claim(&mut stream, &code).await;
        let device_id = ok["payload"]["deviceId"].as_str().unwrap().to_string();

        host.state
            .core()
            .await
            .unwrap()
            .revoke_device(&device_id)
            .await
            .unwrap();

        let removed = next_message(&mut stream).await;
        assert_eq!(removed["type"], "auth-error");
        assert_eq!(removed["payload"]["code"], "DEVICE_REMOVED");
        assert_closed(&mut stream).await;
    }

    #[tokio::test]
    async fn a_connection_that_never_authenticates_is_closed_after_ten_seconds() {
        let host = hosted("ws-timeout").await;
        let mut stream = connect_state(host.state.clone()).await;

        tokio::time::pause();
        tokio::time::advance(AUTH_TIMEOUT + Duration::from_secs(1)).await;

        let timed_out = next_message(&mut stream).await;
        assert_eq!(timed_out["type"], "auth-error");
        assert_eq!(timed_out["payload"]["code"], "AUTH_TIMEOUT");
        assert_closed(&mut stream).await;
    }
}