breakmancer 0.3.0

//! Network and encryption protocol.
//!
//! # Summary
//!
//! The breakmancer protocol consists of a cryptographic handshake
//! followed by a sequence of command/response messages.
//! This implements a secured reverse shell.
//!
//! The protocol follows a client/server model, although in this case the
//! "client" is usually run on what would normally be termed a server
//! (likely a CI/CD server), and
//! the "server" is probably running on the user's development laptop; this is a
//! *reverse* shell, after all. Due to this ambiguity, more specific
//! functional terms will be used here instead.
//!
//! The user first runs `breakmancer listen` to start a session on their
//! development machine; that process will be termed the "Listener". It
//! begins listening on a port and prints a `breakmancer break` command to
//! run on the machine to be debugged.
//!
//! The user then copies this break command into the script that needs
//! debugging. When that line of the script runs, the new breakmancer
//! process reaches back to the Listener, performs a handshake, and waits
//! for further instructions. The new process will be termed the
//! "Breakpoint".
//!
//! The user is then presented with a simple shell-like environment in the
//! Listener session. They can enter commands to be run. These commands
//! are sent to the Breakpoint, which executes them and streams the output
//! back to the Listener for display.
//!
//! # Constraints
//!
//! The Breakpoint may be in a script running in a public environment,
//! such as a CI runner for an open source project. This means it may not
//! always be possible to keep secrets on that side. The protocol will
//! lose some security properties in this scenario, but should not lose
//! all of them.
//!
//! The Breakpoint's output may also not be visible in realtime, so the
//! protocol cannot rely on techniques such as having the user copy
//! session keys from Breakpoint's output to the Listener.
//!
//! # Goals
//!
//! - The Breakpoint must only execute commands specified by the Listener,
//!   without alteration (mutation, replay, reordering, etc.) It must not
//!   be possible for another party to impersonate the Listener, even an
//!   on-path attacker.
//! - The Listener must be assured that all outputs are truly from the
//!   Breakpoint. No party should be able to impersonate the Breakpoint.
//! - Only the Breakpoint should be able to read or infer the contents
//!   of commands. This includes defense against:
//!     - Timing attacks, e.g. inference of text via keystroke timing
//!       during interactive input. (Protocol is currently line-oriented,
//!       though.) **TODO:** Implement buffering and timing jitter?
//!     - Plaintext length leakage, to the extent feasible. **TODO:**
//!       Implement padding.
//! - Confidentiality of traffic must be maintained even if it is stored
//!   and later attacked with a quantum computer.
//! - The Breakpoint may run in a low-confidentiality environment and may
//!   not have secure secrets storage.
//!     - If the setup channel is public (i.e., the setup secret is
//!       exposed to attackers), then the worst outcome should be that the
//!       attacker can impersonate the Breakpoint to the Listener. That
//!       is, the attacker may be able to read the Listener's commands and
//!       return false outputs, but *not* be able to send commands to the
//!       Breakpoint or read the real outputs.
//!     - In the case of a leaked secret *and* later traffic analysis by a
//!       quantum computer, it is expected that the real outputs also lose
//!       confidentiality.
//!
//! # Non-goals
//!
//! - Availability, in the face of an adversary.
//!
//! # Cryptographic protocol
//!
//! At a high level, the Listener and Breakpoint use X25519 to derive two
//! stream pre-keys (one for each direction). The pre-keys are then each
//! hashed with a shared secret to derive the actual ChaCha20-Poly1305
//! stream keys.
//!
//! Details:
//!
//! - The Listener creates an X25519 keypair `listen_pub, listen_priv` and
//!   a shared secret `setup_secret`.
//! - `listen_pub` and `setup_secret` are displayed to the user and
//!   conveyed securely to the Breakpoint.
//! - The Breakpoint creates its own X25519 keypair `break_pub,
//!   break_priv` and sends `break_pub` to the Listener over the network.
//! - Both parties perform X25519 key exchange to derive `pre_rx` and
//!   `pre_tx` receive/transmit pre-keys.
//! - For each pre-key, each party uses a KDF to compute the stream keys:
//!
//!     - The KDF input pieces are:
//!         - ASCII encoding of `breakmancer.streamkey` for domain separation
//!         - ASCII encoding of the protocol version (currently `dev.2`) for
//!           protocol binding
//!         - `listen_pub` and `break_pub` to tie in the public keys
//!         - `pre_rx` or `pre_tx`
//!    - Each KDF input piece is length-prefixed (little-endian 64-bit)
//!    - The full KDF input is fed to BLAKE2b, keyed with `setup_secret`
//!      and with an output size of 32 bytes.
//!
//!   The output is used as a XChaCha20-Poly1305 receive/transmit key.
//!   `rx = KDF(pre_rx)` and `tx = KDF(pre_tx)`.
//! - The parties begin XChaCha20-Poly1305 streams, exchanging headers and
//!   then sending encrypted messages.
//!
//! # Network layer
//!
//! Communication occurs over TCP. Messages are sent length-prefixed
//! (4-byte little-endian).
//!
//! # Insecure message layer
//!
//! Each raw message is a CBOR map with one key (naming the message type)
//! and a type-specific value. There are three types: `BreakpointHello`,
//! `ListenerHello`, and `EncryptedMsg`. Several of these can contain
//! encrypted message data, an XChaCha20-Poly1305 message as bytes.
//!
//! # Secure message layer
//!
//! When an encrypted message is decrypted, it again takes the form of a
//! CBOR map of one key. The types of these are: `BreakpointIntro`,
//! `Command`, `OutputLine`, `Finished`, `TerminateCmd`, and `ExitBreak`.
//!
//! # Application protocol
//!
//! This includes some recap of the cryptographic protocol.
//!
//! Any time there is a protocol error (including a message type that is
//! received in an unexpected context) or network fault (such as a closed
//! TCP connection), that party will drop the connection and start over at
//! the beginning of the handshake.
//!
//! ## Initial, out-of-band setup
//!
//! - The developer runs the `listen` command on their workstation. They
//!   are the "Listener" party. The Listener creates X25519 keypair
//!   `listen_pub, listen_priv` as well as a shared secret `setup_secret`.
//! - The listen command outputs a `break` command to be run in the system
//!   being inspected. When that command runs, it is the "Breakpoint"
//!   party. The command contains `listen_pub` and the address of the
//!   Listener, and the user is told to put `setup_secret` in the
//!   command's environment as securely as possible.
//!
//! ## Handshake
//!
//! - The Breakpoint creates X25519 keypair `break_pub, break_priv`.
//! - The Breakpoint derives its stream keys and its transmit stream's header.
//! - The Breakpoint sends a `BreakpointHello` insecure message containing
//!   `break_pub`, its transmit stream's header, and an encrypted
//!   `BreakpointIntro` message. (The intro just contains information
//!   about *which* breakpoint is running, in case there are multiple.)
//! - The Listener receives the `BreakpointHello` and can now derive its
//!   own stream keys and decrypt and display the intro. It sends a
//!   `ListenerHello` with its own transmit stream's header.
//!
//! At this point, both parties have a secure channel, and all future
//! messages are sent encrypted and wrapped in `EncryptedMsg`-type
//! messages.
//!
//! ## Main loop: Listener
//!
//! Each time the user enters a command. the Listener sends a `Command`
//! message (a command string plus a sequence number that increments for
//! each command). It then switches to waiting for:
//!
//! - `OutputLine` messages bearing lines of stdout or stderr for
//!   display (expecting a matching sequence number.)
//! - A `Finished` message, at which point it increments the command
//!   sequence number and goes back to waiting for user input.
//! - An interrupt from the user, in which case it sends a `TerminateCmd`
//!   (with sequence number) and waits for more output or the command to
//!   finish.
//!
//! Instead of entering a command, the user can also choose to ask the
//! breakpoint to exit and continue script execution (via `ExitBreak`
//! message), or to do that but also exit the listener entirely.
//!
//! ## Main loop: Breakpoint
//!
//! The Breakpoint waits in a loop for `Command` and `ExitBreak`
//! messages.
//!
//! When a `Command` is received, the Breakpoint executes the command and
//! streams back stdout and stderr lines as `OutputLine` messages. If a
//! `TerminateCmd` is received, the Breakpoint checks the sequence number
//! and ends the child process. Either way, once the command finishes, the
//! Breakpoint sends a `Finished` message with the command's exit code and
//! returns to the main loop to wait.
//!
//! If an `ExitBreak` is received in the main loop, the Breakpoint process
//! exits.

