exfiltrate 0.4.0

// SPDX-License-Identifier: MIT OR Apache-2.0
//! WebAssembly worker and WebSocket implementation of the debug server.
//!
//! # How a command gets off the callback
//!
//! Commands used to run inline inside the WebSocket `onmessage` callback, which
//! cost the browser both of the capabilities the native side has: a command
//! could not emit partial output, and a `Cancel` could not even be *received*
//! while one was running. Two obstacles were cited for that, and both are real
//! but neither is load-bearing.
//!
//! **The socket handle is a `JsValue`, so it is neither `Send` nor `Sync`.**
//! True, and it stays true. But the native side does not move its socket
//! either — it moves *bytes* to a single writer. So [`worker_thread`] keeps the
//! handle and performs every write, and everything else submits `Vec<u8>`,
//! which is `Send`, through an [`mpsc`](wasm_lite_std::mpsc) queue.
//!
//! **The `onmessage` callback owns the event loop.** True only while commands
//! run inside it. Each command now gets its own worker, so the callback
//! deserialises, dispatches and returns — which is exactly what makes a `Cancel`
//! frame receivable. Cancellation is then an `AtomicBool` in shared memory,
//! which works across workers because they share linear memory.
//!
//! Spawning a worker from a worker is the obvious trap here, and `wasm_lite`
//! already handles it: worker creation is routed through the main thread on
//! purpose, because Chrome will not start a nested worker while its parent is
//! blocked in `Atomics.wait`.
//!
//! # Events, and the turn that never came
//!
//! There is no idle turn in a browser — nothing runs between messages — so
//! events used to be drained only when a frame happened to arrive. That made
//! `watch` silent in a browser for as long as the client stayed quiet, which is
//! the entire duration of a `watch`. The writer loop now wakes on a short
//! timeout as well as on work, and drains there.
use crate::wire::server::do_command;
use crate::wire::server::reconnect;
use exfiltrate_internal::command::{CommandContext, Response, StreamError};
use exfiltrate_internal::rpc::{Chunk, CommandInvocation, CommandResponse, RPC};
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, LazyLock};
use std::task::{Context, Poll, Waker};
use std::time::Duration;
use wasm_lite::websocket::{BinaryType, CloseEvent, MessageEvent, WebSocket};
use wasm_lite::{Closure, JsValue, console};
use wasm_lite_std::Mutex;
use wasm_lite_std::mpsc;

/// How many commands may run at once in a browser.
///
/// Lower than the native limit of 64 on purpose: a thread here is a Web Worker,
/// which costs far more than a thread does natively. Eight concurrent commands
/// is more than anyone debugging a page runs deliberately, and exceeding it is
/// answered with an error rather than silently queued.
const MAX_CONCURRENT_COMMANDS: usize = 8;

/// How long the writer waits for work before draining events anyway.
///
/// This is the browser's substitute for an idle turn. Short enough that `watch`
/// feels immediate, long enough that an idle page is not doing anything much.
const EVENT_POLL_INTERVAL: Duration = Duration::from_millis(250);

/// One group of frames that must reach the socket without another writer's
/// frames getting between its parts.
type WriteJob = Vec<Vec<u8>>;

impl reconnect::ProxySocket for WebSocket {
    fn ready_state(&self) -> u16 {
        WebSocket::ready_state(self)
    }
}

// `reconnect` restates readyState so it is testable without a browser. These
// used to be `const` assertions against web-sys' constants; wasm_lite reads the
// real ones off the class at runtime instead, so the check moved to a test
// (`ready_states_match_the_browser` below) — which is stronger, because it
// compares against the engine rather than against another copy of the same
// numbers.
const _: () = {
    assert!(reconnect::CONNECTING == WebSocket::CONNECTING);
    assert!(reconnect::OPEN == WebSocket::OPEN);
    assert!(reconnect::CLOSING == WebSocket::CLOSING);
    assert!(reconnect::CLOSED == WebSocket::CLOSED);
};

