breakmancer 0.9.0

//! Breakpoint side of the protocol.

use std::{ops::RangeInclusive, process::Stdio, time::Duration};

use chrono::Utc;
use clap::{Parser, ValueEnum};
use tokio::{
    io::{AsyncBufReadExt, AsyncRead, BufReader},
    process::Command,
    sync::mpsc::{self, Receiver, Sender},
    time,
};

use crate::{
    cli::ConnectVia,
    protocol::{
        self, AuthVersion, BreakpointConnectError, Command as CommandMsg, Connection, ControllerId,
        Finished, IdString, OutputLine, OutputType, SecureMsg, XWingKeypair,
    },
    transport::{tcp::TcpConnectParams, ConnectMethod},
    util,
};

/// Break for debugging.
///
/// Call back to a controller session so that the developer can run
/// commands in this environment. This command and its parameters
/// will be provided by the output of the Controller.
#[derive(Parser, Debug)]
pub struct Cli {
    /// Authentication protocol.
    ///
    /// Only one supported value: `dual-pub`
    #[arg(long, value_enum)]
    auth: AuthVersion,

    /// Transport methods.
    ///
    /// Comma-delimited list of transport methods accepted by the
    /// controller. In this version of breakmancer, the breakpoint
    /// will always pick the first option that it supports.
    ///
    /// Each transport method is a name, a colon, and an (inclusive)
    /// range of acceptable versions. A range can be a single number
    /// or two numbers separated by a hyphen.
    ///
    /// This might look like `tcp:1,xxx:3-4` (with nonexistent transport
    /// method 'xxx' used for sake of example).
    #[arg(long = "via")]
    transport: String,

    /// Callback address and port when using TCP transport.
    ///
    /// IPV6 addresses must be enclosed in `[square brackets]`.
    #[arg(long = "tcp-call", value_name = "HOST:PORT")]
    tcp_call: Option<String>,

    /// Controller's identity.
    ///
    /// Identity of the party initiating the session. This is used in
    /// setting up the encrypted channel, and the exact nature of the
    /// value depends on the protocol version.
    #[arg(required = true, long, short = 'i')]
    controller: String,

    /// Which breakpoint is this?
    ///
    /// Any identifying string to help distinguish between multiple
    /// breakpoints.
    #[arg(long, short = 'w')]
    which: Option<String>,
}

struct Params {
    connect: ConnectMethod,
    controller_id: ControllerId,
    which: Option<String>,
}

/// Parse a string as a version number (small unsigned int), or fail.
fn parse_version_number(num: &str) -> Result<u8, ()> {
    if num.bytes().all(|b| b.is_ascii_digit()) && !num.is_empty() {
        num.parse().map_err(|_err| ())
    } else {
        Err(())
    }
}

/// Parse a string as a closed range, or fail.
///
/// Range can either be in `X` or `X-Y` format, meaning `X..=X` or `X..=Y`.
fn parse_range(range: &str) -> Result<RangeInclusive<u8>, ()> {
    match range.split_once('-') {
        Some((left, right)) => {
            let left = parse_version_number(left).map_err(|_err| ())?;
            let right = parse_version_number(right).map_err(|_err| ())?;
            if left <= right {
                Ok(left..=right)
            } else {
                Err(())
            }
        }
        None => {
            let only = parse_version_number(range).map_err(|_err| ())?;
            Ok(only..=only)
        }
    }
}

/// Select a connection method and version from the provided CLI arg.
fn select_known_connect_method(methods: &str) -> Result<Option<ConnectVia>, String> {
    for option in methods.split(',') {
        let (option_name, version_range) = option
            .split_once(':')
            .ok_or(format!("Missing range delimiter ':' in option: '{option}'"))?;
        let version_range = parse_range(version_range)
            .map_err(|_err| format!("Cannot parse version range in '{option}'"))?;

        // This hardcodes knowledge of available tcp and wormhole
        // versions, which is fine while there's only one of each.
        match ConnectVia::from_str(option_name, false) {
            Ok(ConnectVia::Tcp) => {
                if version_range.contains(&1) {
                    return Ok(Some(ConnectVia::Tcp));
                } else {
                    println!("Skipping {option} because it did not include known versions.");
                    continue;
                }
            }
            Ok(ConnectVia::Wormhole) => {
                if version_range.contains(&1) {
                    return Ok(Some(ConnectVia::Wormhole));
                } else {
                    println!("Skipping {option} because it did not include known versions.");
                    continue;
                }
            }

            Err(_) => {
                println!("Skipping unrecognized transport option '{option_name}'");
                continue;
            }
        }
    }
    Ok(None)
}

