jkipsec 0.1.0

//! IKEv2 cryptographic primitives.
//!
//! For our MVP we hard-code one suite - the modern default that iOS' first
//! proposal selects:
//!
//! - **ENCR**: AES-GCM-16 with 256-bit key (RFC 5282)
//! - **PRF** : HMAC-SHA-256 (output 32 B)
//! - **DH**  : group 19 = NIST P-256 (RFC 5903)
//! - **INTEG**: NONE (AEAD)
//!
//! Everything that depends on the suite (key sizes, ICV length, etc.) lives
//! in the [`Suite`] constants so it's mechanical to add a second suite later.

// Jackson Coxson

use hmac::{Hmac, Mac};
use sha2::Sha256;

/// A negotiable IKEv2 cipher / integrity / PRF / DH suite. Keep small - only
/// the suites we actually want to interop with are listed.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Suite {
    /// Modern iOS default: AES-GCM-16-256 (AEAD) + PRF SHA-256 + ECP-256.
    AesGcm256Sha256Dh19,
    /// Older iOS fallback (iPhone 5 / 8 era): AES-CBC-256 + HMAC-SHA-256-128 +
    /// PRF SHA-256 + ECP-256.
    AesCbc256Sha256Dh19,
}

/// Per-suite static sizes and IDs. Computed once and matched against peer
/// proposals during negotiation; also used by the encrypt/decrypt paths to
/// know how to slice payloads and how big the keymat block is.
#[derive(Debug, Clone, Copy)]
pub struct SuiteParams {
    /// The suite these parameters describe.
    pub suite: Suite,
    /// IKEv2 ENCR transform ID (e.g. 20 = AES-GCM-16, 12 = AES-CBC).
    pub encr_id: u16,
    /// Negotiated AES key length in **bits** (256).
    pub encr_keylen_bits: u16,
    /// AES key length in bytes (32).
    pub encr_key_bytes: usize,
    /// Trailing per-direction key bytes that are *not* the AES key (GCM salt).
    /// 0 for AES-CBC.
    pub encr_salt_bytes: usize,
    /// On-the-wire IV (8 for GCM, 16 for CBC).
    pub encr_iv_bytes: usize,
    /// ICV size on the wire (16 for both - GCM tag and HMAC-SHA-256-128).
    pub encr_icv_bytes: usize,
    /// True for AEAD ciphers (GCM); false for "encrypt-then-MAC" (CBC + HMAC).
    pub aead: bool,
    /// IKEv2 INTEG transform ID (e.g. 12 = HMAC-SHA-256-128). For AEAD this is
    /// `None` (NO_INTEG).
    pub integ_id: Option<u16>,
    /// HMAC key length in bytes. 0 for AEAD.
    pub integ_key_bytes: usize,
    /// PRF output length (32 for SHA-256).
    pub prf_bytes: usize,
}

impl Suite {
    /// Per-suite static sizes and IDs.
    pub const fn params(self) -> SuiteParams {
        match self {
            Suite::AesGcm256Sha256Dh19 => SuiteParams {
                suite: self,
                encr_id: 20, // ENCR_AES_GCM_16
                encr_keylen_bits: 256,
                encr_key_bytes: 32,
                encr_salt_bytes: 4,
                encr_iv_bytes: 8,
                encr_icv_bytes: 16,
                aead: true,
                integ_id: None,
                integ_key_bytes: 0,
                prf_bytes: 32,
            },
            Suite::AesCbc256Sha256Dh19 => SuiteParams {
                suite: self,
                encr_id: 12, // ENCR_AES_CBC
                encr_keylen_bits: 256,
                encr_key_bytes: 32,
                encr_salt_bytes: 0,
                encr_iv_bytes: 16,
                encr_icv_bytes: 16,
                aead: false,
                integ_id: Some(12), // AUTH_HMAC_SHA2_256_128
                integ_key_bytes: 32,
                prf_bytes: 32,
            },
        }
    }

    /// All suites we support, in preferred order. `select_proposal` walks
    /// peer's offers and picks the first one matching any of these.
    pub fn supported() -> &'static [Suite] {
        &[Suite::AesGcm256Sha256Dh19, Suite::AesCbc256Sha256Dh19]
    }
}