pub fn wasm32_go() {
    let thread_result = wasm_lite_std::Builder::new()
        .name("exfiltrate::wasm".to_string())
        .spawn(|| {
            wasm_lite_std::spawn_local(async move {
                let receiver = SEND_WORKER_MESSAGE
                    .1
                    .with_mut_sync(|e| e.take())
                    .expect("no receiver");
                worker_thread(receiver).await;
            });
        });
    match thread_result {
        Ok(_join_handle) => {}
        Err(e) => {
            console::error(&format!("{:?}", e));
            panic!("{:?}", e);
        }
    }
}

/// The event subscription belonging to the browser's current WebSocket connection.
///
/// A browser application has one connection at a time, but not the *same* one
/// for the process's life — the socket drops and reconnects. So this is reset on
/// close rather than being a one-shot: keeping the old handle would leave the
/// previous CLI's subscriptions in force, and the new one would receive topics
/// it never asked for.
static SUBSCRIBER: LazyLock<Mutex<Option<u64>>> = LazyLock::new(|| Mutex::new(None));

/// This connection's event subscription, attaching on first use.
fn subscriber() -> u64 {
    SUBSCRIBER.with_mut_sync(|slot| match slot {
        Some(id) => *id,
        None => {
            let id = crate::events::attach(crate::config_snapshot().event_queue_capacity);
            *slot = Some(id);
            id
        }
    })
}

/// Drops this connection's subscriptions. Called when the socket closes.
fn detach_subscriber() {
    if let Some(id) = SUBSCRIBER.with_mut_sync(|slot| slot.take()) {
        crate::events::detach(id);
    }
}

/// Encodes one message into the frames it occupies on the wire.
///
/// A `CommandResponse` is a frame plus one per attachment, and they must arrive
/// together — hence a job rather than a frame. Doing the split here rather than
/// at each call site is what keeps the attachment rules in one place now that
/// there are two callers: the callback and every command worker.
fn frames_for(mut rpc: RPC) -> Result<WriteJob, String> {
    let mut attachments = Vec::new();
    if let RPC::CommandResponse(ref mut resp) = rpc {
        if resp.response.attachment_count() > exfiltrate_internal::wire::MAX_ATTACHMENTS {
            resp.success = false;
            resp.response = format!(
                "response exceeds the {}-attachment limit",
                exfiltrate_internal::wire::MAX_ATTACHMENTS
            )
            .into();
        }
        attachments = resp.response.split_data();
        resp.num_attachments = u32::try_from(attachments.len())
            .map_err(|_| "response contains too many attachments".to_string())?;
    }

    let mut job =
        vec![rmp_serde::to_vec(&rpc).map_err(|error| format!("cannot encode reply: {error}"))?];
    job.extend(attachments);
    Ok(job)
}

/// Queues one message for the socket worker.
fn enqueue(rpc: RPC) -> Result<(), String> {
    let job = frames_for(rpc)?;
    SEND_WORKER_MESSAGE
        .0
        .send_sync(WorkerMessage::Write(job))
        .map_err(|_| "the socket worker is gone".to_string())
}

/// Writes whatever events have queued for this connection.
///
/// Reads the subscription rather than creating one: this runs on a timer from
/// startup, and attaching a subscriber before anyone has asked for events would
/// mean every page that links this crate holds a queue nothing ever fills.
fn flush_events() {
    let Some(id) = SUBSCRIBER.with_sync(|slot| *slot) else {
        return;
    };
    for event in crate::events::drain(id) {
        if let Err(error) = enqueue(RPC::Event(event)) {
            console::error(&format!("exfiltrate: cannot deliver an event: {error}"));
            return;
        }
    }
}

