//! Canonical Peach Noise state machine.
//!
//! The browser and Android adapters are deliberately thin: all NK/KK
//! transcript processing, DH, AEAD, nonce handling, and transport rekeying
//! happen here through `snow`'s unmodified named patterns.

use snow::{
    Builder, Error, HandshakeState, TransportState,
    params::{CipherChoice, DHChoice, HashChoice, NoiseParams},
    resolvers::{CryptoResolver, DefaultResolver},
    types::{Cipher, Dh, Hash, Random},
};
use subtle::ConstantTimeEq;
use x25519_dalek::{PublicKey, StaticSecret};
use zeroize::{Zeroize, Zeroizing};

const NK: &str = "Noise_NK_25519_ChaChaPoly_BLAKE2s";
const KK: &str = "Noise_KK_25519_ChaChaPoly_BLAKE2s";
pub const NOISE_PROLOGUE: &[u8] = b"peach-sync-v1";
const HASH_LEN: usize = 32;
const TAG_LEN: usize = 16;
const MAX_NOISE_MESSAGE: usize = 65_535;

type CoreResult<T> = Result<T, String>;

enum SessionState {
    Handshake(Option<HandshakeState>),
    Transport(TransportState),
    Closed,
}

pub struct CoreNoiseSession {
    state: SessionState,
    handshake_hash: Option<[u8; HASH_LEN]>,
}

impl CoreNoiseSession {
    pub fn nk_initiator(remote_public_key: &[u8]) -> CoreResult<Self> {
        Self::nk_initiator_inner(remote_public_key, None)
    }

    #[cfg(feature = "test-constructors")]
    #[doc(hidden)]
    pub fn nk_initiator_with_ephemeral_for_test(
        remote_public_key: &[u8],
        ephemeral_private_key: &[u8],
    ) -> CoreResult<Self> {
        Self::nk_initiator_inner(remote_public_key, Some(ephemeral_private_key))
    }

    fn nk_initiator_inner(
        remote_public_key: &[u8],
        ephemeral_private_key: Option<&[u8]>,
    ) -> CoreResult<Self> {
        require_key(remote_public_key, "responder public key")?;
        let params: NoiseParams = NK.parse().map_err(display_error)?;
        let mut builder = core_builder(params)
            .remote_public_key(remote_public_key)
            .map_err(display_error)?;
        if let Some(ephemeral) = ephemeral_private_key {
            require_key(ephemeral, "initiator ephemeral private key")?;
            builder = builder.fixed_ephemeral_key_for_testing_only(ephemeral);
        }
        let handshake = builder.build_initiator().map_err(display_error)?;
        Ok(Self::handshake(handshake))
    }

    pub fn nk_responder(local_private_key: &[u8], local_public_key: &[u8]) -> CoreResult<Self> {
        Self::nk_responder_inner(local_private_key, local_public_key, None)
    }

    #[cfg(feature = "test-constructors")]
    #[doc(hidden)]
    pub fn nk_responder_with_ephemeral_for_test(
        local_private_key: &[u8],
        local_public_key: &[u8],
        ephemeral_private_key: &[u8],
    ) -> CoreResult<Self> {
        Self::nk_responder_inner(
            local_private_key,
            local_public_key,
            Some(ephemeral_private_key),
        )
    }

    fn nk_responder_inner(
        local_private_key: &[u8],
        local_public_key: &[u8],
        ephemeral_private_key: Option<&[u8]>,
    ) -> CoreResult<Self> {
        require_matching_keypair(local_private_key, local_public_key)?;
        let private = Zeroizing::new(local_private_key.to_vec());
        let params: NoiseParams = NK.parse().map_err(display_error)?;
        let mut builder = core_builder(params)
            .local_private_key(private.as_slice())
            .map_err(display_error)?;
        if let Some(ephemeral) = ephemeral_private_key {
            require_key(ephemeral, "responder ephemeral private key")?;
            builder = builder.fixed_ephemeral_key_for_testing_only(ephemeral);
        }
        let handshake = builder.build_responder().map_err(display_error)?;
        Ok(Self::handshake(handshake))
    }

    pub fn kk_initiator(
        local_private_key: &[u8],
        local_public_key: &[u8],
        remote_public_key: &[u8],
    ) -> CoreResult<Self> {
        Self::kk_initiator_inner(local_private_key, local_public_key, remote_public_key, None)
    }

    #[cfg(feature = "test-constructors")]
    #[doc(hidden)]
    pub fn kk_initiator_with_ephemeral_for_test(
        local_private_key: &[u8],
        local_public_key: &[u8],
        remote_public_key: &[u8],
        ephemeral_private_key: &[u8],
    ) -> CoreResult<Self> {
        Self::kk_initiator_inner(
            local_private_key,
            local_public_key,
            remote_public_key,
            Some(ephemeral_private_key),
        )
    }

    fn kk_initiator_inner(
        local_private_key: &[u8],
        local_public_key: &[u8],
        remote_public_key: &[u8],
        ephemeral_private_key: Option<&[u8]>,
    ) -> CoreResult<Self> {
        require_matching_keypair(local_private_key, local_public_key)?;
        require_key(remote_public_key, "responder public key")?;
        let private = Zeroizing::new(local_private_key.to_vec());
        let params: NoiseParams = KK.parse().map_err(display_error)?;
        let mut builder = core_builder(params)
            .local_private_key(private.as_slice())
            .map_err(display_error)?
            .remote_public_key(remote_public_key)
            .map_err(display_error)?;
        if let Some(ephemeral) = ephemeral_private_key {
            require_key(ephemeral, "initiator ephemeral private key")?;
            builder = builder.fixed_ephemeral_key_for_testing_only(ephemeral);
        }
        let handshake = builder.build_initiator().map_err(display_error)?;
        Ok(Self::handshake(handshake))
    }

