use std::{collections::HashMap, sync::Arc};
use tink_core::{utils::wrap_err, Prf, TinkError};
pub struct Set {
pub primary_id: u32,
pub prfs: HashMap<u32, Box<dyn Prf>>,
}
impl Set {
pub fn new(h: &tink_core::keyset::Handle) -> Result<Set, TinkError> {
Set::new_with_key_manager(h, None)
}
fn new_with_key_manager(
h: &tink_core::keyset::Handle,
km: Option<Arc<dyn tink_core::registry::KeyManager>>,
) -> Result<Set, TinkError> {
let ps = h
.primitives_with_key_manager(km)
.map_err(|e| wrap_err("prf::Set: cannot obtain primitive set", e))?;
wrap_prf_set(ps)
}
pub fn compute_primary_prf(
&self,
input: &[u8],
output_length: usize,
) -> Result<Vec<u8>, TinkError> {
let prf = self.prfs.get(&self.primary_id).ok_or_else(|| {
TinkError::new(&format!(
"Could not find primary ID {} in prf.Set",
self.primary_id
))
})?;
prf.compute_prf(input, output_length)
}
}
fn wrap_prf_set(ps: tink_core::primitiveset::PrimitiveSet) -> Result<Set, TinkError> {
let entry = match &ps.primary {
None => return Err("prf::Set: no primary available".into()),
Some(e) => e,
};
match entry.primitive {
tink_core::Primitive::Prf(_) => {}
_ => return Err("prf::Set: not a PRF primitive".into()),
}
let mut set = Set {
primary_id: entry.key_id,
prfs: HashMap::new(),
};
let entries = ps.raw_entries();
if entries.is_empty() {
return Err("Did not find any raw entries".into());
}
if ps.entries.len() != 1 {
return Err("Only raw entries allowed for prf::Set".into());
}
for entry in entries {
let prf = match entry.primitive {
tink_core::Primitive::Prf(prf) => prf,
_ => return Err("prf::Set: not a PRF primitive".into()),
};
set.prfs.insert(entry.key_id, prf);
}
Ok(set)
}