feat: seal messages under batch keys with the wrap in the head and open them through a batch key cache

crates/crypto/src/secret.rs

@@ -68,7 +68,8 @@ impl<const N: usize> Secret<N> {
         Ok(secret)
     }
 
-    pub(crate) fn copy(&self) -> Self {
+    #[must_use]
+    pub fn copy(&self) -> Self {
         let mut out = Self([0; N]);
         out.0.copy_from_slice(&self.0);
         out

crates/keys/src/batch.rs

@@ -0,0 +1,223 @@
+use std::collections::HashMap;
+use std::rc::Rc;
+use std::time::Duration;
+
+use crypto::{BatchKey, Secret, WRAP, seal_batch};
+use rt::Instant;
+
+use crate::Working;
+use crate::cache::{Cache, CacheStats};
+
+const SWEEP: Duration = Duration::from_secs(1);
+
+pub struct Batch {
+    pub id: u64,
+    pub tenant: u32,
+    pub generation: u32,
+    pub day: u32,
+    pub shard: u16,
+    pub wrap: [u8; WRAP],
+    pub key: Rc<BatchKey>,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct BatchSettings {
+    pub seals: u32,
+    pub age: Duration,
+    pub cache: usize,
+}
+
+impl Default for BatchSettings {
+    fn default() -> Self {
+        Self {
+            seals: 1_024,
+            age: Duration::from_secs(60),
+            cache: 65_536,
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub struct BatchStats {
+    pub created: u64,
+    pub full: u64,
+    pub aged: u64,
+    pub day: u64,
+    pub mismatched: u64,
+    pub cache: CacheStats,
+}
+
+struct Opened {
+    batch: Rc<Batch>,
+    seals: u32,
+    since: Instant,
+}
+
+#[derive(Clone)]
+struct Cached {
+    tenant: u32,
+    generation: u32,
+    key: Rc<BatchKey>,
+}
+
+pub struct Batches {
+    settings: BatchSettings,
+    next_sweep: Option<Instant>,
+    open: HashMap<(u32, u32, u16), Opened>,
+    cache: Cache<(u16, u64), Cached>,
+    stats: BatchStats,
+}
+
+impl Batches {
+    pub fn new(settings: BatchSettings) -> Self {
+        Self {
+            settings,
+            next_sweep: None,
+            open: HashMap::new(),
+            cache: Cache::new(settings.cache),
+            stats: BatchStats::default(),
+        }
+    }
+
+    pub fn stats(&self) -> BatchStats {
+        BatchStats {
+            cache: self.cache.stats(),
+            ..self.stats
+        }
+    }
+
+    pub fn open_count(&self) -> usize {
+        self.open.len()
+    }
+
+    pub fn seal(
+        &mut self,
+        working: &Working,
+        shard: u16,
+        first: u64,
+        now: Instant,
+    ) -> Result<Rc<Batch>, crypto::Error> {
+        if self.next_sweep.is_none_or(|next| now >= next) {
+            self.sweep(now);
+        }
+        let (tenant, generation) = (working.keys.tenant(), working.keys.generation());
+        let slot = (tenant, generation, shard);
+        if let Some(opened) = self.open.get_mut(&slot) {
+            if now.saturating_duration_since(opened.since) >= self.settings.age {
+                self.stats.aged += 1;
+            } else if opened.batch.day != working.day {
+                self.stats.day += 1;
+            } else if opened.seals >= self.settings.seals {
+                self.stats.full += 1;
+            } else {
+                opened.seals += 1;
+                return Ok(Rc::clone(&opened.batch));
+            }
+        }
+        let secret = Secret::<32>::random()?;
+        let wrap = seal_batch(
+            &working.today,
+            tenant,
+            generation,
+            working.day,
+            shard,
+            first,
+            &secret,
+        )?;
+        let batch = Rc::new(Batch {
+            id: first,
+            tenant,
+            generation,
+            day: working.day,
+            shard,
+            wrap: wrap.encode(),
+            key: Rc::new(BatchKey::new(&secret)),
+        });
+        self.remember(tenant, generation, shard, first, Rc::clone(&batch.key));
+        self.open.insert(
+            slot,
+            Opened {
+                batch: Rc::clone(&batch),
+                seals: 1,
+                since: now,
+            },
+        );
+        self.stats.created += 1;
+        Ok(batch)
+    }
+
+    pub fn cached(
+        &mut self,
+        tenant: u32,
+        generation: u32,
+        shard: u16,
+        id: u64,
+    ) -> Option<Rc<BatchKey>> {
+        if let Some(opened) = self.open.get(&(tenant, generation, shard))
+            && opened.batch.id == id
+        {
+            return Some(Rc::clone(&opened.batch.key));
+        }
+        let found = self.cache.get(&(shard, id))?;
+        if (found.tenant, found.generation) == (tenant, generation) {
+            Some(found.key)
+        } else {
+            self.stats.mismatched += 1;
+            None
+        }
+    }
+
+    pub fn remember(
+        &mut self,
+        tenant: u32,
+        generation: u32,
+        shard: u16,
+        id: u64,
+        key: Rc<BatchKey>,
+    ) {
+        if let Some(held) = self.cache.peek(&(shard, id))
+            && (held.tenant, held.generation) != (tenant, generation)
+        {
+            self.stats.mismatched += 1;
+            return;
+        }
+        self.cache.put(
+            (shard, id),
+            Cached {
+                tenant,
+                generation,
+                key,
+            },
+        );
+    }
+
+    pub fn sweep(&mut self, now: Instant) {
+        self.next_sweep = Some(now + SWEEP);
+        let age = self.settings.age;
+        let mut closed = Vec::new();
+        self.open.retain(|_, opened| {
+            let young = now.saturating_duration_since(opened.since) < age;
+            if !young {
+                closed.push(Rc::clone(&opened.batch));
+            }
+            young
+        });
+        self.stats.aged += closed.len() as u64;
+        for batch in closed {
+            self.remember(
+                batch.tenant,
+                batch.generation,
+                batch.shard,
+                batch.id,
+                Rc::clone(&batch.key),
+            );
+        }
+    }
+
+    pub fn forget(&mut self, tenant: u32, generation: u32) {
+        self.open
+            .retain(|&(t, g, _), _| (t, g) != (tenant, generation));
+        self.cache
+            .retain(|_, cached| (cached.tenant, cached.generation) != (tenant, generation));
+    }
+}

crates/keys/src/cache.rs

@@ -59,6 +59,11 @@ impl<K: Copy + Eq + Hash, V: Clone> Cache<K, V> {
         }
     }
 
+    pub fn peek(&self, key: &K) -> Option<&V> {
+        let at = *self.index.get(key)?;
+        self.slots[at].as_ref().map(|slot| &slot.value)
+    }
+
     pub fn put(&mut self, key: K, value: V) {
         if let Some(&at) = self.index.get(&key) {
             if let Some(slot) = self.slots[at].as_mut() {

crates/keys/src/lib.rs

@@ -1,4 +1,5 @@
 mod audit;
+mod batch;
 mod cache;
 mod client;
 mod clock;
@@ -12,6 +13,7 @@ mod state;
 mod vault;
 
 pub use audit::{Audit, Op};
+pub use batch::{Batch, BatchSettings, BatchStats, Batches};
 pub use cache::{Cache, CacheStats};
 pub use client::Client;
 pub use clock::{Clock, MIN_GAP, Step};

crates/keys/src/ring.rs

@@ -4,7 +4,7 @@ use std::pin::Pin;
 use std::rc::Rc;
 use std::time::Duration;
 
-use crypto::{Secret, TenantKeys};
+use crypto::{Secret, TenantKeys, Wrap};
 use rt::{Instant, Sender};
 
 use crate::cache::{Cache, CacheStats};
@@ -14,6 +14,7 @@ use crate::{Client, Error};
 pub type Connecting = Pin<Box<dyn Future<Output = Result<Client, Error>>>>;
 pub type Connect = Rc<dyn Fn() -> Connecting>;
 type Waiters = Vec<Sender<Result<Working, Fault>>>;
+type Unwrapping = (Wrap, Vec<Sender<Result<Secret<32>, Fault>>>);
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub struct RingSettings {
@@ -53,6 +54,8 @@ pub enum Fault {
     Unknown,
     Failed,
     Conflict,
+    Gone,
+    Bad,
 }
 
 #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
@@ -65,6 +68,9 @@ pub struct RingStats {
     pub unknown: u64,
     pub unavailable: u64,
     pub failed: u64,
+    pub unwraps: u64,
+    pub gone: u64,
+    pub bad: u64,
     pub tenants: CacheStats,
 }
 
@@ -88,6 +94,7 @@ struct Inner {
     failed: Cell<bool>,
     tenants: RefCell<Cache<(u32, u32), Working>>,
     opening: RefCell<HashMap<(u32, u32), Waiters>>,
+    unwrapping: RefCell<HashMap<(u32, u32, u16, u64), Unwrapping>>,
     stats: Cell<RingStats>,
 }
 
@@ -123,6 +130,7 @@ impl Ring {
                 failed: Cell::new(false),
                 tenants: RefCell::new(Cache::new(settings.tenants)),
                 opening: RefCell::new(HashMap::new()),
+                unwrapping: RefCell::new(HashMap::new()),
                 stats: Cell::new(RingStats::default()),
             }),
         }
@@ -220,6 +228,65 @@ impl Ring {
         }
     }
 
+    pub async fn unwrap(
+        &self,
+        tenant: u32,
+        generation: u32,
+        shard: u16,
+        id: u64,
+        wrap: &Wrap,
+    ) -> Result<Secret<32>, Fault> {
+        if (wrap.generation, wrap.shard) != (generation, shard) {
+            self.count(|s| s.bad += 1);
+            return Err(Fault::Bad);
+        }
+        if self.failed() {
+            self.count(|s| s.refused += 1);
+            return Err(Fault::Failed);
+        }
+        let key = (tenant, generation, shard, id);
+        let (reply, mut answer) = rt::channel();
+        let first = {
+            let mut unwrapping = self.inner.unwrapping.borrow_mut();
+            let (asked, waiting) = unwrapping
+                .entry(key)
+                .or_insert_with(|| (wrap.clone(), Vec::new()));
+            if asked != wrap {
+                drop(unwrapping);
+                self.count(|s| s.bad += 1);
+                return Err(Fault::Bad);
+            }
+            waiting.push(reply);
+            waiting.len() == 1
+        };
+        if first {
+            let ring = self.clone();
+            let request = Request::Unwrap {
+                tenant,
+                batch: id,
+                wrap: wrap.clone(),
+            };
+            rt::spawn(async move {
+                ring.count(|s| s.unwraps += 1);
+                let result = match ring.call(&request).await {
+                    Ok(Answer::Key(secret)) => Ok(secret),
+                    Ok(Answer::Refused(refused)) => Err(ring.fault(None, Some(refused))),
+                    Ok(_) => Err(ring.fault(None, None)),
+                    Err(error) => Err(ring.fault(Some(error), None)),
+                };
+                let waiting = ring.inner.unwrapping.borrow_mut().remove(&key);
+                for waiter in waiting.map(|(_, waiting)| waiting).unwrap_or_default() {
+                    let copy = match &result {
+                        Ok(secret) => Ok(secret.copy()),
+                        Err(fault) => Err(*fault),
+                    };
+                    let _ = waiter.send(copy);
+                }
+            });
+        }
+        answer.recv().await.unwrap_or(Err(Fault::Unavailable))
+    }
+
     pub async fn create(&self, tenant: u32, generation: u32, retention: u16) -> Result<(), Fault> {
         let request = Request::Create {
             tenant,
@@ -288,6 +355,8 @@ impl Ring {
         let fault = match (error, refused) {
             (_, Some(Error::Unknown)) => Fault::Unknown,
             (_, Some(Error::Generation)) => Fault::Conflict,
+            (_, Some(Error::Gone)) => Fault::Gone,
+            (_, Some(Error::Open)) => Fault::Bad,
             (_, Some(Error::Failed | Error::Unreadable))
             | (Some(Error::Failed | Error::Changed), _) => Fault::Failed,
             _ => Fault::Unavailable,
@@ -296,6 +365,8 @@ impl Ring {
             Fault::Unknown | Fault::Conflict => s.unknown += 1,
             Fault::Failed => s.failed += 1,
             Fault::Unavailable => s.unavailable += 1,
+            Fault::Gone => s.gone += 1,
+            Fault::Bad => s.bad += 1,
         });
         fault
     }