    pub fn kk_responder(
        local_private_key: &[u8],
        local_public_key: &[u8],
        remote_public_key: &[u8],
    ) -> CoreResult<Self> {
        Self::kk_responder_inner(local_private_key, local_public_key, remote_public_key, None)
    }

    #[cfg(feature = "test-constructors")]
    #[doc(hidden)]
    pub fn kk_responder_with_ephemeral_for_test(
        local_private_key: &[u8],
        local_public_key: &[u8],
        remote_public_key: &[u8],
        ephemeral_private_key: &[u8],
    ) -> CoreResult<Self> {
        Self::kk_responder_inner(
            local_private_key,
            local_public_key,
            remote_public_key,
            Some(ephemeral_private_key),
        )
    }

    fn kk_responder_inner(
        local_private_key: &[u8],
        local_public_key: &[u8],
        remote_public_key: &[u8],
        ephemeral_private_key: Option<&[u8]>,
    ) -> CoreResult<Self> {
        require_matching_keypair(local_private_key, local_public_key)?;
        require_key(remote_public_key, "initiator public key")?;
        let private = Zeroizing::new(local_private_key.to_vec());
        let params: NoiseParams = KK.parse().map_err(display_error)?;
        let mut builder = core_builder(params)
            .local_private_key(private.as_slice())
            .map_err(display_error)?
            .remote_public_key(remote_public_key)
            .map_err(display_error)?;
        if let Some(ephemeral) = ephemeral_private_key {
            require_key(ephemeral, "responder ephemeral private key")?;
            builder = builder.fixed_ephemeral_key_for_testing_only(ephemeral);
        }
        let handshake = builder.build_responder().map_err(display_error)?;
        Ok(Self::handshake(handshake))
    }

    fn handshake(handshake: HandshakeState) -> Self {
        Self {
            state: SessionState::Handshake(Some(handshake)),
            handshake_hash: None,
        }
    }

    pub fn write_handshake(&mut self, payload: &[u8]) -> CoreResult<Vec<u8>> {
        let handshake = self.handshake_mut()?;
        if !handshake.is_my_turn() {
            return Err("Noise handshake write attempted out of turn".into());
        }
        let mut output = vec![0_u8; MAX_NOISE_MESSAGE];
        let size = handshake
            .write_message(payload, &mut output)
            .map_err(display_error)?;
        output.truncate(size);
        self.finish_if_ready()?;
        Ok(output)
    }

    pub fn read_handshake(&mut self, message: &[u8]) -> CoreResult<Vec<u8>> {
        if message.len() > MAX_NOISE_MESSAGE {
            return Err("Noise handshake message exceeds 65535 bytes".into());
        }
        let handshake = self.handshake_mut()?;
        if handshake.is_my_turn() {
            return Err("Noise handshake read attempted out of turn".into());
        }
        let mut payload = vec![0_u8; MAX_NOISE_MESSAGE];
        let size = handshake
            .read_message(message, &mut payload)
            .map_err(display_error)?;
        payload.truncate(size);
        self.finish_if_ready()?;
        Ok(payload)
    }

    pub fn encrypt(&mut self, plaintext: &[u8]) -> CoreResult<Vec<u8>> {
        if plaintext.len() > MAX_NOISE_MESSAGE - TAG_LEN {
            return Err("Noise plaintext exceeds transport limit".into());
        }
        let transport = self.transport_mut()?;
        let mut output = vec![0_u8; plaintext.len() + TAG_LEN];
        let size = transport
            .write_message(plaintext, &mut output)
            .map_err(display_error)?;
        output.truncate(size);
        Ok(output)
    }

    pub fn decrypt(&mut self, ciphertext: &[u8]) -> CoreResult<Vec<u8>> {
        if ciphertext.len() > MAX_NOISE_MESSAGE {
            return Err("Noise ciphertext exceeds transport limit".into());
        }
        let transport = self.transport_mut()?;
        let mut output = vec![0_u8; ciphertext.len()];
        let size = transport
            .read_message(ciphertext, &mut output)
            .map_err(display_error)?;
        output.truncate(size);
        Ok(output)
    }

    pub fn handshake_hash(&self) -> CoreResult<Vec<u8>> {
        self.handshake_hash
            .map(|hash| hash.to_vec())
            .ok_or_else(|| "Noise handshake is not complete".into())
    }

    pub fn rekey_outgoing(&mut self) -> CoreResult<()> {
        self.transport_mut()?.rekey_outgoing();
        Ok(())
    }

    pub fn rekey_incoming(&mut self) -> CoreResult<()> {
        self.transport_mut()?.rekey_incoming();
        Ok(())
    }

    fn handshake_mut(&mut self) -> CoreResult<&mut HandshakeState> {
        match &mut self.state {
            SessionState::Handshake(Some(handshake)) => Ok(handshake),
            SessionState::Handshake(None) => Err("Noise handshake is transitioning".into()),
            SessionState::Transport(_) => Err("Noise handshake is already complete".into()),
            SessionState::Closed => Err("Noise session is closed".into()),
        }
    }

    fn transport_mut(&mut self) -> CoreResult<&mut TransportState> {
        match &mut self.state {
            SessionState::Transport(transport) => Ok(transport),
            SessionState::Handshake(_) => Err("Noise handshake is not complete".into()),
            SessionState::Closed => Err("Noise session is closed".into()),
        }
    }

