// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use s2n_quic_core::crypto;

header_key!(OneRttHeaderKey);
negotiated_crypto!(OneRttKey, OneRttHeaderKey);

impl crypto::OneRttKey for OneRttKey {
    #[inline]
    fn derive_next_key(&self) -> Self {
        Self(self.0.update())
    }
}

impl crypto::OneRttHeaderKey for OneRttHeaderKey {}

#[cfg(test)]
mod tests {
    use crate::{cipher_suite::TLS_CHACHA20_POLY1305_SHA256, hkdf};
    use hex_literal::hex;
    use s2n_codec::{encoder::scatter, EncoderBuffer};
    use s2n_quic_core::crypto::Key;

    //= https://www.rfc-editor.org/rfc/rfc9001#appendix-A.5
    //# In this example, TLS produces an application write secret from which
    //# a server uses HKDF-Expand-Label to produce four values: a key, an IV,
    //# a header protection key, and the secret that will be used after keys
    //# are updated (this last value is not used further in this example).
    //#
    //# secret
    //#     = 9ac312a7f877468ebe69422748ad00a1
    //#       5443f18203a07d6060f688f30f21632b
    //#
    //# key = HKDF-Expand-Label(secret, "quic key", "", 32)
    //#     = c6d98ff3441c3fe1b2182094f69caa2e
    //#       d4b716b65488960a7a984979fb23e1c8
    //#
    //# iv  = HKDF-Expand-Label(secret, "quic iv", "", 12)
    //#     = e0459b3474bdd0e44a41c144
    //#
    //# hp  = HKDF-Expand-Label(secret, "quic hp", "", 32)
    //#     = 25a282b9e82f06f21f488917a4fc8f1b
    //#       73573685608597d0efcb076b0ab7a7a4
    //#
    //# ku  = HKDF-Expand-Label(secret, "quic ku", "", 32)
    //#     = 1223504755036d556342ee9361d25342
    //#       1a826c9ecdf3c7148684b36b714881f9
    const SECRET: [u8; 32] =
        hex!("9ac312a7f877468ebe69422748ad00a15443f18203a07d6060f688f30f21632b");
    const KU_SECRET: [u8; 32] =
        hex!("1223504755036d556342ee9361d253421a826c9ecdf3c7148684b36b714881f9");

    // Prevent trivial success
    const INVALID_SECRET: [u8; 32] =
        hex!("0000000000000000000000000000000000000000000000000000000000000000");

    /// Return ChaCha20 ciphers because these are the ciphers given in the RFC. The other cipher
    /// implementations don't have RFC values we can test.
    /// This is not exhaustive, but it does show that we are using the KDF and label correctly.
    fn generate_ciphers(
        secret: &[u8],
        next_secret: &[u8],
    ) -> (TLS_CHACHA20_POLY1305_SHA256, TLS_CHACHA20_POLY1305_SHA256) {
        // Create a cipher based on the initial secret
        let key = hkdf::Prk::new_less_safe(hkdf::HKDF_SHA256, secret);
        let cipher = TLS_CHACHA20_POLY1305_SHA256::new(key);

        // Create the cipher after a Key Update has occurred
        let next_cipher = cipher.0.update();

        // Create a cipher based on the expected post-update secret
        let next_key = hkdf::Prk::new_less_safe(hkdf::HKDF_SHA256, next_secret);
        let expected_next_cipher = TLS_CHACHA20_POLY1305_SHA256::new(next_key);

        (next_cipher, expected_next_cipher.0)
    }

    #[test]
    fn test_key_update() {
        let tests = [
            (&SECRET, &KU_SECRET, true),
            (&INVALID_SECRET, &KU_SECRET, false),
        ];

        for (secret, ku_secret, should_match) in tests {
            let (mut next_cipher, mut expected_next_cipher) = generate_ciphers(secret, ku_secret);

            // Encrypt two empty blocks to verify the ciphers are the same
            let mut next_cipher_output = [0; 32];
            let mut expected_cipher_output = [0; 32];

            {
                let next_cipher_output = EncoderBuffer::new(&mut next_cipher_output);
                let mut next_cipher_output = scatter::Buffer::new(next_cipher_output);

                let expected_cipher_output = EncoderBuffer::new(&mut expected_cipher_output);
                let mut expected_cipher_output = scatter::Buffer::new(expected_cipher_output);

                next_cipher
                    .encrypt(0, &[], &mut next_cipher_output)
                    .unwrap();
                expected_next_cipher
                    .encrypt(0, &[], &mut expected_cipher_output)
                    .unwrap();
            }

            if should_match {
                assert_eq!(next_cipher_output, expected_cipher_output);
            } else {
                assert_ne!(next_cipher_output, expected_cipher_output);
            }
        }
    }
}
