dnscrypt 0.2.0

//! Raw `XSalsa20-Poly1305` / `XChaCha20-Poly1305` AEAD and `HSalsa20`/`HChaCha20`
//! key derivation.
//!
//! `DNSCrypt` v2 certificates advertise one of two cipher suites via their ES
//! version field:
//!
//! - `0x0001`: `X25519-XSalsa20Poly1305` (the original `DNSCrypt` construction).
//! - `0x0002`: `X25519-XChacha20Poly1305` (the modern, preferred construction).
//!
//! Both build an AEAD the same way — a stream cipher generates a keystream
//! block, the first 32 bytes become the one-time Poly1305 key, and the rest
//! XORs with the plaintext — differing only in the underlying stream cipher
//! (`XSalsa20` vs `XChaCha20`). [`CipherSuite`] selects between them.

use poly1305::Poly1305;
use poly1305::universal_hash::KeyInit;

/// Which `DNSCrypt` AEAD construction a session uses, selected by the
/// certificate's ES version field.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CipherSuite {
    /// ES version `0x0001`: `X25519-XSalsa20Poly1305`.
    XSalsa20Poly1305,
    /// ES version `0x0002`: `X25519-XChacha20Poly1305`.
    XChaCha20Poly1305,
}

/// Compute `HChaCha20(key, input)` and return the 32-byte derived subkey.
///
/// In the `DNSCrypt` session setup the raw X25519 shared secret is passed through
/// `HChaCha20` with a zero nonce to produce the symmetric session key used for
/// all subsequent query/response encryption.
#[must_use]
pub fn run_hchacha20(key: &[u8; 32], input: &[u8; 16]) -> [u8; 32] {
    use chacha20::{R20, hchacha};
    let out = hchacha::<R20>(&(*key).into(), &(*input).into());
    let mut result = [0u8; 32];
    result.copy_from_slice(&out);
    result
}

/// Compute `HSalsa20(key, input)` and return the 32-byte derived subkey.
///
/// The `XSalsa20Poly1305` analogue of [`run_hchacha20`], used for sessions
/// whose certificate advertises ES version `0x0001`.
#[must_use]
pub fn run_hsalsa20(key: &[u8; 32], input: &[u8; 16]) -> [u8; 32] {
    use salsa20::cipher::consts::U10;
    use salsa20::hsalsa;
    let out = hsalsa::<U10>(&(*key).into(), &(*input).into());
    let mut result = [0u8; 32];
    result.copy_from_slice(&out);
    result
}

/// Fill `keystream` using `suite`'s stream cipher under `key`/`nonce`.
fn apply_keystream(suite: CipherSuite, key: &[u8; 32], nonce: &[u8; 24], keystream: &mut [u8]) {
    match suite {
        CipherSuite::XChaCha20Poly1305 => {
            use chacha20::XChaCha20;
            use chacha20::cipher::{KeyIvInit, StreamCipher};
            XChaCha20::new(key.into(), nonce.into()).apply_keystream(keystream);
        }
        CipherSuite::XSalsa20Poly1305 => {
            use salsa20::XSalsa20;
            use salsa20::cipher::{KeyIvInit, StreamCipher};
            XSalsa20::new(key.into(), nonce.into()).apply_keystream(keystream);
        }
    }
}

/// Encrypt `plaintext` with `suite` and prepend the 16-byte Poly1305 tag.
///
/// The returned buffer layout is `[tag (16 bytes) | ciphertext]`.
#[must_use]
pub fn djb_poly1305_encrypt(
    suite: CipherSuite,
    key: &[u8; 32],
    nonce: &[u8; 24],
    plaintext: &[u8],
) -> Vec<u8> {
    // Generate the full keystream: first 32 bytes become the Poly1305 key,
    // the rest XOR with the plaintext.
    let mut keystream = vec![0u8; 32usize.saturating_add(plaintext.len())];
    apply_keystream(suite, key, nonce, &mut keystream);

    let (poly_key_bytes, keystream_body) = keystream.split_at(32);

    let mut poly_key = [0u8; 32];
    poly_key.copy_from_slice(poly_key_bytes);

    // XOR plaintext with keystream to produce ciphertext.
    let ciphertext: Vec<u8> = plaintext
        .iter()
        .zip(keystream_body.iter())
        .map(|(&p, &k)| p ^ k)
        .collect();

    let tag = Poly1305::new(&poly_key.into()).compute_unpadded(&ciphertext);

    let mut result = Vec::with_capacity(16usize.saturating_add(ciphertext.len()));
    result.extend_from_slice(&tag[..]);
    result.extend_from_slice(&ciphertext);
    result
}