    fn finish_if_ready(&mut self) -> CoreResult<()> {
        let ready = matches!(
            &self.state,
            SessionState::Handshake(Some(handshake)) if handshake.is_handshake_finished()
        );
        if !ready {
            return Ok(());
        }

        let handshake = match std::mem::replace(&mut self.state, SessionState::Closed) {
            SessionState::Handshake(Some(handshake)) => handshake,
            _ => return Err("Noise session transition failed".into()),
        };
        let hash: [u8; HASH_LEN] = handshake
            .get_handshake_hash()
            .try_into()
            .map_err(|_| "Noise handshake hash must be 32 bytes".to_string())?;
        let transport = handshake.into_transport_mode().map_err(display_error)?;
        self.handshake_hash = Some(hash);
        self.state = SessionState::Transport(transport);
        Ok(())
    }
}

impl Drop for CoreNoiseSession {
    fn drop(&mut self) {
        if let Some(hash) = &mut self.handshake_hash {
            hash.zeroize();
        }
        // PB-26-609: snow 0.10 has no zeroize dependency at all, and its
        // default resolver keeps the cipher key inline as
        // `CipherChaChaPoly { key: [u8; 32] }` (snow-0.10.0/src/resolvers/
        // default.rs:152) with no Drop or Zeroize impl. Dropping the
        // TransportState therefore leaves both cipher keys in memory, even
        // though the per-operation ChaCha20Poly1305 it builds does scrub
        // itself. Overwrite both directions through snow's public rekey API
        // first; the session is unusable afterwards by construction.
        if let SessionState::Transport(transport) = &mut self.state {
            transport.rekey_manually(Some(&[0_u8; 32]), Some(&[0_u8; 32]));
        }
        self.state = SessionState::Closed;
    }
}

fn require_key(key: &[u8], label: &str) -> CoreResult<()> {
    if key.len() != 32 {
        return Err(format!("{label} must be 32 bytes"));
    }
    Ok(())
}

fn require_matching_keypair(private_key: &[u8], public_key: &[u8]) -> CoreResult<()> {
    require_key(private_key, "local private key")?;
    require_key(public_key, "local public key")?;
    // Zeroizing so the stack copy is scrubbed: `StaticSecret::from` takes the
    // array by value and zeroizes its own buffer on drop, but a moved-from
    // `[u8; 32]` on the stack is not cleared by Rust.
    let private = Zeroizing::new(<[u8; 32]>::try_from(private_key).expect("length checked"));
    let derived = PublicKey::from(&StaticSecret::from(*private));
    if derived.as_bytes().ct_eq(public_key).unwrap_u8() != 1 {
        return Err("local public key does not match local private key".into());
    }
    Ok(())
}

fn display_error(error: impl std::fmt::Display) -> String {
    error.to_string()
}

fn core_builder(params: NoiseParams) -> Builder<'static> {
    Builder::with_resolver(params, Box::new(RejectZeroResolver))
        .prologue(NOISE_PROLOGUE)
        .expect("fixed Noise prologue is valid")
}

struct RejectZeroResolver;

impl CryptoResolver for RejectZeroResolver {
    fn resolve_rng(&self) -> Option<Box<dyn Random>> {
        DefaultResolver.resolve_rng()
    }

    fn resolve_dh(&self, choice: &DHChoice) -> Option<Box<dyn Dh>> {
        DefaultResolver
            .resolve_dh(choice)
            .map(|inner| Box::new(RejectZeroDh { inner }) as Box<dyn Dh>)
    }

    fn resolve_hash(&self, choice: &HashChoice) -> Option<Box<dyn Hash>> {
        DefaultResolver.resolve_hash(choice)
    }

    fn resolve_cipher(&self, choice: &CipherChoice) -> Option<Box<dyn Cipher>> {
        DefaultResolver.resolve_cipher(choice)
    }
}

struct RejectZeroDh {
    inner: Box<dyn Dh>,
}

/// Small-order X25519 u-coordinates (little-endian) — the points of order
/// 1, 2, 4 and 8 (libhydrogen's contributory-behaviour blocklist). A DH
/// against one of these yields a constant, low-entropy shared secret, so
/// they are rejected as inputs outright (PB-26-014).
const X25519_SMALL_ORDER_PUBKEYS: [[u8; 32]; 9] = [
    [
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00,
    ],
    [
        0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00,
    ],
    [
        0xe0, 0xeb, 0x7a, 0x7c, 0x3b, 0x41, 0xb8, 0xae, 0x16, 0x56, 0xe3, 0xfa, 0xf1, 0x9f, 0xc4,
        0x6a, 0xda, 0x09, 0x8d, 0xeb, 0x9c, 0x32, 0xb1, 0xfd, 0x86, 0x62, 0x05, 0x16, 0x5f, 0x49,
        0xb8, 0x80,
    ],
    [
        0x5f, 0x9c, 0x95, 0xbc, 0xa3, 0x50, 0x8c, 0x24, 0xb1, 0xd0, 0xb1, 0x55, 0x9c, 0x83, 0xef,
        0x5b, 0x04, 0x44, 0x5c, 0xc4, 0x58, 0x1c, 0x8e, 0x86, 0xd8, 0x22, 0x4e, 0xdd, 0xd0, 0x9f,
        0x11, 0x57,
    ],
    [
        0xec, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0x7f,
    ],
    [
        0xed, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0x7f,
    ],
    [
        0xee, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0x7f,
    ],
    [
        0xcd, 0xeb, 0x7a, 0x7c, 0x3b, 0x41, 0xb8, 0xae, 0x16, 0x56, 0xe3, 0xfa, 0xf1, 0x9f, 0xc4,
        0x6a, 0xda, 0x09, 0x8d, 0xeb, 0x9c, 0x32, 0xb1, 0xfd, 0x86, 0x62, 0x05, 0x16, 0x5f, 0x49,
        0xb8, 0x80,
    ],
    [
        0x4c, 0x9c, 0x95, 0xbc, 0xa3, 0x50, 0x8c, 0x24, 0xb1, 0xd0, 0xb1, 0x55, 0x9c, 0x83, 0xef,
        0x5b, 0x04, 0x44, 0x5c, 0xc4, 0x58, 0x1c, 0x8e, 0x86, 0xd8, 0x22, 0x4e, 0xdd, 0xd0, 0x9f,
        0x11, 0xd7,
    ],
];

