From 75affdbbcc0bf1ef694148347788235731b3d9ff Mon Sep 17 00:00:00 2001 From: hamzakargin Date: Tue, 3 Nov 2026 13:29:43 +0400 Subject: [PATCH] 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 | 3 +- crates/keys/src/batch.rs | 223 ++++++++++++++++++++ crates/keys/src/cache.rs | 5 + crates/keys/src/lib.rs | 2 + crates/keys/src/ring.rs | 73 ++++++- crates/keys/tests/batches.rs | 140 +++++++++++++ crates/keys/tests/ring.rs | 83 +++++++- crates/sched/src/keyring.rs | 83 +++++++- crates/sched/src/lib.rs | 4 + crates/sched/src/sealing.rs | 212 +++++++++++++++++++ crates/sched/src/start.rs | 20 +- crates/sched/tests/index.rs | 339 ++++++++++++++++++++++++++++++- crates/sched/tests/node.rs | 1 + crates/sched/tests/start.rs | 72 ++++++- crates/spool/src/intake.rs | 5 +- crates/spool/src/items.rs | 11 +- crates/spool/src/lib.rs | 4 +- crates/spool/tests/intake.rs | 9 +- crates/spool/tests/lifecycle.rs | 1 + crates/spool/tests/limits.rs | 7 + crates/spool/tests/recipients.rs | 1 + 21 files changed, 1262 insertions(+), 36 deletions(-) create mode 100644 crates/keys/src/batch.rs create mode 100644 crates/keys/tests/batches.rs create mode 100644 crates/sched/src/sealing.rs diff --git a/crates/crypto/src/secret.rs b/crates/crypto/src/secret.rs index 1a1bc96..feb39e7 100644 --- a/crates/crypto/src/secret.rs +++ b/crates/crypto/src/secret.rs @@ -68,7 +68,8 @@ impl Secret { 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 diff --git a/crates/keys/src/batch.rs b/crates/keys/src/batch.rs new file mode 100644 index 0000000..30e575b --- /dev/null +++ b/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, +} + +#[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, + seals: u32, + since: Instant, +} + +#[derive(Clone)] +struct Cached { + tenant: u32, + generation: u32, + key: Rc, +} + +pub struct Batches { + settings: BatchSettings, + next_sweep: Option, + 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, 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> { + 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, + ) { + 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)); + } +} diff --git a/crates/keys/src/cache.rs b/crates/keys/src/cache.rs index a6d7529..4d1e47b 100644 --- a/crates/keys/src/cache.rs +++ b/crates/keys/src/cache.rs @@ -59,6 +59,11 @@ impl Cache { } } + 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() { diff --git a/crates/keys/src/lib.rs b/crates/keys/src/lib.rs index a3df4a9..03e0e78 100644 --- a/crates/keys/src/lib.rs +++ b/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}; diff --git a/crates/keys/src/ring.rs b/crates/keys/src/ring.rs index f211985..d48ceda 100644 --- a/crates/keys/src/ring.rs +++ b/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>>>; pub type Connect = Rc Connecting>; type Waiters = Vec>>; +type Unwrapping = (Wrap, Vec, 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, tenants: RefCell>, opening: RefCell>, + unwrapping: RefCell>, stats: Cell, } @@ -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, 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 } diff --git a/crates/keys/tests/batches.rs b/crates/keys/tests/batches.rs new file mode 100644 index 0000000..917f805 --- /dev/null +++ b/crates/keys/tests/batches.rs @@ -0,0 +1,140 @@ +use std::rc::Rc; +use std::time::Duration; + +use crypto::{BatchKey, Secret, Tenant, TenantKeys, Wrap, aad, day_public, open_batch}; +use keys::{BatchSettings, Batches, Working}; +use rt::sim::{Config, SimDriver}; +use rt::{Runtime, now}; + +const DAY: u32 = 20_833; + +fn run(future: impl Future + 'static) -> T { + Runtime::new(Rc::new(SimDriver::new(1, Config::default()))).block_on(future) +} + +fn secret(tenant: u32, day: u32) -> Secret<32> { + let mut bytes = [0; 32]; + bytes[..4].copy_from_slice(&tenant.to_le_bytes()); + bytes[4..8].copy_from_slice(&day.to_le_bytes()); + Secret::new(bytes) +} + +fn working(tenant: u32, generation: u32, day: u32) -> Working { + let root = Tenant::new(&Secret::new([1; 32])); + Working { + keys: TenantKeys::derive(tenant, generation, &root, 1), + day, + oldest: DAY, + today: day_public(&secret(tenant, day)).expect("public"), + tomorrow: day_public(&secret(tenant, day + 1)).expect("public"), + } +} + +fn same_key(a: &BatchKey, b: &BatchKey) -> bool { + let message = [7; 16]; + let aad = aad(4, 1, &message, 0, 1); + let sealed = a.record(&message, 0, 1).seal_one(&aad, b"mail"); + b.record(&message, 0, 1).open_one(&aad, &sealed).is_ok() +} + +#[test] +fn a_batch_rotates_after_its_seals_its_day_and_its_age() { + run(async { + let mut batches = Batches::new(BatchSettings::default()); + let today = working(4, 1, DAY); + let start = now(); + let first = batches.seal(&today, 2, 100, start).expect("seal"); + for n in 1..1_024 { + let again = batches.seal(&today, 2, 100 + n, start).expect("seal"); + assert_eq!(again.id, first.id); + } + let full = batches.seal(&today, 2, 2_000, start).expect("seal"); + assert_eq!(full.id, 2_000); + assert_eq!(batches.stats().full, 1); + let replaced = batches.cached(4, 1, 2, 100).expect("replaced batch cached"); + assert!(same_key(&replaced, &first.key)); + let next_day = batches + .seal(&working(4, 1, DAY + 1), 2, 3_000, start) + .expect("seal"); + assert_eq!((next_day.id, next_day.day), (3_000, DAY + 1)); + assert_eq!(batches.stats().day, 1); + let other_shard = batches.seal(&today, 3, 4_000, start).expect("seal"); + assert_eq!(other_shard.id, 4_000); + assert_eq!(batches.open_count(), 2); + let later = start + Duration::from_secs(60); + batches.sweep(later); + assert_eq!(batches.open_count(), 0); + assert_eq!(batches.stats().aged, 2); + let found = batches.cached(4, 1, 2, 3_000).expect("closed batch cached"); + assert!(same_key(&found, &next_day.key)); + let fresh = batches + .seal(&working(4, 1, DAY + 1), 2, 5_000, later) + .expect("seal"); + assert_eq!(fresh.id, 5_000); + }); +} + +#[test] +fn the_wrap_opens_to_the_batch_key_with_the_day_secret() { + run(async { + let mut batches = Batches::new(BatchSettings::default()); + let batch = batches + .seal(&working(4, 1, DAY), 2, 77, now()) + .expect("seal"); + let wrap = Wrap::decode(&batch.wrap).expect("wrap"); + assert_eq!((wrap.day, wrap.shard, wrap.generation), (DAY, 2, 1)); + let opened = open_batch(&secret(4, DAY), 4, 77, &wrap).expect("open"); + assert!(same_key(&BatchKey::new(&opened), &batch.key)); + assert!(open_batch(&secret(4, DAY), 4, 78, &wrap).is_err()); + assert!(open_batch(&secret(5, DAY), 5, 77, &wrap).is_err()); + }); +} + +#[test] +fn cached_keys_answer_only_their_own_tenant_and_generation() { + run(async { + let mut batches = Batches::new(BatchSettings { + cache: 2, + ..BatchSettings::default() + }); + let batch = batches + .seal(&working(4, 1, DAY), 2, 77, now()) + .expect("seal"); + assert!(batches.cached(4, 1, 2, 77).is_some()); + assert!(batches.cached(5, 1, 2, 77).is_none()); + assert!(batches.cached(4, 2, 2, 77).is_none()); + assert!(batches.cached(4, 1, 3, 77).is_none()); + for id in 1..10 { + batches.remember(9, 1, 2, id, Rc::clone(&batch.key)); + } + assert!(batches.cached(4, 1, 2, 77).is_some(), "open batch evicted"); + batches.remember(5, 1, 6, 900, Rc::clone(&batch.key)); + assert!(batches.cached(4, 1, 6, 900).is_none()); + assert!(batches.stats().mismatched >= 3); + batches.forget(4, 1); + assert!(batches.cached(4, 1, 2, 77).is_none()); + assert_eq!(batches.open_count(), 0); + }); +} + +#[test] +fn a_batch_rotates_on_its_own_age_between_sweeps_and_keeps_its_cache_owner() { + run(async { + let mut batches = Batches::new(BatchSettings::default()); + let today = working(4, 1, DAY); + let start = now(); + let first = batches.seal(&today, 2, 100, start).expect("seal"); + let almost = start + Duration::from_millis(59_900); + assert_eq!(batches.seal(&today, 2, 101, almost).expect("seal").id, 100); + let aged = batches + .seal(&today, 2, 102, start + Duration::from_secs(60)) + .expect("seal"); + assert_eq!(aged.id, 102); + assert_eq!(batches.stats().aged, 1); + batches.remember(5, 1, 2, 100, Rc::clone(&aged.key)); + let kept = batches.cached(4, 1, 2, 100).expect("owner kept"); + assert!(same_key(&kept, &first.key)); + assert!(batches.cached(5, 1, 2, 100).is_none()); + assert!(batches.stats().mismatched >= 2); + }); +} diff --git a/crates/keys/tests/ring.rs b/crates/keys/tests/ring.rs index e164045..a5987e2 100644 --- a/crates/keys/tests/ring.rs +++ b/crates/keys/tests/ring.rs @@ -4,9 +4,12 @@ use std::path::Path; use std::rc::Rc; use std::time::Duration; -use crypto::{Plain, Secret}; +use crypto::{BatchKey, Plain, Secret, Wrap, seal_batch}; use keys::proto::{Answer, Frames, Request}; -use keys::{Audit, Client, Connect, Error, Fault, Ring, RingSettings, Server, Settings, Vault}; +use keys::{ + Audit, BatchSettings, Batches, Client, Connect, Error, Fault, Ring, RingSettings, Server, + Settings, Vault, +}; use rt::sim::{Config, SimDriver}; use rt::{Buf, Runtime, TcpListener, TcpStream, sleep, spawn}; @@ -212,6 +215,82 @@ fn create_sets_retention_adopts_a_retry_and_refuses_another_generation() { }); } +#[test] +fn unwrap_answers_once_per_batch_and_maps_refusals() { + let driver = Rc::new(SimDriver::new(9, quiet())); + run(&driver, async { + let _server = serve("/a", 1).await; + let ring = Ring::new(hosts(&[1]).0, settings()); + ring.create(4, 1, 30).await.expect("create"); + let working = ring.open(4, 1).await.expect("open"); + let mut batches = Batches::new(BatchSettings::default()); + let batch = batches.seal(&working, 2, 77, rt::now()).expect("seal"); + let wrap = Wrap::decode(&batch.wrap).expect("wrap"); + let mut tasks = Vec::new(); + for _ in 0..10 { + let ring = ring.clone(); + let wrap = wrap.clone(); + tasks.push(spawn(async move { ring.unwrap(4, 1, 2, 77, &wrap).await })); + } + let message = [7; 16]; + let aad = crypto::aad(4, 1, &message, 0, 1); + let sealed = batch.key.record(&message, 0, 1).seal_one(&aad, b"mail"); + for task in tasks { + let key = BatchKey::new(&task.await.expect("unwrap")); + let opened = key.record(&message, 0, 1).open_one(&aad, &sealed); + assert_eq!(opened.expect("open"), b"mail"); + } + assert_eq!(ring.stats().unwraps, 1); + let ninety = batches.seal(&working, 3, 90, rt::now()).expect("seal"); + let ninety = Wrap::decode(&ninety.wrap).expect("wrap"); + let mut tampered = ninety.clone(); + tampered.enc[0] ^= 1; + let joined = { + let ring = ring.clone(); + let wrap = ninety.clone(); + spawn(async move { ring.unwrap(4, 1, 3, 90, &wrap).await }) + }; + while ring.stats().unwraps < 2 { + rt::yield_now().await; + } + let other = { + let ring = ring.clone(); + let wrap = ninety.clone(); + spawn(async move { ring.unwrap(5, 1, 3, 90, &wrap).await }) + }; + assert_eq!( + ring.unwrap(4, 1, 3, 90, &tampered).await.err(), + Some(Fault::Bad) + ); + assert!(joined.await.is_ok()); + assert_eq!(other.await.err(), Some(Fault::Unknown)); + assert_eq!( + ring.unwrap(4, 2, 2, 77, &wrap).await.err(), + Some(Fault::Bad) + ); + assert_eq!( + ring.unwrap(4, 1, 3, 77, &wrap).await.err(), + Some(Fault::Bad) + ); + assert_eq!(ring.stats().unwraps, 3); + assert_eq!( + ring.unwrap(4, 1, 2, 78, &wrap).await.err(), + Some(Fault::Bad) + ); + assert_eq!( + ring.unwrap(9, 1, 2, 77, &wrap).await.err(), + Some(Fault::Unknown) + ); + let old = + seal_batch(&working.today, 4, 1, DAY - 5, 2, 80, &Secret::new([3; 32])).expect("wrap"); + assert_eq!( + ring.unwrap(4, 1, 2, 80, &old).await.err(), + Some(Fault::Gone) + ); + assert_eq!(ring.stats().gone, 1); + }); +} + #[test] fn keys_down_is_unavailable_and_retries_back_off() { let driver = Rc::new(SimDriver::new(2, quiet())); diff --git a/crates/sched/src/keyring.rs b/crates/sched/src/keyring.rs index 2450fec..3bce629 100644 --- a/crates/sched/src/keyring.rs +++ b/crates/sched/src/keyring.rs @@ -1,15 +1,29 @@ +use std::cell::RefCell; +use std::rc::Rc; + use control::{Created, KeyError, KeySource, Ready}; -use crypto::TenantKeys; -use keys::{Fault, Ring}; +use crypto::{BatchKey, TenantKeys, Wrap}; +use keys::{BatchSettings, BatchStats, Batches, Fault, Ring}; + +use crate::sealing::{Batched, Keys, Unsealable, Unwrapped}; pub struct Keyring { ring: Ring, retention: u16, + batches: RefCell, } impl Keyring { - pub fn new(ring: Ring, retention: u16) -> Self { - Self { ring, retention } + pub fn new(ring: Ring, retention: u16, batches: BatchSettings) -> Self { + Self { + ring, + retention, + batches: RefCell::new(Batches::new(batches)), + } + } + + pub fn batch_stats(&self) -> BatchStats { + self.batches.borrow().stats() } pub fn ring(&self) -> &Ring { @@ -19,13 +33,23 @@ impl Keyring { fn key_error(fault: Fault) -> KeyError { match fault { - Fault::Unknown => KeyError::Unknown, + Fault::Unknown | Fault::Gone => KeyError::Unknown, Fault::Unavailable => KeyError::Unavailable, - Fault::Failed => KeyError::Failed, + Fault::Failed | Fault::Bad => KeyError::Failed, Fault::Conflict => KeyError::Conflict, } } +fn unsealable(fault: Fault) -> Unsealable { + match fault { + Fault::Unavailable => Unsealable::Unavailable, + Fault::Gone => Unsealable::Gone, + Fault::Unknown | Fault::Conflict => Unsealable::Unknown, + Fault::Bad => Unsealable::Bad, + Fault::Failed => Unsealable::Failed, + } +} + impl KeySource for Keyring { fn get(&self, tenant: u32, generation: u32) -> Option { self.ring @@ -52,3 +76,50 @@ impl KeySource for Keyring { }) } } + +impl Keys for Keyring { + fn batch(&self, tenant: u32, generation: u32, shard: u16, first: u64) -> Batched<'_> { + Box::pin(async move { + let working = self + .ring + .open(tenant, generation) + .await + .map_err(key_error)?; + self.batches + .borrow_mut() + .seal(&working, shard, first, rt::now()) + .map_err(|_| KeyError::Unavailable) + }) + } + + fn batch_key( + &self, + tenant: u32, + generation: u32, + shard: u16, + id: u64, + wrap: Wrap, + ) -> Unwrapped<'_> { + Box::pin(async move { + if let Some(key) = self + .batches + .borrow_mut() + .cached(tenant, generation, shard, id) + { + return Ok(key); + } + let secret = self + .ring + .unwrap(tenant, generation, shard, id, &wrap) + .await + .map_err(unsealable)?; + let mut batches = self.batches.borrow_mut(); + if let Some(key) = batches.cached(tenant, generation, shard, id) { + return Ok(key); + } + let key = Rc::new(BatchKey::new(&secret)); + batches.remember(tenant, generation, shard, id, Rc::clone(&key)); + Ok(key) + }) + } +} diff --git a/crates/sched/src/lib.rs b/crates/sched/src/lib.rs index e81d58e..425e935 100644 --- a/crates/sched/src/lib.rs +++ b/crates/sched/src/lib.rs @@ -7,6 +7,7 @@ mod node; mod normal; mod policy; mod requests; +mod sealing; mod start; mod suppress; mod transport; @@ -26,6 +27,9 @@ pub use requests::{ Begun, DAY, INLINE_IDS, KIND_REQUEST, REQUEST_SHAPE, RequestInputs, Requests, Ticket, retain_request, }; +pub use sealing::{ + Batched, Keys, LABEL_BODY, LABEL_HEAD, LABEL_SLICE, Opener, Sealing, Unsealable, Unwrapped, +}; pub use start::{CONTROL_SHARD, Links, Meshes, Node, Parts, Running, Settings}; pub use suppress::{ Ask, Hash, Inputs, KIND_SUPPRESS, Listed, Page, Reason, SEAL_SUPPRESS, SHAPE, Suppression, diff --git a/crates/sched/src/sealing.rs b/crates/sched/src/sealing.rs new file mode 100644 index 0000000..178d46a --- /dev/null +++ b/crates/sched/src/sealing.rs @@ -0,0 +1,212 @@ +use std::collections::BTreeSet; +use std::pin::Pin; +use std::rc::Rc; + +use control::{KeyError, KeySource}; +use crypto::{ + AAD, BatchKey, ChunkKey, Opener as ChunkOpener, Sealer as ChunkSealer, TAG, Wrap, aad, +}; +use keys::Batch; +use spool::{BODY_ITEM, Body, HEAD_CLEAR, Head, HeadClear, MessageId, Slice, WRAP, slice_item}; +use store::{HEAD_ITEM, INDEX_SIZE, MAX_PAYLOAD}; + +use crate::error::Refusal; +use crate::intake::Draft; +use crate::start::Node; +use crate::transport::Transport; + +pub const LABEL_HEAD: u8 = 1; +pub const LABEL_BODY: u8 = 2; +pub const LABEL_SLICE: u8 = 3; + +const HEAD_ROOM: usize = MAX_PAYLOAD - INDEX_SIZE - HEAD_CLEAR - WRAP - TAG - 4; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Unsealable { + Unavailable, + Gone, + Unknown, + Bad, + Failed, +} + +pub type Batched<'a> = Pin, KeyError>> + 'a>>; +pub type Unwrapped<'a> = Pin, Unsealable>> + 'a>>; + +pub trait Keys: KeySource { + fn batch(&self, tenant: u32, generation: u32, shard: u16, first: u64) -> Batched<'_>; + fn batch_key( + &self, + tenant: u32, + generation: u32, + shard: u16, + id: u64, + wrap: Wrap, + ) -> Unwrapped<'_>; +} + +pub struct Sealing { + tenant: u32, + generation: u32, + id: MessageId, + batch: Rc, + sealed: BTreeSet, +} + +impl Sealing { + pub fn id(&self) -> MessageId { + self.id + } + + pub fn tenant(&self) -> u32 { + self.tenant + } + + pub fn generation(&self) -> u32 { + self.generation + } + + pub fn head( + &mut self, + clear: HeadClear, + expiry: u64, + envelope: &[u8], + headers: &[u8], + ) -> Result { + if envelope.len() + headers.len() > HEAD_ROOM { + return Err(Refusal::Message); + } + let length = u32::try_from(envelope.len()).map_err(|_| Refusal::Message)?; + let (key, aad) = self.key(HEAD_ITEM, LABEL_HEAD)?; + let mut clear_bytes = Vec::with_capacity(4 + envelope.len() + headers.len()); + clear_bytes.extend_from_slice(&length.to_le_bytes()); + clear_bytes.extend_from_slice(envelope); + clear_bytes.extend_from_slice(headers); + Ok(Head { + clear: HeadClear { + batch_key: self.batch.id, + day: self.batch.day, + ..clear + }, + expiry, + wrap: self.batch.wrap, + sealed: key.seal_one(&aad, &clear_bytes), + }) + } + + pub fn inline_body(&mut self, body: &[u8]) -> Result { + let (key, aad) = self.key(BODY_ITEM, LABEL_BODY)?; + Ok(Body::Inline(key.seal_one(&aad, body))) + } + + pub fn body_chunks(&mut self) -> Result { + let (key, aad) = self.key(BODY_ITEM, LABEL_BODY)?; + Ok(ChunkSealer::new(key, &aad)) + } + + pub fn slice(&mut self, n: u32, queues: Vec, recipients: &[u8]) -> Result { + let (key, aad) = self.key(slice_item(n), LABEL_SLICE)?; + Ok(Slice { + queues, + sealed: key.seal_one(&aad, recipients), + }) + } + + fn key(&mut self, item: u32, label: u8) -> Result<(ChunkKey, [u8; AAD]), Refusal> { + if !self.sealed.insert(item) { + return Err(Refusal::Message); + } + let message = self.id.bytes(); + Ok(( + self.batch.key.record(&message, item, label), + aad(self.tenant, self.generation, &message, item, label), + )) + } +} + +pub struct Opener { + tenant: u32, + generation: u32, + message: [u8; 16], + key: Rc, +} + +impl Opener { + pub fn head(&self, sealed: &[u8]) -> Result<(Vec, Vec), Unsealable> { + let (key, aad) = self.key(HEAD_ITEM, LABEL_HEAD); + let mut clear = key.open_one(&aad, sealed).map_err(|_| Unsealable::Bad)?; + let (length, rest) = clear.split_first_chunk::<4>().ok_or(Unsealable::Bad)?; + let length = usize::try_from(u32::from_le_bytes(*length)).map_err(|_| Unsealable::Bad)?; + if length > rest.len() { + return Err(Unsealable::Bad); + } + let headers = clear.split_off(4 + length); + clear.drain(..4); + Ok((clear, headers)) + } + + pub fn inline_body(&self, sealed: &[u8]) -> Result, Unsealable> { + let (key, aad) = self.key(BODY_ITEM, LABEL_BODY); + key.open_one(&aad, sealed).map_err(|_| Unsealable::Bad) + } + + pub fn body_chunks(&self) -> ChunkOpener { + let (key, aad) = self.key(BODY_ITEM, LABEL_BODY); + ChunkOpener::new(key, &aad) + } + + pub fn slice(&self, n: u32, sealed: &[u8]) -> Result, Unsealable> { + let (key, aad) = self.key(slice_item(n), LABEL_SLICE); + key.open_one(&aad, sealed).map_err(|_| Unsealable::Bad) + } + + fn key(&self, item: u32, label: u8) -> (ChunkKey, [u8; AAD]) { + ( + self.key.record(&self.message, item, label), + aad(self.tenant, self.generation, &self.message, item, label), + ) + } +} + +impl Node { + pub async fn sealing(&self, draft: &Draft, shard: u16) -> Result { + let (tenant, generation) = (draft.tenant(), draft.generation()); + let id = self.spool().id(shard).await.map_err(|_| Refusal::Failed)?; + let batch = self + .keys() + .batch(tenant, generation, shard, id.key) + .await + .map_err(|_| Refusal::Keys)?; + Ok(Sealing { + tenant, + generation, + id, + batch, + sealed: BTreeSet::new(), + }) + } + + pub async fn opener( + &self, + tenant: u32, + generation: u32, + id: MessageId, + clear: &HeadClear, + wrap: &[u8; WRAP], + ) -> Result { + let decoded = Wrap::decode(wrap).map_err(|_| Unsealable::Bad)?; + if (decoded.generation, decoded.shard, decoded.day) != (generation, id.shard, clear.day) { + return Err(Unsealable::Bad); + } + let key = self + .keys() + .batch_key(tenant, generation, id.shard, clear.batch_key, decoded) + .await?; + Ok(Opener { + tenant, + generation, + message: id.bytes(), + key, + }) + } +} diff --git a/crates/sched/src/start.rs b/crates/sched/src/start.rs index 1d6ccff..9d7c131 100644 --- a/crates/sched/src/start.rs +++ b/crates/sched/src/start.rs @@ -13,6 +13,7 @@ use crate::clock::Clock; use crate::error::Error; use crate::node::{Config, Note, Sched}; use crate::requests::{KIND_REQUEST, RequestInputs, Requests, retain_request}; +use crate::sealing::Keys; use crate::suppress::{self, Inputs, KIND_SUPPRESS, Suppressions}; use crate::transport::Transport; @@ -70,14 +71,14 @@ pub struct Settings { } pub struct Parts { - pub keys: Rc, + pub keys: Rc, pub transport: T, pub clock: Clock, } pub(crate) struct Inner { pub(crate) settings: Settings, - keys: Rc, + keys: Rc, cores: u16, pub(crate) spool: spool::Spool, pub(crate) sched: Sched, @@ -134,7 +135,7 @@ impl Node { None }; let check = naming(control.as_ref().map(|(loaded, _)| loaded), &settings)?; - let loaded = Queues::load(&core, settings.naming, parts.keys.as_ref()).await?; + let loaded = Queues::load(&core, settings.naming, &*parts.keys).await?; let table = Rc::new(RefCell::new(Table::default())); let clock = parts.clock; retain_current(&mut core, &table, clock); @@ -162,7 +163,7 @@ impl Node { }; let inputs = Inputs { - keys: Rc::clone(&parts.keys), + keys: Rc::clone(&parts.keys) as Rc, changes: changes.clone(), table: Rc::clone(&table), clock, @@ -171,8 +172,11 @@ impl Node { }; let (suppressions, requests, mut index_done) = index_services(links.suppress, inputs).await?; - let (queues, names_done) = - loaded.start(links.names, changes.clone(), Rc::clone(&parts.keys)); + let (queues, names_done) = loaded.start( + links.names, + changes.clone(), + Rc::clone(&parts.keys) as Rc, + ); let (leases, leases_done) = start_leases(links.leases, &table, &spool, settings.naming.shards); @@ -235,6 +239,10 @@ impl Node { &self.inner.spool } + pub(crate) fn keys(&self) -> &dyn Keys { + &*self.inner.keys + } + pub fn sched(&self) -> &Sched { &self.inner.sched } diff --git a/crates/sched/tests/index.rs b/crates/sched/tests/index.rs index b838cd1..1ec3c18 100644 --- a/crates/sched/tests/index.rs +++ b/crates/sched/tests/index.rs @@ -9,17 +9,18 @@ use std::task::Poll; use std::time::Duration; use control::{Created, KeyError, KeySource, Naming, Object, Pool, Ready, State, Table, Tenant}; -use crypto::{Secret, Tenant as Root, TenantKeys}; +use crypto::{BatchKey, Secret, Tenant as Root, TenantKeys, Wrap, day_public, open_batch}; use index::{Entry, Op}; use keys::{ - Audit, Client as KeyClient, Connect, Ring, RingSettings, Server, Settings as KeySettings, Vault, + Audit, BatchSettings, Batches, Client as KeyClient, Connect, Ring, RingSettings, Server, + Settings as KeySettings, Vault, Working, }; use rt::sim::{Config, Disk, Faults, Machine, SimDriver, Stalls, check}; use rt::{JoinHandle, Runtime, TcpListener, TcpStream, sleep, spawn, timeout}; use sched::{ - Answered, Attempt, Begun, Clock, DAY, Draft, Hash, KIND_REQUEST, KIND_SUPPRESS, Keyring, - Meshes, Node, Parts, Reason, Refusal, Running, SHAPE, Settings, Suppression, Ticket, Transport, - normalize, punycode, retain, retain_request, + Answered, Attempt, Batched, Begun, Clock, DAY, Draft, Hash, KIND_REQUEST, KIND_SUPPRESS, + Keyring, Keys, Meshes, Node, Parts, Reason, Refusal, Running, SHAPE, Settings, Suppression, + Ticket, Transport, Unsealable, Unwrapped, normalize, punycode, retain, retain_request, }; use spool::{Body, Head, HeadClear, Message, MessageId, SLICE, Slice}; use store::{Client, Committer, DataFile, Options, SLOT_SIZE, Setup, Tiering, boot}; @@ -34,7 +35,17 @@ const OPTIONS: Options = Options { shards: SHARDS as u32, }; -struct Toy; +struct Toy { + batches: RefCell, +} + +impl Default for Toy { + fn default() -> Self { + Self { + batches: RefCell::new(Batches::new(BatchSettings::default())), + } + } +} const KEYLESS: u32 = 199; const NO_CREATE: u32 = 198; @@ -71,6 +82,56 @@ impl KeySource for Toy { } } +const TOY_DAY: u32 = 20_000; + +fn toy_day_secret(tenant: u32, day: u32) -> Secret<32> { + let mut bytes = [9; 32]; + bytes[..4].copy_from_slice(&tenant.to_le_bytes()); + bytes[4..8].copy_from_slice(&day.to_le_bytes()); + Secret::new(bytes) +} + +impl Keys for Toy { + fn batch(&self, tenant: u32, generation: u32, shard: u16, first: u64) -> Batched<'_> { + let batch = match self.get(tenant, generation) { + Some(keys) => { + let public = |day| day_public(&toy_day_secret(tenant, day)).expect("public"); + let working = Working { + keys, + day: TOY_DAY, + oldest: TOY_DAY, + today: public(TOY_DAY), + tomorrow: public(TOY_DAY + 1), + }; + self.batches + .borrow_mut() + .seal(&working, shard, first, rt::now()) + .map_err(|_| KeyError::Unavailable) + } + None => Err(KeyError::Unavailable), + }; + Box::pin(std::future::ready(batch)) + } + + fn batch_key( + &self, + tenant: u32, + generation: u32, + shard: u16, + id: u64, + wrap: Wrap, + ) -> Unwrapped<'_> { + let key = if (wrap.generation, wrap.shard) == (generation, shard) { + open_batch(&toy_day_secret(tenant, wrap.day), tenant, id, &wrap) + .map(|secret| Rc::new(BatchKey::new(&secret))) + .map_err(|_| Unsealable::Bad) + } else { + Err(Unsealable::Bad) + }; + Box::pin(std::future::ready(key)) + } +} + struct Idle; impl Transport for Idle { @@ -204,10 +265,14 @@ fn keys_up(machine: &Machine) { })); } -fn keyring() -> Rc { +fn keyring() -> Rc { let connect: Connect = Rc::new(|| Box::pin(async { Ok(KeyClient::new(TcpStream::connect(KEYS_AT).await?)) })); - Rc::new(Keyring::new(Ring::new(connect, RingSettings::new(1)), 30)) + Rc::new(Keyring::new( + Ring::new(connect, RingSettings::new(1)), + 30, + BatchSettings::default(), + )) } async fn keys_listening() { @@ -264,7 +329,11 @@ fn try_started( ..Tiering::default() }); } - let keys: Rc = if real_keys { keyring() } else { Rc::new(Toy) }; + let keys: Rc = if real_keys { + keyring() + } else { + Rc::new(Toy::default()) + }; let parts = Parts { keys, transport: Idle, @@ -704,6 +773,257 @@ fn a_different_naming_key_refuses_to_start() { restarted.stop(); } +struct Sealed { + id: MessageId, + clear: HeadClear, + wrap: [u8; spool::WRAP], + head: Vec, + body: Vec, + slice: Vec, +} + +async fn owned_shard(node: &Node) -> u16 { + for shard in 0..SHARDS { + if node.spool().id(shard).await.is_ok() { + return shard; + } + } + panic!("no shard on this core"); +} + +async fn seal_one(node: &Node, shard: u16, n: u8) -> Sealed { + let mut draft = node.draft(1, None).expect("draft"); + let address = format!("s{n}@example.org"); + node.recipient(&mut draft, &address) + .await + .expect("recipient"); + let mut sealing = node.sealing(&draft, shard).await.expect("sealing"); + let clear = HeadClear { + recipients: 1, + accepted: 5, + ..HeadClear::default() + }; + let head = sealing + .head(clear, u64::MAX / 2, &[n; 10], b"Subject: hi\r\n") + .expect("head"); + let body = sealing.inline_body(&[n; 300]).expect("body"); + let slice = sealing + .slice(0, draft.queues().to_vec(), address.as_bytes()) + .expect("slice"); + let sealed = Sealed { + id: sealing.id(), + clear: head.clear, + wrap: head.wrap, + head: head.sealed.clone(), + body: match &body { + Body::Inline(bytes) => bytes.clone(), + Body::Chain(_) => panic!("inline"), + }, + slice: slice.sealed.clone(), + }; + let message = Message { + id: sealing.id(), + tenant: 1, + head, + body, + slices: vec![slice], + }; + node.accept(draft, message).await.expect("accept"); + sealed +} + +async fn opens(node: &Node, sealed: &Sealed, n: u8) { + let opener = node + .opener(1, 1, sealed.id, &sealed.clear, &sealed.wrap) + .await + .expect("opener"); + let (envelope, headers) = opener.head(&sealed.head).expect("head"); + assert_eq!( + (envelope, headers), + (vec![n; 10], b"Subject: hi\r\n".to_vec()) + ); + assert_eq!( + opener.inline_body(&sealed.body).expect("body"), + vec![n; 300] + ); + let address = format!("s{n}@example.org"); + assert_eq!( + opener.slice(0, &sealed.slice).expect("slice"), + address.as_bytes() + ); +} + +#[test] +fn sealed_items_open_only_under_their_own_message_and_item() { + let driver = Rc::new(SimDriver::with_disk(36, Config::default(), created(36))); + let cluster = started(&driver, 2); + let nodes = cluster.nodes.clone(); + cluster.run(async move { + ready(&nodes).await; + let node = &nodes[0]; + let shard = owned_shard(node).await; + let draft = node.draft(1, None).expect("draft"); + let first = node.sealing(&draft, shard).await.expect("sealing"); + let second = node.sealing(&draft, shard).await.expect("sealing"); + assert_ne!(first.id(), second.id()); + let mut once = second; + let clear = HeadClear::default(); + assert!(once.head(clear, 1, b"e", b"h").is_ok()); + assert_eq!( + once.head(clear, 1, b"e", b"h").err(), + Some(Refusal::Message) + ); + assert!(once.inline_body(b"b").is_ok()); + assert!(once.body_chunks().is_err()); + assert!(once.slice(0, vec![1], b"r").is_ok()); + assert_eq!(once.slice(0, vec![1], b"r").err(), Some(Refusal::Message)); + assert!(once.slice(1, vec![1], b"r").is_ok()); + let mut large = node.sealing(&draft, shard).await.expect("sealing"); + let big = vec![0; store::MAX_PAYLOAD]; + assert_eq!( + large.head(clear, 1, b"e", &big).err(), + Some(Refusal::Message) + ); + drop(first); + + let sealed = seal_one(node, shard, 7).await; + opens(node, &sealed, 7).await; + opens(&nodes[1], &sealed, 7).await; + let opener = |tenant, generation, clear: HeadClear| { + let (node, wrap, id) = (node.clone(), sealed.wrap, sealed.id); + async move { + node.opener(tenant, generation, id, &clear, &wrap) + .await + .err() + } + }; + let bad = Some(Unsealable::Bad); + assert_eq!(opener(1, 2, sealed.clear).await, bad); + assert_eq!(opener(2, 1, sealed.clear).await, bad); + let mut day = sealed.clear; + day.day += 1; + assert_eq!(opener(1, 1, day).await, bad); + let mut batch = sealed.clear; + batch.batch_key += 1; + assert_eq!(opener(1, 1, batch).await, bad); + let mut wrap = sealed.wrap; + wrap[50] ^= 1; + assert_eq!( + node.opener(1, 1, sealed.id, &sealed.clear, &wrap) + .await + .err(), + bad + ); + let other = MessageId { + key: sealed.id.key + 1, + ..sealed.id + }; + let moved = node + .opener(1, 1, other, &sealed.clear, &sealed.wrap) + .await + .expect("opener"); + assert_eq!(moved.head(&sealed.head).err(), bad); + let right = node + .opener(1, 1, sealed.id, &sealed.clear, &sealed.wrap) + .await + .expect("opener"); + assert_eq!(right.inline_body(&sealed.head).err(), bad); + assert_eq!(right.slice(1, &sealed.slice).err(), bad); + assert_eq!(right.head(&sealed.body).err(), bad); + }); + cluster.stop(); +} + +#[test] +fn body_chunks_open_in_order_and_a_cut_stream_is_refused() { + let driver = Rc::new(SimDriver::with_disk(39, Config::default(), created(39))); + let cluster = started(&driver, 2); + let nodes = cluster.nodes.clone(); + cluster.run(async move { + ready(&nodes).await; + let node = &nodes[0]; + let shard = owned_shard(node).await; + let draft = node.draft(1, None).expect("draft"); + let clear = HeadClear::default(); + let mut chunked = node.sealing(&draft, shard).await.expect("sealing"); + let head = chunked.head(clear, 1, b"e", b"h").expect("head"); + let mut chunks = chunked.body_chunks().expect("chunks"); + let mut parts = [vec![1; 64], vec![2; 64], vec![3; 10]]; + for part in &mut parts[..2] { + chunks.seal(part).expect("seal"); + } + chunks.seal_last(&mut parts[2]); + let opener = node + .opener(1, 1, chunked.id(), &head.clear, &head.wrap) + .await + .expect("opener"); + let mut cut = parts[..2].to_vec(); + let mut short = opener.body_chunks(); + for part in &mut cut { + short.open(false, part).expect("chunk"); + } + assert!(short.finish().is_err()); + let mut stream = opener.body_chunks(); + let mut whole = Vec::new(); + for (n, part) in parts.iter_mut().enumerate() { + whole.extend_from_slice(stream.open(n == 2, part).expect("chunk")); + } + stream.finish().expect("finish"); + assert_eq!(whole, [vec![1; 64], vec![2; 64], vec![3; 10]].concat()); + }); + cluster.stop(); +} + +fn unwraps(driver: &SimDriver) -> u64 { + let disk = driver.disk(); + let Some(audit) = disk.contents(&Path::new(KEYS_DIR).join("audit")) else { + return 0; + }; + String::from_utf8(audit.to_vec()) + .expect("text") + .lines() + .filter(|line| line.contains(" unwrap 1 1 ")) + .filter_map(|line| line.rsplit(' ').next()?.parse::().ok()) + .sum() +} + +#[test] +fn sealed_messages_open_after_a_restart_with_one_unwrap_per_batch() { + let driver = Rc::new(SimDriver::with_disk(37, Config::default(), created(37))); + let cluster = started_with_keys(&driver, 2); + let nodes = cluster.nodes.clone(); + let sealed = cluster.run(async move { + keys_listening().await; + ready(&nodes).await; + let node = &nodes[0]; + let shard = owned_shard(node).await; + let mut sealed = Vec::new(); + for n in 0..5 { + sealed.push(seal_one(node, shard, n).await); + } + assert!(sealed.iter().all(|s| s.clear.batch_key == sealed[0].id.key)); + for (n, one) in (0..).zip(&sealed) { + opens(node, one, n).await; + } + sleep(Duration::from_millis(300)).await; + sealed + }); + assert_eq!(unwraps(&driver), 0); + drop(cluster); + let driver = Rc::new(SimDriver::with_disk(38, Config::default(), driver.crash())); + let restarted = started_with_keys(&driver, 2); + let nodes = restarted.nodes.clone(); + restarted.run(async move { + wait_for(&nodes, 1, 1, State::Active).await; + for (n, one) in (0..).zip(&sealed) { + opens(&nodes[0], one, n).await; + } + sleep(Duration::from_millis(300)).await; + }); + assert_eq!(unwraps(&driver), 1); + restarted.stop(); +} + #[test] fn listing_pages_through_every_core_in_hash_order() { listing(3, 4); @@ -839,6 +1159,7 @@ fn message(id: MessageId, draft: &Draft) -> Message { ..HeadClear::default() }, expiry: u64::MAX / 2, + wrap: [3; spool::WRAP], sealed: vec![1; 64], }, body: Body::Inline(vec![2; 300]), diff --git a/crates/sched/tests/node.rs b/crates/sched/tests/node.rs index d0bce0f..8280048 100644 --- a/crates/sched/tests/node.rs +++ b/crates/sched/tests/node.rs @@ -196,6 +196,7 @@ fn message(id: MessageId, now: u64, driver: &SimDriver, expiring: bool) -> Messa ..HeadClear::default() }, expiry, + wrap: [3; spool::WRAP], sealed: vec![1; 64], }, body: Body::Inline(vec![2; 300]), diff --git a/crates/sched/tests/start.rs b/crates/sched/tests/start.rs index 179457f..d607176 100644 --- a/crates/sched/tests/start.rs +++ b/crates/sched/tests/start.rs @@ -11,12 +11,13 @@ use control::{ Created, KeyError, KeySource, Limits as TenantLimits, Name, Naming, Object, Pool, Ready, State, Tenant, }; -use crypto::{Secret, Tenant as Root, TenantKeys}; +use crypto::{BatchKey, Secret, Tenant as Root, TenantKeys, Wrap, day_public, open_batch}; +use keys::{BatchSettings, Batches, Working}; use rt::sim::{Config, Disk, Faults, Machine, SimDriver, Stalls, check}; use rt::{JoinHandle, Runtime, sleep, spawn, timeout}; use sched::{ - Answer, Answered, Attempt, Clock, Draft, Kind, Limits, Meshes, Node, Parts, Reason, Refusal, - Running, Settings, Transport, + Answer, Answered, Attempt, Batched, Clock, Draft, Keys, Kind, Limits, Meshes, Node, Parts, + Reason, Refusal, Running, Settings, Transport, Unsealable, Unwrapped, }; use spool::{Body, Head, HeadClear, Message, MessageId, SLICE, Slice}; use store::{Client, DataFile, Options, SLOT_SIZE, Setup, boot}; @@ -124,7 +125,17 @@ fn pool(id: u32) -> Object { }) } -struct Toy; +struct Toy { + batches: RefCell, +} + +impl Default for Toy { + fn default() -> Self { + Self { + batches: RefCell::new(Batches::new(BatchSettings::default())), + } + } +} impl KeySource for Toy { fn get(&self, tenant: u32, generation: u32) -> Option { @@ -145,6 +156,56 @@ impl KeySource for Toy { } } +const TOY_DAY: u32 = 20_000; + +fn toy_day_secret(tenant: u32, day: u32) -> Secret<32> { + let mut bytes = [9; 32]; + bytes[..4].copy_from_slice(&tenant.to_le_bytes()); + bytes[4..8].copy_from_slice(&day.to_le_bytes()); + Secret::new(bytes) +} + +impl Keys for Toy { + fn batch(&self, tenant: u32, generation: u32, shard: u16, first: u64) -> Batched<'_> { + let batch = match self.get(tenant, generation) { + Some(keys) => { + let public = |day| day_public(&toy_day_secret(tenant, day)).expect("public"); + let working = Working { + keys, + day: TOY_DAY, + oldest: TOY_DAY, + today: public(TOY_DAY), + tomorrow: public(TOY_DAY + 1), + }; + self.batches + .borrow_mut() + .seal(&working, shard, first, rt::now()) + .map_err(|_| KeyError::Unavailable) + } + None => Err(KeyError::Unavailable), + }; + Box::pin(std::future::ready(batch)) + } + + fn batch_key( + &self, + tenant: u32, + generation: u32, + shard: u16, + id: u64, + wrap: Wrap, + ) -> Unwrapped<'_> { + let key = if (wrap.generation, wrap.shard) == (generation, shard) { + open_batch(&toy_day_secret(tenant, wrap.day), tenant, id, &wrap) + .map(|secret| Rc::new(BatchKey::new(&secret))) + .map_err(|_| Unsealable::Bad) + } else { + Err(Unsealable::Bad) + }; + Box::pin(std::future::ready(key)) + } +} + struct Simulated { world: Shared, driver: Rc, @@ -347,6 +408,7 @@ fn message(id: MessageId, draft: &Draft, now: u64) -> Message { ..HeadClear::default() }, expiry: u64::MAX / 2, + wrap: [3; spool::WRAP], sealed: vec![1; 64], }, body: Body::Inline(vec![2; 300]), @@ -595,7 +657,7 @@ fn start( .ok()?; let shards = core.shards.clone(); let parts = Parts { - keys: Rc::new(Toy), + keys: Rc::new(Toy::default()), transport, clock: Clock::start(), }; diff --git a/crates/spool/src/intake.rs b/crates/spool/src/intake.rs index 0a6a603..1bfe083 100644 --- a/crates/spool/src/intake.rs +++ b/crates/spool/src/intake.rs @@ -8,7 +8,7 @@ use store::{ use crate::bodies::{Writer, Written}; use crate::error::Error; -use crate::items::{BODY_ITEM, HeadClear, encode_slice, slice_item}; +use crate::items::{BODY_ITEM, HeadClear, WRAP, encode_slice, slice_item}; use crate::limits::{Bounds, Limits, SLICE}; use crate::pending::{self, Pending}; use crate::recipients::Table; @@ -45,6 +45,7 @@ impl MessageId { pub struct Head { pub clear: HeadClear, pub expiry: u64, + pub wrap: [u8; WRAP], pub sealed: Vec, } @@ -304,7 +305,7 @@ impl Spool { head_change, head.clear.encode(), Vec::new(), - head.sealed, + [&head.wrap[..], &head.sealed].concat(), ) .map_err(|_| Error::Limit("header"))?, ); diff --git a/crates/spool/src/items.rs b/crates/spool/src/items.rs index 1a1b05b..8e44ffe 100644 --- a/crates/spool/src/items.rs +++ b/crates/spool/src/items.rs @@ -4,6 +4,7 @@ pub const BODY_ITEM: u32 = 1; pub const SLICE_BASE: u32 = 2; pub const STATE_BASE: u32 = 1 << 16; pub const HEAD_CLEAR: usize = 48; +pub const WRAP: usize = 91; pub fn slice_item(slice: u32) -> u32 { SLICE_BASE + slice @@ -13,6 +14,7 @@ pub fn slice_item(slice: u32) -> u32 { pub struct HeadClear { pub batch_key: u64, pub recipients: u32, + pub day: u32, pub body_size: u64, pub accepted: u64, pub not_before: u64, @@ -25,6 +27,7 @@ impl HeadClear { let mut fields = FieldsMut::new(&mut out); fields.put_u64(0, self.batch_key); fields.put_u32(8, self.recipients); + fields.put_u32(12, self.day); fields.put_u64(16, self.body_size); fields.put_u64(24, self.accepted); fields.put_u64(32, self.not_before); @@ -37,10 +40,11 @@ impl HeadClear { return None; } let fields = Fields::new(clear); - let reserved = fields.u32(12) != 0 || fields.bytes(41, 7).iter().any(|&b| b != 0); + let reserved = fields.bytes(41, 7).iter().any(|&b| b != 0); (!reserved).then(|| Self { batch_key: fields.u64(0), recipients: fields.u32(8), + day: fields.u32(12), body_size: fields.u64(16), accepted: fields.u64(24), not_before: fields.u64(32), @@ -49,6 +53,11 @@ impl HeadClear { } } +pub fn head_parts(region: &[u8]) -> Option<(&[u8; WRAP], &[u8])> { + let (wrap, sealed) = region.split_first_chunk::()?; + Some((wrap, sealed)) +} + pub(crate) fn encode_slice(queues: &[u32], sealed: &[u8]) -> Vec { let start = 4 + 4 * queues.len(); let mut out = vec![0; start + sealed.len()]; diff --git a/crates/spool/src/lib.rs b/crates/spool/src/lib.rs index 2582f49..0b2c3e9 100644 --- a/crates/spool/src/lib.rs +++ b/crates/spool/src/lib.rs @@ -10,8 +10,8 @@ pub use bodies::{KIND_BODY_CHUNK, Written, decode_chunk}; pub use error::Error; pub use intake::{Body, Head, Message, MessageId, Slice, Spool}; pub use items::{ - BODY_ITEM, HEAD_CLEAR, HeadClear, SLICE_BASE, STATE_BASE, STATE_CLEAR, Stamp, StateClear, - slice_item, slice_parts, state_item, + BODY_ITEM, HEAD_CLEAR, HeadClear, SLICE_BASE, STATE_BASE, STATE_CLEAR, Stamp, StateClear, WRAP, + head_parts, slice_item, slice_parts, state_item, }; pub use limits::{Limits, SLICE}; pub use pending::{Open, Pending}; diff --git a/crates/spool/tests/intake.rs b/crates/spool/tests/intake.rs index 1b895e7..56dbac4 100644 --- a/crates/spool/tests/intake.rs +++ b/crates/spool/tests/intake.rs @@ -125,12 +125,14 @@ async fn draft(intake: &Spool, shard: u16, step: u64) -> Result<(Message, Accept batch_key: seed % 5, recipients: u32::try_from(slices.iter().map(|s| s.queues.len()).sum::()) .expect("fits"), + day: u32::try_from(seed % 30_000).expect("fits"), body_size: seed % 100_000, accepted: step, not_before: step + 1, priority: u8::try_from(step % 3).expect("fits"), }, expiry: 1000 + step, + wrap: bytes(seed ^ 1, spool::WRAP).try_into().expect("wrap size"), sealed: bytes(seed, 200 + usize::try_from(seed % 300).expect("fits")), }; let (body, expected) = if step.is_multiple_of(3) { @@ -368,7 +370,12 @@ async fn verify(file: &CoreFile, keeper: &Keeper, model: &Model) { assert_eq!(head.tenant, message.tenant); assert_eq!(HeadClear::decode(&head.clear), Some(message.head.clear)); assert_eq!(head.delivery.due, message.head.expiry); - assert_eq!(read(file, head.location, 7).await, message.head.sealed); + let region = read(file, head.location, 7).await; + let (wrap, sealed) = spool::head_parts(®ion).expect("wrap"); + assert_eq!( + (wrap, sealed), + (&message.head.wrap, message.head.sealed.as_slice()) + ); let body = state.get(store::Key::new(key, BODY_ITEM)).expect("body"); let found = if body.bodies.is_empty() { Expected::Inline(read(file, body.location, 7).await) diff --git a/crates/spool/tests/lifecycle.rs b/crates/spool/tests/lifecycle.rs index 092df6b..1b59303 100644 --- a/crates/spool/tests/lifecycle.rs +++ b/crates/spool/tests/lifecycle.rs @@ -117,6 +117,7 @@ async fn accept( ..HeadClear::default() }, expiry: 1000 + step, + wrap: [3; spool::WRAP], sealed: vec![1; 64], }, body: Body::Inline(vec![2; 300]), diff --git a/crates/spool/tests/limits.rs b/crates/spool/tests/limits.rs index 98cbd9d..a1732b0 100644 --- a/crates/spool/tests/limits.rs +++ b/crates/spool/tests/limits.rs @@ -34,6 +34,7 @@ fn head(recipients: usize) -> Head { ..HeadClear::default() }, expiry: 9, + wrap: [3; spool::WRAP], sealed: vec![1; 100], } } @@ -136,6 +137,12 @@ fn intake_refuses_messages_over_the_limits() { let mut header = message(id, &[3]); header.head.sealed = vec![1; MAX_PAYLOAD]; refused(intake.accept(vec![header]).await, "header"); + let mut edge = message(id, &[3]); + edge.head.sealed = vec![1; MAX_PAYLOAD - INDEX_SIZE - spool::HEAD_CLEAR]; + refused(intake.accept(vec![edge]).await, "header"); + let mut edge = message(id, &[3]); + edge.head.sealed = vec![1; MAX_PAYLOAD - INDEX_SIZE - spool::HEAD_CLEAR - spool::WRAP + 1]; + refused(intake.accept(vec![edge]).await, "header"); let mut inline = message(id, &[3]); inline.body = Body::Inline(vec![3; 1025]); refused(intake.accept(vec![inline]).await, "inline body"); diff --git a/crates/spool/tests/recipients.rs b/crates/spool/tests/recipients.rs index 62c4aed..a83376b 100644 --- a/crates/spool/tests/recipients.rs +++ b/crates/spool/tests/recipients.rs @@ -82,6 +82,7 @@ impl Node { ..HeadClear::default() }, expiry, + wrap: [3; spool::WRAP], sealed: vec![1; 50], }, body: Body::Inline(vec![3; 10]),