breakmancer 0.4.0

//! Controller side of the protocol (run on the developer's laptop).

use std::{
    io::Write,
    net::{IpAddr, Ipv6Addr, SocketAddr, ToSocketAddrs},
    time::Duration,
};

use alkali::asymmetric::kx::x25519blake2b as X25519;
use base64::prelude::{Engine as _, BASE64_STANDARD_NO_PAD};
use clap::Parser;
use indoc::indoc;
use once_cell::sync::Lazy;
use regex::{self, Regex};
use rustyline_async::{Readline, ReadlineEvent};
use tokio::{net::TcpListener, signal};

use crate::{
    protocol::{Command, Connection, ExitBreak, OutputType, SecureMsg, SetupSecret, TerminateCmd},
    sigint, util,
};

/// Start a new session.
///
/// When the controller starts, it will give you a command to run on
/// the breakpoint host. The breakpoint end will try to open a connection
/// to the callback address, at which point a secure handshake
/// will occur and a debugging session can start.
#[derive(Parser, Debug)]
pub struct Cli {
    /// Callback address and port
    ///
    /// IP address or domain name that the developer's machine can be reached at,
    /// and the public-facing port to use. This must not be
    /// obstructed by NAT or firewall.
    ///
    /// IPV6 addresses must be enclosed in `[square brackets]`.
    #[arg(required = true, value_name = "HOST:PORT")]
    callback: String,

    /// Local port to listen on for callback.
    ///
    /// If not specified, defaults to the port from the callback
    /// option. Set this value if you use port-forwarding and have
    /// differing external and internal ports.
    #[arg(long)]
    local_port: Option<u16>,
}

struct Params {
    pub callback: String,
    pub listen: SocketAddr,
}

// Should work for listening on both IPv4 and IPv6.
const UNSPECIFIED_IP: IpAddr = IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0));

fn parse_args(args: &Cli) -> Result<Params, String> {
    // We can't really validate the callback address beyond checking
    // if it's well-formed; it's possible that the breakpoint needs to
    // be given a domain name that won't resolve from the controller's
    // vantage point.
    let (_, callback_port) = util::parse_address(&args.callback)
        .map_err(|err| format!("Invalid callback address: {err}"))?;
    // We'll still use it if it doesn't resolve, but at least warn the user.
    match &args.callback.to_socket_addrs() {
        Ok(addrs) => {
            if addrs.len() == 0 {
                eprintln!("[WARNING] Callback host did not resolve to any IP addresses.");
            }
        }
        Err(err) => eprintln!("[WARNING] Could not resolve callback host to an IP address: {err}"),
    };

    let listen_port = args.local_port.unwrap_or(callback_port);

    Ok(Params {
        callback: args.callback.to_string(),
        listen: SocketAddr::new(UNSPECIFIED_IP, listen_port),
    })
}

/// Print bytes as string (lossy Unicode conversion) and add a newline
/// on the end if there wasn't one already.
fn print_with_trailing_newline(data: &[u8]) {
    print!("{}", String::from_utf8_lossy(data));
    if !data.ends_with("\n".as_bytes()) {
        println!();
    }
}

const CMD_EXIT: &str = "exit";
static CMD_EXIT_PATTERN: Lazy<Regex> =
    Lazy::new(|| Regex::new(&format!("(?i)\\s*{CMD_EXIT}\\s*")).unwrap());

/// What the user did when we asked them for a command at the REPL.
enum ReplInput {
    /// A command the user wants to have run.
    Command(String),
    /// An interrupt of some sort, possibly signaling a desire to
    /// leave the REPL.
    Interrupted,
}

/// Prompt the user for a command at the REPL.
async fn ask_user_next_command(
    cmd_history: &mut Vec<String>,
    sigint_manager: &mut sigint::Manager,
) -> Result<ReplInput, String> {
    // Set up a new readline each time we need one, because
    // otherwise we can't use regular printlns everywhere. (It has
    // to be dropped in order to get the terminal out of raw mode
    // again.) Maybe
    // change to a persistent readline once we switch to a unified
    // logging system or something. The big downside is having to
    // maintain an external history and re-import it repeatedly.
    let (mut readline, rl_writer) = Readline::new(">> ".to_string())
        .map_err(|err| format!("Could not set up input prompt: {err}"))?;
    for cmd in &mut *cmd_history {
        readline.add_history_entry((*cmd).to_string());
    }

    // Suppress the default ^C handler we have set up.
    let suppress_sigint = sigint_manager.suppress();
    let readline_result = readline.readline().await;
    suppress_sigint.release();

    // Prevent dangling prompt
    readline.flush().unwrap();
    // Keep the writer open this long. Otherwise we get `Couldn't read
    // input: line writers closed` and the session keeps restarting.
    drop(rl_writer);

    match readline_result {
        Ok(ReadlineEvent::Line(command)) => Ok(ReplInput::Command(command)),
        Ok(ReadlineEvent::Eof | ReadlineEvent::Interrupted) => Ok(ReplInput::Interrupted),
        Err(err) => Err(format!("Couldn't read input: {err}")),
    }
}

