breakmancer 0.5.0

//! Breakpoint side of the protocol.

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

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

use crate::{
    protocol::{
        self, AuthVersion, BreakpointConnectError, Command as CommandMsg, Connection, Finished,
        OutputLine, OutputType, SecureMsg, XWingKeypair,
    },
    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,

    /// Callback address and port.
    ///
    /// IPV6 addresses must be enclosed in `[square brackets]`.
    #[arg(required = true, long, short = 'c', value_name = "HOST:PORT")]
    callback: 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 {
    callback: String,
    controller_pubkey_digest: String,
    which: Option<String>,
}

fn parse_args(args: &Cli) -> Result<Params, String> {
    util::parse_address(&args.callback)
        .map_err(|err| format!("Invalid callback address: {err}"))?;

    Ok(Params {
        callback: args.callback.to_string(),
        controller_pubkey_digest: args.controller.clone(),
        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> {
    // For audit log purposes, print the (public) connection
    // parameters. This is important in case some of them were passed
    // in as variables.
    log(&format!(
        "Breakpoint will attempt to connect to controller at '{}', expecting controller identity of '{}'.{}",
        args.callback,
        args.controller,
        match &args.which {
            None => String::new(),
            Some(which) => format!(" Breakpoint identity: \"{which}\""),
        },
    ));

    let params = parse_args(args)?;

    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_digest, breakpoint_keypair) = match crate::debug_utils::get_debug_keys() {
            Some(keys) => (
                protocol::controller_digest(&keys.controller_keypair.public_key),
                keys.breakpoint_keypair,
            ),
            None => (params.controller_pubkey_digest.clone(), XWingKeypair::new()),
        };

        let conn = match protocol::breakpoint_open_connection(
            &breakpoint_keypair,
            &params.callback.to_string(),
            &controller_digest,
            params.which.as_ref(),
        )
        .await
        {
            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}")),
        }
    }
}