fn is_small_order_x25519_pubkey(public_key: &[u8]) -> bool {
    // snow 0.10 passes the remote static inside its 56-byte MAXDHLEN buffer
    // (32 key bytes + 24 trailing zeros), so a strict 32-byte try_from always
    // fails and the guard would never fire on a real handshake. Compare the
    // leading 32 bytes whenever a full key is present.
    let Some(key) = public_key.get(..32) else {
        return false;
    };
    let Ok(key) = <[u8; 32]>::try_from(key) else {
        return false;
    };
    X25519_SMALL_ORDER_PUBKEYS
        .iter()
        .any(|low| low.ct_eq(&key).unwrap_u8() == 1)
}

impl Dh for RejectZeroDh {
    fn name(&self) -> &'static str {
        self.inner.name()
    }

    fn pub_len(&self) -> usize {
        self.inner.pub_len()
    }

    fn priv_len(&self) -> usize {
        self.inner.priv_len()
    }

    fn set(&mut self, private_key: &[u8]) {
        self.inner.set(private_key);
    }

    fn generate(&mut self, rng: &mut dyn Random) -> Result<(), Error> {
        self.inner.generate(rng)
    }

    fn pubkey(&self) -> &[u8] {
        self.inner.pubkey()
    }

    fn privkey(&self) -> &[u8] {
        self.inner.privkey()
    }

    fn dh(&self, public_key: &[u8], output: &mut [u8]) -> Result<(), Error> {
        if is_small_order_x25519_pubkey(public_key) {
            return Err(Error::Dh);
        }
        self.inner.dh(public_key, output)?;
        let used = self.inner.dh_len();
        if output[..used].ct_eq(&vec![0_u8; used]).unwrap_u8() == 1 {
            output[..used].zeroize();
            return Err(Error::Dh);
        }
        Ok(())
    }

    fn dh_len(&self) -> usize {
        self.inner.dh_len()
    }
}

#[cfg(target_arch = "wasm32")]
mod wasm {
    use super::CoreNoiseSession;
    use wasm_bindgen::prelude::*;

    fn js_error(message: String) -> JsError {
        JsError::new(&message)
    }

    #[wasm_bindgen]
    pub struct WasmNoiseSession {
        inner: CoreNoiseSession,
    }

    #[wasm_bindgen]
    impl WasmNoiseSession {
        #[wasm_bindgen(js_name = nkInitiator)]
        pub fn nk_initiator(remote_public_key: &[u8]) -> Result<WasmNoiseSession, JsError> {
            Ok(Self {
                inner: CoreNoiseSession::nk_initiator(remote_public_key).map_err(js_error)?,
            })
        }

        #[wasm_bindgen(js_name = nkResponder)]
        pub fn nk_responder(
            local_private_key: &[u8],
            local_public_key: &[u8],
        ) -> Result<WasmNoiseSession, JsError> {
            Ok(Self {
                inner: CoreNoiseSession::nk_responder(local_private_key, local_public_key)
                    .map_err(js_error)?,
            })
        }

        #[cfg(feature = "test-constructors")]
        #[wasm_bindgen(js_name = nkInitiatorWithEphemeralForTest)]
        pub fn nk_initiator_with_ephemeral_for_test(
            remote_public_key: &[u8],
            ephemeral_private_key: &[u8],
        ) -> Result<WasmNoiseSession, JsError> {
            Ok(Self {
                inner: CoreNoiseSession::nk_initiator_with_ephemeral_for_test(
                    remote_public_key,
                    ephemeral_private_key,
                )
                .map_err(js_error)?,
            })
        }

        #[cfg(feature = "test-constructors")]
        #[wasm_bindgen(js_name = nkResponderWithEphemeralForTest)]
        pub fn nk_responder_with_ephemeral_for_test(
            local_private_key: &[u8],
            local_public_key: &[u8],
            ephemeral_private_key: &[u8],
        ) -> Result<WasmNoiseSession, JsError> {
            Ok(Self {
                inner: CoreNoiseSession::nk_responder_with_ephemeral_for_test(
                    local_private_key,
                    local_public_key,
                    ephemeral_private_key,
                )
                .map_err(js_error)?,
            })
        }

        #[wasm_bindgen(js_name = kkInitiator)]
        pub fn kk_initiator(
            local_private_key: &[u8],
            local_public_key: &[u8],
            remote_public_key: &[u8],
        ) -> Result<WasmNoiseSession, JsError> {
            Ok(Self {
                inner: CoreNoiseSession::kk_initiator(
                    local_private_key,
                    local_public_key,
                    remote_public_key,
                )
                .map_err(js_error)?,
            })
        }