/// Tell the breakpoint to exit (continue);
///
/// Consumes the connection, as it will stop working immediately anyhow.
async fn hangup(conn: Connection) -> Result<(), String> {
    conn.tx
        .send(SecureMsg::ExitBreak(ExitBreak {}))
        .await
        .map_err(|err| format!("Could not tell breakpoint to exit: {err}"))?;
    Ok(())
}

/// What the user wants to have happen after the session ends.
enum AfterSession {
    /// Breakpoint should exit, but controller should wait for a new connection.
    Iterate,
    /// Both breakpoint and controller should exit.
    Exit,
}

/// Listen for a connection and perform one shell session.
///
/// The `cmd_history` will be updated with new commands run in this
/// session.
///
/// An Ok return means the session ended normally, without an error
/// (that is, the user asked to end the session). The Ok arm indicates
/// what should happen next.
///
/// The connection is consumed; when the function ends, the connection
/// is dropped, which cancels its tasks.
async fn run_session(
    mut conn: Connection,
    cmd_history: &mut Vec<String>,
    sigint_manager: &mut sigint::Manager,
) -> Result<AfterSession, String> {
    println!("Session ready. You can enter single-line commands. Use `{CMD_EXIT}` to exit.");

    // Each iteration of the loop is a command/response interaction on the REPL.
    let mut cmd_seq = 0;
    loop {
        // Try to get a command to run, or get an interrupt/exit
        // request. Convert these to Some vs. None.
        let command: Option<String> =
            match ask_user_next_command(cmd_history, sigint_manager).await? {
                ReplInput::Command(command) => {
                    if command.is_empty() {
                        continue; // just re-show the prompt
                    }

                    if CMD_EXIT_PATTERN.is_match(&command) {
                        // Another way of exiting the session, besides
                        // ^C/^D. We'll treat it the same way as an
                        // interrupt.
                        None
                    } else {
                        Some(command)
                    }
                }
                ReplInput::Interrupted => None,
            };

        // If there was an interrupt/exit, figure out if they meant it
        // (and return) or whether they actually want to stay in the
        // REPL.
        let Some(command) = command else {
            match ask_user_exit_choice().await {
                ExitChoice::Resume => {
                    // User maybe interrupted by accident, and just
                    // wants to go back to entering commands.
                    continue;
                }
                ExitChoice::Iterate => {
                    hangup(conn).await?;
                    return Ok(AfterSession::Iterate);
                }
                ExitChoice::Exit => {
                    hangup(conn).await?;
                    return Ok(AfterSession::Exit);
                }
            }
        };

        // We actually have a command to run! Save it in the history
        // if it's new and not space-prefixed. (ignorespace,
        // ignoredups)
        if !command.starts_with(' ') && Some(&command) != cmd_history.last() {
            cmd_history.push(command.clone());
        }
        conn.tx
            .send(SecureMsg::Command(Command {
                command: command.clone(),
                seq: cmd_seq,
            }))
            .await
            .map_err(|err| format!("Could not send command to breakpoint: {err}"))?;

        // Once we send the message, we expect zero or more output
        // lines and then an outcome.

        // Whether we're currently trying to kill the remote command
        // (e.g. because it's hanging).
        let mut killing = false;
        loop {
            tokio::select! {
                // Always check the ^C first, to ensure it is handled
                // on a fast stream.
                biased;

                // We turn ^C into a termination command
                _ = signal::ctrl_c() => {
                    if !killing {
                        println!(); // get onto a new line after the ^C
                        println!("Attempting to terminate command...");
                        conn.tx.send(
                            SecureMsg::TerminateCmd(TerminateCmd { seq: cmd_seq })
                        ).await.map_err(|err| format!("Could not send termination command to breakpoint: {err}"))?;
                        killing = true;
                    }
                },

                // But otherwise just stream the output
                msg = conn.rx.recv() => match msg.ok_or("Error reading from breakpoint.")? {
                    SecureMsg::OutputLine(output) => {
                        if output.seq != cmd_seq {
                            return Err(format!(
                                "Received output for command #{} instead of #{}",
                                output.seq, cmd_seq
                            ));
                        }
                        let log_tag = match output.gender {
                            OutputType::Stdout => "out",
                            OutputType::Stderr => "err",
                        };
                        print!("[{log_tag}] ");
                        print_with_trailing_newline(&output.bytes);
                    }
                    SecureMsg::Finished(outcome) => {
                        if outcome.seq != cmd_seq {
                            return Err(format!(
                                "Received completion notification for command #{} instead of #{}",
                                outcome.seq, cmd_seq
                            ));
                        }
                        match outcome.exit_code {
                            Some(status) => println!("[exit: {status}]"),
                            None => println!("[exit: terminated by signal]"),
                        };
                        break;
                    }
                    // TODO: Info about what kind of message was received
                    _ => Err("Received unexpected message type from breakpoint")?,
                },
            };
        }

        cmd_seq += 1;
    }
}

