breakmancer 0.2.0

//! Breakpoint side of the protocol.

use std::{env, 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, OutputLine, OutputType, SecureMsg,
        SetupSecret, SETUP_SECRET_LENGTH,
    },
    util,
};

/// Name of environment variable that contains the callback secret.
const SETUP_SECRET_ENV_NAME: &str = "BM_SETUP_SECRET";

/// Length of callback secret when Base64-encoded as expected.
const SETUP_SECRET_BASE64_LENGTH: usize = 22;

/// Break for debugging.
///
/// Call back to a `listen` session so that the developer can run
/// commands in this environment. This command and its parameters
/// will be provided by the output of the `listen` command.
///
/// `BM_SETUP_SECRET` environment variable must also be set
/// (provided by listen session's output).
#[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,

    /// 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,
    setup_secret: SetupSecret,
    listener_public: X25519::PublicKey,
    which: Option<String>,
}

/// Read callback secret from the environment.
fn read_setup_secret() -> Result<SetupSecret, String> {
    let secret = env::var(SETUP_SECRET_ENV_NAME)
        .map_err(|_| format!("Could not retrieve environment variable {SETUP_SECRET_ENV_NAME}"))?;
    // Checking length before decoding allows for friendlier error messages.
    if secret.len() != SETUP_SECRET_BASE64_LENGTH {
        Err(format!(
            "Callback secret must be {} Base64 chars; got {} chars.",
            SETUP_SECRET_BASE64_LENGTH,
            secret.len()
        ))?;
    }
    let secret_bytes = BASE64_STANDARD_NO_PAD
        .decode(&secret)
        .map_err(|err| format!("Could not parse secret argument as Base64: {err}"))?;
    let secret_bytes: [u8; SETUP_SECRET_LENGTH] = secret_bytes
        .try_into()
        .expect("Callback secret did not decode to the correct number of bytes");
    Ok(SetupSecret::from_bytes(&secret_bytes))
}

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

    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(),
        // The secret is only ever read from the environment, never
        // from an argument, since process arguments can be read by
        // other, unprivileged users on the same host.
        setup_secret: read_setup_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> {
    // 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 connecting to listener at {} and expecting listener identity of {}.{}",
        args.callback,
        args.listener,
        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 {
        // 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 (setup_secret, listener_public, breakpoint_keypair) =
            match crate::debug_utils::get_debug_keys() {
                Some(keys) => (
                    keys.setup_secret,
                    keys.listener_keypair.public_key,
                    keys.breakpoint_keypair,
                ),
                None => (
                    // Annoying to have to clone this...
                    SetupSecret::from_bytes(&params.setup_secret.clone()),
                    params.listener_public,
                    X25519::Keypair::generate().unwrap(),
                ),
            };

        let mut conn = match Connection::new_breakpoint(
            &breakpoint_keypair,
            &params.callback.to_string(),
            &listener_public,
            &setup_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}")),
        }
    }
}

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

    #[test]
    fn setup_secret_size_precomputed() {
        let expected = BASE64_STANDARD_NO_PAD
            .encode([0u8; SETUP_SECRET_LENGTH])
            .len();
        assert_eq!(expected, SETUP_SECRET_BASE64_LENGTH);
    }

    #[test]
    fn setup_secret_size_message() {
        let maybe_secret =
            temp_env::with_var(SETUP_SECRET_ENV_NAME, Some("Abc"), read_setup_secret);
        match maybe_secret {
            Err(msg) => assert_eq!(msg, "Callback secret must be 22 Base64 chars; got 3 chars."),
            Ok(_) => panic!("Expected arg parsing failure"),
        }
    }
}