impl SuiteParams {
    /// Per-direction key+salt blob length (the SK_ei / SK_er chunk).
    pub const fn sk_e_len(&self) -> usize {
        self.encr_key_bytes + self.encr_salt_bytes
    }
    /// Total `prf+(SKEYSEED, ...)` output the IKE_SA needs:
    /// `SK_d | SK_ai | SK_ar | SK_ei | SK_er | SK_pi | SK_pr`.
    pub const fn keymat_len(&self) -> usize {
        self.prf_bytes              // SK_d
            + self.integ_key_bytes * 2  // SK_ai + SK_ar
            + self.sk_e_len() * 2       // SK_ei + SK_er
            + self.prf_bytes * 2 // SK_pi + SK_pr
    }
    /// Total per-direction CHILD_SA keymat: encr key (+ salt) + integ key.
    pub const fn child_per_dir_len(&self) -> usize {
        self.sk_e_len() + self.integ_key_bytes
    }
}

/// Backward-compat shim: old constants used by tests / serializer paths still
/// referring to `crypto::suite::*`. These now reflect the **AES-GCM** suite
/// only - code that needs CBC sizes must call `Suite::params()` instead.
#[allow(missing_docs)]
pub mod suite {
    use super::Suite;
    const P: super::SuiteParams = Suite::AesGcm256Sha256Dh19.params();
    pub const ENCR_KEY_LEN: usize = P.encr_key_bytes;
    pub const ENCR_SALT_LEN: usize = P.encr_salt_bytes;
    pub const ENCR_IV_LEN: usize = P.encr_iv_bytes;
    pub const ENCR_ICV_LEN: usize = P.encr_icv_bytes;
    pub const PRF_KEY_LEN: usize = P.prf_bytes;
    pub const SK_D_LEN: usize = P.prf_bytes;
    pub const SK_P_LEN: usize = P.prf_bytes;
    pub const SK_E_LEN: usize = P.encr_key_bytes + P.encr_salt_bytes;
    pub const KEYMAT_LEN: usize = P.keymat_len();
}

// ------------------------------------------------------------------- PRF (HMAC)

type HmacSha256 = Hmac<Sha256>;

/// `prf(key, data)` per RFC 7296 §2.13 with PRF_HMAC_SHA2_256 - returns 32 bytes.
pub fn prf(key: &[u8], data: &[u8]) -> [u8; 32] {
    // Qualified path: aes_gcm 0.10 still uses crypto-common 0.1, while hmac
    // 0.13 is on 0.2. Two `KeyInit` traits coexist; spell out which one.
    let mut mac =
        <HmacSha256 as hmac::KeyInit>::new_from_slice(key).expect("HMAC accepts any key length");
    mac.update(data);
    mac.finalize().into_bytes().into()
}

/// `prf+(key, data)` per RFC 7296 §2.13. Generates `out_len` bytes by chaining
/// `T_n = prf(key, T_{n-1} | data | n)`.
pub fn prf_plus(key: &[u8], data: &[u8], out_len: usize) -> Vec<u8> {
    let mut out = Vec::with_capacity(out_len);
    let mut t: Vec<u8> = Vec::new();
    let mut counter: u8 = 1;
    while out.len() < out_len {
        let mut input = Vec::with_capacity(t.len() + data.len() + 1);
        input.extend_from_slice(&t);
        input.extend_from_slice(data);
        input.push(counter);
        t = prf(key, &input).to_vec();
        out.extend_from_slice(&t);
        counter = counter.checked_add(1).expect("prf+ exceeds 255 iterations");
    }
    out.truncate(out_len);
    out
}

// ----------------------------------------------------------------- ECDH (P-256)

/// One side of an ECDH P-256 exchange.
pub struct DhEphemeral {
    secret: p256::ecdh::EphemeralSecret,
}

impl DhEphemeral {
    /// Generate a fresh ephemeral keypair using the OS RNG.
    pub fn generate() -> Self {
        // p256 0.13 is still on the rand_core 0.6 trait family, while the
        // rest of the project uses rand 0.10 (rand_core 0.10). We pull in
        // rand_core 0.6 directly (renamed `rand_core_06` in Cargo.toml) so
        // we have an `OsRng` that satisfies p256's bound.
        Self {
            secret: p256::ecdh::EphemeralSecret::random(&mut rand_core_06::OsRng),
        }
    }