/// User's confirmation of how they want to exit the REPL.
enum ExitChoice {
    /// They don't actually want to exit, and want to stay in the REPL.
    Resume,
    /// See [`AfterSession::Iterate`].
    Iterate,
    /// See [`AfterSession::Exit`].
    Exit,
}

/// The user has used an interrupt while the REPL was idle -- what do they want to do next?
async fn ask_user_exit_choice() -> ExitChoice {
    println!(indoc! {
        "You've left the command shell. Do you want to resume, or exit?

         [1] Resume: Go back to issuing commands to the breakpoint.
         [2] Iterate: Exit from breakpoint on remote end, running script normally, but leave controller running with same keys.
         [3] Exit: Exit from breakpoint and controller. For additional debugging you'll need to provision the breakpoint with a new key and secret."
    });
    println!();

    let (mut readline, mut rl_writer) = Readline::new("Choice [1/2/3]: ".to_string()).unwrap();
    let choice = loop {
        match readline.readline().await {
            Ok(ReadlineEvent::Line(choice)) => match choice.as_str() {
                "1" => break ExitChoice::Resume,
                "2" => break ExitChoice::Iterate,
                "3" => break ExitChoice::Exit,
                _ => continue,
            },
            Ok(ReadlineEvent::Eof) => continue,
            Ok(ReadlineEvent::Interrupted) => break ExitChoice::Exit,
            Err(err) => {
                rl_writer
                    .write_all(format!("Failed to read input, exiting: {err}\n").as_bytes())
                    .unwrap();
                break ExitChoice::Exit;
            }
        };
    };
    readline.flush().unwrap(); // prevent dangling prompt
    choice
}

/// Print connection instructions for the user.
fn print_connection_instructions(
    setup_secret: &SetupSecret,
    controller_pubkey: &X25519::PublicKey,
    params: &Params,
) {
    let setup_secret_b64 = &BASE64_STANDARD_NO_PAD.encode(setup_secret.as_slice());
    let controller_pubkey_b64 = &BASE64_STANDARD_NO_PAD.encode(controller_pubkey.as_slice());
    println!(
        indoc! {r#"On the host where you want to set a breakpoint, set the following environment variable in a secure way:

          BM_SETUP_SECRET={}

        (Warning: Anyone who can see the secret can connect to you and pretend to be the breakpoint system, so setting it in an encrypted vault or secret store is recommended. However, they will not be able to impersonate you to the breakpoint system.)

        Then, ensure the breakpoint end runs the following command from the environment you want to inspect:

          breakmancer break -c {} -i {}

        You can add --which="some location" to distinguish between breakpoints.

        (Now listening on {} for callback...)"#},
        setup_secret_b64, &params.callback, &controller_pubkey_b64, &params.listen,
    );
}

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

    // Allow override by debug keys
    let (controller_keypair, setup_secret) = match crate::debug_utils::get_debug_keys() {
        Some(keys) => (keys.controller_keypair, keys.setup_secret),
        None => (
            X25519::Keypair::generate().unwrap(),
            SetupSecret::new_random(),
        ),
    };

    // We need this signal manager for suppressing ^C during certain
    // operations.
    //
    // Except... that's not quite true. We actually *already* suppress
    // ^C at all times, because rustyline-async and tokio's ctrl_c
    // both register a signal handler, and that removes the default
    // action of ^C. So what this manager *actually* does is register
    // yet another handler that exits on ^C, and then provides a way
    // to *specifically* suppress that one handler for certain
    // operations.
    let mut sigint_manager = sigint::Manager::setup();

    let server_socket = TcpListener::bind(params.listen)
        .await
        .map_err(|err| format!("Could not listen on {}: {}", &params.listen, err))?;

    print_connection_instructions(&setup_secret, &controller_keypair.public_key, &params);

    // Share the readline (and its history) across sessions.
    let mut command_history = Vec::<String>::new();

    loop {
        let backoff = Duration::from_secs(10);
        // New connection created (and destroyed) on every loop
        // iteration. This minimizes lingering state to reduce the
        // opportunity for vulnerabilities or other bugs (including
        // leftover in-flight messages).
        let (conn, breakpoint_intro) = match Connection::new_controller(
            &server_socket,
            &controller_keypair,
            &setup_secret,
        )
        .await
        {
            Ok(conn) => conn,
            Err(err) => {
                println!(
                    "Failed to start controller connection; waiting {backoff:?} and trying again: {err}"
                );
                tokio::time::sleep(backoff).await;
                continue;
            }
        };
        if let Some(which) = breakpoint_intro.which {
            println!("Breakpoint: {which}");
        }

        match run_session(conn, &mut command_history, &mut sigint_manager).await {
            Ok(AfterSession::Iterate) => {
                println!("Will wait for new connection...");
            }
            Ok(AfterSession::Exit) => {
                return Ok(());
            }
            Err(err) => {
                println!("Encountered error, restarting connection: {err}");
            }
        }
    }
}