use std::{
    mem,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
};

use alkali::{
    asymmetric::kx::x25519blake2b as X25519,
    hash::generic::blake2b,
    mem::{hardened_buffer, FullAccess},
    symmetric::cipher_stream::{self, DecryptionStream, EncryptionStream},
};
use rand::Fill;
use serde::{Deserialize, Serialize};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{
        tcp::{OwnedReadHalf, OwnedWriteHalf},
        TcpListener, TcpStream,
    },
    sync::mpsc,
};

/// Length of the setup shared secret in bytes.
pub const SETUP_SECRET_LENGTH: usize = 16;

hardened_buffer! {
    pub SetupSecretBuffer(SETUP_SECRET_LENGTH)
}

/// A hardened buffer for safely handling the setup secret.
pub type SetupSecret = SetupSecretBuffer<FullAccess>;

impl SetupSecret {
    /// Create a new shared secret from random bytes.
    pub fn new_random() -> SetupSecret {
        let mut setup_secret = SetupSecret::new_empty().unwrap();
        setup_secret.try_fill(&mut rand::thread_rng()).unwrap();
        setup_secret
    }

    /// Use an existing shared secret that has been read from
    /// somewhere else.
    pub fn from_bytes(bytes: &[u8; SETUP_SECRET_LENGTH]) -> SetupSecret {
        let mut secret = SetupSecret::new_empty().unwrap();
        secret.copy_from_slice(bytes); // does a length check
        secret
    }
}

/// The protocol version that's implemented here.
///
/// (Expecting to later understand multiple protocol versions.)
const VERSION: &str = "dev.2";

/*===================*/
/* Insecure messages */
/*===================*/