        #[wasm_bindgen(js_name = kkResponder)]
        pub fn kk_responder(
            local_private_key: &[u8],
            local_public_key: &[u8],
            remote_public_key: &[u8],
        ) -> Result<WasmNoiseSession, JsError> {
            Ok(Self {
                inner: CoreNoiseSession::kk_responder(
                    local_private_key,
                    local_public_key,
                    remote_public_key,
                )
                .map_err(js_error)?,
            })
        }

        #[cfg(feature = "test-constructors")]
        #[wasm_bindgen(js_name = kkInitiatorWithEphemeralForTest)]
        pub fn kk_initiator_with_ephemeral_for_test(
            local_private_key: &[u8],
            local_public_key: &[u8],
            remote_public_key: &[u8],
            ephemeral_private_key: &[u8],
        ) -> Result<WasmNoiseSession, JsError> {
            Ok(Self {
                inner: CoreNoiseSession::kk_initiator_with_ephemeral_for_test(
                    local_private_key,
                    local_public_key,
                    remote_public_key,
                    ephemeral_private_key,
                )
                .map_err(js_error)?,
            })
        }

        #[cfg(feature = "test-constructors")]
        #[wasm_bindgen(js_name = kkResponderWithEphemeralForTest)]
        pub fn kk_responder_with_ephemeral_for_test(
            local_private_key: &[u8],
            local_public_key: &[u8],
            remote_public_key: &[u8],
            ephemeral_private_key: &[u8],
        ) -> Result<WasmNoiseSession, JsError> {
            Ok(Self {
                inner: CoreNoiseSession::kk_responder_with_ephemeral_for_test(
                    local_private_key,
                    local_public_key,
                    remote_public_key,
                    ephemeral_private_key,
                )
                .map_err(js_error)?,
            })
        }

        #[wasm_bindgen(js_name = writeHandshake)]
        pub fn write_handshake(&mut self, payload: &[u8]) -> Result<Vec<u8>, JsError> {
            self.inner.write_handshake(payload).map_err(js_error)
        }

        #[wasm_bindgen(js_name = readHandshake)]
        pub fn read_handshake(&mut self, message: &[u8]) -> Result<Vec<u8>, JsError> {
            self.inner.read_handshake(message).map_err(js_error)
        }

        pub fn encrypt(&mut self, plaintext: &[u8]) -> Result<Vec<u8>, JsError> {
            self.inner.encrypt(plaintext).map_err(js_error)
        }

        pub fn decrypt(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>, JsError> {
            self.inner.decrypt(ciphertext).map_err(js_error)
        }

        #[wasm_bindgen(getter, js_name = handshakeHash)]
        pub fn handshake_hash(&self) -> Result<Vec<u8>, JsError> {
            self.inner.handshake_hash().map_err(js_error)
        }

        #[wasm_bindgen(js_name = rekeyOutgoing)]
        pub fn rekey_outgoing(&mut self) -> Result<(), JsError> {
            self.inner.rekey_outgoing().map_err(js_error)
        }

        #[wasm_bindgen(js_name = rekeyIncoming)]
        pub fn rekey_incoming(&mut self) -> Result<(), JsError> {
            self.inner.rekey_incoming().map_err(js_error)
        }
    }
}

#[cfg(not(target_arch = "wasm32"))]
mod jni_bridge;

#[cfg(test)]
mod tests {
    use super::{
        CoreNoiseSession, KK, MAX_NOISE_MESSAGE, NK, NOISE_PROLOGUE, RejectZeroResolver,
        X25519_SMALL_ORDER_PUBKEYS, is_small_order_x25519_pubkey,
    };
    use blake2::{
        Blake2sMac256,
        digest::{KeyInit as BlakeKeyInit, Mac},
    };
    use chacha20poly1305::{
        ChaCha20Poly1305, Nonce,
        aead::{Aead, Payload},
    };
    use snow::{Builder, params::NoiseParams};
    use x25519_dalek::{PublicKey, StaticSecret};

    fn keypair(seed: u8) -> ([u8; 32], [u8; 32]) {
        let private = [seed; 32];
        let public = *PublicKey::from(&StaticSecret::from(private)).as_bytes();
        (private, public)
    }

    fn hex(value: &str) -> Vec<u8> {
        assert!(value.len().is_multiple_of(2));
        value
            .as_bytes()
            .chunks_exact(2)
            .map(|pair| u8::from_str_radix(std::str::from_utf8(pair).unwrap(), 16).unwrap())
            .collect()
    }

    fn official_pair(
        pattern: &str,
        initiator_static: Option<&str>,
        initiator_ephemeral: &str,
        initiator_remote_static: &str,
        responder_static: &str,
        responder_ephemeral: &str,
        responder_remote_static: Option<&str>,
    ) -> (CoreNoiseSession, CoreNoiseSession) {
        let prologue = hex("4a6f686e2047616c74");
        let responder_private = hex(responder_static);
        let responder_public = PublicKey::from(&StaticSecret::from(
            <[u8; 32]>::try_from(&responder_private[..]).unwrap(),
        ));
        assert_eq!(
            responder_public.as_bytes(),
            &hex(initiator_remote_static)[..]
        );

        let params: NoiseParams = pattern.parse().unwrap();
        let initiator_ephemeral = hex(initiator_ephemeral);
        let initiator_static = initiator_static.map(hex);
        let responder_ephemeral = hex(responder_ephemeral);
        let responder_remote_static = responder_remote_static.map(hex);
        let mut initiator_builder =
            Builder::with_resolver(params.clone(), Box::new(RejectZeroResolver))
                .prologue(&prologue)
                .unwrap()
                .remote_public_key(responder_public.as_bytes())
                .unwrap()
                .fixed_ephemeral_key_for_testing_only(&initiator_ephemeral);
        if let Some(private) = &initiator_static {
            initiator_builder = initiator_builder.local_private_key(private).unwrap();
        }

        let mut responder_builder = Builder::with_resolver(params, Box::new(RejectZeroResolver))
            .prologue(&prologue)
            .unwrap()
            .local_private_key(&responder_private)
            .unwrap()
            .fixed_ephemeral_key_for_testing_only(&responder_ephemeral);
        if let Some(remote_public) = &responder_remote_static {
            responder_builder = responder_builder.remote_public_key(remote_public).unwrap();
        }

        (
            CoreNoiseSession::handshake(initiator_builder.build_initiator().unwrap()),
            CoreNoiseSession::handshake(responder_builder.build_responder().unwrap()),
        )
    }