/// Verify and decrypt a `suite`-encrypted ciphertext.
///
/// `ciphertext_with_tag` must begin with the 16-byte Poly1305 tag followed
/// by the ciphertext produced by [`djb_poly1305_encrypt`].
///
/// # Errors
///
/// Returns `Err` if the input is shorter than 16 bytes or if the Poly1305
/// tag does not match (constant-time comparison).
pub fn djb_poly1305_decrypt(
    suite: CipherSuite,
    key: &[u8; 32],
    nonce: &[u8; 24],
    ciphertext_with_tag: &[u8],
) -> Result<Vec<u8>, String> {
    if ciphertext_with_tag.len() < 16 {
        return Err("Ciphertext too short".to_string());
    }
    let (tag, ciphertext) = ciphertext_with_tag.split_at(16);

    let mut keystream = vec![0u8; 32usize.saturating_add(ciphertext.len())];
    apply_keystream(suite, key, nonce, &mut keystream);

    let (poly_key_bytes, keystream_body) = keystream.split_at(32);

    let mut poly_key = [0u8; 32];
    poly_key.copy_from_slice(poly_key_bytes);

    let computed_tag = Poly1305::new(&poly_key.into()).compute_unpadded(ciphertext);

    // Constant-time comparison to prevent timing oracle attacks.
    if aws_lc_rs::constant_time::verify_slices_are_equal(&computed_tag[..], tag).is_err() {
        return Err("Poly1305 authentication failed".to_string());
    }

    let plaintext: Vec<u8> = ciphertext
        .iter()
        .zip(keystream_body.iter())
        .map(|(&c, &k)| c ^ k)
        .collect();

    Ok(plaintext)
}

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

    #[test]
    fn test_decrypt_rejects_short_input() {
        let key = [1u8; 32];
        let nonce = [2u8; 24];
        assert!(
            djb_poly1305_decrypt(CipherSuite::XChaCha20Poly1305, &key, &nonce, &[0u8; 15]).is_err()
        );
        assert!(djb_poly1305_decrypt(CipherSuite::XChaCha20Poly1305, &key, &nonce, &[]).is_err());
    }

    #[test]
    fn test_run_hchacha20_deterministic() {
        let key = [3u8; 32];
        let input = [4u8; 16];
        assert_eq!(run_hchacha20(&key, &input), run_hchacha20(&key, &input));
        assert_ne!(
            run_hchacha20(&key, &input),
            run_hchacha20(&[5u8; 32], &input)
        );
    }

    #[test]
    fn test_run_hsalsa20_deterministic() {
        let key = [3u8; 32];
        let input = [4u8; 16];
        assert_eq!(run_hsalsa20(&key, &input), run_hsalsa20(&key, &input));
        assert_ne!(run_hsalsa20(&key, &input), run_hsalsa20(&[5u8; 32], &input));
        assert_ne!(run_hsalsa20(&key, &input), run_hchacha20(&key, &input));
    }

    #[test]
    fn test_encrypt_empty_plaintext_roundtrip() {
        for suite in [
            CipherSuite::XChaCha20Poly1305,
            CipherSuite::XSalsa20Poly1305,
        ] {
            let key = [6u8; 32];
            let nonce = [7u8; 24];
            let encrypted = djb_poly1305_encrypt(suite, &key, &nonce, &[]);
            assert_eq!(encrypted.len(), 16);
            assert_eq!(
                djb_poly1305_decrypt(suite, &key, &nonce, &encrypted).unwrap(),
                Vec::<u8>::new()
            );
        }
    }
}