/// Initial message sent from breakpoint to listener.
///
/// Negotiates protocol, delivers the other public key, starts one
/// half of secure channel, and provides breakpoint's details.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
struct BreakpointHello {
    /// Protocol to use.
    ///
    /// The breakpoint will have been provisioned with a list of
    /// protocols acceptable to the listener, and this is the protocol
    /// the breakpoint has chosen to use.
    pub protocol: String,
    /// The breakpoint's ephemeral X25519 public key.
    #[serde(with = "serde_bytes")]
    pub breakpoint_pub_key: X25519::PublicKey,
    #[serde(with = "serde_bytes")]
    /// Starts the secure channel from the breakpoint.
    ///
    /// This is an `XChaCha20poly1305` header for `cipher_stream`.
    pub cipher_header: [u8; cipher_stream::HEADER_LENGTH],
    /// Encrypted [`BreakpointIntro`], the first `SecureMsg` in the stream.
    ///
    /// This must be decrypted so that cipher stream synchronization is maintained.
    #[serde(with = "serde_bytes")]
    pub encrypted_intro: Vec<u8>,
}

/// Listener's first message back to the breakpoint.
///
/// Starts the other half of the secure channel.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
struct ListenerHello {
    /// Starts the secure channel from the listener.
    ///
    /// `XChaCha20poly1305` header for `cipher_stream`.
    #[serde(with = "serde_bytes")]
    pub cipher_header: [u8; cipher_stream::HEADER_LENGTH],
}

/// Contains an encrypted [`SecureMsg`].
///
/// A stream of these messages forms the basis of the secure channel.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
struct EncryptedMsg {
    #[serde(with = "serde_bytes")]
    pub encrypted: Vec<u8>,
}

/// A message sent insecurely over the network.
///
/// These messages do not have confidentiality or integrity protection
/// except where they specifically contain encrypted fields.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
enum InsecureMsg {
    BreakpointHello(BreakpointHello),
    ListenerHello(ListenerHello),
    EncryptedMsg(EncryptedMsg),
}

/*=================*/
/* Secure messages */
/*=================*/

/// Embedded in the [`BreakpointHello`], this tells the listener about
/// the breakpoint.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct BreakpointIntro {
    /// Which breakpoint is calling back (if option is in use.)
    pub which: Option<String>,
}

/// A command sent by the listener for the breakpoint to execute.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct Command {
    /// Code to execute on the breakpoint.
    pub command: String,
    /// Sequence ID for this command.
    pub seq: u32,
}

/// One line of output from a command, sent by breakpoint to listener.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct OutputLine {
    #[serde(with = "serde_bytes")]
    pub bytes: Vec<u8>,
    pub gender: OutputType,
    /// Must match the most recent [`Command::seq`].
    pub seq: u32,
}

#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
pub enum OutputType {
    Stdout,
    Stderr,
}

/// Final results of a command, sent by breakpoint to listener.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct Finished {
    /// This is a `std::process::ExitStatus`
    pub exit_code: Option<i32>,
    /// Must match the most recent [`Command::seq`].
    pub seq: u32,
}

/// Sent by listener to tell breakpoint to terminate a command that is
/// in-process.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct TerminateCmd {
    /// Sequence number of the command we're going to try to terminate.
    pub seq: u32,
}

/// Listener tells the breakpoint to exit and resume normal script
/// execution.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct ExitBreak {}

/// All message types sent over the secure channel.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub enum SecureMsg {
    BreakpointIntro(BreakpointIntro),
    Command(Command),
    OutputLine(OutputLine),
    Finished(Finished),
    TerminateCmd(TerminateCmd),
    /// This variant is special-cased to end the connection when sent or received.
    ExitBreak(ExitBreak),
}

/*============*/

/// Labels for the parties in the protocol.
enum Party {
    /// The process on the developer's computer, initiating the
    /// session, waiting for a callback. Acts as a server,
    /// network-wise.
    Listener,
    /// The process on the remote end, calling back to the
    /// listener. Acts as a client, network-wise.
    Breakpoint,
}

/// Bytes used for KDF input length-prefixing.
const KDF_INPUT_LEN_BYTES: usize = 8;

/// Length-prefix a collection of KDF inputs.
///
/// The aim here is to avoid canonicalization attacks.
///
/// The input is a collection of inputs to a KDF. The output is the
/// concatenation of length/input pairs. The length of each input is
/// encoded as a little-endian 64-bit integer.
///
/// Accordingly, inputs may only be of length 0 to 2^64 - 1.
fn length_prefixed_kdf_input(inputs: &Vec<&[u8]>) -> Vec<u8> {
    let mut out = Vec::new();
    for &input in inputs {
        let length = input.len();
        assert!(
            u64::try_from(length).is_ok(),
            "Length of KDF input was greater than allowed ({length} > 2^64 - 1)",
        );
        let len_prefix = &length.to_le_bytes()[..KDF_INPUT_LEN_BYTES];
        out.extend(len_prefix);
        out.extend(input);
    }
    out
}