/// Turn CLI args into run parameters.
fn parse_args(args: &Cli) -> Result<Params, String> {
    let connect_type = select_known_connect_method(&args.transport)?.ok_or(
        "No recognized transport option was provided. (Version mismatch between controller and breakpoint?)".to_string()
    )?;
    let connect = match connect_type {
        ConnectVia::Tcp => {
            let callback_addr = args
                .tcp_call
                .as_ref()
                .ok_or(String::from("TCP callback address was not specified."))?;
            // TODO: Move into TcpConnectParams constructor?
            util::parse_address(callback_addr)
                .map_err(|err| format!("Invalid callback address: {err}"))?;
            ConnectMethod::Tcp(TcpConnectParams {
                callback_addr: callback_addr.clone(),
                local_port: None,
            })
        }
        ConnectVia::Wormhole => ConnectMethod::Wormhole,
    };

    Ok(Params {
        connect,
        controller_id: ControllerId::from_received_string(&args.controller)
            .map_err(|err| format!("Not a valid verification string: {err}"))?,
        which: args.which.clone(),
    })
}

fn log(msg: &str) {
    let timestamp = Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true);
    println!("[{timestamp}] {msg}");
}

/// Start a task that reads from the subprocess stdout or
/// stderr and sends lines to a channel for sending over the network.
fn spool_output<T>(source: T, output_type: OutputType, sink: Sender<SecureMsg>, seq: u32)
where
    T: AsyncRead + Send + Unpin + 'static,
{
    // Cancellation: This task is dropped immediately and detaches,
    // but will terminate when the other end of the channel closes.
    tokio::spawn(async move {
        let mut reader = BufReader::new(source);
        loop {
            let mut buf = Vec::new();
            let num_read = reader.read_until(0x0A, &mut buf).await.unwrap();
            if num_read == 0 {
                return;
            }

            let out_msg = SecureMsg::OutputLine(OutputLine {
                bytes: buf,
                gender: output_type.clone(),
                seq,
            });
            if sink.send(out_msg).await.is_err() {
                // Message transmitter has closed early (child
                // process has probably been killed).
                return;
            }
        }
    });
}