    /// Public share encoded as `X || Y`, 64 bytes - the wire format for IKEv2
    /// group 19 (RFC 5903 §6).
    pub fn public_share(&self) -> [u8; 64] {
        use p256::elliptic_curve::sec1::ToEncodedPoint;
        let pk = self.secret.public_key();
        let encoded = pk.to_encoded_point(false); // uncompressed: 0x04 || X || Y
        let bytes = encoded.as_bytes();
        debug_assert_eq!(bytes.len(), 65);
        let mut out = [0u8; 64];
        out.copy_from_slice(&bytes[1..]);
        out
    }

    /// Compute the shared secret `g^ir` from the peer's `X || Y` (64 bytes).
    /// Returns the 32-byte X coordinate of the resulting point (RFC 5903 §7).
    pub fn diffie_hellman(self, peer: &[u8]) -> Result<[u8; 32], CryptoError> {
        if peer.len() != 64 {
            return Err(CryptoError::InvalidPeerShare);
        }
        let mut uncompressed = [0u8; 65];
        uncompressed[0] = 0x04;
        uncompressed[1..].copy_from_slice(peer);
        let pk = p256::PublicKey::from_sec1_bytes(&uncompressed)
            .map_err(|_| CryptoError::InvalidPeerShare)?;
        let shared = self.secret.diffie_hellman(&pk);
        let bytes = shared.raw_secret_bytes();
        let mut out = [0u8; 32];
        out.copy_from_slice(bytes);
        Ok(out)
    }
}

// ------------------------------------------------------------------- AES-GCM

use aes_gcm::{
    Aes256Gcm,
    aead::{AeadInOut, KeyInit, Nonce, Tag, inout::InOutBuf},
};

/// Encrypt `plaintext` in place, producing the ICV.
///
/// `key` is the 32-byte AES key (the encryption portion of `SK_ei`/`SK_er`),
/// `salt` is the 4-byte GCM salt (the trailing portion of `SK_ei`/`SK_er`),
/// `iv` is the 8-byte IV that goes on the wire inside the SK payload, and
/// `aad` is the IKEv2 header + SK header (everything authenticated).
pub fn aes_gcm_seal(
    key: &[u8],
    salt: &[u8],
    iv: &[u8],
    aad: &[u8],
    buf: &mut Vec<u8>,
) -> Result<[u8; 16], CryptoError> {
    if key.len() != 32 || salt.len() != 4 || iv.len() != 8 {
        return Err(CryptoError::AeadFailed);
    }
    let cipher = Aes256Gcm::new_from_slice(key).expect("32-byte key");
    let mut nonce_bytes = [0u8; 12];
    nonce_bytes[..4].copy_from_slice(salt);
    nonce_bytes[4..].copy_from_slice(iv);
    let nonce = Nonce::<Aes256Gcm>::from(nonce_bytes);
    let tag = cipher
        .encrypt_inout_detached(&nonce, aad, InOutBuf::from(buf.as_mut_slice()))
        .map_err(|_| CryptoError::AeadFailed)?;
    Ok(tag.into())
}

/// Decrypt-and-verify in place. On success `buf` contains the plaintext.
pub fn aes_gcm_open(
    key: &[u8],
    salt: &[u8],
    iv: &[u8],
    aad: &[u8],
    buf: &mut Vec<u8>,
    tag: &[u8],
) -> Result<(), CryptoError> {
    if key.len() != 32 || salt.len() != 4 || iv.len() != 8 || tag.len() != 16 {
        return Err(CryptoError::AeadFailed);
    }
    let cipher = Aes256Gcm::new_from_slice(key).expect("32-byte key");
    let mut nonce_bytes = [0u8; 12];
    nonce_bytes[..4].copy_from_slice(salt);
    nonce_bytes[4..].copy_from_slice(iv);
    let nonce = Nonce::<Aes256Gcm>::from(nonce_bytes);
    let mut tag_bytes = [0u8; 16];
    tag_bytes.copy_from_slice(tag);
    let tag = Tag::<Aes256Gcm>::from(tag_bytes);
    cipher
        .decrypt_inout_detached(&nonce, aad, InOutBuf::from(buf.as_mut_slice()), &tag)
        .map_err(|_| CryptoError::AeadFailed)
}