/// Make a one-direction stream key for either listener or breakpoint.
///
/// The stream pre-key is one half of the result of an X25519 key
/// exchange, either the transmit or receive key. It is combined with
/// the shared secret `setup_secret` to form a new XChaCha20-Poly1305
/// for that party and stream direction.
fn make_stream_key(
    stream_prekey: &[u8; X25519::SESSION_KEY_LENGTH],
    pubkey_server: &X25519::PublicKey,
    pubkey_client: &X25519::PublicKey,
    setup_secret: &SetupSecret,
) -> cipher_stream::xchacha20poly1305::Key<FullAccess> {
    let kdf_in_raw = length_prefixed_kdf_input(&vec![
        "breakmancer.streamkey".as_bytes(),
        VERSION.as_bytes(),
        pubkey_server,
        pubkey_client,
        stream_prekey,
    ]);
    let key_bytes = blake2b::hash_custom_to_vec(
        &kdf_in_raw,
        Some(setup_secret.as_slice()),
        cipher_stream::xchacha20poly1305::KEY_LENGTH,
    )
    .unwrap();
    cipher_stream::xchacha20poly1305::Key::<FullAccess>::try_from(key_bytes.as_slice()).unwrap()
}

/// Create the pair of transmit/receive keys that will be used for
/// secure communication.
fn derive_stream_keys(
    my_role: &Party,
    my_keypair: &X25519::Keypair,
    other_public: &X25519::PublicKey,
    setup_secret: &SetupSecret,
) -> (
    cipher_stream::xchacha20poly1305::Key<FullAccess>,
    cipher_stream::xchacha20poly1305::Key<FullAccess>,
) {
    let prekeys = match my_role {
        Party::Breakpoint => my_keypair.client_keys(other_public),
        Party::Listener => my_keypair.server_keys(other_public),
    };
    let (tx_prekey, rx_prekey) = prekeys.unwrap();
    let (pubkey_server, pubkey_client) = match my_role {
        Party::Breakpoint => (other_public, &my_keypair.public_key),
        Party::Listener => (&my_keypair.public_key, other_public),
    };
    let tx = make_stream_key(&tx_prekey, pubkey_server, pubkey_client, setup_secret);
    let rx = make_stream_key(&rx_prekey, pubkey_server, pubkey_client, setup_secret);
    (tx, rx)
}

/// Serialize data to CBOR, at any layer of protocol.
///
/// This standardizes our approach to data serialization and encoding.
fn serialize<D>(data: &D) -> Vec<u8>
where
    D: Serialize,
{
    // Note that serde_cbor serializes structs as maps of strings to
    // values by default. This is important—if the other party
    // includes additional, unexpected keys, we can skip over and
    // ignore them.
    serde_cbor::to_vec(&data).unwrap()
}

/// Deserialize data from CBOR, at any layer of protocol.
///
/// This standardizes our approach to data serialization and encoding.
fn deserialize<'a, T>(bytes: &'a [u8]) -> Result<T, String>
where
    T: Deserialize<'a>,
{
    serde_cbor::from_slice(bytes).map_err(|err| format!("{err}"))
}

/// Encrypt a [`SecureMsg`] for transport.
///
/// This must be done immediately before sending the message in
/// order to ensure that the cipher state stays in sync.
fn encrypt_message(msg: &SecureMsg, cipher_tx: &mut EncryptionStream) -> EncryptedMsg {
    let msg_bytes = serialize(msg);
    let msg_slice = msg_bytes.as_slice();
    let mut ciphered = vec![0u8; msg_slice.len() + cipher_stream::OVERHEAD_LENGTH];
    cipher_tx.encrypt(msg_slice, None, &mut ciphered).unwrap();
    EncryptedMsg {
        encrypted: ciphered,
    }
}

/// Decrypt an encrypted message.
///
/// This must be done immediately after receiving the message in
/// order to ensure that the cipher state stays in sync.
fn decrypt_message(
    msg: &EncryptedMsg,
    cipher_rx: &mut DecryptionStream,
) -> Result<SecureMsg, String> {
    let ciphered = msg.encrypted.as_slice();
    let mut data = vec![0u8; ciphered.len() - cipher_stream::OVERHEAD_LENGTH];
    cipher_rx
        .decrypt(ciphered, None, &mut data)
        .map_err(|err| format!("Unable to decrypt message: {err}"))?;
    deserialize(&data).map_err(|err| format!("Could not parse message data: {err}"))
}

/// Our network frames will be length-prefixed by a 4-byte integer.
const FRAME_HEADER_SIZE: u32 = 4;
/// Max frame body size that can be accepted given [`FRAME_HEADER_SIZE`].
const FRAME_DATA_MAX_BYTES: usize = 2usize.pow(FRAME_HEADER_SIZE * 8);

