exfiltrate 0.2.4

// SPDX-License-Identifier: MIT OR Apache-2.0
//! WebAssembly worker and WebSocket implementation of the debug server.
use crate::wire::server::do_command;
use crate::wire::server::reconnect;
use exfiltrate_internal::rpc::RPC;
use std::sync::{Arc, LazyLock};
use std::time::Duration;
use wasm_lite::websocket::{BinaryType, CloseEvent, MessageEvent, WebSocket};
use wasm_lite::{Closure, JsValue, console};
use wasm_lite_std::Mutex;

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);
        }
    }
}

fn handle_msg(data: &[u8]) -> Result<RPC, String> {
    //parse as RPC
    match rmp_serde::from_slice(data) {
        Ok(msg) => {
            let msg: RPC = msg;
            match msg {
                RPC::Command(command) => {
                    let reply = do_command(command);
                    Ok(RPC::CommandResponse(reply))
                }
                RPC::CommandResponse(r) => Err(format!("Expected command, got: {:?}", r)),
                _ => Err("Unknown RPC variant received".to_string()),
            }
        }
        Err(e) => Err(format!("{:?}", e)),
    }
}

/// 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(receiver: continue_stream::Receiver<WorkerMessage>) {
    console::log("thread started");

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

    loop {
        let r = receiver.receive().await;

        match r {
            Some(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));
                        }
                    }
                }
            }
            None => {
                console::log("receiver closed, exiting thread");
                break;
            }
        }
    }
}

enum WorkerMessage {
    Reconnect,
}

#[allow(clippy::type_complexity)]
static SEND_WORKER_MESSAGE: LazyLock<(
    continue_stream::Sender<WorkerMessage>,
    Mutex<Option<continue_stream::Receiver<WorkerMessage>>>,
)> = LazyLock::new(|| {
    let (s, r) = continue_stream::continuation();
    (s, Mutex::new(Some(r)))
});

/// 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.
struct OneShot<T> {
    c: Arc<Mutex<Option<r#continue::Sender<T>>>>,
}

impl<T> OneShot<T> {
    /// Creates a new one-shot sender.
    fn new(sender: r#continue::Sender<T>) -> Self {
        OneShot {
            c: Arc::new(Mutex::new(Some(sender))),
        }
    }

    /// 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) {
        if let Some(sender) = self.c.with_mut_sync(|l| l.take()) {
            sender.send(value);
        }
    }
}

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

const WEB_ADDR: &str = "ws://localhost:1338";

async fn create_web_socket() -> Result<WebSocket, String> {
    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) = r#continue::continuation::<Result<(), String>>();
    let func_sender = OneShot::new(func_sender);
    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.
        wasm_lite_std::spawn_local(async move {
            let _ = debug_ws_handshake(WEB_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| {
        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));
        SEND_WORKER_MESSAGE.0.send(WorkerMessage::Reconnect);
    });
    ws.set_onclose(Some(onclose_callback.as_js_value()));
    onclose_callback.forget(); //leak the closure

    let move_ws = ws.clone();
    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;
        };

        match handle_msg(&data) {
            Ok(mut rpc) => {
                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();
                    let Ok(count) = u32::try_from(attachments.len()) else {
                        console::error("Response contains too many attachments");
                        return;
                    };
                    resp.num_attachments = count;
                }

                // Serialize the main message (now small)
                let msgpack_reply = match rmp_serde::to_vec(&rpc) {
                    Ok(reply) => reply,
                    Err(error) => {
                        console::error(&format!("Error serializing response: {error}"));
                        return;
                    }
                };
                if let Err(error) = move_ws.send_bytes(&msgpack_reply) {
                    console::error(&format!("Error sending response: {error:?}"));
                    return;
                }

                // Send attachments
                for attachment in attachments {
                    if let Err(error) = move_ws.send_bytes(&attachment) {
                        console::error(&format!("Error sending attachment: {error:?}"));
                        return;
                    }
                }
            }
            Err(e) => {
                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::*;

    /// 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
            ]
        );
    }
}