/// Execute the command and send its results to the controller.
async fn execute_and_stream_results(
    conn: &mut Connection,
    command_msg: &CommandMsg,
) -> Result<(), String> {
    let mut child = Command::new("bash")
        .arg("-c")
        .arg(&command_msg.command)
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
        .unwrap();
    let child_out = child.stdout.take().unwrap();
    let child_err = child.stderr.take().unwrap();

    // We'll read stdout and stderr from the process in separate
    // threads and send them a line at a time into this channel; a
    // loop will receive both and actually send them over the wire.
    //
    // `Connection` itself can't be used by the threads, since the
    // cipher_stream structs are !Send.
    let (stdout_producer, mut msg_consumer): (Sender<SecureMsg>, Receiver<SecureMsg>) =
        mpsc::channel(32);
    let stderr_producer = stdout_producer.clone();

    let seq = command_msg.seq;

    // Each of these threads reads stdout or stderr until EOF and then
    // exits (and drops the Sender).
    //
    // TODO: Consider whether this is the ideal approach for reading
    // stdout and stderr. There's no inherent way to keep them in
    // sync. Reading each one as fast as possible may help, but if one
    // is way more active than the other then it may get desynced. If
    // the child process only emits full lines to one at a time, then
    // keeping a channel with a buffer capacity of 1 would be ideal --
    // but it might make a different situation worse.
    spool_output(child_out, OutputType::Stdout, stdout_producer, seq);
    spool_output(child_err, OutputType::Stderr, stderr_producer, seq);

    // Limit how fast we spool out the output. This is important for
    // not flooding the connection in case a termination command comes
    // down the line.
    let mut interval = time::interval(Duration::from_millis(10));
    interval.set_missed_tick_behavior(time::MissedTickBehavior::Delay);
    loop {
        tokio::select! {
            // Always check the termination branch first, to ensure it
            // is handled on a fast stream.
            biased;

            // The controller might ask us to terminate the command.
            controller_msg = conn.recv() => match controller_msg? {
                None => {
                    // Connection has been lost, so kill the child
                    // process before exiting (since the controller will
                    // no longer have control over it.)
                    log("[DEBUG] Attempting to terminate the command due to lost connection...");
                    child.kill().await.map_err(|err| format!(
                        "Failed to kill child process after connection was lost: {err}"
                    ))?;
                    Err("Lost connection while running command.")?;
                },
                Some(SecureMsg::TerminateCmd(term)) => {
                    if term.seq != seq {
                        Err(format!(
                            "Controller requested termination of command #{} but we're running command #{}.",
                            term.seq, seq
                        ))?;
                    }
                    log("[DEBUG] Attempting to terminate the command at the controller's request...");
                    child.kill().await.map_err(|err| format!(
                        "Failed to kill child process at controller's request: {err}"
                    ))?;
                    log("Done.");
                    break;
                },
                Some(wrong_variant) => {
                    Err(format!(
                        "Received unexpected message type while running command: {}",
                        Into::<&str>::into(wrong_variant)
                    ))?;
                },
            },

            // The main stdout/stderr-sending branch.
            return_msg_maybe = msg_consumer.recv() => match return_msg_maybe {
                Some(msg) => conn.send(msg)
                    .await
                    .map_err(|err| format!("Could not send command output to controller: {err}"))?,
                None => break,
            },
        }

        interval.tick().await;
    }

    // TODO: Check whether this is correct -- should the wait happen
    // earlier? Or are we guaranteed the process has already exited
    // anyhow, since stdout and stderr have reached EOF?
    let exit_status = child.wait_with_output().await.unwrap();
    // Some debug code to see if this situation ever happens.
    let extra_stdout = exit_status.stdout.len();
    let extra_stderr = exit_status.stderr.len();
    if extra_stdout > 0 || extra_stderr > 0 {
        eprintln!(
            "Possible bug: There was still data in stdout ({extra_stdout} bytes) or stderr ({extra_stderr} bytes) by the time we called wait on the process."
        );
    }

    // Send exit code and let controller know the process finished.
    conn.send(SecureMsg::Finished(Finished {
        exit_code: exit_status.status.code(),
        seq,
    }))
    .await
    .map_err(|err| format!("Could not send command exit code to controller: {err}"))
}

/// Run one session to completion.
///
/// An Ok return indicates the user asked for the connection to end
/// and the breakpoint to exit.
///
/// The connection is consumed; when the function ends, the connection
/// is dropped, which cancels its tasks.
async fn run_session(mut conn: Connection) -> Result<(), String> {
    log("Waiting for first command...");

    loop {
        let command_msg = match conn
            .recv()
            .await?
            .ok_or("Controller no longer sending messages")?
        {
            SecureMsg::Command(msg) => msg,
            SecureMsg::ExitBreak(_) => {
                eprintln!("Controller asked for normal execution to resume; exiting breakpoint.");
                return Ok(());
            }
            wrong_variant => Err(format!(
                "Received unexpected message type: {}",
                Into::<&str>::into(wrong_variant)
            ))?,
        };
        log(&format!("Running command: {}", command_msg.command));

        execute_and_stream_results(&mut conn, &command_msg).await?;
    }
}

