attrkey 0.1.0-alpha.1

use argon2::{Algorithm, Argon2, Version};
use digest::{
    Digest, HashMarker, OutputSizeUser,
    block_buffer::Eager,
    core_api::{
        BlockSizeUser, BufferKindUser, CoreProxy, FixedOutputCore, UpdateCore,
    },
    typenum::{IsLess, Le, NonZero, U256},
};
use rayon::prelude::*;
use std::marker::PhantomData;
use zeroize::Zeroize;

use crate::errors::ArrkeyError;
use crate::keyspace::Keyspace;
use crate::params::Params;

/// Represents a collection of attribute values and that can be hardened into
/// a `Keyspace` using memory-hard key derivation.
#[derive(Clone, Debug)]
pub struct Attributes<'a, D> {
    arr: Vec<Vec<u8>>,
    argon2_context: Argon2<'a>,
    context_string: Option<String>,
    constraint: u8,
    _digest: PhantomData<D>,
}

impl<'a, D> Attributes<'a, 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 collection of `Attributes`.
    ///
    /// Requires at least two attribute values, and a set of valid parameters.
    /// Returns an error when fewer than two attributes are provided.
    pub fn new(arr: Vec<Vec<u8>>, params: Params) -> Result<Self, ArrkeyError> {
        if arr.len() <= 1 {
            return Err(ArrkeyError::AttributesTooFew);
        }

        let argon2_params = argon2::Params::new(
            params.m_cost,
            params.t_cost,
            params.p_cost,
            Some(32),
        )
        .expect("password hashing params already validated");

        let argon2_context =
            Argon2::new(Algorithm::Argon2id, Version::V0x13, argon2_params);

        Ok(Self {
            arr,
            argon2_context,
            context_string: params.context,
            constraint: params.constraint,
            _digest: PhantomData,
        })
    }

    /// Harden the attributes into a cryptographic [Keyspace].
    ///
    /// Applies a memory-hard key derivation function
    /// ([Argon2id](argon2::Argon2)) to harden all attributes in parallel,
    /// using per-attribute salts derived from a shared base salt. All hardened
    /// attributes are used to produce a keyspace, from which attribute-based
    /// keys can be selected.
    ///
    /// Returns a [Keyspace] of the hardened attributes on success, or an error
    /// when the memory-hard key derivation process fails.
    pub fn harden(&self, salt: &[u8]) -> Result<Keyspace<D>, ArrkeyError> {
        let mut arr_salts = self.derive_salts(salt);

        let hardened: Vec<[u8; 32]> = self
            .arr
            .par_iter()
            .zip(arr_salts.par_iter())
            .map(|(pwd, salt)| -> Result<[u8; 32], ArrkeyError> {
                let mut tmp = [0u8; 32];
                self.argon2_context
                    .hash_password_into(pwd, salt, &mut tmp)
                    .map_err(|_| ArrkeyError::HardenFail)?;
                Ok(tmp)
            })
            .collect::<Result<Vec<_>, _>>()?;
        arr_salts.zeroize();

        Ok(Keyspace::new(hardened, self.constraint))
    }

    /// Return the number of attributes in this collection.
    pub fn len(&self) -> usize {
        self.arr.len()
    }

    fn derive_salts(&self, base_salt: &[u8]) -> Vec<Vec<u8>> {
        (0..self.arr.len())
            .map(|i| {
                let mut hasher = D::new();

                hasher.update(base_salt);
                if let Some(ctx) = &self.context_string {
                    hasher.update(ctx.as_bytes());
                }
                hasher.update(&i.to_le_bytes());

                hasher.finalize().to_vec()
            })
            .collect()
    }
}

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

    const ATTR_1: &'static [u8; 18] = b"Pigs on the Wing I";
    const ATTR_2: &'static [u8; 4] = b"Dogs";
    const ATTR_3: &'static [u8; 27] = b"Pigs (Three Different Ones)";
    const ATTR_4: &'static [u8; 5] = b"Sheep";
    const ATTR_5: &'static [u8; 19] = b"Pigs on the Wing II";

    fn get_arr() -> Vec<Vec<u8>> {
        vec![
            ATTR_1.to_vec(),
            ATTR_2.to_vec(),
            ATTR_3.to_vec(),
            ATTR_4.to_vec(),
            ATTR_5.to_vec(),
        ]
    }

    #[test]
    fn new_produces_attributes_collection() {
        let arr = get_arr();

        // Valid attribute values and params must yield Ok result
        let result = Attributes::<Sha256>::new(arr, Params::default());
        assert!(result.is_ok());
    }

    #[test]
    fn new_too_few() {
        let arr_empty: Vec<Vec<u8>> = vec![];
        let arr_one: Vec<Vec<u8>> = vec![ATTR_1.to_vec()];

        // Too few attributes (zero attributes) must result in an error
        let result = Attributes::<Sha256>::new(arr_empty, Params::default());
        assert_eq!(result.unwrap_err(), ArrkeyError::AttributesTooFew);

        // Too few attributes (one attribute) must result in an error
        let result = Attributes::<Sha256>::new(arr_one, Params::default());
        assert_eq!(result.unwrap_err(), ArrkeyError::AttributesTooFew);
    }

    #[test]
    fn harden_produces_keyspace() {
        let arr = get_arr();
        let salt = b"Animals";

        // Should construct collection of attributes
        let result = Attributes::<Sha256>::new(arr, Params::default());
        assert!(result.is_ok());
        let attributes = result.unwrap();

        // Hardening from attribute collection must produce keyspace
        let result = attributes.harden(salt);
        assert!(result.is_ok());
    }

    #[test]
    fn harden_fails_argon2_failure() {
        let huge_attribute = vec![0u8; (u32::MAX as usize) + 1];
        let arr = vec![ATTR_1.to_vec(), huge_attribute];
        let salt = b"Animals";

        // Should construct collection of attributes
        let result = Attributes::<Sha256>::new(arr, Params::default());
        assert!(result.is_ok());
        let attributes = result.unwrap();

        // Hardening must fail for attribute too large for Argon2
        let result = attributes.harden(salt);
        assert_eq!(result.unwrap_err(), ArrkeyError::HardenFail);
    }

    #[test]
    fn len_returns_collection_length() {
        let arr = get_arr();
        let arr_len = arr.len();

        // Should construct collection of attributes
        let result = Attributes::<Sha256>::new(arr, Params::default());
        assert!(result.is_ok());
        let attributes = result.unwrap();

        // Returned length must match original number of attributes
        let len = attributes.len();
        assert_eq!(len, arr_len);
    }
}