Skip to main content

vitaminc_aead/
seal.rs

1//! Sealing and opening leaves: the cipher traits, and the traversals that
2//! take a tree apart into leaves and put it back together.
3//!
4//! A cipher does two things, seal one leaf and open one leaf. Everything
5//! about containers (element and entry AAD, the empty markers, duplicate
6//! keys, what counts as sealed) lives in
7//! [`EncryptSink`](crate::EncryptSink) and [`walk`](crate::walk), so a
8//! second cipher cannot diverge on it.
9//!
10//! [`Cipher`] is synchronous. [`AsyncCipher`] is for anything that has to
11//! be awaited: a remote key provider, a KMS, a cloud HSM. Its `seal_many`
12//! and `open_many` take every leaf of one document at once, so a provider
13//! with a batch API makes one round trip per document. The walk itself is
14//! never awaited: encryption plans a tree of [`Pending`] leaves first, and
15//! decryption opens every leaf before the sink sees anything.
16//!
17//! Both traits are re-exported from the crate root.
18
19use core::future::Future;
20use std::collections::HashSet;
21
22use futures_util::stream::{self, StreamExt, TryStreamExt};
23use vitaminc_protected::Protected;
24
25use crate::frame::PassthroughFrame;
26use crate::{CipherText, Context};
27
28/// A ciphertext tree whose passthrough subtrees are frames.
29pub type Node<L> = CipherText<L, PassthroughFrame>;
30
31/// A leaf the walk has planned but nobody has sealed: its tagged plaintext
32/// and the AAD derived for its place in the tree.
33pub struct Pending {
34    pub plaintext: Protected<Vec<u8>>,
35    pub aad: Context<'static>,
36}
37
38/// A leaf after opening: its tagged plaintext.
39pub type Opened = Protected<Vec<u8>>;
40
41/// Seal and open one leaf, synchronously.
42///
43/// The `aad` is the caller-visible context for the leaf's position. The
44/// cipher binds the wire version into it itself
45/// ([`Context::for_leaf`]) at the AEAD boundary.
46pub trait Cipher {
47    /// The sealed leaf.
48    type Leaf;
49    /// Kept opaque: see [`Unspecified`](crate::Unspecified).
50    type Error;
51
52    fn seal(
53        &self,
54        plaintext: Protected<Vec<u8>>,
55        aad: &Context<'_>,
56    ) -> Result<Self::Leaf, Self::Error>;
57
58    fn open(&self, leaf: Self::Leaf, aad: &Context<'_>) -> Result<Protected<Vec<u8>>, Self::Error>;
59}
60
61/// Seal and open leaves through something that has to be awaited.
62///
63/// `seal_many` and `open_many` take every leaf of one document in walk
64/// order. The defaults seal each leaf concurrently, at most
65/// [`DEFAULT_CONCURRENCY`] at a time; a provider with a batch API
66/// overrides them to make one round trip per document. Either way the
67/// cipher must return exactly one leaf per leaf it was given, in the same
68/// order, or the document is refused.
69///
70/// `Sync`, because an async cipher is shared across tasks.
71pub trait AsyncCipher: Sync {
72    type Leaf: Send;
73    type Error: Send;
74
75    fn seal(
76        &self,
77        plaintext: Protected<Vec<u8>>,
78        aad: Context<'static>,
79    ) -> impl Future<Output = Result<Self::Leaf, Self::Error>> + Send;
80
81    fn open(
82        &self,
83        leaf: Self::Leaf,
84        aad: Context<'static>,
85    ) -> impl Future<Output = Result<Protected<Vec<u8>>, Self::Error>> + Send;
86
87    /// Every pending leaf of one document, in walk order.
88    fn seal_many(
89        &self,
90        leaves: Vec<Pending>,
91    ) -> impl Future<Output = Result<Vec<Self::Leaf>, Self::Error>> + Send {
92        stream::iter(leaves)
93            .map(|p| self.seal(p.plaintext, p.aad))
94            .buffered(DEFAULT_CONCURRENCY)
95            .try_collect()
96    }
97
98    /// Every sealed leaf of one document with the AAD its position derives,
99    /// in walk order.
100    fn open_many(
101        &self,
102        leaves: Vec<(Self::Leaf, Context<'static>)>,
103    ) -> impl Future<Output = Result<Vec<Protected<Vec<u8>>>, Self::Error>> + Send {
104        stream::iter(leaves)
105            .map(|(leaf, aad)| self.open(leaf, aad))
106            .buffered(DEFAULT_CONCURRENCY)
107            .try_collect()
108    }
109}
110
111/// How many leaves the default `seal_many` and `open_many` have in flight
112/// at once, so a large document does not send every leaf to a remote
113/// provider at the same moment.
114pub const DEFAULT_CONCURRENCY: usize = 32;
115
116/// A cipher returned a different number of leaves than it was given.
117#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
118pub enum LeafCount {
119    #[error("too few leaves")]
120    TooFew,
121    #[error("too many leaves")]
122    TooMany,
123}
124
125/// Seal every pending leaf, one at a time.
126pub fn seal_all<C: Cipher>(
127    cipher: &C,
128    node: Node<Pending>,
129) -> Result<Node<C::Leaf>, SealError<C::Error>> {
130    let mut pending = Vec::new();
131    let shape = split(node, &mut pending);
132    let leaves = pending
133        .into_iter()
134        .map(|p| cipher.seal(p.plaintext, &p.aad))
135        .collect::<Result<Vec<_>, _>>()
136        .map_err(SealError::Cipher)?;
137    fill(shape, leaves).map_err(SealError::Count)
138}
139
140/// Seal every pending leaf in one `seal_many` call.
141///
142/// A tree with no leaves (a top-level passthrough) makes no call at all,
143/// as the sync path and the per-leaf defaults would not, so a batching
144/// provider is never asked for an empty batch.
145pub async fn seal_all_async<C: AsyncCipher>(
146    cipher: &C,
147    node: Node<Pending>,
148) -> Result<Node<C::Leaf>, SealError<C::Error>> {
149    let mut pending = Vec::new();
150    let shape = split(node, &mut pending);
151    let leaves = if pending.is_empty() {
152        Vec::new()
153    } else {
154        cipher.seal_many(pending).await.map_err(SealError::Cipher)?
155    };
156    fill(shape, leaves).map_err(SealError::Count)
157}
158
159/// Why sealing failed.
160///
161/// A cipher's own error is printed in the message rather than chained as
162/// the `source()`, so a cipher error type need not implement `Error`.
163#[derive(Debug, thiserror::Error)]
164pub enum SealError<E> {
165    #[error("cipher failed: {0}")]
166    Cipher(E),
167    #[error("cipher returned {0}")]
168    Count(LeafCount),
169}
170
171/// Open every leaf, one at a time.
172///
173/// The tree's structure is checked before any leaf is opened, so a
174/// malformed tree costs no cipher calls. A cipher's failure is returned as
175/// it is: for an AEAD that is wrong key, wrong AAD or tampering, for a
176/// remote provider it may also be an outage, and the cipher's error type
177/// says which. A leaf count mismatch is a cipher bug.
178pub fn open_all<'a, C: Cipher>(
179    cipher: &C,
180    node: Node<C::Leaf>,
181    aad: &Context<'a>,
182) -> Result<Node<Opened>, OpenError<C::Error>> {
183    check_shape(&node)?;
184    let mut sealed = Vec::new();
185    let shape = split_with_aad(node, aad, &mut sealed);
186    let opened = sealed
187        .into_iter()
188        .map(|(leaf, aad)| cipher.open(leaf, &aad))
189        .collect::<Result<Vec<_>, _>>()
190        .map_err(OpenError::Cipher)?;
191    fill(shape, opened).map_err(OpenError::Count)
192}
193
194/// Open every leaf in one `open_many` call, after the same structural
195/// check as [`open_all`]. A tree with no leaves makes no call.
196pub async fn open_all_async<'a, C: AsyncCipher>(
197    cipher: &C,
198    node: Node<C::Leaf>,
199    aad: &Context<'a>,
200) -> Result<Node<Opened>, OpenError<C::Error>> {
201    check_shape(&node)?;
202    let mut sealed = Vec::new();
203    let shape = split_with_aad(node, aad, &mut sealed);
204    let opened = if sealed.is_empty() {
205        Vec::new()
206    } else {
207        cipher.open_many(sealed).await.map_err(OpenError::Cipher)?
208    };
209    fill(shape, opened).map_err(OpenError::Count)
210}
211
212/// Why opening failed.
213#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
214pub enum OpenError<E> {
215    /// The cipher refused a leaf.
216    #[error("cipher failed: {0}")]
217    Cipher(E),
218    /// The tree breaks a container rule: a container with nothing sealed
219    /// in it, a repeated map key, or a passthrough that is not a frame.
220    #[error("malformed ciphertext")]
221    Malformed,
222    /// The cipher returned the wrong number of leaves.
223    #[error("cipher returned {0}")]
224    Count(LeafCount),
225}
226
227/// The container rules `EncryptSink` enforced, checked on the sealed tree
228/// so a malformed one is refused before any leaf is opened. `walk`
229/// checks them again, since it accepts opened trees from anywhere.
230fn check_shape<L, E>(node: &Node<L>) -> Result<(), OpenError<E>> {
231    let sealed_any = |mut nodes: core::slice::Iter<'_, Node<L>>| {
232        nodes.any(|n| !matches!(n, CipherText::Passthrough(_)))
233    };
234    match node {
235        CipherText::Sequence(items) => {
236            if !sealed_any(items.iter()) {
237                return Err(OpenError::Malformed);
238            }
239            items.iter().try_for_each(check_shape)
240        }
241        CipherText::Map(entries) => {
242            let mut seen = HashSet::new();
243            if !entries.iter().all(|(k, _)| seen.insert(k.as_str()))
244                || !entries
245                    .iter()
246                    .any(|(_, n)| !matches!(n, CipherText::Passthrough(_)))
247            {
248                return Err(OpenError::Malformed);
249            }
250            entries.iter().try_for_each(|(_, n)| check_shape(n))
251        }
252        CipherText::Passthrough(frame) => {
253            frame.parse().map(|_| ()).map_err(|_| OpenError::Malformed)
254        }
255        CipherText::Single(_) | CipherText::EmptySequence(_) | CipherText::EmptyMap(_) => Ok(()),
256    }
257}
258
259/// A tree's shape with its leaves taken out, in walk order.
260fn split<L>(node: Node<L>, leaves: &mut Vec<L>) -> Node<()> {
261    match node {
262        CipherText::Single(leaf) => {
263            leaves.push(leaf);
264            CipherText::Single(())
265        }
266        CipherText::Sequence(items) => {
267            CipherText::Sequence(items.into_iter().map(|n| split(n, leaves)).collect())
268        }
269        CipherText::EmptySequence(leaf) => {
270            leaves.push(leaf);
271            CipherText::EmptySequence(())
272        }
273        CipherText::Map(entries) => CipherText::Map(
274            entries
275                .into_iter()
276                .map(|(k, n)| (k, split(n, leaves)))
277                .collect(),
278        ),
279        CipherText::EmptyMap(leaf) => {
280            leaves.push(leaf);
281            CipherText::EmptyMap(())
282        }
283        CipherText::Passthrough(p) => CipherText::Passthrough(p),
284    }
285}
286
287/// A tree's shape with its leaves taken out, each paired with the AAD its
288/// position derives. The derivations are the ones `EncryptSink` used.
289fn split_with_aad<L>(
290    node: Node<L>,
291    aad: &Context<'_>,
292    leaves: &mut Vec<(L, Context<'static>)>,
293) -> Node<()> {
294    match node {
295        CipherText::Single(leaf) => {
296            leaves.push((leaf, aad.clone().into_owned()));
297            CipherText::Single(())
298        }
299        CipherText::Sequence(items) => {
300            let element_aad = aad.for_sequence_element();
301            CipherText::Sequence(
302                items
303                    .into_iter()
304                    .map(|n| split_with_aad(n, &element_aad, leaves))
305                    .collect(),
306            )
307        }
308        CipherText::EmptySequence(leaf) => {
309            leaves.push((leaf, aad.for_empty_sequence()));
310            CipherText::EmptySequence(())
311        }
312        CipherText::Map(entries) => CipherText::Map(
313            entries
314                .into_iter()
315                .map(|(k, n)| {
316                    let shape = split_with_aad(n, &aad.for_map_entry(&k), leaves);
317                    (k, shape)
318                })
319                .collect(),
320        ),
321        CipherText::EmptyMap(leaf) => {
322            leaves.push((leaf, aad.for_empty_map()));
323            CipherText::EmptyMap(())
324        }
325        CipherText::Passthrough(p) => CipherText::Passthrough(p),
326    }
327}
328
329/// Put leaves back into a shape, in walk order.
330fn fill<L>(shape: Node<()>, leaves: Vec<L>) -> Result<Node<L>, LeafCount> {
331    let mut leaves = leaves.into_iter();
332    let node = fill_from(shape, &mut leaves).ok_or(LeafCount::TooFew)?;
333    if leaves.next().is_some() {
334        return Err(LeafCount::TooMany);
335    }
336    Ok(node)
337}
338
339fn fill_from<L>(shape: Node<()>, leaves: &mut impl Iterator<Item = L>) -> Option<Node<L>> {
340    Some(match shape {
341        CipherText::Single(()) => CipherText::Single(leaves.next()?),
342        CipherText::Sequence(items) => CipherText::Sequence(
343            items
344                .into_iter()
345                .map(|n| fill_from(n, leaves))
346                .collect::<Option<_>>()?,
347        ),
348        CipherText::EmptySequence(()) => CipherText::EmptySequence(leaves.next()?),
349        CipherText::Map(entries) => CipherText::Map(
350            entries
351                .into_iter()
352                .map(|(k, n)| fill_from(n, leaves).map(|n| (k, n)))
353                .collect::<Option<_>>()?,
354        ),
355        CipherText::EmptyMap(()) => CipherText::EmptyMap(leaves.next()?),
356        CipherText::Passthrough(p) => CipherText::Passthrough(p),
357    })
358}
359
360/// Every leaf of a tree, in walk order, with the AAD its position derives.
361/// For tests and for ciphers that inspect a plan.
362pub fn leaves_with_aad<L>(node: Node<L>, aad: &Context<'_>) -> Vec<(L, Context<'static>)> {
363    let mut out = Vec::new();
364    split_with_aad(node, aad, &mut out);
365    out
366}
367
368/// A cipher that stores the plaintext and AAD in the clear, for tests of
369/// the traversals and the sink. Opening checks the AAD matches.
370#[cfg(test)]
371pub(crate) mod mock {
372    use super::*;
373    use core::sync::atomic::{AtomicUsize, Ordering};
374    use futures_util::future::try_join_all;
375
376    #[derive(Debug, Clone, PartialEq, Eq)]
377    pub struct Leaf {
378        pub aad: Vec<u8>,
379        pub plaintext: Vec<u8>,
380    }
381
382    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
383    pub struct Refused;
384
385    #[derive(Default)]
386    pub struct Mock;
387
388    impl Cipher for Mock {
389        type Leaf = Leaf;
390        type Error = Refused;
391
392        fn seal(&self, plaintext: Protected<Vec<u8>>, aad: &Context<'_>) -> Result<Leaf, Refused> {
393            Ok(Leaf {
394                aad: aad.as_bytes().to_vec(),
395                plaintext: vitaminc_protected::Controlled::risky_unwrap(plaintext),
396            })
397        }
398
399        fn open(&self, leaf: Leaf, aad: &Context<'_>) -> Result<Protected<Vec<u8>>, Refused> {
400            if leaf.aad != aad.as_bytes() {
401                return Err(Refused);
402            }
403            Ok(Protected::new(leaf.plaintext))
404        }
405    }
406
407    /// The mock behind an await, counting calls. `batch` decides whether
408    /// `seal_many`/`open_many` are one call or the per-leaf defaults.
409    pub struct Remote {
410        pub calls: AtomicUsize,
411        pub batch: bool,
412    }
413
414    impl Remote {
415        pub fn new(batch: bool) -> Self {
416            Self {
417                calls: AtomicUsize::new(0),
418                batch,
419            }
420        }
421
422        pub fn calls(&self) -> usize {
423            self.calls.load(Ordering::SeqCst)
424        }
425    }
426
427    impl AsyncCipher for Remote {
428        type Leaf = Leaf;
429        type Error = Refused;
430
431        async fn seal(
432            &self,
433            plaintext: Protected<Vec<u8>>,
434            aad: Context<'static>,
435        ) -> Result<Leaf, Refused> {
436            self.calls.fetch_add(1, Ordering::SeqCst);
437            Cipher::seal(&Mock, plaintext, &aad)
438        }
439
440        async fn open(
441            &self,
442            leaf: Leaf,
443            aad: Context<'static>,
444        ) -> Result<Protected<Vec<u8>>, Refused> {
445            self.calls.fetch_add(1, Ordering::SeqCst);
446            Cipher::open(&Mock, leaf, &aad)
447        }
448
449        async fn seal_many(&self, leaves: Vec<Pending>) -> Result<Vec<Leaf>, Refused> {
450            if !self.batch {
451                return try_join_all(leaves.into_iter().map(|p| self.seal(p.plaintext, p.aad)))
452                    .await;
453            }
454            self.calls.fetch_add(1, Ordering::SeqCst);
455            leaves
456                .into_iter()
457                .map(|p| Cipher::seal(&Mock, p.plaintext, &p.aad))
458                .collect()
459        }
460
461        async fn open_many(
462            &self,
463            leaves: Vec<(Leaf, Context<'static>)>,
464        ) -> Result<Vec<Protected<Vec<u8>>>, Refused> {
465            if !self.batch {
466                return try_join_all(leaves.into_iter().map(|(l, aad)| self.open(l, aad))).await;
467            }
468            self.calls.fetch_add(1, Ordering::SeqCst);
469            leaves
470                .into_iter()
471                .map(|(l, aad)| Cipher::open(&Mock, l, &aad))
472                .collect()
473        }
474    }
475}
476
477#[cfg(test)]
478mod tests {
479    use super::mock::{Leaf, Mock, Refused, Remote};
480    use super::*;
481    use crate::frame::PassthroughFrame;
482    use crate::plaintext::record::{record, Event, Owned};
483    use crate::{decrypt_into, encrypt, encrypt_async, plan, walk, EncryptError, Int, IntoAad};
484    use core::sync::atomic::{AtomicUsize, Ordering};
485    use futures_executor::block_on;
486    use std::collections::BTreeMap;
487    use vitaminc_protected::Controlled;
488
489    fn sample() -> BTreeMap<&'static str, Vec<Option<u32>>> {
490        let mut m = BTreeMap::new();
491        m.insert("a", vec![Some(1), None, Some(3)]);
492        m.insert("b", vec![]);
493        m.insert("c", vec![Some(4)]);
494        m
495    }
496
497    #[test]
498    fn leaves_come_out_in_walk_order_with_their_own_aad() {
499        let planned = plan::<_, Refused, _>(sample(), "ctx").expect("plans");
500        let leaves = leaves_with_aad(planned, &"ctx".into_aad());
501        // a: three elements, b: one empty marker, c: one element.
502        assert_eq!(leaves.len(), 5);
503        let a = "ctx".into_aad().for_map_entry("a").for_sequence_element();
504        assert_eq!(leaves[0].1, a);
505        assert_eq!(leaves[1].1, a);
506        assert_eq!(leaves[2].1, a);
507        assert_eq!(
508            leaves[3].1,
509            "ctx".into_aad().for_map_entry("b").for_empty_sequence()
510        );
511        assert_eq!(
512            leaves[4].1,
513            "ctx".into_aad().for_map_entry("c").for_sequence_element()
514        );
515        // The marker is empty; the absent `u32` is a typed null, as long
516        // as the `u32` leaves beside it.
517        assert!(leaves[3].0.plaintext.risky_ref().is_empty());
518        assert_eq!(leaves[1].0.plaintext.risky_ref(), &[0x15, 0x06, 0, 0, 0]);
519        assert_eq!(leaves[0].0.plaintext.risky_ref().len(), 5);
520        // And the plan's AAD is what the sealed leaf is bound under.
521        let sealed = encrypt(sample(), &Mock, "ctx").expect("encrypts");
522        let sealed = leaves_with_aad(sealed, &"ctx".into_aad());
523        assert!(sealed.iter().all(|(leaf, aad)| leaf.aad == aad.as_bytes()));
524    }
525
526    #[test]
527    fn open_then_walk_is_the_original() {
528        let sealed = encrypt(sample(), &Mock, "ctx").expect("encrypts");
529        let opened = open_all(&Mock, sealed, &"ctx".into_aad()).expect("opens");
530        let mut log = Vec::new();
531        walk(opened, crate::plaintext::record::Record(&mut log)).expect("walks");
532        assert_eq!(log, record(sample()).expect("records"));
533    }
534
535    #[test]
536    fn a_wrong_aad_fails_to_open() {
537        let sealed = encrypt(sample(), &Mock, "ctx").expect("encrypts");
538        assert_eq!(
539            open_all(&Mock, sealed, &"other".into_aad()).err(),
540            Some(OpenError::Cipher(Refused))
541        );
542    }
543
544    /// Counts opens, so a test can see that a malformed tree costs none.
545    struct Counting(core::cell::Cell<usize>);
546
547    impl Cipher for Counting {
548        type Leaf = Leaf;
549        type Error = Refused;
550        fn seal(&self, plaintext: Protected<Vec<u8>>, aad: &Context<'_>) -> Result<Leaf, Refused> {
551            Mock.seal(plaintext, aad)
552        }
553        fn open(&self, leaf: Leaf, aad: &Context<'_>) -> Result<Protected<Vec<u8>>, Refused> {
554            self.0.set(self.0.get() + 1);
555            Mock.open(leaf, aad)
556        }
557    }
558
559    #[test]
560    fn a_malformed_tree_is_refused_before_any_leaf_is_opened() {
561        let leaf = || match encrypt(1u16, &Mock, "ctx").expect("encrypts") {
562            CipherText::Single(leaf) => leaf,
563            _ => unreachable!(),
564        };
565        let passthrough = || CipherText::Passthrough(PassthroughFrame(vec![0x00]));
566        let malformed: [Node<Leaf>; 5] = [
567            CipherText::Map(vec![
568                ("k".into(), CipherText::Single(leaf())),
569                ("k".into(), CipherText::Single(leaf())),
570            ]),
571            CipherText::Sequence(vec![passthrough(), passthrough()]),
572            CipherText::Sequence(vec![]),
573            CipherText::Sequence(vec![
574                CipherText::Single(leaf()),
575                CipherText::Passthrough(PassthroughFrame(vec![0xEF])),
576            ]),
577            // A map whose value is a sequence with nothing sealed in it.
578            CipherText::Map(vec![
579                ("a".into(), CipherText::Single(leaf())),
580                ("b".into(), CipherText::Sequence(vec![passthrough()])),
581            ]),
582        ];
583        for node in malformed {
584            let cipher = Counting(core::cell::Cell::new(0));
585            assert_eq!(
586                open_all(&cipher, node, &"ctx".into_aad()).err(),
587                Some(OpenError::Malformed)
588            );
589            assert_eq!(cipher.0.get(), 0);
590        }
591    }
592
593    #[test]
594    fn a_provider_error_on_open_is_not_reported_as_tampering() {
595        struct Down;
596        impl AsyncCipher for Down {
597            type Leaf = Leaf;
598            type Error = &'static str;
599            async fn seal(
600                &self,
601                p: Protected<Vec<u8>>,
602                aad: Context<'static>,
603            ) -> Result<Leaf, &'static str> {
604                Mock.seal(p, &aad).map_err(|_| "refused")
605            }
606            async fn open(
607                &self,
608                _: Leaf,
609                _: Context<'static>,
610            ) -> Result<Protected<Vec<u8>>, &'static str> {
611                Err("kms unavailable")
612            }
613        }
614        let sealed = encrypt(34u32, &Mock, ()).expect("encrypts");
615        let mut log = Vec::new();
616        assert!(matches!(
617            block_on(crate::decrypt_into_async(
618                &Down,
619                sealed,
620                (),
621                crate::plaintext::record::Record(&mut log)
622            )),
623            Err(crate::DecryptError::Cipher("kms unavailable"))
624        ));
625    }
626
627    /// Tracks how many leaves are in flight at once through the default
628    /// `seal_many`/`open_many`. Each call yields once, so the stream has to
629    /// hold several open to make progress.
630    #[derive(Default)]
631    struct Concurrency {
632        now: AtomicUsize,
633        max: AtomicUsize,
634    }
635
636    struct YieldOnce(bool);
637
638    impl Future for YieldOnce {
639        type Output = ();
640        fn poll(
641            mut self: core::pin::Pin<&mut Self>,
642            cx: &mut core::task::Context<'_>,
643        ) -> core::task::Poll<()> {
644            if self.0 {
645                return core::task::Poll::Ready(());
646            }
647            self.0 = true;
648            cx.waker().wake_by_ref();
649            core::task::Poll::Pending
650        }
651    }
652
653    impl Concurrency {
654        async fn one<T>(&self, f: impl FnOnce() -> T) -> T {
655            let now = self.now.fetch_add(1, Ordering::SeqCst) + 1;
656            self.max.fetch_max(now, Ordering::SeqCst);
657            YieldOnce(false).await;
658            self.now.fetch_sub(1, Ordering::SeqCst);
659            f()
660        }
661    }
662
663    impl AsyncCipher for Concurrency {
664        type Leaf = Leaf;
665        type Error = Refused;
666        async fn seal(
667            &self,
668            p: Protected<Vec<u8>>,
669            aad: Context<'static>,
670        ) -> Result<Leaf, Refused> {
671            self.one(|| Mock.seal(p, &aad)).await
672        }
673        async fn open(
674            &self,
675            l: Leaf,
676            aad: Context<'static>,
677        ) -> Result<Protected<Vec<u8>>, Refused> {
678            self.one(|| Mock.open(l, &aad)).await
679        }
680    }
681
682    /// Each call waits longer than the one after it, so leaves finish in
683    /// the reverse of the order they started.
684    struct LastFirst(AtomicUsize);
685
686    impl LastFirst {
687        async fn one<T>(&self, f: impl FnOnce() -> T) -> T {
688            let wait = 64usize.saturating_sub(self.0.fetch_add(1, Ordering::SeqCst));
689            for _ in 0..wait {
690                YieldOnce(false).await;
691            }
692            f()
693        }
694    }
695
696    impl AsyncCipher for LastFirst {
697        type Leaf = Leaf;
698        type Error = Refused;
699        async fn seal(
700            &self,
701            p: Protected<Vec<u8>>,
702            aad: Context<'static>,
703        ) -> Result<Leaf, Refused> {
704            self.one(|| Mock.seal(p, &aad)).await
705        }
706        async fn open(
707            &self,
708            l: Leaf,
709            aad: Context<'static>,
710        ) -> Result<Protected<Vec<u8>>, Refused> {
711            self.one(|| Mock.open(l, &aad)).await
712        }
713    }
714
715    #[test]
716    fn the_default_batch_keeps_walk_order_when_leaves_finish_out_of_order() {
717        let values = vec![1u32, 2, 3, 4];
718        let cipher = LastFirst(AtomicUsize::new(0));
719        let sealed = block_on(encrypt_async(values.clone(), &cipher, "ctx")).expect("encrypts");
720        let opened = open_all(&Mock, sealed, &"ctx".into_aad()).expect("opens");
721        let mut log = Vec::new();
722        walk(opened, crate::plaintext::record::Record(&mut log)).expect("walks");
723        assert_eq!(log, record(values.clone()).expect("records"));
724
725        let sealed = encrypt(values.clone(), &Mock, "ctx").expect("encrypts");
726        let mut log = Vec::new();
727        block_on(crate::decrypt_into_async(
728            &cipher,
729            sealed,
730            "ctx",
731            crate::plaintext::record::Record(&mut log),
732        ))
733        .expect("decrypts");
734        assert_eq!(log, record(values).expect("records"));
735    }
736
737    #[test]
738    fn a_malformed_tree_costs_an_async_cipher_no_round_trip() {
739        let leaf = || match encrypt(1u16, &Mock, "ctx").expect("encrypts") {
740            CipherText::Single(leaf) => leaf,
741            _ => unreachable!(),
742        };
743        let duplicate_keys: Node<Leaf> = CipherText::Map(vec![
744            ("k".into(), CipherText::Single(leaf())),
745            ("k".into(), CipherText::Single(leaf())),
746        ]);
747        let remote = Remote::new(true);
748        assert_eq!(
749            block_on(open_all_async(&remote, duplicate_keys, &"ctx".into_aad())).err(),
750            Some(OpenError::Malformed)
751        );
752        assert_eq!(remote.calls(), 0);
753    }
754
755    #[test]
756    fn seal_and_open_errors_say_what_failed() {
757        assert_eq!(LeafCount::TooFew.to_string(), "too few leaves");
758        assert_eq!(LeafCount::TooMany.to_string(), "too many leaves");
759        assert_eq!(
760            SealError::<&str>::Cipher("down").to_string(),
761            "cipher failed: down"
762        );
763        assert_eq!(
764            SealError::<&str>::Count(LeafCount::TooMany).to_string(),
765            "cipher returned too many leaves"
766        );
767        assert_eq!(
768            OpenError::<&str>::Cipher("down").to_string(),
769            "cipher failed: down"
770        );
771        assert_eq!(
772            OpenError::<&str>::Malformed.to_string(),
773            "malformed ciphertext"
774        );
775        assert_eq!(
776            OpenError::<&str>::Count(LeafCount::TooFew).to_string(),
777            "cipher returned too few leaves"
778        );
779        fn is_error<E: core::error::Error>(_: &E) {}
780        is_error(&SealError::<&str>::Cipher("down"));
781        is_error(&OpenError::<&str>::Malformed);
782    }
783
784    #[test]
785    fn the_default_batch_bounds_how_many_leaves_are_in_flight() {
786        let cipher = Concurrency::default();
787        let many: Vec<u32> = (0..200).collect();
788        let sealed = block_on(encrypt_async(many.clone(), &cipher, "ctx")).expect("encrypts");
789        let max = cipher.max.load(Ordering::SeqCst);
790        assert!(max > 1 && max <= DEFAULT_CONCURRENCY, "{max}");
791        let mut log = Vec::new();
792        block_on(crate::decrypt_into_async(
793            &cipher,
794            sealed,
795            "ctx",
796            crate::plaintext::record::Record(&mut log),
797        ))
798        .expect("decrypts");
799        assert!(cipher.max.load(Ordering::SeqCst) <= DEFAULT_CONCURRENCY);
800        assert_eq!(log, record(many).expect("records"));
801    }
802
803    #[test]
804    fn encrypt_async_walks_before_returning_a_send_future() {
805        // A source a JavaScript binding might hand over: not `Send`.
806        #[derive(Clone)]
807        struct Local(std::rc::Rc<u32>);
808        impl crate::Encrypt for Local {
809            fn visit<S: crate::Sink>(self, sink: S) -> Result<S::Ok, S::Error> {
810                (*self.0).visit(sink)
811            }
812        }
813        fn assert_send<T: Send>(t: T) -> T {
814            t
815        }
816        let remote = Remote::new(true);
817        let future = assert_send(encrypt_async(Local(std::rc::Rc::new(34)), &remote, "ctx"));
818        assert_eq!(remote.calls(), 0);
819        block_on(future).expect("encrypts");
820        assert_eq!(remote.calls(), 1);
821    }
822
823    #[test]
824    fn a_cipher_that_miscounts_is_refused() {
825        struct DropsOne;
826        impl Cipher for DropsOne {
827            type Leaf = Leaf;
828            type Error = Refused;
829            fn seal(
830                &self,
831                plaintext: Protected<Vec<u8>>,
832                aad: &Context<'_>,
833            ) -> Result<Leaf, Refused> {
834                Mock.seal(plaintext, aad)
835            }
836            fn open(&self, leaf: Leaf, aad: &Context<'_>) -> Result<Protected<Vec<u8>>, Refused> {
837                Mock.open(leaf, aad)
838            }
839        }
840        struct Extra(Remote);
841        impl AsyncCipher for Extra {
842            type Leaf = Leaf;
843            type Error = Refused;
844            async fn seal(
845                &self,
846                plaintext: Protected<Vec<u8>>,
847                aad: Context<'static>,
848            ) -> Result<Leaf, Refused> {
849                self.0.seal(plaintext, aad).await
850            }
851            async fn open(
852                &self,
853                leaf: Leaf,
854                aad: Context<'static>,
855            ) -> Result<Protected<Vec<u8>>, Refused> {
856                self.0.open(leaf, aad).await
857            }
858            async fn seal_many(&self, leaves: Vec<Pending>) -> Result<Vec<Leaf>, Refused> {
859                let mut out = self.0.seal_many(leaves).await?;
860                out.push(out[0].clone());
861                Ok(out)
862            }
863            async fn open_many(
864                &self,
865                mut leaves: Vec<(Leaf, Context<'static>)>,
866            ) -> Result<Vec<Protected<Vec<u8>>>, Refused> {
867                leaves.pop();
868                self.0.open_many(leaves).await
869            }
870        }
871        let extra = Extra(Remote::new(true));
872        assert!(matches!(
873            block_on(encrypt_async(sample(), &extra, "ctx")),
874            Err(EncryptError::Count(LeafCount::TooMany))
875        ));
876        let sealed = encrypt(sample(), &DropsOne, "ctx").expect("encrypts");
877        assert_eq!(
878            block_on(open_all_async(&extra, sealed, &"ctx".into_aad())).err(),
879            Some(OpenError::<Refused>::Count(LeafCount::TooFew))
880        );
881    }
882
883    #[test]
884    fn a_batching_cipher_makes_one_round_trip_each_way() {
885        let remote = Remote::new(true);
886        let sealed = block_on(encrypt_async(sample(), &remote, "ctx")).expect("encrypts");
887        assert_eq!(remote.calls(), 1);
888        let mut log = Vec::new();
889        block_on(crate::decrypt_into_async(
890            &remote,
891            sealed,
892            "ctx",
893            crate::plaintext::record::Record(&mut log),
894        ))
895        .expect("decrypts");
896        assert_eq!(remote.calls(), 2);
897        assert_eq!(log, record(sample()).expect("records"));
898    }
899
900    #[test]
901    fn a_document_with_no_leaves_makes_no_cipher_call() {
902        // A top-level passthrough seals nothing, so a batching provider is
903        // not asked for an empty batch either way.
904        struct Pass;
905        impl crate::Encrypt for Pass {
906            fn visit<S: crate::Sink>(self, sink: S) -> Result<S::Ok, S::Error> {
907                sink.passthrough("public")
908            }
909        }
910        let remote = Remote::new(true);
911        let sealed = block_on(encrypt_async(Pass, &remote, "ctx")).expect("encrypts");
912        assert!(matches!(sealed, CipherText::Passthrough(_)));
913        let mut log = Vec::new();
914        block_on(crate::decrypt_into_async(
915            &remote,
916            sealed,
917            "ctx",
918            crate::plaintext::record::Record(&mut log),
919        ))
920        .expect("decrypts");
921        assert_eq!(remote.calls(), 0);
922        assert_eq!(log, record(Pass).expect("records"));
923    }
924
925    #[test]
926    fn without_a_batch_api_each_leaf_is_its_own_call() {
927        let remote = Remote::new(false);
928        let sealed = block_on(encrypt_async(sample(), &remote, "ctx")).expect("encrypts");
929        assert_eq!(remote.calls(), 5);
930        let mut log = Vec::new();
931        block_on(crate::decrypt_into_async(
932            &remote,
933            sealed,
934            "ctx",
935            crate::plaintext::record::Record(&mut log),
936        ))
937        .expect("decrypts");
938        assert_eq!(remote.calls(), 10);
939    }
940
941    #[test]
942    fn sync_and_async_ciphers_read_each_others_leaves() {
943        let remote = Remote::new(true);
944        let sealed = block_on(encrypt_async(34u32, &remote, ())).expect("encrypts");
945        let mut log = Vec::new();
946        decrypt_into(
947            &Mock,
948            sealed,
949            (),
950            crate::plaintext::record::Record(&mut log),
951        )
952        .expect("decrypts");
953        assert_eq!(log, [Event::Scalar(None, Owned::Int(Int::U32(34)))]);
954    }
955
956    #[test]
957    fn a_cipher_failure_fails_the_document() {
958        struct Down;
959        impl AsyncCipher for Down {
960            type Leaf = Leaf;
961            type Error = &'static str;
962            async fn seal(
963                &self,
964                _: Protected<Vec<u8>>,
965                _: Context<'static>,
966            ) -> Result<Leaf, &'static str> {
967                Err("key provider unavailable")
968            }
969            async fn open(
970                &self,
971                _: Leaf,
972                _: Context<'static>,
973            ) -> Result<Protected<Vec<u8>>, &'static str> {
974                Err("key provider unavailable")
975            }
976        }
977        assert!(matches!(
978            block_on(encrypt_async(34u32, &Down, ())),
979            Err(EncryptError::Cipher("key provider unavailable"))
980        ));
981    }
982}