attrkey 0.1.0-alpha.2

use annihilation::AnnihlKey;
use hkdf::{Hkdf, hmac::digest::OutputSizeUser};
use rand_core::{Rng, SeedableRng};
use sha2::Sha256;
use std::marker::PhantomData;
use zeroize::{Zeroize, ZeroizeOnDrop};

use crate::errors::AttrkeyError;

const MIN_CONSTRAINT: u8 = 1;

/// A `Keyspace` represents a collection of hardened attribute values from
/// which attribute-based cryptographic keys can be derived.
#[derive(Clone, Zeroize, ZeroizeOnDrop)]
pub struct Keyspace<R> {
    arr: Vec<Vec<u8>>,
    constraint: u8,
    context: Vec<u8>,
    _rng: PhantomData<R>,
}

impl<R> Keyspace<R>
where
    R: SeedableRng + Rng + Sync,
    R::Seed: AsMut<[u8]> + Default + Zeroize,
{
    /// Construct a new `Keyspace` from a collection of hardened attributes,
    /// proof-of-work constraint, and optional context.
    pub fn new(
        arr: Vec<Vec<u8>>,
        constraint: u8,
        context: Option<&[u8]>,
    ) -> Result<Self, AttrkeyError> {
        if constraint < MIN_CONSTRAINT {
            return Err(AttrkeyError::ConstraintTooSmall);
        }

        Ok(Self {
            arr,
            constraint,
            context: context.map(|c| c.to_vec()).unwrap_or_default(),
            _rng: PhantomData,
        })
    }

    /// Derive a cryptographic key from a selection of attribute indices.
    ///
    /// Partitions the keyspace into selected and remaining attributes, then
    /// computes an [annihilation key](annihilation::AnnihlKey::to_annihilation)
    /// from the partition, used as a seed for RNG expansion.
    ///
    /// Returns the derived key on success (default length = 32), or an error
    /// when the selection is invalid or partition annihilation fails.
    pub fn derive_key(
        &self,
        selection: &[usize],
        output_len: Option<usize>,
    ) -> Result<Vec<u8>, AttrkeyError> {
        let seed = self.derive_seed(selection)?;

        let mut rng = R::from_seed(seed);

        let dst_len = output_len.unwrap_or(32);
        let mut dst = vec![0u8; dst_len];
        rng.fill_bytes(&mut dst);
        Ok(dst)
    }

    /// Derive a cryptographic key from a selection of attribute indices into a
    /// provided buffer.
    ///
    /// Partitions the keyspace into selected and remaining attributes, then
    /// computes an [annihilation key](annihilation::AnnihlKey::to_annihilation)
    /// from the partition, used as a seed for RNG expansion.
    ///
    /// Fills the provided buffer with the derived bytes on success, or an error
    /// when the selection is invalid or partition annihilation fails.
    pub fn derive_key_into(
        &self,
        selection: &[usize],
        dst: &mut [u8],
    ) -> Result<(), AttrkeyError> {
        let seed = self.derive_seed(selection)?;

        let mut rng = R::from_seed(seed);

        rng.fill_bytes(dst);
        Ok(())
    }

    /// Extract an [annihilative pair](annihilation::AnnihlKey::new_pair) from
    /// an attribute selection.
    ///
    /// Partitions attributes by the selection indices, concatenating selected
    /// attributes as key material and remaining attributes as antikey
    /// material, then mines for a pair that satisfies the keyspace's
    /// proof-of-work constraint.
    ///
    /// Returns the annihilative pair on success, or an error when the
    /// selection is empty, selects all attributes, or contains out-of-bound
    /// indices.
    pub fn extract_pair(
        &self,
        selection: &[usize],
    ) -> Result<(AnnihlKey, AnnihlKey), AttrkeyError> {
        if selection.is_empty() {
            return Err(AttrkeyError::IndicesTooFew);
        }
        if selection.len() >= self.arr.len() {
            return Err(AttrkeyError::IndicesTooMany);
        }
        if selection.iter().any(|&index| index >= self.arr.len()) {
            return Err(AttrkeyError::IndexOutOfBounds);
        }

        let (selected, remaining): (Vec<_>, Vec<_>) = self
            .arr
            .iter()
            .enumerate()
            .partition(|(index, _)| selection.contains(index));

        let mut ikm: Vec<u8> = selected
            .into_iter()
            .flat_map(|(_, arr)| arr.iter().copied())
            .collect();
        let mut iam: Vec<u8> = remaining
            .into_iter()
            .flat_map(|(_, arr)| arr.iter().copied())
            .collect();

        let pair = AnnihlKey::new_pair(&ikm, &iam, self.constraint);
        ikm.zeroize();
        iam.zeroize();

        Ok(pair)
    }

    fn derive_seed(
        &self,
        selection: &[usize],
    ) -> Result<R::Seed, AttrkeyError> {
        let (key, antikey) = self.extract_pair(selection)?;

        let mut prk = key
            .into_annihilation(antikey)
            .map_err(|_| AttrkeyError::Annihilation)?;

        let seed_len = R::Seed::default().as_mut().len();
        let max_okm_len = 255 * <Sha256 as OutputSizeUser>::output_size();
        if seed_len > max_okm_len {
            prk.zeroize();
            return Err(AttrkeyError::SeedTooLong);
        }

        let hk = Hkdf::<Sha256>::from_prk(&prk)
            .expect("annihilation PRK is always 32 bytes");

        let mut tmp = vec![0u8; seed_len];
        hk.expand(&self.context, &mut tmp)
            .expect("seed length is within HKDF-SHA256 bounds");

        let mut seed = R::Seed::default();
        seed.as_mut().copy_from_slice(&tmp);
        prk.zeroize();
        tmp.zeroize();

        Ok(seed)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chacha20::ChaCha20Rng;
    use rand_core::TryRng;
    use scrypt::{Params, Scrypt};

    use crate::attributes::Attributes;

    /// Testing Keyspace type.
    type TestKeyspace = Keyspace<ChaCha20Rng>;

    /// Seed type for [LargeSeedRng] exceeding HKDF-SHA256's
    /// maximum OKM length (255 * 32).
    #[derive(Clone)]
    struct LargeSeed([u8; 8161]);

    impl Default for LargeSeed {
        fn default() -> Self {
            Self([0u8; 8161])
        }
    }

    impl AsRef<[u8]> for LargeSeed {
        fn as_ref(&self) -> &[u8] {
            &self.0
        }
    }

    impl AsMut<[u8]> for LargeSeed {
        fn as_mut(&mut self) -> &mut [u8] {
            &mut self.0
        }
    }

    impl Zeroize for LargeSeed {
        fn zeroize(&mut self) {
            self.0.iter_mut().for_each(|b| *b = 0);
        }
    }

    /// Triggers the [SeedTooLong](AttrkeyError::SeedTooLong) error path
    /// for [Keyspace] methods.
    struct LargeSeedRng;

    impl SeedableRng for LargeSeedRng {
        type Seed = LargeSeed;

        fn from_seed(_seed: Self::Seed) -> Self {
            LargeSeedRng
        }
    }

    impl TryRng for LargeSeedRng {
        type Error = core::convert::Infallible;

        fn try_fill_bytes(
            &mut self,
            _dst: &mut [u8],
        ) -> Result<(), Self::Error> {
            Ok(())
        }

        fn try_next_u32(&mut self) -> Result<u32, Self::Error> {
            Ok(0)
        }

        fn try_next_u64(&mut self) -> Result<u64, Self::Error> {
            Ok(0)
        }
    }

    fn get_arr() -> Vec<Vec<u8>> {
        vec![
            b"Pigs on the Wing I".to_vec(),
            b"Dogs".to_vec(),
            b"Pigs (Three Different Ones)".to_vec(),
            b"Sheep".to_vec(),
            b"Pigs on the Wing II".to_vec(),
        ]
    }

    fn produce_keyspace() -> TestKeyspace {
        let params = Params::new(4, 8, 1).unwrap();
        let scrypt = Scrypt::new_with_params(params);

        Attributes::<Scrypt, ChaCha20Rng>::new(get_arr(), scrypt)
            .expect("should construct collection")
            .harden(4, None)
            .expect("hardening should produce keyspace")
    }

    fn produce_large_seed_keyspace() -> Keyspace<LargeSeedRng> {
        let params = Params::new(4, 8, 1).unwrap();
        let scrypt = Scrypt::new_with_params(params);

        Attributes::<Scrypt, LargeSeedRng>::new(get_arr(), scrypt)
            .expect("should construct collection")
            .harden(4, None)
            .expect("hardening should produce keyspace")
    }

    #[test]
    fn new_constructs_keyspace() {
        let result = TestKeyspace::new(get_arr(), 4, None);
        assert!(result.is_ok());
    }

    #[test]
    fn new_with_context_constructs_keyspace() {
        let result = TestKeyspace::new(get_arr(), 4, Some(b"Animals"));
        assert!(result.is_ok());
    }

    #[test]
    fn new_fails_constraint_too_small() {
        let result = TestKeyspace::new(get_arr(), 0, None);
        assert!(matches!(result, Err(AttrkeyError::ConstraintTooSmall)));
    }

    #[test]
    fn derive_key_produces_key() {
        let keyspace = produce_keyspace();

        // Key derivation must succeed for valid attribute selection
        let result = keyspace.derive_key(&[1, 3], None);
        assert!(result.is_ok());

        // Derived key must default to 32 byes when length is None
        assert_eq!(result.unwrap().len(), 32);
    }

    #[test]
    fn derive_key_with_output_len_produces_key() {
        let keyspace = produce_keyspace();

        // Key derivation must succeed for valid attribute selection
        let result = keyspace.derive_key(&[1, 3], Some(48));
        assert!(result.is_ok());

        // Derived key length must match provided output length
        assert_eq!(result.unwrap().len(), 48);
    }

    #[test]
    fn derive_key_fails_too_few_indices() {
        let keyspace = produce_keyspace();

        // Empty selection must result in an error
        let result = keyspace.derive_key(&[], None);
        assert_eq!(result, Err(AttrkeyError::IndicesTooFew));
    }

    #[test]
    fn derive_key_fails_too_many_indices() {
        let keyspace = produce_keyspace();
        let selection: Vec<usize> = (0..keyspace.arr.len()).collect();

        // Selecting all attributes must result in an error
        let result = keyspace.derive_key(&selection, None);
        assert_eq!(result, Err(AttrkeyError::IndicesTooMany));
    }

    #[test]
    fn derive_key_fails_out_of_bounds() {
        let keyspace = produce_keyspace();
        let selection: Vec<usize> = [0, keyspace.arr.len() + 1].to_vec();

        // Out of bounds index in selection must result in an error
        let result = keyspace.derive_key(&selection, None);
        assert_eq!(result, Err(AttrkeyError::IndexOutOfBounds));
    }

    #[test]
    fn derive_key_fails_seed_too_long() {
        let keyspace = produce_large_seed_keyspace();

        // RNG seed exceeding HKDF-SHA256 OKM limit must result in an error
        let result = keyspace.derive_key(&[1, 3], None);
        assert_eq!(result, Err(AttrkeyError::SeedTooLong));
    }

    #[test]
    fn derive_key_into_fills_buffer() {
        let keyspace = produce_keyspace();
        let mut dst = [0u8; 48];

        // Key derivation must succeed for valid attribute selection
        let result = keyspace.derive_key_into(&[1, 3], &mut dst);
        assert!(result.is_ok());
    }

    #[test]
    fn derive_key_into_fails_too_few_indices() {
        let keyspace = produce_keyspace();
        let mut dst = [0u8; 48];

        // Empty selection must result in an error
        let result = keyspace.derive_key_into(&[], &mut dst);
        assert_eq!(result, Err(AttrkeyError::IndicesTooFew));
    }

    #[test]
    fn derive_key_into_fails_too_many_indices() {
        let keyspace = produce_keyspace();
        let mut dst = [0u8; 48];
        let selection: Vec<usize> = (0..keyspace.arr.len()).collect();

        // Selecting all attributes must result in an error
        let result = keyspace.derive_key_into(&selection, &mut dst);
        assert_eq!(result, Err(AttrkeyError::IndicesTooMany));
    }

    #[test]
    fn derive_key_into_fails_out_of_bounds() {
        let keyspace = produce_keyspace();
        let mut dst = [0u8; 48];
        let selection: Vec<usize> = [0, keyspace.arr.len() + 1].to_vec();

        // Out of bounds index in selection must result in an error
        let result = keyspace.derive_key_into(&selection, &mut dst);
        assert_eq!(result, Err(AttrkeyError::IndexOutOfBounds));
    }

    #[test]
    fn derive_key_into_fails_seed_too_long() {
        let keyspace = produce_large_seed_keyspace();
        let mut dst = [0u8; 48];

        // RNG seed exceeding HKDF-SHA256 OKM limit must result in an error
        let result = keyspace.derive_key_into(&[1, 3], &mut dst);
        assert_eq!(result, Err(AttrkeyError::SeedTooLong));
    }

    #[test]
    fn extract_pair_produces_pair() {
        let keyspace = produce_keyspace();

        // Valid selection must return an annihilative pair
        let result = keyspace.extract_pair(&[1, 3]);
        assert!(result.is_ok());
    }

    #[test]
    fn extract_pair_fails_empty_selection() {
        let keyspace = produce_keyspace();

        // Empty selection must result in an error
        let result = keyspace.extract_pair(&[]);
        assert_eq!(result, Err(AttrkeyError::IndicesTooFew));
    }

    #[test]
    fn extract_pair_fails_too_many_indices() {
        let keyspace = produce_keyspace();
        let selection: Vec<usize> = (0..keyspace.arr.len()).collect();

        // Selecting all attributes must result in an error
        let result = keyspace.extract_pair(&selection);
        assert_eq!(result, Err(AttrkeyError::IndicesTooMany));
    }

    #[test]
    fn extract_pair_fails_out_of_bounds() {
        let keyspace = produce_keyspace();
        let selection: Vec<usize> = [0, keyspace.arr.len() + 1].to_vec();

        // Out of bounds index in selection must result in an error
        let result = keyspace.extract_pair(&selection);
        assert_eq!(result, Err(AttrkeyError::IndexOutOfBounds));
    }
}