/// Runs one command on its own worker and queues its chunks and final response.
///
/// Returns false when the command could not be started, which the caller
/// answers rather than swallowing.
fn start_command(command: CommandInvocation) -> bool {
    let reply_id = command.reply_id;
    let token = Arc::new(AtomicBool::new(false));
    let admitted = RUNNING.with_mut_sync(|running| {
        if running.len() >= MAX_CONCURRENT_COMMANDS {
            return false;
        }
        running.insert(reply_id, token.clone());
        true
    });
    if !admitted {
        return false;
    }

    let name = command.name.clone();
    let spawned = wasm_lite_std::spawn_named(format!("exfiltrate::command {name}"), move || {
        let seq = Arc::new(AtomicU64::new(0));
        // The sink holds no `JsValue` — it serialises and hands over bytes, and
        // the socket worker is the only thing that touches the socket.
        let context = CommandContext::new(
            token,
            Arc::new(move |payload: Response| {
                let chunk = RPC::Chunk(Chunk {
                    reply_id,
                    seq: seq.fetch_add(1, Ordering::Relaxed),
                    payload,
                });
                enqueue(chunk).map_err(StreamError::Disconnected)
            }),
        );

        let response = do_command(command, &context);
        RUNNING.with_mut_sync(|running| running.remove(&reply_id));
        if let Err(error) = enqueue(RPC::CommandResponse(response)) {
            console::error(&format!(
                "exfiltrate: cannot deliver reply {reply_id}: {error}"
            ));
        }
        // Anything the command emitted while it ran goes out with it.
        flush_events();
    });
    if spawned.is_err() {
        RUNNING.with_mut_sync(|running| running.remove(&reply_id));
        return false;
    }
    true
}

/// Handles one inbound frame.
///
/// Everything here is meant to be quick, because this runs on the callback that
/// owns the worker's event loop: a `Command` is handed to its own worker and a
/// `Cancel` only sets a flag. That promptness is what makes a `Cancel` arriving
/// mid-command receivable at all.
fn handle_msg(data: &[u8]) -> Result<(), String> {
    let msg: RPC = rmp_serde::from_slice(data).map_err(|error| {
        format!(
            "could not parse a message from the client: {error}. This usually means the CLI \
             and the linked library are different versions."
        )
    })?;

    let result = match msg {
        RPC::Hello(peer) => {
            let local = crate::config_snapshot().build_info();
            if !peer.is_compatible() {
                console::error(&peer.skew_message(&local));
            }
            enqueue(RPC::Hello(local))
        }
        RPC::Command(command) => {
            let reply_id = command.reply_id;
            if start_command(command) {
                Ok(())
            } else {
                enqueue(RPC::CommandResponse(CommandResponse::new(
                    false,
                    format!(
                        "this connection already has {MAX_CONCURRENT_COMMANDS} commands in \
                         flight, or a worker could not be started; wait for one to finish"
                    )
                    .into(),
                    reply_id,
                )))
            }
        }
        RPC::Cancel(cancel) => {
            // Shared linear memory is what makes this work across workers: the
            // flag the command polls is the flag set here.
            let known = RUNNING.with_sync(|running| {
                running
                    .get(&cancel.reply_id)
                    .map(|token| token.store(true, Ordering::Relaxed))
                    .is_some()
            });
            if !known {
                console::log(&format!(
                    "exfiltrate: cancel for reply {} arrived after the command finished",
                    cancel.reply_id
                ));
            }
            Ok(())
        }
        RPC::Subscribe(subscription) => {
            if let Err(error) = crate::events::subscribe(subscriber(), &subscription.topic) {
                console::error(&format!("exfiltrate: subscribe refused: {error}"));
            }
            Ok(())
        }
        RPC::Unsubscribe(subscription) => {
            if let Err(error) = crate::events::unsubscribe(subscriber(), &subscription.topic) {
                console::error(&format!("exfiltrate: unsubscribe refused: {error}"));
            }
            Ok(())
        }
        RPC::CommandResponse(_) | RPC::Chunk(_) | RPC::Event(_) => {
            Err("the client sent a message only a server may send".to_string())
        }
        _ => {
            // `RPC` is non-exhaustive; a newer client may know variants this
            // build does not. Ignoring one keeps future additions from being a
            // flag day.
            console::log("exfiltrate: ignoring an RPC variant this build does not know");
            Ok(())
        }
    };
    result?;
    flush_events();
    Ok(())
}