/// Send a message in the clear.
///
/// Each frame is a little-endian length header ([`FRAME_HEADER_SIZE`]
/// bytes), and then the `clear_msg` serialized as CBOR.
///
/// This is not cancellation safe.
async fn send_clear_frame<W>(stream: &mut W, clear_msg: &InsecureMsg) -> Result<(), String>
where
    W: AsyncWriteExt + Unpin,
{
    let data = serialize(clear_msg);
    let data_len = data.len();
    if data_len > FRAME_DATA_MAX_BYTES {
        return Err(String::from(
            "Cannot send frame of more than {data_len} bytes",
        ));
    }
    let len_bytes = &data_len.to_le_bytes()[..FRAME_HEADER_SIZE as usize];

    let mut frame = Vec::<u8>::from(len_bytes);
    frame.extend_from_slice(&data);
    stream
        .write_all(&frame)
        .await
        .map_err(|err| format!("Failed to send frame: {err}"))?;
    Ok(())
}

/// Receive a message in the clear.
///
/// This is not cancellation safe.
async fn receive_clear_frame<R>(stream: &mut R) -> Result<InsecureMsg, String>
where
    R: AsyncReadExt + Unpin,
{
    let mut len_header = [0u8; mem::size_of::<usize>()];
    stream
        .read_exact(&mut len_header[..FRAME_HEADER_SIZE as usize])
        .await
        .map_err(|err| format!("Failed to receive frame length: {err}"))?;
    let mut data = vec![0u8; usize::from_le_bytes(len_header)];
    stream
        .read_exact(&mut data)
        .await
        .map_err(|err| format!("Failed to receive frame body: {err}"))?;

    deserialize(&data).map_err(|err| format!("Could not parse frame body: {err}"))
}

/// Newtype wrapper for cipher stream so we can mark it as Send.
struct EncryptionStreamWrapper(EncryptionStream);
// SAFETY: Alkali's maintainer has said this should be safe to do.
// Copying from <https://github.com/tom25519/alkali/issues/5>:
//
// « This doesn't happen automatically because we use Sodium's
// `sodium_malloc` to store the stream state on the heap, rather than
// just storing it on the stack or putting it in a `Box` through Rust,
// and so we have to maintain a raw pointer to this state for the
// lifetime of the struct. This is not strictly necessary at all, but
// with access to the state an attacker can decrypt/forge messages in
// the stream, so we treat it as equivalent to a key in terms of
// sensitivity, and we always protect keys in this way in this
// library. So `unsafe` will be necessary to mark these structs as
// `Send`/`Sync`. »
//
// TODO: Once this is resolved in alkali, drop these wrappers.
unsafe impl Send for EncryptionStreamWrapper {}
/// Newtype wrapper for cipher stream so we can mark it as Send.
struct DecryptionStreamWrapper(DecryptionStream);
// SAFETY: See description for EncryptionStreamWrapper.
unsafe impl Send for DecryptionStreamWrapper {}

/// The transmission side of the connection.
struct ConnectionTx {
    /// Write half of the TCP stream.
    tcp_write: OwnedWriteHalf,
    /// Cipher stream for encrypting messages to be sent over TCP.
    cipher_tx: EncryptionStreamWrapper,
    /// Indicates when a clean shutdown is supposed to be
    /// occurring.
    shutting_down: Arc<AtomicBool>,
}

impl ConnectionTx {
    pub fn from(
        tcp_write: OwnedWriteHalf,
        cipher_tx: EncryptionStream,
        shutting_down: Arc<AtomicBool>,
    ) -> ConnectionTx {
        ConnectionTx {
            tcp_write,
            cipher_tx: EncryptionStreamWrapper(cipher_tx),
            shutting_down,
        }
    }

    fn log_error(&self, msg: &str) {
        if !self.shutting_down.load(Ordering::SeqCst) {
            eprintln!("[ERROR] {msg}");
        }
    }

    /// Send insecure message.
    async fn send_insecure(&mut self, msg: &InsecureMsg) -> Result<(), String> {
        send_clear_frame(&mut self.tcp_write, msg).await
    }

    /// Encrypt a [`SecureMsg`] for transport.
    ///
    /// This must be done immediately before sending the message in
    /// order to ensure that the cipher state stays in sync.
    fn encrypt_message(&mut self, msg: &SecureMsg) -> EncryptedMsg {
        encrypt_message(msg, &mut self.cipher_tx.0)
    }

    /// Send a secure message.
    async fn send_secure(&mut self, msg: &SecureMsg) -> Result<(), String> {
        let msg = self.encrypt_message(msg);
        self.send_insecure(&InsecureMsg::EncryptedMsg(msg)).await
    }

    /// Spawns the secure message sender thread.
    ///
    /// The use of channels allows us to avoid the cancel-unsafety of
    /// send and receive.
    pub fn spawn_tx(mut self) -> mpsc::Sender<SecureMsg> {
        let (send_producer, mut send_consumer) = mpsc::channel::<SecureMsg>(200);
        tokio::spawn(async move {
            loop {
                let Some(msg) = send_consumer.recv().await else {
                    self.log_error(
                        "Internal error getting new message to send, shutting down connection.",
                    );
                    break;
                };

                if let Err(err) = self.send_secure(&msg).await {
                    self.log_error(&format!(
                        "Failed to send message to network, shutting down connection: {err}"
                    ));
                    break;
                }

                if let SecureMsg::ExitBreak(_) = msg {
                    self.shutting_down.store(true, Ordering::SeqCst);
                };
            }

            if let Err(err) = self.tcp_write.shutdown().await {
                self.log_error(&format!("Couldn't do clean hangup: {err}"));
            }
        });
        send_producer
    }
}

