exfiltrate 0.3.0

// SPDX-License-Identifier: MIT OR Apache-2.0
//! Tests for the command server.

use super::*;
use exfiltrate_internal::args::{ArgKind, ArgSpec};
use exfiltrate_internal::command::Command;
// Imported here rather than relying on the module's own imports, which are
// native-only: these tests run on wasm32 too.
use std::sync::atomic::AtomicBool;
use std::sync::{Arc, Mutex};

static STREAMING_ARGS: &[ArgSpec] = &[ArgSpec::flag("count", "how many chunks", ArgKind::Integer)];

struct Streaming;

impl Command for Streaming {
    fn name(&self) -> &'static str {
        "server_test_streaming"
    }
    fn short_description(&self) -> &'static str {
        "emits chunks"
    }
    fn full_description(&self) -> &'static str {
        "emits chunks"
    }
    fn args(&self) -> &'static [ArgSpec] {
        STREAMING_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> {
        for index in 0..3 {
            context.check_cancelled()?;
            let _ = context.emit(format!("chunk {index}"));
        }
        Ok("done".into())
    }
}

fn invoke(name: &str, context: &CommandContext) -> CommandResponse {
    do_command(
        CommandInvocation::new(name.to_string(), Vec::new(), 1),
        context,
    )
}

// Only reachable from the native-only panic-isolation test below.
#[cfg(not(target_arch = "wasm32"))]
struct Panicking;

#[cfg(not(target_arch = "wasm32"))]
impl Command for Panicking {
    fn name(&self) -> &'static str {
        "server_test_panicking"
    }
    fn short_description(&self) -> &'static str {
        "panics on purpose"
    }
    fn full_description(&self) -> &'static str {
        "panics on purpose"
    }
    fn execute(&self, _args: Vec<String>) -> Result<Response, Response> {
        panic!("boom")
    }
}

// Native only: `.cargo/config.toml` builds wasm32 with `panic_abort`, so
// invoking a panicking command there would abort the test binary rather
// than produce a failed response. That asymmetry is the point of the test —
// see `crate::panics` — and it cannot be exercised on the target that lacks
// the behaviour.
#[cfg(not(target_arch = "wasm32"))]
#[test]
fn a_panicking_command_fails_the_request_rather_than_the_connection() {
    crate::try_add_command(Panicking).ok();
    let previous = std::panic::take_hook();
    std::panic::set_hook(Box::new(|_| {}));
    let response = invoke("server_test_panicking", &CommandContext::detached());
    std::panic::set_hook(previous);

    assert!(!response.success);
    let message = response.response.to_string();
    assert!(message.contains("server_test_panicking"), "{message}");
    assert!(message.contains("boom"), "{message}");
}

#[test]
fn an_unknown_command_points_at_the_list() {
    let response = invoke("server_test_no_such_command", &CommandContext::detached());
    assert!(!response.success);
    let message = response.response.to_string();
    assert!(message.contains("command not found"), "{message}");
    assert!(message.contains("exfiltrate list"), "{message}");
}

#[test]
fn a_streaming_command_emits_chunks_before_its_final_response() {
    crate::try_add_command(Streaming).ok();
    let chunks = Arc::new(Mutex::new(Vec::new()));
    let sink_chunks = chunks.clone();
    let context = CommandContext::new(
        Arc::new(AtomicBool::new(false)),
        Arc::new(move |payload: Response| {
            sink_chunks.lock().unwrap().push(payload.to_string());
            Ok(())
        }),
    );

    let response = invoke("server_test_streaming", &context);
    assert!(response.success);
    assert_eq!(response.response.to_string(), "done");
    assert_eq!(*chunks.lock().unwrap(), ["chunk 0", "chunk 1", "chunk 2"]);
}

#[test]
fn a_cancelled_command_stops_and_says_so() {
    crate::try_add_command(Streaming).ok();
    let cancelled = Arc::new(AtomicBool::new(true));
    let context = CommandContext::new(cancelled, Arc::new(|_| Ok(())));
    let response = invoke("server_test_streaming", &context);
    assert!(!response.success);
    assert!(
        response.response.to_string().contains("cancelled"),
        "{}",
        response.response
    );
}

#[cfg(not(target_arch = "wasm32"))]
#[test]
fn a_reachable_address_invents_a_token_rather_than_refusing_to_listen() {
    use exfiltrate_internal::transport::Address;

    let public = Address::parse("0.0.0.0:1337").unwrap();
    let Credential::Generated(token) = credential_for(&public, None).unwrap() else {
        panic!("binding every interface with no credential is the exposure this prevents")
    };
    assert!(!token.is_empty());
    // Two servers must not share one, which is the whole reason it is
    // generated rather than derived from anything.
    let Credential::Generated(second) = credential_for(&public, None).unwrap() else {
        panic!("expected a generated token")
    };
    assert_ne!(token, second);
}

#[cfg(not(target_arch = "wasm32"))]
#[test]
fn a_configured_token_is_honoured_everywhere_and_a_local_address_needs_none() {
    use exfiltrate_internal::transport::Address;

    // Even where the address would not have required one: setting a token
    // and having it quietly ignored would be the surprise.
    assert_eq!(
        credential_for(
            &Address::parse("127.0.0.1:1337").unwrap(),
            Some("TOKEN".to_string())
        )
        .unwrap(),
        Credential::Configured("TOKEN".to_string())
    );
    assert_eq!(
        credential_for(&Address::parse("127.0.0.1:1337").unwrap(), None).unwrap(),
        Credential::None
    );
    // The filesystem is the access control for these two, so neither needs
    // a token to be safe to bind.
    #[cfg(unix)]
    {
        assert_eq!(
            credential_for(&Address::parse("unix:/tmp/x.sock").unwrap(), None).unwrap(),
            Credential::None
        );
        assert_eq!(
            credential_for(&Address::parse("fd:3").unwrap(), None).unwrap(),
            Credential::None
        );
    }
}

#[test]
fn a_registered_command_is_found_by_name_rather_than_by_position() {
    crate::commands::register_commands(&crate::Config::default());
    let response = invoke("list", &CommandContext::detached());
    assert!(response.success);
}