breakmancer 0.1.1

//! Breakpoint side of the protocol.

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

use alkali::asymmetric::kx::x25519blake2b as X25519;
use base64::prelude::{Engine as _, BASE64_STANDARD_NO_PAD};
use chrono::Utc;
use clap::Parser;
use tokio::{
    io::{AsyncBufReadExt, AsyncRead, BufReader},
    process::Command,
    sync::mpsc::{self, Receiver, Sender},
    time::{self, sleep},
};

use crate::{
    protocol::{
        Command as CommandMsg, Connection, Finished, OobSecret, OutputLine, OutputType, SecureMsg,
    },
    util,
};

#[derive(Parser, Debug)]
pub struct Cli {
    /// Callback address and port.
    ///
    /// IPV6 addresses must be enclosed in `[square brackets]`.
    #[arg(required = true, long, short = 'c', value_name = "HOST:PORT")]
    callback: String,

    /// Secret for encrypted channel.
    ///
    /// It is strongly advised to pass this via environment variable,
    /// as process arguments are generally visible to all users on a
    /// system.
    ///
    /// The nature of this secret depends on the protocol version.
    #[arg(required = true, long, env = "BM_CALLBACK_SECRET")]
    secret: String,

    /// Listener'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')]
    listener: 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,
    oob_secret: OobSecret,
    listener_public: X25519::PublicKey,
    which: Option<String>,
}

fn parse_args(args: &Cli) -> Result<Params, String> {
    util::parse_address(&args.callback)
        .map_err(|err| format!("Invalid callback address: {err}"))?;
    let oob_secret_bytes = BASE64_STANDARD_NO_PAD
        .decode(&args.secret)
        .map_err(|err| format!("Could not parse secret argument as Base64: {err}"))?;
    let oob_secret = OobSecret::from_bytes(&oob_secret_bytes);
    let listen_pub = BASE64_STANDARD_NO_PAD
        .decode(&args.listener)
        .map_err(|err| format!("Could not parse listener argument as Base64: {err}"))?;
    let listen_pub = <[u8; X25519::PUBLIC_KEY_LENGTH]>::try_from(listen_pub).map_err(|_| {
        format!(
            "Listener identity must be {} bytes of data",
            X25519::PUBLIC_KEY_LENGTH
        )
    })?;

    Ok(Params {
        callback: args.callback.to_string(),
        oob_secret,
        listener_public: listen_pub,
        which: args.which.clone(),
    })
}

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

// Spins up a thread that reads from the subprocess stdout or
// stderr and sends lines to a channel for sending over the wire.
fn spool_output<T>(source: T, output_type: OutputType, sink: Sender<SecureMsg>, seq: u32)
where
    T: AsyncRead + Send + Unpin + 'static,
{
    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 listener.
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 listener might ask us to terminate the command.
            listener_msg = conn.rx.recv() => match listener_msg {
                None => {
                    // Connection has been lost, so kill the child
                    // process before exiting (since the listener 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!(
                            "Listener requested termination of command #{} but we're running command #{}.",
                            term.seq, seq
                        ))?;
                    }
                    log("[DEBUG] Attempting to terminate the command at the listener's request...");
                    child.kill().await.map_err(|err| format!(
                        "Failed to kill child process at listener's request: {err}"
                    ))?;
                    log("Done.");
                    break;
                },
                Some(_) => {
                    // TODO: More informative (include type of message)
                    Err("Unexpected message from listener while running command")?;
                },
            },

            // The main stdout/stderr-sending branch.
            return_msg_maybe = msg_consumer.recv() => match return_msg_maybe {
                Some(msg) => conn.tx.send(msg)
                    .await
                    .map_err(|err| format!("Could not send command output to listener: {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 listener know the process finished.
    conn.tx
        .send(SecureMsg::Finished(Finished {
            exit_code: exit_status.status.code(),
            seq,
        }))
        .await
        .map_err(|err| format!("Could not send command exit code to listener: {err}"))
}

/// What to do after the connection ends.
struct AfterDisconnect {
    reconnect: bool,
}

async fn run_session(conn: &mut Connection) -> Result<AfterDisconnect, String> {
    log("Waiting for first command...");

    loop {
        let command_msg = match conn
            .rx
            .recv()
            .await
            .ok_or("Listener no longer sending messages")?
        {
            SecureMsg::Command(msg) => msg,
            SecureMsg::ExitBreak(_) => {
                eprintln!("Listener asked for normal execution to resume; exiting breakpoint.");
                return Ok(AfterDisconnect { reconnect: false });
            }
            // TODO: Info about what kind of message was received
            _ => Err("Received unexpected message from listener")?,
        };
        log(&format!("Running command: {}", command_msg.command));

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

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

    let backoff = time::Duration::from_secs(10);
    loop {
        // Allow override by debug keys
        //
        // Note that we create a new breakpoint keypair on every
        // iteration, just on the principle of clearing as much stale
        // state as possible.
        let (oob_secret, listener_public, breakpoint_keypair) =
            match crate::debug_utils::get_debug_keys() {
                Some(keys) => (
                    keys.oob_secret,
                    keys.listener_keypair.public_key,
                    keys.breakpoint_keypair,
                ),
                None => (
                    // TODO: Terrible, do something smarter.
                    OobSecret::from_bytes(params.oob_secret.as_slice()),
                    params.listener_public,
                    X25519::Keypair::generate().unwrap(),
                ),
            };

        let mut conn = match Connection::new_breakpoint(
            &breakpoint_keypair,
            &params.callback.to_string(),
            &listener_public,
            &oob_secret,
            &params.which,
        )
        .await
        {
            Ok(conn) => conn,
            Err(err) => {
                log(&format!(
                    "Failed to connect to listener; waiting {backoff:?} and trying again: {err}"
                ));
                sleep(backoff).await;
                continue;
            }
        };

        match run_session(&mut conn).await {
            Ok(after) => {
                // This branch is always taken, which is a little odd,
                // but probably better to have it explicit.
                if !after.reconnect {
                    return Ok(());
                }
            }
            Err(err) => log(&format!("Encountered error, restarting connection: {err}")),
        }
    }
}