/// The receiving side of the connection.
struct ConnectionRx {
    /// Read half of the TCP stream.
    tcp_read: OwnedReadHalf,
    /// Cipher stream for decrypting messages received over TCP.
    cipher_rx: DecryptionStreamWrapper,
    /// Indicates when a clean shutdown is supposed to be
    /// occurring.
    shutting_down: Arc<AtomicBool>,
}

impl ConnectionRx {
    pub fn from(
        tcp_read: OwnedReadHalf,
        cipher_rx: DecryptionStream,
        shutting_down: Arc<AtomicBool>,
    ) -> ConnectionRx {
        ConnectionRx {
            tcp_read,
            cipher_rx: DecryptionStreamWrapper(cipher_rx),
            shutting_down,
        }
    }

    fn log_error(&self, msg: &str) {
        if !self.shutting_down.load(Ordering::SeqCst) {
            eprintln!("[ERROR] {msg}");
        }
    }

    /// Receive insecure message.
    async fn receive_insecure(&mut self) -> Result<InsecureMsg, String> {
        receive_clear_frame(&mut self.tcp_read).await
    }

    /// Decrypt an encrypted message.
    ///
    /// This must be done immediately after receiving the message in
    /// order to ensure that the cipher state stays in sync.
    fn decrypt_message(&mut self, msg: &EncryptedMsg) -> Result<SecureMsg, String> {
        decrypt_message(msg, &mut self.cipher_rx.0)
    }

    /// Receive a secure message.
    async fn receive_secure(&mut self) -> Result<SecureMsg, String> {
        match self.receive_insecure().await? {
            InsecureMsg::EncryptedMsg(enc) => self.decrypt_message(&enc),
            clear_msg => Err(format!("Unexpected insecure message type: {clear_msg:?}")),
        }
    }

    /// Spawns the secure message receiver thread.
    ///
    /// The use of channels allows us to avoid the cancel-unsafety of
    /// send and receive.
    pub fn spawn_rx(mut self) -> mpsc::Receiver<SecureMsg> {
        let (recv_producer, recv_consumer) = mpsc::channel::<SecureMsg>(200);
        tokio::spawn(async move {
            loop {
                let msg = match self.receive_secure().await {
                    Ok(msg) => msg,
                    Err(err) => {
                        self.log_error(&format!(
                            "Couldn't receive message from network, shutting down connection: {err}"
                        ));
                        break;
                    }
                };

                if let SecureMsg::ExitBreak(_) = msg {
                    self.shutting_down.store(true, Ordering::SeqCst);
                };

                if let Err(err) = recv_producer.send(msg).await {
                    self.log_error(&format!(
                        "Internal error passing a received message around; shutting down: {err}"
                    ));
                    break;
                }
            }
        });
        recv_consumer
    }
}

/// A connection to the other party (breakpoint or listener).
pub struct Connection {
    /// Receives messages from the other party, until connection drops
    /// or there is a protocol error.
    pub rx: mpsc::Receiver<SecureMsg>,
    /// Sends messages to the other party, until connection drops or
    /// there is a protocol error.
    pub tx: mpsc::Sender<SecureMsg>,
    /// Indicates when a clean shutdown is supposed to be
    /// occurring.
    shutting_down: Arc<AtomicBool>,
}

impl Connection {
    /// Create a new listener connection, completing once the
    /// breakpoint has connected and completed the handshake.
    ///
    /// Resolves to a pair of channels that serve as an abstraction
    /// over the secure channel, as well as the breakpoint's
    /// self-introduction.
    pub async fn new_listener(
        server_socket: &TcpListener,
        listener_keypair: &X25519::Keypair,
        setup_secret: &SetupSecret,
    ) -> Result<(Connection, BreakpointIntro), String> {
        // Wait for a new connection
        let (stream, client_addr) = server_socket
            .accept()
            .await
            .map_err(|err| format!("Couldn't get new connection: {err}"))?;
        println!("Connection from {client_addr:?}");

        let (mut tcp_read, mut tcp_write) = stream.into_split();

        // Receive Breakpoint's public key, stream start, and introduction
        let breakpoint_hello = match receive_clear_frame(&mut tcp_read).await {
            Ok(InsecureMsg::BreakpointHello(hello)) => Ok(hello),
            Ok(_) => Err("Expected a breakpoint hello.".to_owned()),
            Err(err) => Err(format!(
                "Couldn't receive data when waiting for a breakpoint hello: {err}"
            )
            .to_owned()),
        }?;

        // Now we can use X25519 to compute the stream pre-keys for each
        // party, and then mix them with the setup secret to create
        // the session stream keys.
        let (tx, rx) = derive_stream_keys(
            &Party::Listener,
            listener_keypair,
            &breakpoint_hello.breakpoint_pub_key,
            setup_secret,
        );

        // The breakpoint has already sent over its cipher stream header.
        let mut rx_stream = DecryptionStream::new(&rx, &breakpoint_hello.cipher_header).unwrap();

        // Now we can learn about the breakpoint.
        let maybe_intro = decrypt_message(
            &EncryptedMsg {
                encrypted: breakpoint_hello.encrypted_intro,
            },
            &mut rx_stream,
        )
        .map_err(|err| format!("Could not decrypt breakpoint introduction: {err}"))?;
        let breakpoint_intro = match maybe_intro {
            SecureMsg::BreakpointIntro(intro) => intro,
            _ => Err("Unexpected message type for breakpoint intro.")?,
        };

        // And the listener can respond with a hello of its own,
        // carrying the cipher stream header for that direction.
        let tx_stream = EncryptionStream::new(&tx).unwrap();
        let tx_header = tx_stream.get_header();
        send_clear_frame(
            &mut tcp_write,
            &InsecureMsg::ListenerHello(ListenerHello {
                cipher_header: tx_header,
            }),
        )
        .await
        .map_err(|err| format!("Couldn't send encrypted stream header: {err}"))?;

        let shutting_down = Arc::new(AtomicBool::from(false));

        Ok((
            Connection {
                rx: ConnectionRx::from(tcp_read, rx_stream, shutting_down.clone()).spawn_rx(),
                tx: ConnectionTx::from(tcp_write, tx_stream, shutting_down.clone()).spawn_tx(),
                shutting_down,
            },
            breakpoint_intro,
        ))
    }