    fn assert_cacophony_vector(
        mut initiator: CoreNoiseSession,
        mut responder: CoreNoiseSession,
        handshake_hash: &str,
        messages: &[(&str, &str)],
    ) {
        for (index, (payload, expected_ciphertext)) in messages.iter().enumerate() {
            let plaintext = hex(payload);
            let expected = hex(expected_ciphertext);
            if index < 2 {
                let (sender, receiver) = if index % 2 == 0 {
                    (&mut initiator, &mut responder)
                } else {
                    (&mut responder, &mut initiator)
                };
                let actual = sender.write_handshake(&plaintext).unwrap();
                assert_eq!(actual, expected, "handshake message {index}");
                assert_eq!(
                    receiver.read_handshake(&actual).unwrap(),
                    plaintext,
                    "handshake payload {index}",
                );
            } else {
                let (sender, receiver) = if index % 2 == 0 {
                    (&mut initiator, &mut responder)
                } else {
                    (&mut responder, &mut initiator)
                };
                let actual = sender.encrypt(&plaintext).unwrap();
                assert_eq!(actual, expected, "transport message {index}");
                assert_eq!(
                    receiver.decrypt(&actual).unwrap(),
                    plaintext,
                    "transport payload {index}",
                );
            }
        }
        let expected_hash = hex(handshake_hash);
        assert_eq!(initiator.handshake_hash().unwrap(), expected_hash);
        assert_eq!(responder.handshake_hash().unwrap(), expected_hash);
    }

    fn complete(initiator: &mut CoreNoiseSession, responder: &mut CoreNoiseSession) {
        let message1 = initiator.write_handshake(&[]).unwrap();
        assert!(responder.read_handshake(&message1).unwrap().is_empty());
        let message2 = responder.write_handshake(&[]).unwrap();
        assert!(initiator.read_handshake(&message2).unwrap().is_empty());
        assert_eq!(
            initiator.handshake_hash().unwrap(),
            responder.handshake_hash().unwrap()
        );
    }

    #[test]
    fn mismatched_prologue_rejects_handshake() {
        let (private, _) = keypair(7);
        let remote = keypair(8).1;
        let params: NoiseParams = NK.parse().unwrap();
        let mut initiator = Builder::with_resolver(params.clone(), Box::new(RejectZeroResolver))
            .prologue(b"peach-sync-old")
            .unwrap()
            .remote_public_key(&remote)
            .unwrap()
            .fixed_ephemeral_key_for_testing_only(&[1; 32])
            .build_initiator()
            .unwrap();
        let mut responder = Builder::with_resolver(params, Box::new(RejectZeroResolver))
            .prologue(NOISE_PROLOGUE)
            .unwrap()
            .local_private_key(&private)
            .unwrap()
            .fixed_ephemeral_key_for_testing_only(&[2; 32])
            .build_responder()
            .unwrap();
        let mut message = [0_u8; MAX_NOISE_MESSAGE];
        let size = initiator.write_message(&[], &mut message).unwrap();
        assert!(responder.read_message(&message[..size], &mut []).is_err());
    }

    #[test]
    fn nk_round_trip_uses_one_state_machine() {
        let (responder_private, responder_public) = keypair(7);
        let mut initiator = CoreNoiseSession::nk_initiator(&responder_public).unwrap();
        let mut responder =
            CoreNoiseSession::nk_responder(&responder_private, &responder_public).unwrap();
        complete(&mut initiator, &mut responder);

        let ciphertext = initiator.encrypt(b"peach").unwrap();
        assert_eq!(responder.decrypt(&ciphertext).unwrap(), b"peach");
        let response = responder.encrypt(b"sync").unwrap();
        assert_eq!(initiator.decrypt(&response).unwrap(), b"sync");
    }

    // PB-26-014: contributory behaviour — a small-order remote static must
    // fail the handshake at the DH step instead of contributing a constant
    // low-entropy secret to the keying material. Drive ALL nine entries:
    // seven have an all-zero DH output (the pre-existing check catches them)
    // but two produce non-zero output and only this guard rejects them.
    #[test]
    fn small_order_remote_static_fails_handshake() {
        let (local_private, local_public) = keypair(11);
        for low in X25519_SMALL_ORDER_PUBKEYS.iter() {
            let mut initiator =
                CoreNoiseSession::kk_initiator(&local_private, &local_public, low).unwrap();
            assert!(
                initiator.write_handshake(&[]).is_err(),
                "handshake against small-order static must fail"
            );
            assert!(is_small_order_x25519_pubkey(low));
        }
        let (_, real_public) = keypair(9);
        assert!(!is_small_order_x25519_pubkey(&real_public));
    }

