breakmancer 0.1.1

//! Network and encryption protocol.
//!
//! Summary: The Listener and Breakpoint use X25519 to derive two stream
//! pre-keys, then combine each with an out-of-band shared secret to derive the
//! real ChaCha20-Poly1305 stream keys. (One stream in each direction.)
//!
//! Initial setup:
//!
//! - The developer runs the `listen` command on their workstation. They are the
//!   "Listener" party.
//! - Listener creates X25519 keypair `listen_pub, listen_priv` as well as a shared
//!   secret `prekey_oob`.
//! - 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 user is told to put `prekey_oob` in the
//!   command's environment as securely as possible.
//! - Breakpoint creates X25519 keypair `break_pub, break_priv`.
//! - Breakpoint performs X25519 key exchange to derive the pair of directional
//!   stream keys `pre_rx, pre_tx`. Each of these is hashed with BLAKE2b to 32
//!   bytes, keyed with `prekey_oob`. These are now the actual
//!   `rx, tx` stream keys.
//! - Breakpoint creates an XChaCha20-Poly1305 stream for transmission to Listener.
//! - Breakpoint sends `break_pub` to Listener in the clear in a `BreakpointHello` message,
//!   which contains a map of `protocol` (string), `breakpoint_pub_key` (32 bytes),
//!   `cipher_header` (24 bytes) for the XChaCha20-Poly1305 stream it will transmit,
//!   and `encrypted_intro` (bytes) as an additional, encrypted payload.
//!   (Protocol version is just `"dev"` for now; this will allow software version
//!   slop between the two sides in the future.)
//! - Listener receives the hello and is able to perform key exchange. It is able
//!   to decrypt the `encrypted_intro` to learn about the breakpoint's location.
//! - Listener creates its own XChaCha20-Poly1305 stream and sends the header to
//!   the breakpoint in a `ListenerHello` with a `cipher_header` field.
//!
//! At this point, we begin the application-level protocol, exchanging
//! `EncryptedMsg` messages containing commands and output lines.
//!
//! Any time there is a protocol error or network fault, both parties
//! start over at the step where Breakpoint generates and sends a keypair.
use std::mem::size_of;

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

/// Length of the out-of-band shared secret in bytes.
pub const OOB_SECRET_LENGTH: usize = 16;

hardened_buffer! {
    pub OutOfBandSharedSecret(OOB_SECRET_LENGTH)
}

/// A hardened buffer for safely handling the out-of-band shared
/// secret.
pub type OobSecret = OutOfBandSharedSecret<FullAccess>;

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

    /// Use an existing shared secret that has been read from
    /// somewhere else.
    pub fn from_bytes(bytes: &[u8]) -> OobSecret {
        // TODO was observed to panic
        let mut secret = OobSecret::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";

/*===================*/
/* 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),
    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,
}

/// 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 `oob_secret` to form a new XChaCha20-Poly1305
/// for that party and stream direction.
fn make_stream_key(
    stream_prekey: &[u8],
    oob_secret: &OobSecret,
) -> cipher_stream::xchacha20poly1305::Key<FullAccess> {
    let context = "cipher_stream key from KX prekey and OOB secret".as_bytes();
    let key_bytes = blake2b::hash_custom_to_vec(
        &[context, b"|", stream_prekey].concat(),
        Some(oob_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,
    oob_secret: &OobSecret,
) -> (
    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 tx = make_stream_key(tx_prekey.as_slice(), oob_secret);
    let rx = make_stream_key(rx_prekey.as_slice(), oob_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; 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,
}

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

    /// 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 {
                if let Some(msg) = send_consumer.recv().await {
                    if let Err(err) = self.send_secure(&msg).await {
                        eprintln!(
                            "[DEBUG] Failed to send message, shutting down connection: {err}"
                        );
                        break;
                    }
                } else {
                    eprintln!(
                        "[DEBUG] Failed to get new message to send, shutting down connection."
                    );
                    break;
                }
            }

            if let Err(err) = self.tcp_write.shutdown().await {
                eprintln!("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,
}

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

    /// 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) => {
                        eprintln!(
                            "[DEBUG] Couldn't receive message, shutting down connection: {err}"
                        );
                        break;
                    }
                };
                if let Err(err) = recv_producer.send(msg).await {
                    eprintln!("[DEBUG] Couldn't consume new message, 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>,
}

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,
        oob_secret: &OobSecret,
    ) -> 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 out-of-band secret to create
        // the session stream keys.
        let (tx, rx) = derive_stream_keys(
            &Party::Listener,
            listener_keypair,
            &breakpoint_hello.breakpoint_pub_key,
            oob_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}"))?;

        Ok((
            Connection {
                rx: ConnectionRx::from(tcp_read, rx_stream).spawn_rx(),
                tx: ConnectionTx::from(tcp_write, tx_stream).spawn_tx(),
            },
            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,
        oob_secret: &OobSecret,
        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,
            oob_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}"))?;

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

    /// Attempt to close the connection gracefully.
    pub fn hangup(self) {
        // Dropping the Connection drops the [tx] Sender, which means
        // the spawn_tx loop gets a None and sends a shutdown. A bit
        // janky, but it works?
        std::mem::drop(self);
    }
}

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

    use super::*;

    /// 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:?}"
            );
        }
    }
}