    /// Create a new breakpoint connection, completing once the
    /// listener has responded and completed the handshake.
    ///
    /// Resolves to a pair of channels that serve as an abstraction
    /// over the secure channel.
    pub async fn new_breakpoint(
        breakpoint_keypair: &X25519::Keypair,
        listener_addr: &String,
        listener_public: &X25519::PublicKey,
        setup_secret: &SetupSecret,
        which: &Option<String>,
    ) -> Result<Connection, String> {
        let stream = TcpStream::connect(listener_addr)
            .await
            .map_err(|err| format!("Unable to connect to {listener_addr}: {err}"))?;
        match stream.peer_addr() {
            Ok(connected) => eprintln!("Connecting to listener at {connected}"),
            Err(err) => eprintln!("Unable to get address of listener connection: {err}"),
        };

        let (mut tcp_read, mut tcp_write) = stream.into_split();

        // The breakpoint already has all the key material for the
        // secure channel.
        let (tx, rx) = derive_stream_keys(
            &Party::Breakpoint,
            breakpoint_keypair,
            listener_public,
            setup_secret,
        );
        // It can't receive encrypted messages until it gets the
        // rx_header, but it can create the tx_header now.
        let mut tx_stream = EncryptionStream::new(&tx).unwrap();
        let tx_header = tx_stream.get_header();

        // Start by sending over the breakpoint public key, tx_header,
        // and encrypted introduction.
        let intro = SecureMsg::BreakpointIntro(BreakpointIntro {
            which: which.clone(),
        });
        let hello = BreakpointHello {
            protocol: VERSION.to_owned(),
            breakpoint_pub_key: breakpoint_keypair.public_key,
            cipher_header: tx_header,
            encrypted_intro: encrypt_message(&intro, &mut tx_stream).encrypted,
        };
        send_clear_frame(&mut tcp_write, &InsecureMsg::BreakpointHello(hello))
            .await
            .map_err(|err| format!("Unable to send hello: {err}"))?;

        let listener_hello = match receive_clear_frame(&mut tcp_read).await {
            Ok(InsecureMsg::ListenerHello(hello)) => Ok(hello),
            Ok(_) => Err("Expected a listener cipher header.".to_owned()),
            Err(err) => Err(format!(
                "Couldn't receive data when waiting for a listener cipher header: {err}"
            )
            .to_owned()),
        }?;
        let rx_stream = DecryptionStream::new(&rx, &listener_hello.cipher_header)
            .map_err(|err| format!("Unable to start decryption stream: {err}"))?;

        let shutting_down = Arc::new(AtomicBool::from(false));

        Ok(Connection {
            rx: ConnectionRx::from(tcp_read, rx_stream, shutting_down.clone()).spawn_rx(),
            tx: ConnectionTx::from(tcp_write, tx_stream, shutting_down.clone()).spawn_tx(),
            shutting_down,
        })
    }
}

impl Drop for Connection {
    /// Attempt to close the connection gracefully.
    fn drop(&mut self) {
        // Dropping the Connection drops the [tx] Sender, which means
        // the spawn_tx loop gets a None and sends a shutdown. So all
        // of that works automatically, but we do need to give both
        // halves of the connection a heads-up that the TCP connection
        // is about to stop working and that most errors should be
        // suppressed.
        self.shutting_down.store(true, Ordering::SeqCst);
    }
}

#[cfg(test)]
mod tests {
    use crate::test_utils::{all_variant_mapping, partial_hexlify};

    use super::*;

    /// Ensure hardcoded length allows a u64.
    #[test]
    fn kdf_inputs() {
        assert_eq!(KDF_INPUT_LEN_BYTES * 8, 64);
    }

