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::sync::atomic::{AtomicBool, AtomicU64, Ordering};
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;
use wasm_lite_std::mpsc;
const MAX_CONCURRENT_COMMANDS: usize = 8;
const EVENT_POLL_INTERVAL: Duration = Duration::from_millis(250);
type WriteJob = Vec<Vec<u8>>;
impl reconnect::ProxySocket for WebSocket {
fn ready_state(&self) -> u16 {
WebSocket::ready_state(self)
}
}
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);
}
}
}
static SUBSCRIBER: LazyLock<Mutex<Option<u64>>> = LazyLock::new(|| Mutex::new(None));
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
}
})
}
fn detach_subscriber() {
if let Some(id) = SUBSCRIBER.with_mut_sync(|slot| slot.take()) {
crate::events::detach(id);
}
}
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)
}
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())
}
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;
}
}
}
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));
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}"
));
}
flush_events();
});
if spawned.is_err() {
RUNNING.with_mut_sync(|running| running.remove(&reply_id));
return false;
}
true
}
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) => {
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())
}
_ => {
console::log("exfiltrate: ignoring an RPC variant this build does not know");
Ok(())
}
};
result?;
flush_events();
Ok(())
}
pub async fn debug_ws_handshake(ws_url: &str) -> Result<(), JsValue> {
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");
opts.set_mode("cors");
match wasm_lite::fetch::fetch(&http_url, &opts).await {
Err(err) => {
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(())
}
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 {
let deadline = wasm_lite_std::time::Instant::now() + EVENT_POLL_INTERVAL;
match receiver.recv_async_timeout(deadline).await {
Ok(WorkerMessage::Reconnect) => {
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));
}
}
}
}
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(),
Err(_) => {
console::log("receiver closed, exiting thread");
break;
}
}
}
}
enum WorkerMessage {
Reconnect,
Write(WriteJob),
}
#[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)))
});
static RUNNING: LazyLock<Mutex<HashMap<u32, Arc<AtomicBool>>>> =
LazyLock::new(|| Mutex::new(HashMap::new()));
struct OneShot<T> {
c: Arc<Mutex<Option<r#continue::Sender<T>>>>,
}
impl<T> OneShot<T> {
fn new(sender: r#continue::Sender<T>) -> Self {
OneShot {
c: Arc::new(Mutex::new(Some(sender))),
}
}
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),
}
}
}
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,
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();
let move_func_sender = func_sender.clone();
let onerror_callback = Closure::new_with_arg(move |_event| {
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();
let onclose_callback = Closure::new_with_arg(move |event| {
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();
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();
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;
#[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");
}
static NO_ARGS: &[ArgSpec] = &[];
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())
}
}
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();
}
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}"),
}
}
}
#[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");
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");
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));
}
}