attrkey 0.1.0-alpha.1

use annihilation::AnnihlKey;
use digest::{
    Digest, HashMarker, OutputSizeUser,
    block_buffer::Eager,
    core_api::{
        BlockSizeUser, BufferKindUser, CoreProxy, FixedOutputCore, UpdateCore,
    },
    typenum::{IsLess, Le, NonZero, U256},
};
use rand_chacha::ChaCha8Rng;
use rand_core::{Rng, SeedableRng};
use std::marker::PhantomData;
use zeroize::Zeroize;

use crate::errors::ArrkeyError;

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

impl<D> Keyspace<D>
where
    D: Digest + CoreProxy + OutputSizeUser + Sync,
    D::Core: Sync
        + HashMarker
        + UpdateCore
        + FixedOutputCore
        + BufferKindUser<BufferKind = Eager>
        + Default
        + Clone
        + BlockSizeUser,
    <D::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
    Le<<D::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
{
    /// Construct a new `Keyspace` from a collection of hardened attributes and
    /// a proof-of-work constraint.
    pub fn new(arr: Vec<[u8; 32]>, constraint: u8) -> Self {
        Self {
            arr,
            constraint,
            _digest: PhantomData,
        }
    }

    /// Derive a cryptographic key from a selection of attribute indices.
    ///
    /// Partitions the keyspace into selected and remaining attributes, then
    /// creates an [annihilative pair](annihilation::AnnihlKey::new_pair) from
    /// the partition. Computes the pair's annihilation key, and uses it as a
    /// seed for ChaCha8 CSPRNG expansion, defaulting to a 32 byte output
    /// length unless otherwise specified.
    ///
    /// Returns the derived key on success, or an error when the selection is
    /// invalid or annihilation fails.
    pub fn derive_key(
        &self,
        selection: &[usize],
        output_len: Option<usize>,
    ) -> Result<Vec<u8>, ArrkeyError> {
        let (key, antikey) = self.extract_pair(selection)?;

        let mut seed = key
            .into_annihilation(antikey)
            .map_err(|_| ArrkeyError::Annihilation)?;
        let mut rng = ChaCha8Rng::from_seed(seed);
        seed.zeroize();

        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
    /// creates an [annihilative pair](annihilation::AnnihlKey::new_pair) from
    /// the partition. Computes the pair's annihilation key, and uses it as a
    /// seed to fill the destination buffer with ChaCha8 CSPRNG output.
    ///
    /// Returns the derived key on success, or an error when the selection is
    /// invalid or annihilation fails.
    pub fn derive_key_into(
        &self,
        selection: &[usize],
        dst: &mut [u8],
    ) -> Result<(), ArrkeyError> {
        let (key, antikey) = self.extract_pair(selection)?;

        let mut seed = key
            .into_annihilation(antikey)
            .map_err(|_| ArrkeyError::Annihilation)?;
        let mut rng = ChaCha8Rng::from_seed(seed);
        seed.zeroize();

        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), ArrkeyError> {
        if selection.is_empty() {
            return Err(ArrkeyError::IndicesTooFew);
        }
        if selection.len() >= self.arr.len() {
            return Err(ArrkeyError::IndicesTooMany);
        }
        if selection.iter().any(|&index| index >= self.arr.len()) {
            return Err(ArrkeyError::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).collect();
        let mut iam: Vec<u8> =
            remaining.into_iter().flat_map(|(_, arr)| *arr).collect();

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

        Ok(pair)
    }
}

#[cfg(test)]
mod tests {
    use sha2::Sha256;

    use crate::{Attributes, Params};

    use super::*;

    fn produce_keyspace() -> Keyspace<Sha256> {
        let arr = 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(),
        ];

        let result = Attributes::<Sha256>::new(arr, Params::default());
        let attributes = result.unwrap();

        let salt = b"Animals";
        let result = attributes.harden(salt);
        result.unwrap()
    }

    #[test]
    fn new_constructs_keyspace() {
        let produced_keyspace = produce_keyspace();
        let arr = produced_keyspace.arr;
        let constraint = produced_keyspace.constraint;

        // Constructed keyspace must be the same as the produced
        let keyspace = Keyspace::<Sha256>::new(arr.clone(), constraint);
        assert_eq!(arr, keyspace.arr);
    }

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

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

    #[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(ArrkeyError::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(ArrkeyError::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(ArrkeyError::IndexOutOfBounds));
    }

    #[test]
    fn derive_key_into_succeeds_valid_selection() {
        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(ArrkeyError::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(ArrkeyError::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(ArrkeyError::IndexOutOfBounds));
    }

    #[test]
    fn extract_pair_succeeds_valid_selection() {
        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(ArrkeyError::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(ArrkeyError::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(ArrkeyError::IndexOutOfBounds));
    }
}