/// Debug a WebSocket URL by hitting it over HTTP.
///
/// A WebSocket that fails to connect tells script nothing about why — the spec
/// fires a bare `error` event on purpose. Fetching the same URL over HTTP is
/// the one way to get a status code and a body out of the server, so this
/// exists to turn "it didn't connect" into something diagnosable.
pub async fn debug_ws_handshake(ws_url: &str) -> Result<(), JsValue> {
    // Convert ws:// -> http:// and wss:// -> https://
    let http_url = if let Some(rest) = ws_url.strip_prefix("ws://") {
        format!("http://{rest}")
    } else if let Some(rest) = ws_url.strip_prefix("wss://") {
        format!("https://{rest}")
    } else {
        ws_url.to_string()
    };

    let opts = wasm_lite::fetch::RequestInit::new();
    opts.set_method("GET");
    // CORS mode is usually fine; tweak if you know you're same-origin
    opts.set_mode("cors");

    match wasm_lite::fetch::fetch(&http_url, &opts).await {
        Err(err) => {
            // `Display`, not `Debug`: this is a `TypeError` whose message is
            // the only diagnostic the browser offers.
            console::log(&format!("fetch error: {err}"));
        }
        Ok(response) => {
            console::log(&format!("Fetch status: {}", response.status()));
            match response.text().await {
                Ok(body) => console::log(&body),
                Err(e) => console::log(&format!("could not read body: {e}")),
            }
        }
    }
    Ok(())
}

/// Main worker thread function that manages WebSocket connections.
///
/// This function runs in a dedicated thread and:
/// - Handles connection requests
/// - Manages the WebSocket lifecycle
/// - Routes messages between the WebSocket and the proxy system
///
/// # Arguments
///
/// * `receiver` - Channel for receiving control messages
async fn worker_thread(mut receiver: mpsc::Receiver<WorkerMessage>) {
    console::log("thread started");

    let mut socket: Option<WebSocket> = None;
    let _ = SEND_WORKER_MESSAGE.0.send_sync(WorkerMessage::Reconnect);

    loop {
        // The timeout is what gives a browser the idle turn it does not
        // otherwise have. Without it, an event emitted while the client sat
        // quietly in `watch` would wait for the client to say something first.
        let deadline = wasm_lite_std::time::Instant::now() + EVENT_POLL_INTERVAL;
        match receiver.recv_async_timeout(deadline).await {
            Ok(WorkerMessage::Reconnect) => {
                // Discards a socket that can no longer carry traffic, so a
                // dropped connection reconnects rather than wedging. See
                // `super::reconnect`, where both rules are tested.
                if reconnect::needs_connect(&mut socket) {
                    console::log("WebSocket: connecting...");

                    let attempt = create_web_socket().await;
                    match reconnect::store_attempt(&mut socket, attempt) {
                        Ok(()) => {
                            console::log("WebSocket created successfully");
                        }
                        Err(e) => {
                            console::log(&format!("Failed to create WebSocket: {:?}", e));
                        }
                    }
                }
            }
            // The only place the socket handle is ever used for writing. That
            // is what lets a chunk sink on another worker exist at all: it
            // hands over bytes, never the `JsValue`.
            Ok(WorkerMessage::Write(job)) => {
                let Some(ws) = &socket else {
                    console::log("exfiltrate: dropping a reply; there is no socket");
                    continue;
                };
                for frame in job {
                    if let Err(error) = ws.send_bytes(&frame) {
                        console::error(&format!("exfiltrate: write failed: {error:?}"));
                        break;
                    }
                }
            }
            Err(mpsc::RecvTimeoutError::Timeout) => flush_events(),
            // `RecvTimeoutError` is non-exhaustive, so this arm is the
            // catch-all as well as the disconnect case. Anything a future
            // version adds is a reason to stop, not to spin.
            Err(_) => {
                console::log("receiver closed, exiting thread");
                break;
            }
        }
    }
}