// -------------------------------------------------------------- AES-CBC

use aes::Aes256;
use aes::cipher::{BlockModeDecrypt, BlockModeEncrypt, KeyIvInit, block_padding::NoPadding};

type Aes256CbcEnc = cbc::Encryptor<Aes256>;
type Aes256CbcDec = cbc::Decryptor<Aes256>;

/// Encrypt `plaintext` with AES-256-CBC and return the ciphertext. Caller is
/// responsible for ensuring `plaintext.len()` is a multiple of 16 (block size)
/// - IKEv2's RFC 4303 §2.4 padding scheme guarantees this at the higher layer.
pub fn aes_cbc_256_encrypt(
    key: &[u8; 32],
    iv: &[u8; 16],
    plaintext: &[u8],
) -> Result<Vec<u8>, CryptoError> {
    if !plaintext.len().is_multiple_of(16) {
        return Err(CryptoError::AeadFailed);
    }
    let cipher = Aes256CbcEnc::new(key.into(), iv.into());
    // Use NoPadding so the caller's 1..=N+pad_len trailer (per RFC 4303 §2.4)
    // is preserved verbatim - the cbc crate would otherwise PKCS#7-pad on top.
    let mut buf = plaintext.to_vec();
    let len = buf.len();
    cipher
        .encrypt_padded::<NoPadding>(&mut buf, len)
        .map_err(|_| CryptoError::AeadFailed)?;
    Ok(buf)
}

/// Decrypt AES-256-CBC ciphertext.
pub fn aes_cbc_256_decrypt(
    key: &[u8; 32],
    iv: &[u8; 16],
    ciphertext: &[u8],
) -> Result<Vec<u8>, CryptoError> {
    if !ciphertext.len().is_multiple_of(16) {
        return Err(CryptoError::AeadFailed);
    }
    let cipher = Aes256CbcDec::new(key.into(), iv.into());
    let mut buf = ciphertext.to_vec();
    cipher
        .decrypt_padded::<NoPadding>(&mut buf)
        .map_err(|_| CryptoError::AeadFailed)?;
    Ok(buf)
}

/// HMAC-SHA-256 truncated to 128 bits - IKEv2 INTEG transform 12.
pub fn hmac_sha256_128(key: &[u8], data: &[u8]) -> [u8; 16] {
    let full = prf(key, data);
    let mut out = [0u8; 16];
    out.copy_from_slice(&full[..16]);
    out
}

// -------------------------------------------------------------- Auth Key

/// Derive the storable IKEv2 auth key from a raw pre-shared key.
///
/// Returns `prf(PSK, "Key Pad for IKEv2")` - a 32-byte one-way derivation
/// that is all jkipsec needs to verify a client's AUTH payload. Store this
/// in your database instead of the plaintext PSK; the raw secret can be
/// discarded after this call.
///
/// See RFC 7296 §2.15.
pub fn derive_auth_key(psk: &[u8]) -> [u8; 32] {
    prf(psk, b"Key Pad for IKEv2")
}

// ------------------------------------------------------------------- Errors

/// Errors returned by the crypto primitives.
#[derive(Debug, thiserror::Error)]
pub enum CryptoError {
    /// The peer's DH share was malformed (wrong length or not a valid point).
    #[error("invalid peer DH share")]
    InvalidPeerShare,
    /// Authenticated encryption / decryption failed (tag mismatch, bad key
    /// length, or unsupported plaintext alignment for AES-CBC).
    #[error("AEAD encryption or decryption failed")]
    AeadFailed,
}

