elaeis4/quicktap/src/noise/handshake/response.rs

//! `Noise_IKpsk2` handshake response packet

use x25519_dalek::PublicKey;

use crate::noise::error::NoiseError;
use crate::noise::handshake::{HandshakeState, needs_cookie};
use crate::qcrypto::aead::{qcrypto_aead, qcrypto_aead_decrypt};
use crate::qcrypto::hashes::{qcrypto_hash_twice, qcrypto_mac};
use crate::qcrypto::hkdf::qcrypto_hkdf;
use crate::qcrypto::LABEL_MAC1;
use crate::qcrypto::pki::{qcrypto_dh_generate_longterm, qcrypto_dh_longterm};

struct HandshakeResponseRaw {
    sender: [u8; 4],
    receiver: [u8; 4],
    ephemeral: [u8; 32],
    empty: [u8; 16],
    mac1: [u8; 16],
    mac2: [u8; 16]
}
impl HandshakeResponseRaw {
    fn to_bytes(&self, session: &mut HandshakeState) -> [u8; 92] {
        let mut output_array = [0u8; 92];

        output_array[0] = 2u8;
        output_array[4..8].copy_from_slice(&self.sender);
        output_array[8..12].copy_from_slice(&self.receiver);
        output_array[12..44].copy_from_slice(&self.ephemeral);
        output_array[44..60].copy_from_slice(&self.empty);

        let mac1: [u8; 16] = qcrypto_mac(&qcrypto_hash_twice(LABEL_MAC1.as_bytes(), session.s_pub_i.as_bytes()), &output_array[..60]);

        output_array[60..76].copy_from_slice(&mac1);

        let mac2 = if needs_cookie(session) { qcrypto_mac(&session.cookies[session.cookies.len() - 1].cookie, &output_array[..76]) } else { [0u8; 16] };

        output_array[76..92].copy_from_slice(&mac2);

        output_array
    }

    #[allow(clippy::unwrap_used)] // Only used for type conversions in known safe ways
    fn from_bytes(packet: [u8; 92]) -> Self {
        Self {
            sender: packet[4..8].try_into().unwrap(),
            receiver: packet[8..12].try_into().unwrap(),
            ephemeral: packet[12..44].try_into().unwrap(),
            empty: packet[44..60].try_into().unwrap(),
            mac1: packet[60..76].try_into().unwrap(),
            mac2: packet[76..92].try_into().unwrap()
        }
    }
}

/// Creates a handshake response packet using the current active handshake session.
/// # Errors
/// This function will error if an encryption step is unsuccessful
/// # Panics
/// This function, while containing unwraps, will never panic.
#[allow(clippy::unwrap_used)] // Used for known safe type conversions only
pub fn generate_handshake_response(session: &mut HandshakeState) -> Result<[u8; 92], NoiseError> {
    let eph_keypair = qcrypto_dh_generate_longterm();

    let mut msg = HandshakeResponseRaw {
        sender: [0u8; 4],
        receiver: [0u8; 4],
        ephemeral: [0u8; 32],
        empty: [0u8; 16],
        mac1: [0u8; 16],
        mac2: [0u8; 16]
    };

    msg.receiver = session.i_i.to_le_bytes();
    msg.sender = session.i_r.to_le_bytes();

    session.ck = qcrypto_hkdf::<1>(&session.ck, eph_keypair.1.as_bytes())[0];

    msg.ephemeral = eph_keypair.1.to_bytes();

    session.h = qcrypto_hash_twice(&session.h, &msg.ephemeral);

    session.ck = qcrypto_hkdf::<1>(&session.ck, qcrypto_dh_longterm(&eph_keypair.0, &session.e_pub_i).as_bytes())[0];
    session.ck = qcrypto_hkdf::<1>(&session.ck, qcrypto_dh_longterm(&eph_keypair.0, &session.s_pub_i).as_bytes())[0];

    let cr_t_k = qcrypto_hkdf::<3>(&session.ck, &session.q);
    session.ck = cr_t_k[0];
    let t = cr_t_k[1];
    let k = cr_t_k[2];

    session.h = qcrypto_hash_twice(&session.h, &t);

    println!("{:?} {:?} {:?}", k, 0, session.h);
    msg.empty = match qcrypto_aead(&k, 0, &[], &session.h) {
        Ok(s) => s.try_into().unwrap(),
        Err(e) => return Err(NoiseError::ChaCha20Error(e))
    };

    session.h = qcrypto_hash_twice(&session.h, &msg.empty);

    Ok(msg.to_bytes(session))
}

/// Decrypts a handshake response packet using the current active handshake session.
/// # Errors
/// This function will error if a decryption step is unsuccessful
/// # Panics
/// This function, while containing unwraps, will never panic.
pub fn parse_handshake_response(session: &mut HandshakeState, packet: [u8; 92]) -> Result<(), NoiseError> {
    let msg = HandshakeResponseRaw::from_bytes(packet);

    let e_pub_r = PublicKey::from(msg.ephemeral);

    let mut ck = qcrypto_hkdf::<1>(&session.ck, e_pub_r.as_bytes())[0];
    let mut h = qcrypto_hash_twice(&session.h, &msg.ephemeral);

    ck = qcrypto_hkdf::<1>(&ck, qcrypto_dh_longterm(&session.e_priv_me, &e_pub_r).as_bytes())[0];
    ck = qcrypto_hkdf::<1>(&ck, qcrypto_dh_longterm(session.s_priv_me, &e_pub_r).as_bytes())[0];

    let cr_t_k = qcrypto_hkdf::<3>(&ck, &session.q);
    ck = cr_t_k[0];
    let t = cr_t_k[1];
    let k = cr_t_k[2];

    h = qcrypto_hash_twice(&h, &t);

    let empty = match qcrypto_aead_decrypt(&k, 0, &msg.empty, &h) {
        Ok(s) => s,
        Err(e) => return Err(NoiseError::ChaCha20Error(e))
    };

    if !empty.is_empty() {
        return Err(NoiseError::PacketUnauthenticated)
    }

    h = qcrypto_hash_twice(&h, &msg.empty);

    let mac1: [u8; 16] = qcrypto_mac(&qcrypto_hash_twice(LABEL_MAC1.as_bytes(), session.s_pub_i.as_bytes()), &packet[..60]);
    let mac2 = if needs_cookie(session) { qcrypto_mac(&session.cookies[session.cookies.len() - 1].cookie, &packet[..76]) } else { [0u8; 16] };

    if mac1 != msg.mac1 || mac2 != msg.mac2 {
        return Err(NoiseError::PacketUnauthenticated)
    }

    if msg.receiver != session.i_i.to_le_bytes() {
        return Err(NoiseError::UnrelatedHandshakePacket)
    }

    session.e_pub_r = e_pub_r;
    session.ck = ck;
    session.h = h;

    Ok(())
}