breakmancer 0.5.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 start` to start a session on their
//! development machine; that process will be termed the "Controller". 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 (we'll call this the "Breakpoint") reaches back to the Controller,
//! performs a handshake, and waits for further instructions.
//!
//! The user is then presented with a simple shell-like environment in the
//! Controller 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 Controller 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 dual-pub
//! protocol therefore does not rely on provisioning a private key or
//! a shared secret on the Breakpoint side.
//!
//! Instead, dual-pub requires that the Breakpoint's output is visible
//! more or less in realtime. This allows the Breakpoint to generate
//! an ephemeral keypair and print the public key; the user can then
//! copy this (or rather a digest of it) to the Controller.
//!
//! # Goals
//!
//! - The Breakpoint must only execute commands specified by the Controller,
//!   without alteration (mutation, replay, reordering, etc.) It must not
//!   be possible for another party to impersonate the Controller, even an
//!   on-path attacker.
//! - The Controller 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.
//!
//! # Non-goals
//!
//! - Availability, in the face of an adversary.
//!
//! # Cryptographic protocol
//!
//! At a high level, the Controller and Breakpoint use X-Wing to send
//! each other secrets during an initial handshake, then combine those
//! to generate a pair of XChaCha20-Poly1305 stream keys for all
//! subsequent messages.
//!
//! Both parties generate their X-Wing keypairs at startup, and the
//! user is enlisted to copy digests of the public keys in both
//! directions in order to establish trust.
//!
//! Details:
//!
//! - The Controller creates an X-Wing keypair `ctrl_pub`,
//!   `ctrl_priv` as well as a digest `ctrl_pub_digest`.
//!     - The digest is a BLAKE2b digest with `breakmancer.dual-pub.pubkey-digest.controller`
//!       as the `key` param, and 16 byte digest size. The digest is then
//!       truncated to 12 bytes and Base58-encoded (using the Bitcoin alphabet).
//!     - The role of the digest is to allow provisioning each party with
//!       the other's public key in a trusted manner without requiring a user to
//!       copy and paste a kilobyte+ of text.
//! - The Breakpoint does the same: `break_pub`, `break_priv`, and
//!   `break_pub_digest` (with `breakpoint` instead of `controller` in the
//!   `key` param).
//! - The digests are communicated to the counterparties out-of-band.
//! - During the handshake over the network:
//!     - Each party sends the other their full public key. The parties then
//!       compute the digest for the received key and compare it to the trusted
//!       one received out-of-band.
//!     - Each party uses X-Wing to encapsulate a secret for the other. The
//!       ciphertexts are shared across the network and decapsulated. The
//!       resulting shared secrets are `setup_secret_part1` (sent by
//!       the Controller) and `setup_secret_part2` (sent by the Breakpoint).
//!     - Each party now computes the stream keys:
//!         - Concatenate `ctrl_pub`, `break_pub`, `setup_secret_part1`,
//!           `setup_secret_part2`.
//!         - Feed this into BLAKE2b with a key of `breakmancer.dual-pub.stream-keys`
//!           and a digest length of 64.
//!         - Split the digest into two 32-byte values.
//!         - These are now the XChaCha20-Poly1305 stream keys
//!           `ctrl_stream_key` and `break_stream_key`. The first is used by the
//!           Controller to send messages, the latter by the Breakpoint.
//!     - Each party initializes a libsodium `cipher_stream` with their respective
//!       transmit key and sends the initializing header to the other party.
//!     - The handshake is now complete.
//! - All subsequent messages are encrypted using `cipher_stream`.
//!
//! # 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 five types: `BreakpointHello`,
//! `ControllerHello`, `BreakpointSetupFinalize`, `ControllerSetupFinalize`,
//! and `EncryptedMsg`. Several of these can contain
//! encrypted message data, a `cipher_stream` 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), network fault (such as a closed
//! TCP connection), or cryptography failure (such as a public key that
//! doesn't match the expected digest) that party will drop the
//! connection and start over at the beginning of the handshake.
//!
//! ## Initial, out-of-band setup
//!
//! - The developer runs the `start` command on their workstation. This
//!   is the "Controller" process. The Controller creates a fresh keypair.
//! - The Controller outputs a `break` command to be run in the system
//!   being inspected. When that command runs, it is the "Breakpoint"
//!   party. The command contains `ctrl_pub_digest` and the address of the
//!   Controller.
//! - When the Breakpoint starts, it also creates a fresh keypair. It then
//!   connects to the Controller and prints `break_pub_digest` to stdout.
//!
//! ## Handshake
//!
//! Due to the need to exchange public keys before secrets can be exchanged,
//! the full handshake requires four messages:
//!
//! - The Breakpoint sends a `BreakpointHello` containing `protocol`
//!   with value (`"dual-pub"`), `break_pub`, and `expected_controller_digest`.
//!     - The Controller can reject the connection early if the digest
//!       doesn't match. This is *not* a security measure; it only
//!       allows early, non-interactive detection of a handshake that
//!       is doomed to failure.
//!     - A rejection here may be indicated by sending a `HandshakeRejected`
//!       message giving the reason (and indicating whether the failure is
//!       permanent, which it would be at this stage). Or, the controller
//!       could simply close the connection.
//! - The Controller asks for `break_pub_digest` as input and validates
//!   `break_pub` against this. It then sends a `ControllerHello` containing
//!   `ctrl_pub` and `setup_secret_part1`.
//!     - A digest mismatch here could again result in a `HandshakeRejected`
//!       message, although this one should be non-permanent.
//! - The Breakpoint validates `ctrl_pub`, decapsulates the secret, derives the
//!   stream keys, and responds with a `BreakpointSetupFinalize` containing
//!   `setup_secret_part2`, a `cipher_header` for its transmit stream, and an
//!   `encrypted_intro` (encrypted message containing a `BreakpointIntro`).
//!     - The intro just contains information about *which* breakpoint is
//!       running, in case there are multiple.
//! - The Controller decapsulates the second secret and derives stream keys as
//!   well, then sends a `ControllerSetupFinalize` with its own `cipher_header`.
//!
//! At this point, both parties have a secure channel, and all future
//! messages are sent encrypted and wrapped in `EncryptedMsg`-type
//! messages.
//!
//! Any message not in this strict order is an error and requires starting a
//! fresh connection.
//!
//! ## Main loop: Controller
//!
//! Each time the user enters a command. the Controller 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 controller 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::{
    hash::generic::blake2b,
    mem::FullAccess,
    symmetric::cipher_stream::{self, xchacha20poly1305, DecryptionStream, EncryptionStream},
};
use clap::ValueEnum;
use kem::{Decapsulate, Encapsulate};
use rust_base58::ToBase58;
use serde::{Deserialize, Serialize};
use strum_macros::IntoStaticStr;
use subtle::ConstantTimeEq;
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{
        tcp::{OwnedReadHalf, OwnedWriteHalf},
        TcpListener, TcpStream,
    },
    sync::mpsc::{self, error::SendError},
};

/// Authentication protocol version that a set of sessions will use.
#[derive(ValueEnum, Clone, Debug)]
pub enum AuthVersion {
    /// Dual public key auth.
    ///
    /// Authenticates both Controller to Breakpoint, and vice versa,
    /// with each party generating public keys. The keys are X-Wing
    /// (X25519 + ML-KEM) to provide post-quantum security. Each party
    /// provides a verification string to the other party, out of
    /// band, in order to validate the public keys.
    ///
    /// This requires that the user be able to observe streaming logs
    /// from the Breakpoint, which is how the Breakpoint communicates
    /// its verification string. (The Controller's verification string
    /// is provided to the Breakpoint in its invocation in the script.)
    DualPub,
}

/// Byte count in verifier digest.
///
/// Digest entropy chosen to be sufficient that it would take an
/// interactive attacker 1000 years to compute a matching key if they
/// could make 1 trillion guesses per second simultaneously on 1
/// million computers. (In practice, keys should normally live hours
/// to days.)
///
/// That would be about 2^95 guesses; the closest number of bytes
/// would be 12.
const VERIFIER_DIGEST_BYTES: usize = 12;

/// Wrapper representing an X-Wing keypair.
#[derive(Clone)]
pub struct XWingKeypair {
    pub public_key: x_wing::EncapsulationKey,
    pub private_key: x_wing::DecapsulationKey,
}

impl XWingKeypair {
    /// Create a new, random keypair.
    pub fn new() -> XWingKeypair {
        let (decaps, encaps) = x_wing::generate_key_pair_from_os_rng();
        XWingKeypair {
            public_key: encaps,
            private_key: decaps,
        }
    }

    /// Derive the keypair from the private key.
    pub fn from_decaps(decaps: &x_wing::DecapsulationKey) -> XWingKeypair {
        XWingKeypair {
            public_key: decaps.encapsulation_key(),
            private_key: decaps.clone(),
        }
    }
}

/// Compute the verification digest string for the controller.
///
/// This is a Base58 digest of the controller's public key (along with
/// some domain separation).
pub fn controller_digest(public_key: &x_wing::EncapsulationKey) -> String {
    pubkey_digest(public_key, "controller")
}

/// Compute the verification digest string for the breakpoint.
///
/// This is a Base58 digest of the breakpoint's public key (along with
/// some domain separation).
pub fn breakpoint_digest(public_key: &x_wing::EncapsulationKey) -> String {
    pubkey_digest(public_key, "breakpoint")
}

/// Internal function for computing either the controller or breakpoint digest.
fn pubkey_digest(public_key: &x_wing::EncapsulationKey, party: &str) -> String {
    let pub_raw = public_key.as_bytes();
    // Domain separation provided here, via the `key` parameter
    let domain = format!("breakmancer.dual-pub.pubkey-digest.{party}");
    // Alkali/libsodium doesn't allow a direct output of 12 bytes, so
    // generate the smallest allowed digest size and then truncate it.
    let digest = &blake2b::hash_custom_to_vec(&pub_raw, Some(domain.as_bytes()), 16).unwrap()
        [..VERIFIER_DIGEST_BYTES];
    digest.to_base58()
}

/// Verify a key digest by comparing it to a trusted value.
///
/// `trusted_digest` is the key digest that was received out of band
/// and is trusted. `received_key_digest` is the output of either
/// [`controller_digest`] or [`breakpoint_digest`].
///
/// The comparison is performed in constant time.
///
/// Some cleanup is performed -- whitespace is trimmed, and lookalike
/// characters that are not part of the Base58 alphabet are changed to
/// the expected in-alphabet character.
fn verify_digest(trusted_digest: &str, received_key_digest: &str) -> bool {
    // If someone typed this, automatically correct Base58
    // lookalikes. Also, trim whitespace.
    let trusted_digest = trusted_digest.replace(['l', 'I'], "1").trim().to_string();

    trusted_digest
        .as_bytes()
        .ct_eq(received_key_digest.as_bytes())
        .into()
}

/// The protocol version that's implemented here.
///
/// (Expecting to later understand multiple protocol versions.)
const PROTO_DUAL_PUB: &str = "dual-pub";

//************************//
//                        //
//  Unencrypted messages  //
//                        //
//************************//

/// Message sent any time during handshake to tell counterparty why
/// the handshake is being rejected and that the connection will be
/// shut down.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
struct HandshakeRejected {
    /// Human-readable reason for rejection.
    reason: String,
    /// Whether the other party should consider this a permanent
    /// rejection (and should not re-attempt a connection with the
    /// same parameters).
    permanent: bool,
}

/// Initial message sent from breakpoint to controller.
///
/// Confirms protocol and delivers breakpoint's public key.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
struct BreakpointHello {
    /// Protocol to use.
    ///
    /// The breakpoint will have been provisioned with a list of
    /// protocols acceptable to the controller, and this is the protocol
    /// the breakpoint has chosen to use.
    protocol: String,

    /// The breakpoint's public key.
    #[serde(with = "serde_bytes")]
    break_pub: [u8; x_wing::ENCAPSULATION_KEY_SIZE],

    /// The verification string that the breakpoint expects from the
    /// controller. (Intended only for early cooperative termination,
    /// not for security.)
    expected_controller_digest: String,
}

/// Controller's first message back to the breakpoint.
///
/// Delivers the controller's full public key and its contribution to
/// the setup secret.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
struct ControllerHello {
    /// The controller's public key.
    #[serde(with = "serde_bytes")]
    ctrl_pub: [u8; x_wing::ENCAPSULATION_KEY_SIZE],

    /// The controller's half of the setup secret, encapsulated for
    /// the breakpoint's public key.
    #[serde(with = "serde_bytes")]
    setup_secret_part1: [u8; x_wing::CIPHERTEXT_SIZE],
}

/// Breakpoint's second message.
///
/// This completes setup for the breakpoint->controller direction.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
struct BreakpointSetupFinalize {
    /// The breakpoint's half of the setup secret, encapsulated for
    /// the controller's public key.
    #[serde(with = "serde_bytes")]
    setup_secret_part2: [u8; x_wing::CIPHERTEXT_SIZE],

    /// Starts the secure channel from the breakpoint.
    ///
    /// This is an `XChaCha20poly1305` header for `cipher_stream`.
    #[serde(with = "serde_bytes")]
    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.
    encrypted_intro: Vec<u8>,
}

/// Controller's second message.
///
/// This completes setup for the controller->breakpoint direction.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
struct ControllerSetupFinalize {
    /// Starts the secure channel from the controller.
    ///
    /// This is an `XChaCha20poly1305` header for `cipher_stream`.
    #[serde(with = "serde_bytes")]
    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")]
    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.
#[allow(clippy::large_enum_variant)]
#[derive(Debug, PartialEq, Serialize, Deserialize, IntoStaticStr)]
enum InsecureMsg {
    HandshakeRejected(HandshakeRejected),
    BreakpointHello(BreakpointHello),
    ControllerHello(ControllerHello),
    BreakpointSetupFinalize(BreakpointSetupFinalize),
    ControllerSetupFinalize(ControllerSetupFinalize),
    EncryptedMsg(EncryptedMsg),
}

//***************************************//
//                                       //
//  Secure messages (all are encrypted)  //
//                                       //
//***************************************//

/// Embedded in the [`BreakpointHello`], this tells the controller 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 controller 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 controller.
#[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 controller.
#[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 controller 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,
}

/// Controller 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, IntoStaticStr)]
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),
}

//*****************************************//
//                                         //
//  Message encryption and key derivation  //
//                                         //
//*****************************************//

/// Create the pair of keys that will be used for secure
/// communication.
///
/// Accepts the public keys of both parties and the secrets that each
/// party has encapsulated for the other via the setup messages.
///
/// The first returned key is for the [`cipher_stream`] of messages from
/// the controller to the breakpoint; the second is for the other
/// direction.
fn derive_stream_keys_dual_pub(
    pubkey_controller: &x_wing::EncapsulationKey,
    pubkey_breakpoint: &x_wing::EncapsulationKey,
    setup_secret_controller: &x_wing::SharedSecret,
    setup_secret_breakpoint: &x_wing::SharedSecret,
) -> (
    xchacha20poly1305::Key<FullAccess>,
    xchacha20poly1305::Key<FullAccess>,
) {
    let mut in_raw = Vec::new();
    in_raw.extend(&pubkey_controller.as_bytes());
    in_raw.extend(&pubkey_breakpoint.as_bytes());
    in_raw.extend(*setup_secret_controller);
    in_raw.extend(*setup_secret_breakpoint);

    let key_len = xchacha20poly1305::KEY_LENGTH;
    // Hash input is public keys followed by setup secrets. Key is a
    // domain separator.
    let key_bytes = blake2b::hash_custom_to_vec(
        &in_raw,
        Some("breakmancer.dual-pub.stream-keys".as_bytes()),
        key_len * 2,
    )
    .unwrap();

    // Split output into two halves, as stream keys
    let controller_stream_key =
        xchacha20poly1305::Key::<_>::try_from(&key_bytes[0..key_len]).unwrap();
    let breakpoint_stream_key =
        xchacha20poly1305::Key::<_>::try_from(&key_bytes[key_len..]).unwrap();

    (controller_stream_key, breakpoint_stream_key)
}

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

//****************************************//
//                                        //
//  Secured connection abstraction layer  //
//                                        //
//****************************************//

/// 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.
    pub fn spawn_tx(mut self) -> mpsc::Sender<SecureMsg> {
        let (send_producer, mut send_consumer) = mpsc::channel::<SecureMsg>(200);
        // Cancellation: This task handle is dropped, but execution
        // continues (on tokio or other detach-on-drop runtimes) until
        // the Connection drops (or there is an error), at which point
        // the loop exits.
        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 {
                    // Let all parts of this end of the connection
                    // know that a shutdown is pending, and that
                    // errors are expected.
                    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.
    pub fn spawn_rx(mut self) -> mpsc::Receiver<SecureMsg> {
        let (recv_producer, recv_consumer) = mpsc::channel::<SecureMsg>(200);
        // Cancellation:  This task  handle is dropped,  but execution
        // continues (on tokio or other detach-on-drop runtimes) until
        // the Connection drops (or there is an error), at which point
        // the loop exits.
        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
    }
}

//*******************************//
//                               //
//  Controller connection setup  //
//                               //
//*******************************//

/// Create a new controller connection and complete when the
/// breakpoint's hello has been received.
///
/// This does not complete the connection setup; the caller is
/// responsible for either dropping the returned
/// [`ControllerReceivedHello`] or calling its
/// `verify_breakpoint_and_complete_setup` method.
pub async fn controller_open_connection(
    controller_keypair: XWingKeypair,
    server_socket: &TcpListener,
) -> Result<ControllerReceivedHello, 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
    let breakpoint_hello = match receive_clear_frame(&mut tcp_read).await {
        Ok(InsecureMsg::BreakpointHello(hello)) => hello,
        Ok(wrong) => Err(format!(
            "Unexpected message type when waiting for BreakpointHello: {}",
            Into::<&str>::into(wrong)
        ))?,
        Err(err) => Err(format!(
            "Couldn't receive data when waiting for a BreakpointHello: {err}"
        ))?,
    };
    if breakpoint_hello.protocol != PROTO_DUAL_PUB {
        Err(format!(
            "Unexpected protocol field from breakpoint: {}",
            breakpoint_hello.protocol
        ))?;
    }
    let breakpoint_pub_key: x_wing::EncapsulationKey = (&breakpoint_hello.break_pub).into();

    // The Breakpoint gives us an opportunity to close the connection
    // early on verification mismatch instead of requiring an
    // interactive step.
    let ctrl_digest = controller_digest(&controller_keypair.public_key);
    if !verify_digest(&ctrl_digest, &breakpoint_hello.expected_controller_digest) {
        let err_msg = format!(
            "Breakpoint sent wrong expected controller verification string, \
             exiting early. (Breakpoint had outdated connection string?) \
             Expected = {}, sent = {}.",
            ctrl_digest, breakpoint_hello.expected_controller_digest
        );
        send_clear_frame(
            &mut tcp_write,
            &InsecureMsg::HandshakeRejected(HandshakeRejected {
                reason: String::from(
                    "Controller doesn't match expected verification string sent by breakpoint.",
                ),
                permanent: true,
            }),
        )
        .await
        .map_err(|err| format!("Unable to send rejection message. Rejection was: {err}"))?;
        Err(err_msg)?;
    }

    Ok(ControllerReceivedHello {
        controller_keypair,
        tcp_read,
        tcp_write,
        breakpoint_pub_unverified: breakpoint_pub_key,
    })
}

/// State where the controller has received a hello from the
/// breakpoint.
///
/// This represents a partially set-up connection requiring some
/// interaction with the user before setup can complete.
pub struct ControllerReceivedHello {
    controller_keypair: XWingKeypair,
    tcp_read: OwnedReadHalf,
    tcp_write: OwnedWriteHalf,
    /// Unverified public key for a breakpoint, received over the
    /// network.
    breakpoint_pub_unverified: x_wing::EncapsulationKey,
}

impl ControllerReceivedHello {
    /// Given a breakpoint verification digest received out-of-band,
    /// finish the setup process.
    ///
    /// If this completes successfully, returns the secured connection
    /// as well as the decrypted breakpoint introduction message.
    pub async fn verify_breakpoint_and_complete_setup(
        mut self,
        expected_breakpoint_digest: &str,
    ) -> Result<(Connection, BreakpointIntro), String> {
        let key_digest = breakpoint_digest(&self.breakpoint_pub_unverified);
        if !verify_digest(expected_breakpoint_digest, &key_digest) {
            let err_msg = "Breakpoint verification string did not match. \
                 Is it possible that you have multiple breakpoints attempting \
                 to connect here at the same time and entered the verification \
                 string from the wrong one?";
            send_clear_frame(
                &mut self.tcp_write,
                &InsecureMsg::HandshakeRejected(HandshakeRejected {
                    reason: String::from(
                        "Breakpoint sent a public key that doesn't match \
                         the verification string the user entered.",
                    ),
                    permanent: false, // could be a mis-paste or a typo
                }),
            )
            .await
            .map_err(|err| format!("Unable to send rejection message. Rejection was: {err}"))?;
            return Err(err_msg.to_string());
        }
        let breakpoint_pubkey = self.breakpoint_pub_unverified; // now verified!

        // We can now trust the breakpoint's public key, which means
        // we can send them our half of the setup secret. They'll also
        // need the controller's public key so they can send us
        // *their* portion of the setup secret.

        let (controller_secret_encap, controller_secret) = breakpoint_pubkey
            .encapsulate(&mut rand::rngs::OsRng)
            .map_err(|err| format!("Failed to create controller secret: {err}"))?;

        send_clear_frame(
            &mut self.tcp_write,
            &InsecureMsg::ControllerHello(ControllerHello {
                ctrl_pub: self.controller_keypair.public_key.as_bytes(),
                setup_secret_part1: controller_secret_encap.as_bytes(),
            }),
        )
        .await
        .map_err(|err| format!("Couldn't respond to breakpoint's hello: {err}"))?;

        // The breakpoint now can send us its own half of the setup
        // secret, after which both parties will have enough
        // information to compute the stream secrets. So the
        // breakpoint also sends us the beginning of its stream,
        // including a first message.

        let breakpoint_finalize = match receive_clear_frame(&mut self.tcp_read).await {
            Ok(InsecureMsg::BreakpointSetupFinalize(finalize)) => finalize,
            Ok(wrong) => Err(format!(
                "Unexpected message type when waiting for BreakpointSetupFinalize: {}",
                Into::<&str>::into(wrong)
            ))?,
            Err(err) => Err(format!(
                "Couldn't receive data when waiting for a BreakpointSetupFinalize: {err}"
            ))?,
        };

        let breakpoint_secret_encap: x_wing::Ciphertext =
            (&breakpoint_finalize.setup_secret_part2).into();
        let breakpoint_secret = self
            .controller_keypair
            .private_key
            .decapsulate(&breakpoint_secret_encap)
            .map_err(|err| {
                format!("Unable to decrypt breakpoint's half of the setup secret: {err}")
            })?;

        let (controller_stream_key, breakpoint_stream_key) = derive_stream_keys_dual_pub(
            &self.controller_keypair.public_key,
            &breakpoint_pubkey,
            &controller_secret,
            &breakpoint_secret,
        );

        // The breakpoint has already included a message for us! We'll
        // set up the receive stream and decrypt it now, before
        // setting up the transmit stream.

        let mut rx_stream =
            DecryptionStream::new(&breakpoint_stream_key, &breakpoint_finalize.cipher_header)
                .map_err(|err| format!("Unable to start decryption stream: {err}"))?;

        let maybe_intro = decrypt_message(
            &EncryptedMsg {
                encrypted: breakpoint_finalize.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.")?,
        };

        // The last step is setting up the transmit stream and sending
        // the start of it to the breakpoint.

        let tx_stream = EncryptionStream::new(&controller_stream_key).unwrap();
        let tx_header = tx_stream.get_header();

        send_clear_frame(
            &mut self.tcp_write,
            &InsecureMsg::ControllerSetupFinalize(ControllerSetupFinalize {
                cipher_header: tx_header,
            }),
        )
        .await
        .map_err(|err| format!("Unable to send controller setup finalization: {err}"))?;

        // Now both parties should have a fully secured connection.

        let shutting_down = Arc::new(AtomicBool::from(false));
        let secured_conn = Connection {
            rx: ConnectionRx::from(self.tcp_read, rx_stream, shutting_down.clone()).spawn_rx(),
            tx: ConnectionTx::from(self.tcp_write, tx_stream, shutting_down.clone()).spawn_tx(),
            shutting_down,
        };

        Ok((secured_conn, breakpoint_intro))
    }
}

//*******************************//
//                               //
//  Breakpoint connection setup  //
//                               //
//*******************************//

/// Error occurring during breakpoint's connection setup.
pub struct BreakpointConnectError {
    /// Error message string.
    pub message: String,
    /// Whether this should be considered a permanent error,
    /// indicating that the same connection parameters should not be
    /// retried.
    pub permanent: bool,
}

/// By default, an error shouldn't be considered a permanent error.
impl From<String> for BreakpointConnectError {
    fn from(message: String) -> Self {
        Self {
            message,
            permanent: false,
        }
    }
}

/// Create a new breakpoint connection, completing once the
/// controller has responded and completed the handshake.
///
/// Resolves to a pair of channels that serve as an abstraction
/// over the secure channel.
#[allow(clippy::too_many_lines)]
pub async fn breakpoint_open_connection(
    breakpoint_keypair: &XWingKeypair,
    controller_addr: &String,
    controller_expected_digest: &str,
    which: Option<&String>,
) -> Result<Connection, BreakpointConnectError> {
    let stream = TcpStream::connect(controller_addr)
        .await
        .map_err(|err| format!("Unable to connect to {controller_addr}: {err}"))?;
    match stream.peer_addr() {
        Ok(connected) => eprintln!("Connecting to controller at {connected}"),
        Err(err) => eprintln!("Unable to get address of controller connection: {err}"),
    }

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

    // Send a hello to the controller with our public key. This will
    // need to be verified with the help of the user.

    send_clear_frame(
        &mut tcp_write,
        &InsecureMsg::BreakpointHello(BreakpointHello {
            protocol: PROTO_DUAL_PUB.to_owned(),
            break_pub: breakpoint_keypair.public_key.as_bytes(),
            expected_controller_digest: controller_expected_digest.to_owned(),
        }),
    )
    .await
    .map_err(|err| format!("Unable to send hello: {err}"))?;

    let breakpoint_digest = breakpoint_digest(&breakpoint_keypair.public_key);
    // TODO: Move this printing into the breakpoint module?
    println!("This is the verification string for the breakpoint:");
    println!();
    println!("    {breakpoint_digest}");
    println!();
    println!("Copy that string into the breakmancer controller when prompted.");

    // Once the controller has verified the breakpoint's public key,
    // it responds with its own public key (which we'll need to
    // verify) as well as its half of the setup secret.

    let controller_hello = match receive_clear_frame(&mut tcp_read).await {
        Ok(InsecureMsg::ControllerHello(hello)) => hello,
        Ok(InsecureMsg::HandshakeRejected(HandshakeRejected { reason, permanent })) => {
            Err(BreakpointConnectError {
                message: format!("Controller rejected connection: {reason}"),
                permanent,
            })?
        }
        Ok(wrong) => Err(format!(
            "Unexpected message type when waiting for ControllerHello: {}",
            Into::<&str>::into(wrong)
        ))?,
        Err(err) => Err(format!(
            "Couldn't receive data when waiting for a ControllerHello: {err}"
        ))?,
    };

    let unverified_controller_pubkey: x_wing::EncapsulationKey =
        (&controller_hello.ctrl_pub).into();
    let received_digest = controller_digest(&unverified_controller_pubkey);
    if !verify_digest(controller_expected_digest, &received_digest) {
        Err("Controller verification string did not match. \
             Is it possible that the controller has since been \
             restarted? (This would mean it is using a new set of keys.)"
            .to_string())?;
    }
    let controller_pubkey = unverified_controller_pubkey; // now verified!

    let controller_secret_encap: x_wing::Ciphertext = (&controller_hello.setup_secret_part1).into();
    let controller_secret = breakpoint_keypair
        .private_key
        .decapsulate(&controller_secret_encap)
        .map_err(|err| format!("Unable to decrypt controller's half of the setup secret: {err}"))?;

    // We can now trust the controller's public key, which means we
    // can send them our half of the setup secret and derive the full
    // shared secret for ourselves, and that in turn means we can go
    // ahead and include an encrypted introduction message.

    let (breakpoint_secret_encap, breakpoint_secret) = controller_pubkey
        .encapsulate(&mut rand::rngs::OsRng)
        .map_err(|err| format!("Failed to create breakpoint secret: {err}"))?;

    let (controller_stream_key, breakpoint_stream_key) = derive_stream_keys_dual_pub(
        &controller_pubkey,
        &breakpoint_keypair.public_key,
        &controller_secret,
        &breakpoint_secret,
    );

    let mut tx_stream = EncryptionStream::new(&breakpoint_stream_key).unwrap();
    let tx_header = tx_stream.get_header();

    let intro = &SecureMsg::BreakpointIntro(BreakpointIntro {
        which: which.cloned(),
    });
    send_clear_frame(
        &mut tcp_write,
        &InsecureMsg::BreakpointSetupFinalize(BreakpointSetupFinalize {
            setup_secret_part2: breakpoint_secret_encap.as_bytes(),
            cipher_header: tx_header,
            encrypted_intro: encrypt_message(intro, &mut tx_stream).encrypted,
        }),
    )
    .await
    .map_err(|err| format!("Couldn't send breakpoint setup finalize: {err}"))?;

    // Finally, we just need the controller's setup finalize message,
    // which will contain the stream header in the other direction.

    let controller_finalize = match receive_clear_frame(&mut tcp_read).await {
        Ok(InsecureMsg::ControllerSetupFinalize(setup)) => setup,
        Ok(InsecureMsg::HandshakeRejected(HandshakeRejected { reason, permanent })) => {
            Err(BreakpointConnectError {
                message: format!("Controller rejected connection: {reason}"),
                permanent,
            })?
        }
        Ok(wrong) => Err(format!(
            "Unexpected message type when waiting for ControllerSetupFinalize: {}",
            Into::<&str>::into(wrong)
        ))?,
        Err(err) => Err(format!(
            "Couldn't receive data when waiting for a ControllerSetupFinalize: {err}"
        ))?,
    };
    let rx_stream =
        DecryptionStream::new(&controller_stream_key, &controller_finalize.cipher_header)
            .map_err(|err| format!("Unable to start decryption stream: {err}"))?;

    // Now both parties should have a fully secured connection.

    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,
    })
}

//*************************************//
//                                     //
//  Secured connection (either party)  //
//                                     //
//*************************************//

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

impl Connection {
    /// Receive a message over the secured channel.
    pub async fn recv(&mut self) -> Option<SecureMsg> {
        self.rx.recv().await
    }

    /// Send a message over the secured channel.
    pub async fn send(&self, msg: SecureMsg) -> Result<(), SendError<SecureMsg>> {
        self.tx.send(msg).await
    }
}

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::*;

    /// Pinning test for digests.
    #[test]
    fn pubkey_digest() {
        let privkey = x_wing::DecapsulationKey::from([7u8; 32]);
        let pubkey = privkey.encapsulation_key();
        let digest_c = controller_digest(&pubkey);
        let digest_b = breakpoint_digest(&pubkey);

        // Note that Base58 does not produce the same output length
        // for a given input length -- it is content-dependent.
        assert_eq!(&digest_c, "rfFydfomQ4546wuW");
        assert_eq!(&digest_b, "3YVBK2JSiKM7Y5i4i");
    }

    /// Pinning test for key derivation
    #[test]
    fn stream_key_pinning() {
        let ctrl_pair = x_wing::DecapsulationKey::from([20u8; 32]);
        let break_pair = x_wing::DecapsulationKey::from([21u8; 32]);
        let ctrl_secret = x_wing::SharedSecret::from([22u8; 32]);
        let break_secret = x_wing::SharedSecret::from([23u8; 32]);
        let (ctrl_stream_key, break_stream_key) = derive_stream_keys_dual_pub(
            &ctrl_pair.encapsulation_key(),
            &break_pair.encapsulation_key(),
            &ctrl_secret,
            &break_secret,
        );
        assert_eq!(
            hex::encode(ctrl_stream_key.as_slice()),
            "cab71a3c070baa772ece7eb5f108032527bc3e9d6d51839ec4b394cb3afd175f"
        );
        assert_eq!(
            hex::encode(break_stream_key.as_slice()),
            "bc1aeb4d8668dbc18f7fadb2d9f9fa8d3487996016dbbb8795acebaffcfa5b0c"
        );
    }

    /// 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"),
                break_pub: [5u8; 1216],
                expected_controller_digest: String::from("abc"),
            }),
        )
        .await
        .unwrap();

        let expected = [
            // Length header, 1292 byte body.
            b"\x0D\x05\x00\x00" as &[u8],
            b"\xA1", // 1 kv map
            b"o",
            b"BreakpointHello",
            b"\xA3", // 3 kv map
            b"h",
            b"protocol",
            b"d",
            b"test",
            b"i",
            b"break_pub",
            b"Y",
            b"\x04\xC0",
            &[b'\x05'; 1216],
            b"x\x1A",
            b"expected_controller_digest",
            b"c",
            b"abc",
        ]
        .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 => Vec<u8>:

            InsecureMsg::HandshakeRejected =
                (
                    HandshakeRejected {
                        reason: String::from("not feeling it"),
                        permanent: true,
                    }
                )
                => [
                    b"\xA1qHandshakeRejected\xA2" as &[u8],
                    b"freasonnnot feeling it",
                    b"ipermanent\xF5",
                ].concat(),

            InsecureMsg::BreakpointHello =
                (
                    BreakpointHello {
                        protocol: String::from("SOMEPROTO"),
                        break_pub: [6u8; 1216],
                        expected_controller_digest: String::from("cdigest"),
                    }
                )
                => [
                    b"\xA1oBreakpointHello\xA3" as &[u8],
                    b"hprotocoliSOMEPROTO",
                    b"ibreak_pubY\x04\xC0", &[b'\x06'; 1216],
                    b"x\x1Aexpected_controller_digestgcdigest",
                ].concat(),

            InsecureMsg::ControllerHello =
                (
                    ControllerHello {
                        ctrl_pub: [6u8; 1216],
                        setup_secret_part1: [4u8; 1120],
                    }
                )
                => [
                    b"\xA1oControllerHello\xA2" as &[u8],
                    b"hctrl_pubY\x04\xC0", &[b'\x06'; 1216],
                    b"rsetup_secret_part1Y\x04\x60", &[b'\x04'; 1120],
                ].concat(),

            InsecureMsg::BreakpointSetupFinalize =
                (
                    BreakpointSetupFinalize {
                        setup_secret_part2: [5u8; 1120],
                        cipher_header: [8u8; 24],
                        encrypted_intro: [7u8; 10].to_vec(),
                    }
                )
                => [
                    b"\xA1wBreakpointSetupFinalize\xA3" as &[u8],
                    b"rsetup_secret_part2Y\x04\x60", &[b'\x05'; 1120],
                    b"mcipher_headerX\x18", &[b'\x08'; 24],
                    b"oencrypted_intro\x8A", &[b'\x07'; 10],
                ].concat(),

            InsecureMsg::ControllerSetupFinalize =
                (
                    ControllerSetupFinalize {
                        cipher_header: [9u8; 24],
                    }
                )
                => [
                    b"\xA1wControllerSetupFinalize\xA1" as &[u8],
                    b"mcipher_headerX\x18", &[b'\x09'; 24],
                ].concat(),

            InsecureMsg::EncryptedMsg =
                (
                    EncryptedMsg {
                        encrypted: [9u8; 3].to_vec(),
                    }
                )
                => b"\xA1lEncryptedMsg\xA1iencrypted\x43\x09\x09\x09".to_vec(),
        );
        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:?}"
            );
        }
    }
}