use core::future::Future;
use core::marker::PhantomData;
use vitaminc_aead::plaintext::{Error, Path, Step};
use vitaminc_aead::seal::{seal_all, seal_all_async, AsyncCipher, Cipher};
use vitaminc_aead::shape::{self, Locate, MapRules, SeqRules, Settle};
use vitaminc_aead::{
Encrypt, EncryptError, EncryptSink, FixedType, IntoAad, Logical, MapSink, Node, Pending,
Physical, Refusal, SeqSink, Sink, SourceError, Tee,
};
use vitaminc_protected::OpaqueDebug;
use crate::scalar::{normalize, NormalizeError, Normalized};
pub trait Normalize: Copy {
type Term;
const INDEX: &'static str;
fn normalize(
&self,
ty: Option<Logical>,
value: Physical<'_>,
) -> Result<Self::Term, NormalizeError>;
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Equality;
impl Normalize for Equality {
type Term = Normalized;
const INDEX: &'static str = "equality";
fn normalize(
&self,
ty: Option<Logical>,
value: Physical<'_>,
) -> Result<Normalized, NormalizeError> {
normalize(ty, value)
}
}
#[derive(OpaqueDebug)]
pub struct Term<T = Normalized> {
#[non_sensitive]
pub path: Path,
pub normalized: T,
}
#[derive(Clone, OpaqueDebug, thiserror::Error)]
#[error("no {index} term: {error}")]
pub struct TermError {
#[non_sensitive]
index: &'static str,
#[non_sensitive]
path: Path,
#[non_sensitive]
#[source]
error: NormalizeError,
}
impl TermError {
pub fn index(&self) -> &'static str {
self.index
}
pub fn path(&self) -> &Path {
&self.path
}
pub fn error(&self) -> NormalizeError {
self.error
}
}
#[derive(Debug, thiserror::Error)]
#[error(transparent)]
pub struct TermsError(SourceError);
impl TermsError {
pub fn term(&self) -> Option<&TermError> {
self.0.inner().downcast_ref()
}
pub fn refusal(&self) -> Option<&Refusal> {
self.0.inner().downcast_ref()
}
pub fn source_error(&self) -> &SourceError {
&self.0
}
}
impl Error for TermsError {
fn source(error: SourceError) -> Self {
Self(error)
}
}
pub struct Terms<'t, E, N: Normalize = Equality> {
out: &'t mut Vec<Term<N::Term>>,
keys: Keys<'t>,
at: Place<'t>,
rules: Option<NodeRules<'t>>,
by: N,
_error: PhantomData<E>,
}
struct NodeRules<'t> {
depth: usize,
settle: Option<Settle<'t>>,
}
impl NodeRules<'_> {
fn sealed(self) {
if let Some(settle) = self.settle {
settle.sealed();
}
}
}
struct Held<'t, R> {
node: NodeRules<'t>,
rules: R,
}
enum Keys<'t> {
Own(String),
Lent(&'t mut String),
}
impl Keys<'_> {
fn get(&mut self) -> &mut String {
match self {
Keys::Own(keys) => keys,
Keys::Lent(keys) => keys,
}
}
fn as_str(&self) -> &str {
match self {
Keys::Own(keys) => keys,
Keys::Lent(keys) => keys,
}
}
}
#[derive(Clone, Copy)]
enum Place<'t> {
Root,
Element {
parent: &'t Place<'t>,
index: usize,
},
Entry {
parent: &'t Place<'t>,
key: (usize, usize),
},
}
impl Place<'_> {
fn locate(&self, keys: &str, steps: &mut Vec<Step>) {
match *self {
Place::Root => {}
Place::Element { parent, index } => {
parent.locate(keys, steps);
steps.push(Step::Index(index));
}
Place::Entry {
parent,
key: (start, end),
} => {
parent.locate(keys, steps);
steps.extend(keys.get(start..end).map(|key| Step::Key(key.to_owned())));
}
}
}
fn path(&self, keys: &str) -> Path {
let mut steps = Vec::new();
self.locate(keys, &mut steps);
Path::new(steps)
}
}
struct At<'a> {
place: &'a Place<'a>,
keys: &'a str,
}
impl Locate for At<'_> {
fn locate(&self, steps: &mut Vec<Step>) {
self.place.locate(self.keys, steps);
}
}
fn refused<E: Error>(refusal: Refusal) -> E {
E::source(SourceError::new(refusal))
}
impl<'t, E> Terms<'t, E> {
pub fn new(out: &'t mut Vec<Term>) -> Self {
Self::by(Equality, out)
}
fn beside_encryption(out: &'t mut Vec<Term>) -> Self {
Self {
rules: None,
..Self::new(out)
}
}
}
impl<'t, E, N: Normalize> Terms<'t, E, N> {
pub fn by(by: N, out: &'t mut Vec<Term<N::Term>>) -> Self {
Self {
out,
keys: Keys::Own(String::new()),
at: Place::Root,
rules: Some(NodeRules {
depth: 0,
settle: None,
}),
by,
_error: PhantomData,
}
}
fn at(&self) -> At<'_> {
At {
place: &self.at,
keys: self.keys.as_str(),
}
}
}
impl<'t, E: Error, N: Normalize> Sink for Terms<'t, E, N> {
type Ok = ();
type Error = E;
type Seq = TermsSeq<'t, E, N>;
type Map = TermsMap<'t, E, N>;
fn scalar(self, ty: Option<Logical>, value: Physical<'_>) -> Result<(), E> {
if let Some(rules) = &self.rules {
shape::scalar(rules.depth, ty, &value, &self.at()).map_err(refused)?;
}
match self.by.normalize(ty, value) {
Ok(normalized) => {
let path = self.at.path(self.keys.as_str());
self.out.push(Term { path, normalized });
}
Err(NormalizeError::NoDomain) if matches!(ty, None | Some(Logical::Undefined)) => {}
Err(error) => {
return Err(E::source(SourceError::new(TermError {
index: N::INDEX,
path: self.at.path(self.keys.as_str()),
error,
})))
}
}
if let Some(rules) = self.rules {
rules.sealed();
}
Ok(())
}
fn typed_null(self, _: FixedType) -> Result<(), E> {
self.scalar(None, Physical::Null)
}
fn seq(self, _: Option<usize>) -> Result<Self::Seq, E> {
let rules = match self.rules {
Some(node) => {
shape::container(
node.depth,
&At {
place: &self.at,
keys: self.keys.as_str(),
},
)
.map_err(refused)?;
Some(Held {
node,
rules: SeqRules::new(),
})
}
None => None,
};
Ok(TermsSeq {
mark: self.keys.as_str().len(),
out: self.out,
keys: self.keys,
at: self.at,
next: 0,
rules,
by: self.by,
_error: PhantomData,
})
}
fn map(self, _: Option<usize>) -> Result<Self::Map, E> {
let rules = match self.rules {
Some(node) => {
shape::container(
node.depth,
&At {
place: &self.at,
keys: self.keys.as_str(),
},
)
.map_err(refused)?;
Some(Held {
node,
rules: MapRules::new(),
})
}
None => None,
};
Ok(TermsMap {
mark: self.keys.as_str().len(),
out: self.out,
keys: self.keys,
at: self.at,
rules,
by: self.by,
_error: PhantomData,
})
}
fn passthrough<P: Encrypt + Clone>(self, value: P) -> Result<(), E> {
if let Some(rules) = self.rules {
let at = At {
place: &self.at,
keys: self.keys.as_str(),
};
shape::frame_passthrough(value, rules.depth, &at, refused)?;
if let Some(settle) = rules.settle {
settle.passthrough();
}
}
Ok(())
}
}
pub struct TermsSeq<'t, E, N: Normalize = Equality> {
out: &'t mut Vec<Term<N::Term>>,
keys: Keys<'t>,
mark: usize,
at: Place<'t>,
next: usize,
rules: Option<Held<'t, SeqRules>>,
by: N,
_error: PhantomData<E>,
}
impl<E: Error, N: Normalize> SeqSink for TermsSeq<'_, E, N> {
type Ok = ();
type Error = E;
type Element<'s>
= Terms<'s, E, N>
where
Self: 's;
fn element(&mut self) -> Result<Terms<'_, E, N>, E> {
let TermsSeq {
out,
keys,
mark,
at,
next,
rules,
by,
..
} = self;
let keys = keys.get();
let rules = match rules {
Some(held) => Some(NodeRules {
depth: held.node.depth + 1,
settle: Some(
held.rules
.element(&At { place: at, keys })
.map_err(refused)?,
),
}),
None => None,
};
keys.truncate(*mark);
let index = *next;
*next += 1;
Ok(Terms {
out,
keys: Keys::Lent(keys),
at: Place::Element { parent: at, index },
rules,
by: *by,
_error: PhantomData,
})
}
fn end(self) -> Result<(), E> {
if let Some(held) = self.rules {
let at = At {
place: &self.at,
keys: self.keys.as_str(),
};
held.rules.end(&at).map_err(refused)?;
held.node.sealed();
}
Ok(())
}
}
pub struct TermsMap<'t, E, N: Normalize = Equality> {
out: &'t mut Vec<Term<N::Term>>,
keys: Keys<'t>,
mark: usize,
at: Place<'t>,
rules: Option<Held<'t, MapRules>>,
by: N,
_error: PhantomData<E>,
}
impl<E: Error, N: Normalize> MapSink for TermsMap<'_, E, N> {
type Ok = ();
type Error = E;
type Entry<'s>
= Terms<'s, E, N>
where
Self: 's;
fn entry(&mut self, key: &str) -> Result<Terms<'_, E, N>, E> {
let TermsMap {
out,
keys,
mark,
at,
rules,
by,
..
} = self;
let keys = keys.get();
let rules = match rules {
Some(held) => Some(NodeRules {
depth: held.node.depth + 1,
settle: Some(
held.rules
.entry(key, &At { place: at, keys })
.map_err(refused)?,
),
}),
None => None,
};
keys.truncate(*mark);
keys.push_str(key);
let key = (*mark, keys.len());
Ok(Terms {
out,
keys: Keys::Lent(keys),
at: Place::Entry { parent: at, key },
rules,
by: *by,
_error: PhantomData,
})
}
fn end(self) -> Result<(), E> {
if let Some(held) = self.rules {
let at = At {
place: &self.at,
keys: self.keys.as_str(),
};
held.rules.end(&at).map_err(refused)?;
held.node.sealed();
}
Ok(())
}
}
pub fn terms<P: Encrypt>(value: P) -> Result<Vec<Term>, TermsError> {
terms_by(Equality, value)
}
pub fn terms_by<N, P>(by: N, value: P) -> Result<Vec<Term<N::Term>>, TermsError>
where
N: Normalize,
P: Encrypt,
{
let mut out = Vec::new();
value.visit(Terms::by(by, &mut out))?;
Ok(out)
}
pub type WithTerms<L> = (Node<L>, Vec<Term>);
pub fn plan_with_terms<'a, P, E, A>(value: P, aad: A) -> Result<WithTerms<Pending>, EncryptError<E>>
where
P: Encrypt,
A: IntoAad<'a>,
{
let mut terms = Vec::new();
let (planned, ()) = value.visit(Tee(
EncryptSink::root(aad),
Terms::beside_encryption(&mut terms),
))?;
Ok((planned, terms))
}
pub fn encrypt_with_terms<'a, P, C, A>(
value: P,
cipher: &C,
aad: A,
) -> Result<WithTerms<C::Leaf>, EncryptError<C::Error>>
where
P: Encrypt,
C: Cipher,
A: IntoAad<'a>,
{
let (planned, terms) = plan_with_terms(value, aad)?;
Ok((seal_all(cipher, planned)?, terms))
}
pub fn encrypt_with_terms_async<'a, 'c, P, C, A>(
value: P,
cipher: &'c C,
aad: A,
) -> impl Future<Output = Result<WithTerms<C::Leaf>, EncryptError<C::Error>>> + Send + 'c
where
P: Encrypt,
C: AsyncCipher,
A: IntoAad<'a>,
{
let planned = plan_with_terms(value, aad);
async move {
let (planned, terms) = planned?;
Ok((seal_all_async(cipher, planned).await?, terms))
}
}