/// What the socket worker is asked to do.
enum WorkerMessage {
    /// Open a socket, if the current one can no longer carry traffic.
    Reconnect,
    /// Put these frames on the wire, in this order, with nothing between them.
    Write(WriteJob),
}

/// The queue into the socket worker.
///
/// Multi-producer, which is the point: every command worker holds a sender, and
/// so does the `onmessage` callback. Only the receiving end — and therefore only
/// one thread — ever touches the `JsValue`.
#[allow(clippy::type_complexity)]
static SEND_WORKER_MESSAGE: LazyLock<(
    mpsc::Sender<WorkerMessage>,
    Mutex<Option<mpsc::Receiver<WorkerMessage>>>,
)> = LazyLock::new(|| {
    let (sender, receiver) = mpsc::channel();
    (sender, Mutex::new(Some(receiver)))
});

/// The cancellation flags of the commands running right now.
static RUNNING: LazyLock<Mutex<HashMap<u32, Arc<AtomicBool>>>> =
    LazyLock::new(|| Mutex::new(HashMap::new()));

/// The shared half of [`OneShot`] and [`OneShotFuture`].
struct OneShotState<T> {
    value: Option<T>,
    waker: Option<Waker>,
    /// Set once a value has been accepted, so later callbacks are no-ops
    /// rather than overwriting a result the receiver may already have taken.
    sent: bool,
}

/// A one-shot sender that can only send a value once.
///
/// This is used for sending completion signals from WebSocket
/// event handlers back to the async context. It ensures that
/// only the first event (either success or error) is processed.
///
/// Hand-rolled rather than `r#continue::continuation()`. That dependency put
/// `exfiltrate` and `continue` in a package cycle -- `continue`'s own
/// `exfiltrate` feature depends back on this crate -- which Cargo rejects
/// outright, so `continue` could never expose its continuation registry to
/// anyone (Metropolis/continue#4). Forty lines are cheaper than that, and this
/// crate's premise is being small and runtime-free anyway.
struct OneShot<T> {
    state: Arc<Mutex<OneShotState<T>>>,
}

/// Resolves when its [`OneShot`] sends, or never if it is dropped first.
///
/// "Never" is correct here: the WebSocket fires `onopen` or `onerror` for every
/// connection attempt, so the only way to reach that state is for the browser
/// to abandon the socket without either -- in which case there is nothing to
/// report and nothing to retry against.
struct OneShotFuture<T> {
    state: Arc<Mutex<OneShotState<T>>>,
}

fn one_shot<T>() -> (OneShot<T>, OneShotFuture<T>) {
    let state = Arc::new(Mutex::new(OneShotState {
        value: None,
        waker: None,
        sent: false,
    }));
    (
        OneShot {
            state: state.clone(),
        },
        OneShotFuture { state },
    )
}

impl<T> OneShot<T> {
    /// Sends a value if not already sent.
    ///
    /// This method is idempotent - subsequent calls after the first
    /// successful send will be no-ops.
    fn send_if_needed(&self, value: T) {
        // The waker is taken out under the lock and woken outside it: waking
        // can poll inline, and polling reacquires this same lock.
        let waker = self.state.with_mut_sync(|state| {
            if state.sent {
                return None;
            }
            state.sent = true;
            state.value = Some(value);
            state.waker.take()
        });
        if let Some(waker) = waker {
            waker.wake();
        }
    }
}

impl<T> Clone for OneShot<T> {
    fn clone(&self) -> Self {
        OneShot {
            state: Arc::clone(&self.state),
        }
    }
}

impl<T> Future for OneShotFuture<T> {
    type Output = T;

    fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<T> {
        self.state.with_mut_sync(|state| match state.value.take() {
            Some(value) => Poll::Ready(value),
            None => {
                state.waker = Some(context.waker().clone());
                Poll::Pending
            }
        })
    }
}

/// The WebSocket URL the browser application dials.
///
/// A browser has no environment variables, so unlike the native side there is
/// nothing to read at startup: [`Config::addr`](crate::Config::addr) is the only
/// way to point a wasm build at a proxy that is not on the default port.
fn web_addr() -> String {
    crate::config_snapshot()
        .addr
        .filter(|addr| !addr.is_empty())
        .unwrap_or_else(|| exfiltrate_internal::wire::WEB_ADDR.to_string())
}

