orbit_auth/
table.rs

1//! Bounded current auth state with typed protection. Reads copy fixed atomic
2//! metadata; writes commit one inactive bank under a kernel-released lock.
3//! Authentication does not prevent deletion, rollback or a compromised key holder.
4use std::collections::HashMap;
5use std::marker::PhantomData;
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::sync::{Arc, LazyLock, Mutex, Weak};
8
9use chacha20poly1305::aead::AeadInOut;
10use chacha20poly1305::{KeyInit, Tag, XChaCha20Poly1305, XNonce};
11use orbit_core::Fleet;
12#[cfg(unix)]
13use orbit_core::shm::{ShmRegion, ring_segment_name};
14use zeroize::Zeroizing;
15
16use crate::{
17    Blake3State, EncryptedState, Error, Result, StateProtection,
18    UnprotectedState,
19};
20
21/// Current session/replay state. Principal values live only in process memory.
22pub const AUTH_STATE_KIND: u8 = 198;
23pub const AUTH_STATE_CAPACITY: usize = orbit_core::compile::usize_from_env(
24    option_env!("ORBIT_AUTH_STATE_CAPACITY"),
25    1024,
26);
27const WORDS: usize = 12;
28// Largest encoding: nonce (24), record (96), AEAD tag (16).
29const BANK_WORDS: usize = 17;
30type Wire = [u64; BANK_WORDS];
31#[cfg(unix)]
32const HEADER: usize = 64;
33#[cfg(unix)]
34const MAGIC: &[u8; 8] = b"OAUTH004";
35
36/// Separate fleet-scoped protection key; never stored in the shared mapping.
37pub(crate) struct RecordProtection<P: StateProtection> {
38    key: Zeroizing<[u8; 32]>,
39    policy: PhantomData<P>,
40}
41impl<P: StateProtection> RecordProtection<P> {
42    pub(crate) fn new(key: Zeroizing<[u8; 32]>) -> Self {
43        Self { key, policy: PhantomData }
44    }
45    fn context(
46        index: usize,
47        revision: u64,
48    ) -> [u8; 51] {
49        let mut aad = [0; 51];
50        aad[..26].copy_from_slice(b"orbit-auth/v4/state-record");
51        aad[26] = AUTH_STATE_KIND;
52        aad[27..35].copy_from_slice(&P::ID.to_le_bytes());
53        aad[35..43].copy_from_slice(&(index as u64).to_le_bytes());
54        aad[43..].copy_from_slice(&revision.to_le_bytes());
55        aad
56    }
57    fn tag(
58        &self,
59        aad: &[u8],
60        record: &[u8],
61    ) -> blake3::Hash {
62        let mut mac = Zeroizing::new(blake3::Hasher::new_keyed(&self.key));
63        mac.update(aad);
64        mac.update(record);
65        mac.finalize()
66    }
67    fn encode(
68        &self,
69        index: usize,
70        revision: u64,
71        record: Record,
72    ) -> Result<Wire> {
73        let mut bytes = Zeroizing::new([0; BANK_WORDS * 8]);
74        let mut plaintext = Zeroizing::new([0; WORDS * 8]);
75        for (dst, word) in
76            plaintext.as_chunks_mut::<8>().0.iter_mut().zip(record.words())
77        {
78            dst.copy_from_slice(&word.to_le_bytes());
79        }
80        let aad = Self::context(index, revision);
81        match P::ID {
82            Blake3State::ID => {
83                bytes[..96].copy_from_slice(plaintext.as_ref());
84                bytes[96..128].copy_from_slice(
85                    self.tag(&aad, plaintext.as_ref()).as_bytes(),
86                );
87            }
88            EncryptedState::ID => {
89                let mut nonce = [0; 24];
90                getrandom::fill(&mut nonce).map_err(|_| Error::Randomness)?;
91                let cipher = XChaCha20Poly1305::new((&*self.key).into());
92                let tag = cipher
93                    .encrypt_inout_detached(
94                        &XNonce::from(nonce),
95                        &aad,
96                        (&mut plaintext[..]).into(),
97                    )
98                    .map_err(|_| Error::PolicyUnavailable)?;
99                bytes[..24].copy_from_slice(&nonce);
100                bytes[24..120].copy_from_slice(plaintext.as_ref());
101                bytes[120..].copy_from_slice(&tag);
102            }
103            UnprotectedState::ID => {
104                bytes[..96].copy_from_slice(plaintext.as_ref())
105            }
106            _ => unreachable!("sealed policy"),
107        }
108        Ok(std::array::from_fn(|i| {
109            u64::from_le_bytes(bytes[i * 8..i * 8 + 8].try_into().unwrap())
110        }))
111    }
112    fn decode(
113        &self,
114        index: usize,
115        revision: u64,
116        wire: &Wire,
117    ) -> Result<Record> {
118        let mut bytes = Zeroizing::new([0; BANK_WORDS * 8]);
119        for (dst, word) in bytes.as_chunks_mut::<8>().0.iter_mut().zip(wire) {
120            dst.copy_from_slice(&word.to_le_bytes());
121        }
122        let aad = Self::context(index, revision);
123        let mut plaintext = Zeroizing::new([0; WORDS * 8]);
124        match P::ID {
125            Blake3State::ID => {
126                let expected = self.tag(&aad, &bytes[..96]);
127                let actual = blake3::Hash::from_bytes(
128                    bytes[96..128].try_into().unwrap(),
129                );
130                // Hash::eq is constant-time; do not compare raw tag byte arrays.
131                if expected != actual || bytes[128..].iter().any(|&v| v != 0) {
132                    return Err(Error::PolicyUnavailable);
133                }
134                plaintext.copy_from_slice(&bytes[..96]);
135            }
136            EncryptedState::ID => {
137                plaintext.copy_from_slice(&bytes[24..120]);
138                let cipher = XChaCha20Poly1305::new((&*self.key).into());
139                let nonce =
140                    XNonce::try_from(&bytes[..24]).expect("fixed nonce length");
141                let tag =
142                    Tag::try_from(&bytes[120..]).expect("fixed tag length");
143                cipher
144                    .decrypt_inout_detached(
145                        &nonce,
146                        &aad,
147                        (&mut plaintext[..]).into(),
148                        &tag,
149                    )
150                    .map_err(|_| Error::PolicyUnavailable)?;
151            }
152            UnprotectedState::ID => {
153                if bytes[96..].iter().any(|&v| v != 0) {
154                    return Err(Error::PolicyUnavailable);
155                }
156                plaintext.copy_from_slice(&bytes[..96]);
157            }
158            _ => unreachable!("sealed policy"),
159        }
160        Record::from_words(std::array::from_fn(|i| {
161            u64::from_le_bytes(plaintext[i * 8..i * 8 + 8].try_into().unwrap())
162        }))
163    }
164}
165
166pub(crate) const SESSION: u64 = 1;
167pub(crate) const REVOKED: u64 = 2;
168pub(crate) const CONSUMED: u64 = 3;
169
170#[derive(Clone, Copy, Debug, PartialEq, Eq)]
171pub(crate) struct StoreKey(pub [u64; 4]);
172
173#[derive(Clone, Copy, Debug, PartialEq, Eq)]
174pub(crate) struct Record {
175    pub(crate) key: StoreKey,
176    pub(crate) subject: [u64; 4],
177    pub(crate) created_at: u64,
178    pub(crate) expires_at: u64,
179    pub(crate) generation: u64,
180    pub(crate) status: u64,
181}
182impl Record {
183    fn words(self) -> [u64; WORDS] {
184        [
185            self.key.0[0],
186            self.key.0[1],
187            self.key.0[2],
188            self.key.0[3],
189            self.subject[0],
190            self.subject[1],
191            self.subject[2],
192            self.subject[3],
193            self.created_at,
194            self.expires_at,
195            self.generation,
196            self.status,
197        ]
198    }
199    fn from_words(words: [u64; WORDS]) -> Result<Self> {
200        let record = Self {
201            key: StoreKey(words[..4].try_into().expect("four words")),
202            subject: words[4..8].try_into().expect("four words"),
203            created_at: words[8],
204            expires_at: words[9],
205            generation: words[10],
206            status: words[11],
207        };
208        if !matches!(record.status, SESSION | REVOKED | CONSUMED)
209            || record.created_at >= record.expires_at
210            || (record.status == CONSUMED
211                && (record.subject != [0; 4] || record.generation != 0))
212        {
213            return Err(Error::PolicyUnavailable);
214        }
215        Ok(record)
216    }
217}
218
219#[derive(Clone, Copy, Debug)]
220pub(crate) struct Snapshot {
221    pub(crate) index: usize,
222    pub(crate) revision: u64,
223    pub(crate) record: Record,
224    wire: Wire,
225}
226
227enum Read {
228    Empty,
229    Value(u64, Wire),
230    Changed,
231}
232
233#[repr(C, align(64))]
234pub(crate) struct Slot {
235    // Never wraps or resets: cached slot references cannot survive reuse (ABA).
236    revision: AtomicU64,
237    banks: [[AtomicU64; BANK_WORDS]; 2],
238}
239impl Slot {
240    fn empty() -> Self {
241        Self {
242            revision: AtomicU64::new(0),
243            banks: std::array::from_fn(|_| {
244                std::array::from_fn(|_| AtomicU64::new(0))
245            }),
246        }
247    }
248
249    fn read_once(&self) -> Result<Read> {
250        let before = self.revision.load(Ordering::SeqCst);
251        if before == 0 {
252            return Ok(Read::Empty);
253        }
254        let words: [u64; BANK_WORDS] = std::array::from_fn(|i| {
255            self.banks[(before & 1) as usize][i].load(Ordering::SeqCst)
256        });
257        // Every bank word and revision operation participates in the SC total
258        // order. Even if two writers cycle back to this bank during the read,
259        // the non-wrapping revision changes and the partial copy is discarded.
260        if self.revision.load(Ordering::SeqCst) != before {
261            return Ok(Read::Changed);
262        }
263        Ok(Read::Value(before, words))
264    }
265
266    /// Caller holds Table::write's exclusive writer lock.
267    pub(crate) fn store<P: StateProtection>(
268        &self,
269        record: Record,
270        index: usize,
271        mac: &RecordProtection<P>,
272    ) -> Result<u64> {
273        Record::from_words(record.words())?;
274        let revision = self
275            .revision
276            .load(Ordering::SeqCst)
277            .checked_add(1)
278            .ok_or(Error::PolicyUnavailable)?;
279        let wire = mac.encode(index, revision, record)?;
280        let bank = &self.banks[(revision & 1) as usize];
281        for (dst, value) in bank.iter().zip(wire) {
282            dst.store(value, Ordering::SeqCst);
283        }
284        self.revision.store(revision, Ordering::SeqCst);
285        Ok(revision)
286    }
287}
288
289pub(crate) struct Memory {
290    _fleet: Arc<Fleet>,
291    protection: u64,
292    slots: Vec<Slot>,
293    writer: Mutex<()>,
294}
295type MemoryRegistry = HashMap<usize, Weak<Memory>>;
296static MEMORY: LazyLock<Mutex<MemoryRegistry>> =
297    LazyLock::new(|| Mutex::new(HashMap::new()));
298
299pub(crate) enum Table {
300    Memory(Arc<Memory>),
301    #[cfg(unix)]
302    Shm(ShmRegion),
303}
304impl Table {
305    pub(crate) fn new(
306        fleet: Arc<Fleet>,
307        protection: u64,
308    ) -> Result<Self> {
309        if fleet.is_shm() {
310            #[cfg(unix)]
311            {
312                return Self::shm(
313                    &ring_segment_name(fleet.name(), AUTH_STATE_KIND),
314                    protection,
315                );
316            }
317            #[cfg(not(unix))]
318            {
319                return Err(Error::PolicyUnavailable);
320            }
321        }
322        let mut registry =
323            MEMORY.lock().map_err(|_| Error::PolicyUnavailable)?;
324        registry.retain(|_, value| value.strong_count() > 0);
325        let table = registry.entry(Arc::as_ptr(&fleet) as usize).or_default();
326        if let Some(memory) = table.upgrade() {
327            if memory.protection != protection {
328                return Err(Error::IncompatibleProtection);
329            }
330            return Ok(Self::Memory(memory));
331        }
332        let memory = Arc::new(Memory {
333            _fleet: fleet,
334            protection,
335            slots: (0..AUTH_STATE_CAPACITY).map(|_| Slot::empty()).collect(),
336            writer: Mutex::new(()),
337        });
338        *table = Arc::downgrade(&memory);
339        Ok(Self::Memory(memory))
340    }
341
342    #[cfg(unix)]
343    fn shm(
344        name: &str,
345        protection: u64,
346    ) -> Result<Self> {
347        let size = HEADER + AUTH_STATE_CAPACITY * size_of::<Slot>();
348        let (region, _initialization) =
349            ShmRegion::open_or_create_locked(name, size)
350                .map_err(|_| Error::PolicyUnavailable)?;
351        let mut expected = [0u8; HEADER];
352        expected[..8].copy_from_slice(MAGIC);
353        for (i, value) in [
354            AUTH_STATE_CAPACITY,
355            size_of::<Slot>(),
356            BANK_WORDS,
357            usize::from(cfg!(target_endian = "little")),
358        ]
359        .into_iter()
360        .enumerate()
361        {
362            expected[8 + i * 8..16 + i * 8]
363                .copy_from_slice(&(value as u64).to_le_bytes());
364        }
365        expected[40..48].copy_from_slice(&protection.to_le_bytes());
366        // SAFETY: mmap is page aligned and covers all slots. The initialization
367        // lock excludes openers; published fields are thereafter only atomic.
368        unsafe {
369            if region.created() {
370                std::ptr::write_bytes(region.as_ptr(), 0, size);
371                std::ptr::copy_nonoverlapping(
372                    expected.as_ptr(),
373                    region.as_ptr(),
374                    HEADER,
375                );
376            } else {
377                let actual =
378                    std::slice::from_raw_parts(region.as_ptr(), HEADER);
379                if actual[..40] != expected[..40]
380                    || actual[48..] != expected[48..]
381                {
382                    return Err(Error::IncompatibleLayout);
383                }
384                if actual[40..48] != expected[40..48] {
385                    return Err(Error::IncompatibleProtection);
386                }
387            }
388        }
389        Ok(Self::Shm(region))
390    }
391
392    fn slots(&self) -> &[Slot] {
393        match self {
394            Self::Memory(memory) => &memory.slots,
395            #[cfg(unix)]
396            Self::Shm(region) => {
397                // SAFETY: validated fixed geometry and alignment; region lives
398                // as long as self. All peer-visible fields are atomic and there
399                // are no mutable Rust references into the shared mapping.
400                unsafe {
401                    std::slice::from_raw_parts(
402                        region.as_ptr().add(HEADER).cast::<Slot>(),
403                        AUTH_STATE_CAPACITY,
404                    )
405                }
406            }
407        }
408    }
409
410    /// Protected policies compare all bytes with the verified local copy. The
411    /// unprotected policy trusts writers to publish a revision for every change.
412    pub(crate) fn unchanged<P: StateProtection>(
413        &self,
414        snapshot: &Snapshot,
415    ) -> Result<bool> {
416        let Some(slot) = self.slots().get(snapshot.index) else {
417            return Ok(false);
418        };
419        if P::ID == UnprotectedState::ID {
420            return Ok(
421                slot.revision.load(Ordering::SeqCst) == snapshot.revision
422            );
423        }
424        Ok(matches!(slot.read_once()?, Read::Value(revision, wire)
425            if revision == snapshot.revision && wire == snapshot.wire))
426    }
427
428    pub(crate) fn lookup<P: StateProtection>(
429        &self,
430        key: StoreKey,
431        mac: &RecordProtection<P>,
432    ) -> Result<Option<Snapshot>> {
433        let result = lookup(self.slots(), key, mac)?;
434        if result.raced {
435            // A raced bank copy takes one synchronized read. No spin/poll loop,
436            // no deadline, and no acceptance based on a torn snapshot.
437            self.write(|slots| find(slots, key, mac))
438        } else {
439            Ok(result.snapshot)
440        }
441    }
442
443    pub(crate) fn write<T>(
444        &self,
445        f: impl FnOnce(&[Slot]) -> Result<T>,
446    ) -> Result<T> {
447        match self {
448            Self::Memory(memory) => {
449                let _lock = memory
450                    .writer
451                    .lock()
452                    .map_err(|_| Error::PolicyUnavailable)?;
453                f(&memory.slots)
454            }
455            #[cfg(unix)]
456            Self::Shm(region) => {
457                let _lock = region
458                    .lock_exclusive()
459                    .map_err(|_| Error::PolicyUnavailable)?;
460                f(self.slots())
461            }
462        }
463    }
464}
465
466struct Lookup {
467    snapshot: Option<Snapshot>,
468    raced: bool,
469}
470impl Lookup {
471    fn stable(snapshot: Option<Snapshot>) -> Self {
472        Self { snapshot, raced: false }
473    }
474    fn changed() -> Self {
475        Self { snapshot: None, raced: true }
476    }
477}
478fn lookup<P: StateProtection>(
479    slots: &[Slot],
480    key: StoreKey,
481    mac: &RecordProtection<P>,
482) -> Result<Lookup> {
483    for offset in 0..slots.len() {
484        let index = (key.0[0] as usize).wrapping_add(offset) % slots.len();
485        match slots[index].read_once()? {
486            Read::Empty => return Ok(Lookup::stable(None)),
487            Read::Value(revision, wire) => {
488                let record = mac.decode(index, revision, &wire)?;
489                if record.key == key {
490                    return Ok(Lookup::stable(Some(Snapshot {
491                        index,
492                        revision,
493                        record,
494                        wire,
495                    })));
496                }
497            }
498            Read::Changed => return Ok(Lookup::changed()),
499        }
500    }
501    Ok(Lookup::stable(None))
502}
503
504/// Used only while the writer lock is held; no bank can change here.
505pub(crate) fn find<P: StateProtection>(
506    slots: &[Slot],
507    key: StoreKey,
508    mac: &RecordProtection<P>,
509) -> Result<Option<Snapshot>> {
510    let result = lookup(slots, key, mac)?;
511    if result.raced {
512        Err(Error::PolicyUnavailable)
513    } else {
514        Ok(result.snapshot)
515    }
516}
517
518/// Enumerate one subject while the existing exclusive writer lock is held.
519/// Authenticate all occupied slots before making a bulk update.
520pub(crate) fn subject_sessions<P: StateProtection>(
521    slots: &[Slot],
522    subject: [u64; 4],
523    now: u64,
524    mac: &RecordProtection<P>,
525) -> Result<Vec<Snapshot>> {
526    let mut sessions = Vec::new();
527    for (index, slot) in slots.iter().enumerate() {
528        match slot.read_once()? {
529            Read::Empty => {}
530            Read::Changed => return Err(Error::PolicyUnavailable),
531            Read::Value(revision, wire) => {
532                let record = mac.decode(index, revision, &wire)?;
533                if record.status == SESSION
534                    && record.subject == subject
535                    && now < record.expires_at
536                {
537                    if now < record.created_at {
538                        return Err(Error::NotYetValid);
539                    }
540                    sessions.push(Snapshot { index, revision, record, wire });
541                }
542            }
543        }
544    }
545    Ok(sessions)
546}
547
548/// Caller has checked absence and holds the writer lock. Reclaims only expired
549/// state; no tombstone is removed early to admit another record.
550pub(crate) fn insert<P: StateProtection>(
551    slots: &[Slot],
552    record: Record,
553    now: u64,
554    mac: &RecordProtection<P>,
555) -> Result<()> {
556    for offset in 0..slots.len() {
557        let index =
558            (record.key.0[0] as usize).wrapping_add(offset) % slots.len();
559        let reusable = match slots[index].read_once()? {
560            Read::Empty => true,
561            Read::Value(revision, wire) => {
562                let old = mac.decode(index, revision, &wire)?;
563                now >= old.expires_at
564            }
565            Read::Changed => return Err(Error::PolicyUnavailable),
566        };
567        if reusable {
568            slots[index].store(record, index, mac)?;
569            return Ok(());
570        }
571    }
572    Err(Error::StateFull)
573}
574
575const _: () = assert!(
576    AUTH_STATE_CAPACITY > 0
577        && AUTH_STATE_CAPACITY
578            <= (isize::MAX as usize - 64) / size_of::<Slot>()
579);
580const _: () = assert!(size_of::<Slot>() == 320 && align_of::<Slot>() == 64);
581
582#[cfg(test)]
583#[path = "../tests/unit/table.rs"]
584mod tests;
orbit_auth/table.rs

orbit_auth/
table.rs
