use crate::utils::{wrap_err, TinkError};
use std::{
collections::{hash_map, HashMap},
convert::TryFrom,
};
#[derive(Clone)]
pub struct Entry {
pub key_id: crate::KeyId,
pub primitive: crate::Primitive,
pub prefix: Vec<u8>,
pub prefix_type: tink_proto::OutputPrefixType,
pub status: tink_proto::KeyStatusType,
}
impl Entry {
fn new(
key_id: crate::KeyId,
p: crate::Primitive,
prefix: &[u8],
prefix_type: tink_proto::OutputPrefixType,
status: tink_proto::KeyStatusType,
) -> Self {
Entry {
key_id,
primitive: p,
prefix: prefix.to_vec(),
prefix_type,
status,
}
}
}
#[derive(Clone, Default)]
pub struct PrimitiveSet {
pub primary: Option<Entry>,
pub entries: HashMap<Vec<u8>, Vec<Entry>>,
}
impl PrimitiveSet {
pub fn new() -> Self {
PrimitiveSet {
primary: None,
entries: HashMap::new(),
}
}
pub fn raw_entries(&self) -> Vec<Entry> {
self.entries_for_prefix(&crate::cryptofmt::RAW_PREFIX)
}
pub fn entries_for_prefix(&self, prefix: &[u8]) -> Vec<Entry> {
match self.entries.get(prefix) {
Some(v) => v.clone(),
None => Vec::new(),
}
}
pub fn add(
&mut self,
p: crate::Primitive,
key: &tink_proto::keyset::Key,
) -> Result<Entry, TinkError> {
if key.status != tink_proto::KeyStatusType::Enabled as i32 {
return Err("The key must be ENABLED".into());
}
let prefix =
crate::cryptofmt::output_prefix(key).map_err(|e| wrap_err("primitiveset", e))?;
let entry = Entry::new(
key.key_id,
p,
&prefix,
tink_proto::OutputPrefixType::try_from(key.output_prefix_type)
.map_err(|_e| TinkError::new("invalid key prefix type"))?,
tink_proto::KeyStatusType::try_from(key.status)
.map_err(|_e| TinkError::new("invalid key status"))?,
);
let retval = entry.clone();
match self.entries.entry(prefix) {
hash_map::Entry::Occupied(mut oe) => oe.get_mut().push(entry),
hash_map::Entry::Vacant(ve) => {
ve.insert(vec![entry]);
}
};
Ok(retval)
}
}
pub struct TypedEntry<P: From<crate::Primitive>> {
pub key_id: crate::KeyId,
pub primitive: P,
pub prefix: Vec<u8>,
pub prefix_type: tink_proto::OutputPrefixType,
pub status: tink_proto::KeyStatusType,
}
impl<P: From<crate::Primitive>> From<Entry> for TypedEntry<P> {
fn from(entry: Entry) -> Self {
Self {
key_id: entry.key_id,
primitive: entry.primitive.into(),
prefix: entry.prefix,
prefix_type: entry.prefix_type,
status: entry.status,
}
}
}
pub struct TypedPrimitiveSet<P: From<crate::Primitive>> {
pub primary: Option<TypedEntry<P>>,
pub entries: HashMap<Vec<u8>, Vec<TypedEntry<P>>>,
}
impl<P: From<crate::Primitive>> TypedPrimitiveSet<P> {
pub fn raw_entries(&self) -> Option<&Vec<TypedEntry<P>>> {
self.entries_for_prefix(&crate::cryptofmt::RAW_PREFIX)
}
pub fn entries_for_prefix(&self, prefix: &[u8]) -> Option<&Vec<TypedEntry<P>>> {
self.entries.get(prefix)
}
}
impl<T> Clone for TypedPrimitiveSet<T>
where
TypedEntry<T>: Clone,
T: From<crate::Primitive>,
{
fn clone(&self) -> Self {
Self {
primary: self.primary.as_ref().cloned(),
entries: self.entries.clone(),
}
}
}
impl<P: From<crate::Primitive>> From<PrimitiveSet> for TypedPrimitiveSet<P> {
fn from(ps: PrimitiveSet) -> Self {
Self {
primary: ps.primary.map(|e| e.into()),
entries: ps
.entries
.into_iter()
.map(|(k, v)| (k, v.into_iter().map(TypedEntry::<P>::from).collect()))
.collect(),
}
}
}
impl Clone for TypedEntry<Box<dyn crate::Aead>> {
fn clone(&self) -> Self {
Self {
key_id: self.key_id,
primitive: self.primitive.box_clone(),
prefix: self.prefix.clone(),
prefix_type: self.prefix_type,
status: self.status,
}
}
}
impl Clone for TypedEntry<Box<dyn crate::DeterministicAead>> {
fn clone(&self) -> Self {
Self {
key_id: self.key_id,
primitive: self.primitive.box_clone(),
prefix: self.prefix.clone(),
prefix_type: self.prefix_type,
status: self.status,
}
}
}
impl Clone for TypedEntry<Box<dyn crate::HybridDecrypt>> {
fn clone(&self) -> Self {
Self {
key_id: self.key_id,
primitive: self.primitive.box_clone(),
prefix: self.prefix.clone(),
prefix_type: self.prefix_type,
status: self.status,
}
}
}
impl Clone for TypedEntry<Box<dyn crate::HybridEncrypt>> {
fn clone(&self) -> Self {
Self {
key_id: self.key_id,
primitive: self.primitive.box_clone(),
prefix: self.prefix.clone(),
prefix_type: self.prefix_type,
status: self.status,
}
}
}
impl Clone for TypedEntry<Box<dyn crate::Mac>> {
fn clone(&self) -> Self {
Self {
key_id: self.key_id,
primitive: self.primitive.box_clone(),
prefix: self.prefix.clone(),
prefix_type: self.prefix_type,
status: self.status,
}
}
}
impl Clone for TypedEntry<Box<dyn crate::Signer>> {
fn clone(&self) -> Self {
Self {
key_id: self.key_id,
primitive: self.primitive.box_clone(),
prefix: self.prefix.clone(),
prefix_type: self.prefix_type,
status: self.status,
}
}
}
impl Clone for TypedEntry<Box<dyn crate::StreamingAead>> {
fn clone(&self) -> Self {
Self {
key_id: self.key_id,
primitive: self.primitive.box_clone(),
prefix: self.prefix.clone(),
prefix_type: self.prefix_type,
status: self.status,
}
}
}
impl Clone for TypedEntry<Box<dyn crate::Verifier>> {
fn clone(&self) -> Self {
Self {
key_id: self.key_id,
primitive: self.primitive.box_clone(),
prefix: self.prefix.clone(),
prefix_type: self.prefix_type,
status: self.status,
}
}
}