//! Raw `XChaCha20-Poly1305` AEAD and `HChaCha20` key derivation.
//!
//! These primitives implement the `DNSCrypt` v2 `"XSalsa20Poly1305"` wire format,
//! which `DNSCrypt` labels "X25519-XChacha20Poly1305".  The construction is:
//!
//! 1. Derive a 64-byte keystream block from the key and nonce using `XChaCha20`.
//! 2. Use the first 32 bytes as the Poly1305 one-time key.
//! 3. Encrypt/decrypt the payload using the remaining keystream bytes.
//! 4. Authenticate the ciphertext with Poly1305.

use chacha20::XChaCha20;
use chacha20::cipher::{KeyIvInit, StreamCipher};
use poly1305::Poly1305;
use poly1305::universal_hash::KeyInit;
use subtle::ConstantTimeEq;

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

/// Encrypt `plaintext` with `XChaCha20-Poly1305` and prepend the 16-byte tag.
///
/// The returned buffer layout is `[tag (16 bytes) | ciphertext]`.
#[must_use]
pub fn xchacha20_djb_poly1305_encrypt(
    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.
    #[cfg(feature = "zeroize")]
    let mut keystream = zeroize::Zeroizing::new(vec![0u8; 32 + plaintext.len()]);
    #[cfg(not(feature = "zeroize"))]
    let mut keystream = vec![0u8; 32 + plaintext.len()];

    XChaCha20::new(key.into(), nonce.into()).apply_keystream(&mut keystream);

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

    #[cfg(feature = "zeroize")]
    let mut poly_key = zeroize::Zeroizing::new([0u8; 32]);
    #[cfg(not(feature = "zeroize"))]
    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();

    #[cfg(feature = "zeroize")]
    let tag = Poly1305::new(&(*poly_key).into()).compute_unpadded(&ciphertext);
    #[cfg(not(feature = "zeroize"))]
    let tag = Poly1305::new(&poly_key.into()).compute_unpadded(&ciphertext);

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

/// Verify and decrypt a XChaCha20-Poly1305 ciphertext.
///
/// `ciphertext_with_tag` must begin with the 16-byte Poly1305 tag followed
/// by the ciphertext produced by [`xchacha20_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 xchacha20_djb_poly1305_decrypt(
    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);

    #[cfg(feature = "zeroize")]
    let mut keystream = zeroize::Zeroizing::new(vec![0u8; 32 + ciphertext.len()]);
    #[cfg(not(feature = "zeroize"))]
    let mut keystream = vec![0u8; 32 + ciphertext.len()];

    XChaCha20::new(key.into(), nonce.into()).apply_keystream(&mut keystream);

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

    #[cfg(feature = "zeroize")]
    let mut poly_key = zeroize::Zeroizing::new([0u8; 32]);
    #[cfg(not(feature = "zeroize"))]
    let mut poly_key = [0u8; 32];

    poly_key.copy_from_slice(poly_key_bytes);

    #[cfg(feature = "zeroize")]
    let computed_tag = Poly1305::new(&(*poly_key).into()).compute_unpadded(ciphertext);
    #[cfg(not(feature = "zeroize"))]
    let computed_tag = Poly1305::new(&poly_key.into()).compute_unpadded(ciphertext);

    // Constant-time comparison to prevent timing oracle attacks.
    if computed_tag[..].ct_eq(tag).unwrap_u8() == 0 {
        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!(xchacha20_djb_poly1305_decrypt(&key, &nonce, &[0u8; 15]).is_err());
        assert!(xchacha20_djb_poly1305_decrypt(&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_encrypt_empty_plaintext_roundtrip() {
        let key = [6u8; 32];
        let nonce = [7u8; 24];
        let encrypted = xchacha20_djb_poly1305_encrypt(&key, &nonce, &[]);
        assert_eq!(encrypted.len(), 16);
        assert_eq!(
            xchacha20_djb_poly1305_decrypt(&key, &nonce, &encrypted).unwrap(),
            Vec::<u8>::new()
        );
    }
}