    // snow hands the guard the remote static inside its 56-byte MAXDHLEN
    // buffer (32 key bytes + 24 trailing zeros) — the guard must compare the
    // leading 32 bytes, not require exactly 32 (PB-26-014 regression).
    #[test]
    fn small_order_guard_handles_snow_padded_buffers() {
        for low in X25519_SMALL_ORDER_PUBKEYS.iter() {
            let mut padded = [0u8; 56];
            padded[..32].copy_from_slice(low);
            assert!(is_small_order_x25519_pubkey(&padded));
        }
        let (_, real_public) = keypair(13);
        let mut padded = [0u8; 56];
        padded[..32].copy_from_slice(&real_public);
        assert!(!is_small_order_x25519_pubkey(&padded));
    }

    #[test]
    fn kk_round_trip_authenticates_both_static_keys() {
        let (initiator_private, initiator_public) = keypair(11);
        let (responder_private, responder_public) = keypair(19);
        let mut initiator = CoreNoiseSession::kk_initiator(
            &initiator_private,
            &initiator_public,
            &responder_public,
        )
        .unwrap();
        let mut responder = CoreNoiseSession::kk_responder(
            &responder_private,
            &responder_public,
            &initiator_public,
        )
        .unwrap();
        complete(&mut initiator, &mut responder);
    }

    #[test]
    fn mismatched_local_keypair_is_rejected() {
        let (private, _) = keypair(2);
        let (_, wrong_public) = keypair(3);
        let result = CoreNoiseSession::nk_responder(&private, &wrong_public);
        assert!(matches!(result, Err(error) if error.contains("does not match")));
    }

    #[test]
    fn tampered_transport_frame_is_rejected() {
        let (responder_private, responder_public) = keypair(23);
        let mut initiator = CoreNoiseSession::nk_initiator(&responder_public).unwrap();
        let mut responder =
            CoreNoiseSession::nk_responder(&responder_private, &responder_public).unwrap();
        complete(&mut initiator, &mut responder);
        let mut ciphertext = initiator.encrypt(b"vault").unwrap();
        ciphertext[0] ^= 0x80;
        assert!(responder.decrypt(&ciphertext).is_err());
    }

    #[test]
    fn cacophony_nk_vector_matches_shared_core() {
        let pair = official_pair(
            NK,
            None,
            "893e28b9dc6ca8d611ab664754b8ceb7bac5117349a4439a6b0569da977c464a",
            "31e0303fd6418d2f8c0e78b91f22e8caed0fbe48656dcf4767e4834f701b8f62",
            "4a3acbfdb163dec651dfa3194dece676d437029c62a408b4c5ea9114246e4893",
            "bbdb4cdbd309f1a1f2e1456967fe288cadd6f712d65dc7b7793d5e63da6b375b",
            None,
        );
        assert_cacophony_vector(
            pair.0,
            pair.1,
            "d7244d974066aae2376f7ba5534f60a6e4e82cd7c9751e226cae3928e6b49f14",
            &[
                (
                    "4c756477696720766f6e204d69736573",
                    "ca35def5ae56cec33dc2036731ab14896bc4c75dbb07a61f879f8e3afa4c794454ae7612d1724af42adb130160a9a94e67b5b169b4e00c189f6467cd17eb7cad",
                ),
                (
                    "4d757272617920526f746862617264",
                    "95ebc60d2b1fa672c1f46a8aa265ef51bfe38e7ccb39ec5be34069f144808843986a5c929337e337ac8b4a074af12ab9f76318a5f18c8b599a443af07383ce",
                ),
                (
                    "462e20412e20486179656b",
                    "550027c7a5d450017bcb5e12b8253b1c53fd2213aeda84891d5f95",
                ),
                (
                    "4361726c204d656e676572",
                    "dfbce0c38210ccee35e830aca9dd8b8b3997b933e75bfc8864b759",
                ),
            ],
        );
    }

    #[test]
    fn cacophony_kk_vector_matches_shared_core() {
        let pair = official_pair(
            KK,
            Some("e61ef9919cde45dd5f82166404bd08e38bceb5dfdfded0a34c8df7ed542214d1"),
            "893e28b9dc6ca8d611ab664754b8ceb7bac5117349a4439a6b0569da977c464a",
            "31e0303fd6418d2f8c0e78b91f22e8caed0fbe48656dcf4767e4834f701b8f62",
            "4a3acbfdb163dec651dfa3194dece676d437029c62a408b4c5ea9114246e4893",
            "bbdb4cdbd309f1a1f2e1456967fe288cadd6f712d65dc7b7793d5e63da6b375b",
            Some("6bc3822a2aa7f4e6981d6538692b3cdf3e6df9eea6ed269eb41d93c22757b75a"),
        );
        assert_cacophony_vector(
            pair.0,
            pair.1,
            "1362b8627a00907ce11e558aba8ce7cbca88e83f0e84ce7db5159b1c3e25ab59",
            &[
                (
                    "4c756477696720766f6e204d69736573",
                    "ca35def5ae56cec33dc2036731ab14896bc4c75dbb07a61f879f8e3afa4c7944266a5f53784aa3becb0f7485c2759c328937867a4cbaafef07422b0725e098be",
                ),
                (
                    "4d757272617920526f746862617264",
                    "95ebc60d2b1fa672c1f46a8aa265ef51bfe38e7ccb39ec5be34069f144808843008aeea5d76d6abcbab87a18502c8a8352d9933ac11e2a7d228038d721e31e",
                ),
                (
                    "462e20412e20486179656b",
                    "5f92113edf78c3e56e6d67201f5f9e0c8f2930c3e1ffb64ede0358",
                ),
                (
                    "4361726c204d656e676572",
                    "30ebbd9cdcef7f40d99c8cd11e880dac28f5c9e5032c1059b3b56a",
                ),
            ],
        );
    }