pub async fn run(args: &Cli) -> Result<(), String> {
    let params = parse_args(args)?;

    // For audit log purposes, print the (public) connection
    // parameters. This is important in case some of them were passed
    // in as variables.
    let describe_call = match &params.connect {
        ConnectMethod::Tcp(params) => format!("at '{}'", params.callback_addr),
        ConnectMethod::Wormhole => String::from("via magic-wormhole"),
    };
    log(&format!(
        "Breakpoint will attempt to connect to controller {}, expecting controller identity of '{}'.{}",
        describe_call,
        params.controller_id.id_string(),
        match &params.which {
            None => String::new(),
            Some(which) => format!(" Breakpoint identity: \"{which}\""),
        },
    ));

    let backoff = time::Duration::from_secs(10);
    loop {
        // Create a new keypair on each iteration, on the principle of
        // clearing as much stale state as possible.
        //
        // Allow override by debug keys.
        let (controller_id, breakpoint_keypair) = match crate::debug_utils::get_debug_keys() {
            Some(keys) => (
                ControllerId::from_key(&keys.controller_keypair.public_key),
                keys.breakpoint_keypair,
            ),
            None => (params.controller_id.clone(), XWingKeypair::new()),
        };

        let caller = params.connect.new_caller(&controller_id);

        // Awkward. Maybe a try-block will clean this up once that stabilizes.
        let maybe_conn = {
            // Start setting up the connection
            match protocol::breakpoint_open_connection(
                &breakpoint_keypair,
                &caller,
                &controller_id,
                params.which.clone(),
            )
            .await
            {
                Ok(half_conn) => {
                    // In user-facing messaging we use the more
                    // verbose "verification" language; internally we
                    // use "ID" because honestly it's just shorter.
                    println!("This is the verification string for the breakpoint:");
                    println!();
                    println!("    {}", half_conn.breakpoint_id.id_string());
                    println!();
                    println!("Copy that string into the breakmancer controller when prompted.");

                    // Finish setting up the connection.
                    half_conn.continue_after_id_printed().await
                }
                Err(err) => Err(err),
            }
        };

        let conn = match maybe_conn {
            Ok(conn) => conn,
            Err(BreakpointConnectError { message, permanent }) => {
                log(&format!("Failed to connect to controller: {message}"));
                if permanent {
                    return Err("Not attempting a new connection.".to_string());
                }

                log(&format!("Waiting {backoff:?} before trying again.",));
                time::sleep(backoff).await;
                continue;
            }
        };

        match run_session(conn).await {
            Ok(()) => {
                // User asked for breakpoint to resume (so now we do a clean exit).
                return Ok(());
            }
            Err(err) => log(&format!("Encountered error, restarting connection: {err}")),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parsing_transport_version() {
        parse_version_number("").unwrap_err();
        parse_version_number("+1").unwrap_err();
        parse_version_number("-1").unwrap_err();
        parse_version_number("5a").unwrap_err();

        assert_eq!(parse_version_number("1").unwrap(), 1);
        assert_eq!(parse_version_number("203").unwrap(), 203);
    }

    #[test]
    fn parsing_transport_range() {
        parse_range("").unwrap_err();
        parse_range("1-").unwrap_err();
        parse_range("-2").unwrap_err();
        parse_range("1:2").unwrap_err();
        parse_range("20-3").unwrap_err(); // reversed

        assert_eq!(parse_range("10").unwrap(), 10..=10);
        assert_eq!(parse_range("10-10").unwrap(), 10..=10);
        assert_eq!(parse_range("1-4").unwrap(), 1..=4);
    }

    #[test]
    fn parsing_transport_options() {
        assert_eq!(
            select_known_connect_method(""),
            Err("Missing range delimiter ':' in option: ''".to_string())
        );
        assert_eq!(
            select_known_connect_method("tcp-1"),
            Err("Missing range delimiter ':' in option: 'tcp-1'".to_string())
        );
        assert_eq!(
            select_known_connect_method("tcp:x"),
            Err("Cannot parse version range in 'tcp:x'".to_string())
        );

        assert_eq!(
            select_known_connect_method("tcp:1"),
            Ok(Some(ConnectVia::Tcp))
        );
        assert_eq!(
            select_known_connect_method("xxx:1,tcp:1,wormhole:55-66"),
            Ok(Some(ConnectVia::Tcp))
        );
        assert_eq!(select_known_connect_method("wormhole:55-66"), Ok(None));
        assert_eq!(
            select_known_connect_method("wormhole:1-66"),
            Ok(Some(ConnectVia::Wormhole))
        );
    }
}