    /// Pinning test for network frame format.
    #[tokio::test]
    async fn network_frame() {
        let mut written = Vec::<u8>::new();
        send_clear_frame(
            &mut written,
            &InsecureMsg::BreakpointHello(BreakpointHello {
                protocol: String::from("test"),
                breakpoint_pub_key: [5u8; 32],
                cipher_header: [1u8; 24],
                encrypted_intro: [3u8; 6].to_vec(),
            }),
        )
        .await
        .unwrap();

        let expected = [
            // Length header, 148 byte body.
            b"\x94\x00\x00\x00" as &[u8],

            b"\xA1", // 1 kv map
            b"o", b"BreakpointHello",
            b"\xA4", // 4 kv map
            b"h", b"protocol",
            b"d", b"test",
            b"r", b"breakpoint_pub_key",
            b"X", b"\x20", b"\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05\x05",
            b"m", b"cipher_header",
            b"X\x18", b"\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01",
            b"o", b"encrypted_intro",
            b"F", b"\x03\x03\x03\x03\x03\x03",
        ].concat();
        assert_eq!(
            expected,
            written,
            "expected {} but got {}",
            partial_hexlify(&expected),
            partial_hexlify(&written)
        );
    }

    /// Test [`InsecureMsg`] de/serialization (including field names)
    /// using pinned values.
    #[tokio::test]
    async fn insecure_message_format_pinning() {
        let examples = all_variant_mapping!(
            InsecureMsg => &[u8]:

            InsecureMsg::BreakpointHello =
                (
                    BreakpointHello {
                        protocol: String::from("someproto"),
                        breakpoint_pub_key: [6u8; 32],
                        cipher_header: [7u8; 24],
                        encrypted_intro: [3u8; 6].to_vec(),
                    }
                )
                => b"\xA1oBreakpointHello\xA4hprotocolisomeprotorbreakpoint_pub_keyX\x20\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06\x06mcipher_headerX\x18\x07\x07\x07\x07\x07\x07\x07\x07\x07\x07\x07\x07\x07\x07\x07\x07\x07\x07\x07\x07\x07\x07\x07\x07oencrypted_intro\x46\x03\x03\x03\x03\x03\x03",

            InsecureMsg::ListenerHello =
                (
                    ListenerHello {
                        cipher_header: [8u8; 24],
                    }
                )
                => b"\xA1mListenerHello\xA1mcipher_headerX\x18\x08\x08\x08\x08\x08\x08\x08\x08\x08\x08\x08\x08\x08\x08\x08\x08\x08\x08\x08\x08\x08\x08\x08\x08",

            InsecureMsg::EncryptedMsg =
                (
                    EncryptedMsg {
                        encrypted: [9u8; 3].to_vec(),
                    }
                )
                => b"\xA1lEncryptedMsg\xA1iencrypted\x43\x09\x09\x09",
        );
        for (variant_value, serialized) in examples {
            assert_eq!(
                serialized,
                serialize(&variant_value),
                "Serialized bytes did not match expected value for input {variant_value:?}"
            );
            let deserialized = deserialize(serialized)
                .map_err(|err| format!("Could not deserialize: {err}; expected {variant_value:?}"))
                .unwrap();
            assert_eq!(
                variant_value, deserialized,
                "Deserialized value did not match expected value {variant_value:?}"
            );
        }
    }

    /// Test [`SecureMsg`] de/serialization (including field names) using
    /// pinned values.
    #[tokio::test]
    async fn secure_message_format_pinning() {
        let examples = all_variant_mapping!(
            SecureMsg => &[u8]:

            SecureMsg::BreakpointIntro =
                (
                    BreakpointIntro {
                        which: Some("loop".to_owned()),
                    }
                )
                => b"\xA1oBreakpointIntro\xA1ewhichdloop",

            SecureMsg::Command =
                (
                    Command {
                        command: "pwd".to_owned(),
                        seq: 5,
                    }
                )
                => b"\xA1gCommand\xA2gcommandcpwdcseq\x05",

            SecureMsg::OutputLine =
                (
                    OutputLine {
                        bytes: b"hello world!".to_vec(),
                        gender: OutputType::Stdout,
                        seq: 7,
                    }
                )
                => b"\xA1jOutputLine\xA3ebytes\x4Chello world!fgenderfStdoutcseq\x07",

            SecureMsg::Finished =
                (
                    Finished {
                        exit_code: Some(1),
                        seq: 9,
                    }
                )
                => b"\xA1hFinished\xA2iexit_code\x01cseq\x09",

            SecureMsg::TerminateCmd =
                (
                    TerminateCmd { seq: 11 }
                )
                => b"\xA1lTerminateCmd\xA1cseq\x0b",

            SecureMsg::ExitBreak =
                (ExitBreak {})
                => b"\xA1iExitBreak\xA0",
        );
        for (variant_value, serialized) in examples {
            assert_eq!(
                serialized,
                serialize(&variant_value),
                "Serialized bytes did not match expected value for input {variant_value:?}"
            );
            let deserialized = deserialize(serialized)
                .map_err(|err| format!("Could not deserialize: {err}; expected {variant_value:?}"))
                .unwrap();
            assert_eq!(
                variant_value, deserialized,
                "Deserialized value did not match expected value {variant_value:?}"
            );
        }
    }
}