    #[test]
    fn rfc_8439_chacha20_poly1305_aead_vector_matches_snow_cipher() {
        let key = hex("808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f");
        let nonce = hex("070000004041424344454647");
        let aad = hex("50515253c0c1c2c3c4c5c6c7");
        let plaintext = hex(concat!(
            "4c616469657320616e642047656e746c656d656e206f662074686520636c6173",
            "73206f66202739393a204966204920636f756c64206f6666657220796f75206f",
            "6e6c79206f6e652074697020666f7220746865206675747572652c2073756e73",
            "637265656e20776f756c642062652069742e",
        ));
        let expected = hex(concat!(
            "d31a8d34648e60db7b86afbc53ef7ec2a4aded51296e08fea9e2b5a736ee62d6",
            "3dbea45e8ca9671282fafb69da92728b1a71de0a9e060b2905d6a5b67ecd3b36",
            "92ddbd7f2d778b8c9803aee328091b58fab324e4fad675945585808b4831d7bc",
            "3ff4def08e4b7a9de576d26586cec64b6116",
            "1ae10b594f09e26a7e902ecbd0600691",
        ));
        let cipher = ChaCha20Poly1305::new_from_slice(&key).unwrap();
        assert_eq!(
            cipher
                .encrypt(
                    Nonce::from_slice(&nonce),
                    Payload {
                        msg: &plaintext,
                        aad: &aad,
                    },
                )
                .unwrap(),
            expected,
        );
    }

    /// Source of the item starting at `marker`, cut at the next top-level item
    /// so the search cannot drift into a later function or into these tests.
    fn source_after(marker: &str) -> String {
        let source = include_str!("lib.rs");
        let after = source.split(marker).nth(1).expect("marker is present");
        let end = after
            .find("\nfn ")
            .or_else(|| after.find("\nimpl "))
            .unwrap_or(after.len());
        after[..end].to_string()
    }

    // PB-26-608: `require_matching_keypair` copies the private key onto the
    // stack. Rust does not scrub a moved-from `[u8; 32]`, and stack residue
    // cannot be observed portably, so this pins the wrapper instead.
    #[test]
    fn pb_26_608_private_key_copy_is_zeroized() {
        let body = source_after("fn require_matching_keypair");
        assert!(
            body.contains("Zeroizing::new"),
            "the private-key copy must be Zeroizing-wrapped"
        );
    }

    // PB-26-609: proves the mechanism `CoreNoiseSession::drop` relies on —
    // rekeying both directions with zeros really does replace the stored
    // ChaChaPoly key rather than adding a second one.
    #[test]
    fn pb_26_609_zero_rekey_replaces_the_transport_cipher_key() {
        let params: NoiseParams = NK.parse().unwrap();
        let responder_static = keypair(9);
        let mut handshake_initiator = Builder::new(params.clone())
            .remote_public_key(&responder_static.1)
            .unwrap()
            .build_initiator()
            .unwrap();
        let mut handshake_responder = Builder::new(params)
            .local_private_key(&responder_static.0)
            .unwrap()
            .build_responder()
            .unwrap();

        let mut wire = [0_u8; 1024];
        let mut out = [0_u8; 1024];
        let len = handshake_initiator.write_message(&[], &mut wire).unwrap();
        handshake_responder
            .read_message(&wire[..len], &mut out)
            .unwrap();
        let len = handshake_responder.write_message(&[], &mut wire).unwrap();
        handshake_initiator
            .read_message(&wire[..len], &mut out)
            .unwrap();

        let mut initiator = handshake_initiator
            .into_transport_mode()
            .expect("handshake must be finished");
        let mut responder = handshake_responder
            .into_transport_mode()
            .expect("handshake must be finished");

        let mut frame = [0_u8; 64];
        let len = initiator.write_message(b"secret", &mut frame).unwrap();
        let frame = frame[..len].to_vec();

        initiator.rekey_manually(Some(&[0_u8; 32]), Some(&[0_u8; 32]));

        // A zeroed key cannot decrypt the frame the previous key produced, so
        // the old key material is gone from the live state.
        assert!(initiator.read_message(&frame, &mut [0_u8; 64]).is_err());
        // And a peer holding the real key cannot read what the rekeyed side
        // writes next, so the replacement is real, not additive.
        let mut next = [0_u8; 64];
        let len = initiator.write_message(b"after", &mut next).unwrap();
        assert!(
            responder
                .read_message(&next[..len], &mut [0_u8; 64])
                .is_err()
        );
    }

    // PB-26-609: pins that drop performs the zero rekey. Heap residue cannot be
    // observed portably, so this asserts the call site rather than the bytes.
    #[test]
    fn pb_26_609_drop_zeroizes_transport_cipher_keys() {
        let body = source_after("impl Drop for CoreNoiseSession");
        assert!(
            body.contains("rekey_manually"),
            "drop must overwrite both transport cipher keys before closing"
        );
    }

    #[test]
    fn blake2_reference_keyed_vector_matches_snow_hash_dependency() {
        let key: Vec<u8> = (0_u8..32).collect();
        let mut mac = <Blake2sMac256 as BlakeKeyInit>::new_from_slice(&key).unwrap();
        Mac::update(&mut mac, &[]);
        assert_eq!(
            Mac::finalize(mac).into_bytes().as_slice(),
            &hex("48a8997da407876b3d79c0d92325ad3b89cbb754d86ab71aee047ad345fd2c49"),
        );
    }
}