// ------------------------------------------------------------------- Tests

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn aes_cbc_round_trip() {
        let key = [0x42u8; 32];
        let iv = [0x11u8; 16];
        // 32 bytes - already block-aligned.
        let plaintext = b"hello, ipsec, this is a 32-byteX";
        assert_eq!(plaintext.len() % 16, 0);
        let ct = aes_cbc_256_encrypt(&key, &iv, plaintext).unwrap();
        assert_eq!(ct.len(), plaintext.len());
        assert_ne!(&ct[..], &plaintext[..]);
        let pt = aes_cbc_256_decrypt(&key, &iv, &ct).unwrap();
        assert_eq!(&pt[..], &plaintext[..]);
    }

    #[test]
    fn aes_cbc_rejects_unaligned() {
        let key = [0u8; 32];
        let iv = [0u8; 16];
        let unaligned = [0u8; 17];
        assert!(aes_cbc_256_encrypt(&key, &iv, &unaligned).is_err());
        assert!(aes_cbc_256_decrypt(&key, &iv, &unaligned).is_err());
    }

    #[test]
    fn hmac_sha256_128_truncates_to_16() {
        let key = b"some-key";
        let data = b"some-data";
        let full = prf(key, data);
        let trunc = hmac_sha256_128(key, data);
        assert_eq!(&trunc[..], &full[..16]);
    }

    /// RFC 4868 §2.7.2.1 test vector for HMAC-SHA-256 ("Hi There" / 0x0b * 20).
    #[test]
    fn hmac_known_vector() {
        let key = [0x0bu8; 20];
        let data = b"Hi There";
        let out = prf(&key, data);
        let expected = [
            0xb0, 0x34, 0x4c, 0x61, 0xd8, 0xdb, 0x38, 0x53, 0x5c, 0xa8, 0xaf, 0xce, 0xaf, 0x0b,
            0xf1, 0x2b, 0x88, 0x1d, 0xc2, 0x00, 0xc9, 0x83, 0x3d, 0xa7, 0x26, 0xe9, 0x37, 0x6c,
            0x2e, 0x32, 0xcf, 0xf7,
        ];
        assert_eq!(out, expected);
    }

    /// prf+ length and chaining: T1 must equal `prf(K, S | 0x01)`, T2 must equal
    /// `prf(K, T1 | S | 0x02)`, etc.
    #[test]
    fn prf_plus_chains() {
        let key = b"some-prf-key";
        let seed = b"some-seed-data";
        let out = prf_plus(key, seed, 100);
        assert_eq!(out.len(), 100);

        let mut expected_t1 = Vec::from(&seed[..]);
        expected_t1.push(0x01);
        let t1 = prf(key, &expected_t1);
        assert_eq!(&out[..32], &t1[..]);

        let mut t2_input = Vec::from(&t1[..]);
        t2_input.extend_from_slice(seed);
        t2_input.push(0x02);
        let t2 = prf(key, &t2_input);
        assert_eq!(&out[32..64], &t2[..]);
    }

    /// Round-trip ECDH: two locally-generated keypairs must agree.
    #[test]
    fn ecdh_round_trip() {
        let a = DhEphemeral::generate();
        let b = DhEphemeral::generate();
        let a_pub = a.public_share();
        let b_pub = b.public_share();
        let ab = a.diffie_hellman(&b_pub).unwrap();
        let ba = b.diffie_hellman(&a_pub).unwrap();
        assert_eq!(ab, ba);
    }

    /// Round-trip AES-GCM: encrypt then decrypt.
    #[test]
    fn aes_gcm_round_trip() {
        let key = [0x42u8; 32];
        let salt = [0x11u8; 4];
        let iv = [0x22u8; 8];
        let aad = b"associated-data";
        let plaintext = b"hello, ipsec";
        let mut buf = plaintext.to_vec();
        let tag = aes_gcm_seal(&key, &salt, &iv, aad, &mut buf).unwrap();
        assert_ne!(&buf[..], &plaintext[..]);
        aes_gcm_open(&key, &salt, &iv, aad, &mut buf, &tag).unwrap();
        assert_eq!(&buf[..], &plaintext[..]);
    }

    /// Tampering with the AAD must fail decryption.
    #[test]
    fn aes_gcm_aad_tamper_fails() {
        let key = [0x42u8; 32];
        let salt = [0x11u8; 4];
        let iv = [0x22u8; 8];
        let mut buf = b"hello".to_vec();
        let tag = aes_gcm_seal(&key, &salt, &iv, b"good-aad", &mut buf).unwrap();
        assert!(aes_gcm_open(&key, &salt, &iv, b"BAD-aad", &mut buf, &tag).is_err());
    }
}