async fn create_web_socket() -> Result<WebSocket, String> {
    let web_addr = web_addr();
    let ws = match WebSocket::new(&web_addr) {
        Ok(ws) => ws,
        // Only a malformed URL reaches here; a server that is not there is
        // reported through the error/close events below.
        Err(e) => return Err(format!("{e}")),
    };

    let (func_sender, func_fut) = one_shot::<Result<(), String>>();
    ws.set_binary_type(BinaryType::ArrayBuffer);

    let move_func_sender = func_sender.clone();
    let onopen_callback = Closure::new_with_arg(move |_event| {
        console::log("WebSocket opened!");
        move_func_sender.send_if_needed(Ok(()));
    });
    ws.set_onopen(Some(onopen_callback.as_js_value()));
    onopen_callback.forget(); //leak the closure

    let move_func_sender = func_sender.clone();
    let onerror_callback = Closure::new_with_arg(move |_event| {
        // The error event carries no detail — the spec fires a bare `Event` —
        // so the only way to learn anything is to ask the server over HTTP.
        let probe_addr = web_addr.clone();
        wasm_lite_std::spawn_local(async move {
            let _ = debug_ws_handshake(&probe_addr).await;
        });
        console::log("Websocket error");
        move_func_sender.send_if_needed(Err("Cannot connect to server".to_string()));
    });
    ws.set_onerror(Some(onerror_callback.as_js_value()));
    onerror_callback.forget(); //leak the closure

    let onclose_callback = Closure::new_with_arg(move |event| {
        // The next connection gets a fresh subscription; carrying this one over
        // would deliver topics the new client never subscribed to.
        detach_subscriber();
        let event = CloseEvent::from_js(event);
        console::log(&format!(
            "WebSocket closed: code {} clean {} {}",
            event.code(),
            event.was_clean(),
            event.reason()
        ));
        wasm_lite_std::sleep(Duration::from_secs(10));
        let _ = SEND_WORKER_MESSAGE.0.send_sync(WorkerMessage::Reconnect);
    });
    ws.set_onclose(Some(onclose_callback.as_js_value()));
    onclose_callback.forget(); //leak the closure

    // Deliberately does not capture the socket. Replies go through the queue to
    // the one task that owns the handle, which is what lets a command run
    // somewhere else and still answer.
    let onmessage_callback = Closure::new_with_arg(move |event| {
        let event = MessageEvent::from_js(event);
        let Some(data) = event.data_bytes() else {
            console::log("Received non-binary message");
            return;
        };
        if let Err(e) = handle_msg(&data) {
            console::error(&format!("Error handling message: {}", e));
        }
    });
    ws.set_onmessage(Some(onmessage_callback.as_js_value()));
    onmessage_callback.forget(); //leak the closure

    let f = func_fut.await;
    f.map(|_| ws)
}

#[cfg(test)]
mod tests {
    use super::*;
    use exfiltrate_internal::args::ArgSpec;
    use exfiltrate_internal::command::Command;
    use wasm_lite_std::time::Instant;

    /// The `readyState` values `reconnect` restates must be the ones the engine
    /// actually uses.
    ///
    /// This was a `const` assertion against web-sys' constants. It cannot be
    /// one any more — wasm_lite reads them off the class — and that is an
    /// improvement: a `const` assert only ever compared two hard-coded copies
    /// of the same four numbers, while this compares against the browser.
    #[wasm_lite::wasm_lite_test]
    fn ready_states_match_the_browser() {
        assert_eq!(
            WebSocket::browser_ready_states(),
            [
                reconnect::CONNECTING,
                reconnect::OPEN,
                reconnect::CLOSING,
                reconnect::CLOSED
            ]
        );
    }

    #[wasm_lite::wasm_lite_test]
    fn a_response_and_its_attachments_are_one_job_with_a_matching_count() {
        let file = exfiltrate_internal::command::FileInfo {
            proposed_extension: "txt".to_string(),
            remark: None,
            contents: b"hello".to_vec(),
        };
        let job = frames_for(RPC::CommandResponse(CommandResponse::new(
            true,
            Response::Files(vec![file]),
            1,
        )))
        .unwrap();

        assert_eq!(job.len(), 2, "one response frame and one attachment");
        let RPC::CommandResponse(decoded) = rmp_serde::from_slice::<RPC>(&job[0]).unwrap() else {
            panic!("the first frame is the response")
        };
        assert_eq!(decoded.num_attachments, 1);
        assert_eq!(job[1], b"hello");
    }

    /// The one-shot replaced `r#continue::continuation()` to break a package
    /// cycle (Metropolis/continue#4), so its three behaviours are asserted here
    /// rather than inherited: a value already in the slot, a wake for a value
    /// that arrives after the poll, and a second send being ignored.
    ///
    /// Driven by hand rather than with `block_on`. An executor would make this
    /// a test of the executor's wake plumbing as much as the one-shot's, and
    /// `block_on` is `Atomics.wait` -- so a sender scheduled onto the same
    /// thread's event loop deadlocks against it, which is exactly what the
    /// first draft of this test did.
    #[wasm_lite::wasm_lite_test]
    fn one_shot_delivers_once() {
        struct CountWakes(AtomicU64);
        impl std::task::Wake for CountWakes {
            fn wake(self: Arc<Self>) {
                self.0.fetch_add(1, Ordering::Relaxed);
            }
        }

        let wakes = Arc::new(CountWakes(AtomicU64::new(0)));
        let waker = Waker::from(wakes.clone());
        let mut context = Context::from_waker(&waker);

        // Already in the slot: ready on the first poll, and nothing to wake.
        let (sender, mut future) = one_shot::<u32>();
        sender.send_if_needed(1);
        assert_eq!(
            Pin::new(&mut future).poll(&mut context),
            Poll::Ready(1),
            "a value sent before the first poll is waiting in the slot"
        );
        assert_eq!(wakes.0.load(Ordering::Relaxed), 0, "nobody was parked");

        // Arrives after the poll parked a waker -- the case the WebSocket
        // callbacks actually hit.
        let (sender, mut future) = one_shot::<u32>();
        assert_eq!(
            Pin::new(&mut future).poll(&mut context),
            Poll::Pending,
            "an empty slot parks"
        );
        sender.send_if_needed(2);
        assert_eq!(
            wakes.0.load(Ordering::Relaxed),
            1,
            "sending must wake the parked receiver, or it never polls again"
        );
        assert_eq!(Pin::new(&mut future).poll(&mut context), Poll::Ready(2));

        // `onopen` and `onerror` can both fire; only the first counts.
        let (sender, mut future) = one_shot::<u32>();
        sender.send_if_needed(3);
        sender.send_if_needed(4);
        assert_eq!(
            Pin::new(&mut future).poll(&mut context),
            Poll::Ready(3),
            "the second send must not overwrite the first"
        );
    }

    static NO_ARGS: &[ArgSpec] = &[];

    /// Emits three chunks and then finishes, so a test can tell partial output
    /// from a final response by their order on the queue.
    struct Streaming;

    impl Command for Streaming {
        fn name(&self) -> &'static str {
            "wasm_test_streaming"
        }
        fn short_description(&self) -> &'static str {
            "emits chunks"
        }
        fn full_description(&self) -> &'static str {
            "emits chunks"
        }
        fn args(&self) -> &'static [ArgSpec] {
            NO_ARGS
        }
        fn execute(&self, args: Vec<String>) -> Result<Response, Response> {
            self.execute_with(args, &CommandContext::detached())
        }
        fn execute_with(
            &self,
            _args: Vec<String>,
            context: &CommandContext,
        ) -> Result<Response, Response> {
            assert!(
                context.supports_streaming(),
                "a browser command must get a live context"
            );
            for index in 0..3 {
                context.emit(format!("chunk {index}")).expect("emit failed");
            }
            Ok("done".into())
        }
    }

    /// Runs until cancelled. Nothing else can end it, so a test that finishes
    /// has proved the `Cancel` arrived while this was still running.
    struct Forever;

    impl Command for Forever {
        fn name(&self) -> &'static str {
            "wasm_test_forever"
        }
        fn short_description(&self) -> &'static str {
            "runs until cancelled"
        }
        fn full_description(&self) -> &'static str {
            "runs until cancelled"
        }
        fn args(&self) -> &'static [ArgSpec] {
            NO_ARGS
        }
        fn execute(&self, args: Vec<String>) -> Result<Response, Response> {
            self.execute_with(args, &CommandContext::detached())
        }
        fn execute_with(
            &self,
            _args: Vec<String>,
            context: &CommandContext,
        ) -> Result<Response, Response> {
            loop {
                context.check_cancelled()?;
                wasm_lite_std::sleep(Duration::from_millis(5));
            }
        }
    }

    fn deliver(rpc: RPC) {
        handle_msg(&rmp_serde::to_vec(&rpc).unwrap()).unwrap();
    }

    /// The next message the socket worker would have written.
    fn next_rpc(receiver: &mpsc::Receiver<WorkerMessage>) -> RPC {
        loop {
            let deadline = Instant::now() + Duration::from_secs(20);
            match receiver.recv_sync_timeout(deadline) {
                Ok(WorkerMessage::Write(job)) => {
                    return rmp_serde::from_slice(&job[0]).expect("undecodable frame");
                }
                Ok(WorkerMessage::Reconnect) => continue,
                Err(error) => panic!("nothing was written: {error}"),
            }
        }
    }

    /// Everything the rework is for, in one test.
    ///
    /// One test rather than three because it takes the socket worker's receiver,
    /// which there is only one of. `(worker)` because it blocks waiting on that
    /// receiver, and blocking the *main* thread would stop the browser from
    /// starting the workers these commands run on.
    #[wasm_lite::wasm_lite_test(worker)]
    fn a_browser_command_streams_and_can_be_cancelled_while_it_runs() {
        crate::try_add_command(Streaming).ok();
        crate::try_add_command(Forever).ok();
        let receiver = SEND_WORKER_MESSAGE
            .1
            .with_mut_sync(|slot| slot.take())
            .expect("no test may have taken the receiver already");

        // Streaming: the chunks must arrive *before* the response, which is the
        // capability a detached context could not offer at all.
        deliver(RPC::Command(CommandInvocation::new(
            "wasm_test_streaming".to_string(),
            Vec::new(),
            1,
        )));
        let mut chunks = Vec::new();
        let response = loop {
            match next_rpc(&receiver) {
                RPC::Chunk(chunk) => {
                    assert_eq!(chunk.reply_id, 1);
                    chunks.push(chunk.payload.to_string());
                }
                RPC::CommandResponse(response) => break response,
                _ => {}
            }
        };
        assert_eq!(chunks, ["chunk 0", "chunk 1", "chunk 2"]);
        assert!(response.success, "{}", response.response);
        assert_eq!(response.response.to_string(), "done");

        // Cancellation: `wasm_test_forever` ends no other way, so reaching the
        // assertion at all proves the callback took a `Cancel` while a command
        // was running — which is what running inline made impossible.
        deliver(RPC::Command(CommandInvocation::new(
            "wasm_test_forever".to_string(),
            Vec::new(),
            2,
        )));
        deliver(RPC::Cancel(exfiltrate_internal::rpc::Cancel {
            reply_id: 2,
        }));
        let cancelled = loop {
            if let RPC::CommandResponse(response) = next_rpc(&receiver)
                && response.reply_id == 2
            {
                break response;
            }
        };
        assert!(!cancelled.success);
        assert!(
            cancelled.response.to_string().contains("cancelled"),
            "{}",
            cancelled.response
        );

        SEND_WORKER_MESSAGE
            .1
            .with_mut_sync(|slot| *slot = Some(receiver));
    }
}