link-assistant-router 0.66.0

//! Token persistence backends.
//!
//! Issued-token state can be persisted as real Links Notation text, a native
//! file-mapped doublets graph, or both.
//!
//! This module provides:
//!
//! - The [`TokenStore`] trait — a small async-ish (sync, for now) API for
//!   listing, persisting, and revoking [`TokenRecord`]s.
//! - [`MemoryTokenStore`] — an in-memory implementation, used in tests and
//!   for [`StoragePolicy::Memory`].
//! - [`TextTokenStore`] — persists records through `lino-objects-codec` at
//!   `<data_dir>/tokens.lino`.
//! - [`BinaryTokenStore`] — persists the same `Type → SubType → Value` graph
//!   in a `doublets` store backed by `platform-mem` at
//!   `<data_dir>/tokens.bin`.
//! - [`DualTokenStore`] — fans writes out to two stores (typically text +
//!   binary) and reads from the *first* store, falling back to the second
//!   on miss. This is the default when [`StoragePolicy::Both`] is set.
//! - [`build_token_store`] — factory that picks the right combination of
//!   stores based on the [`StoragePolicy`] from configuration.
//!
//! All persistence operations are best-effort: failures are surfaced as
//! [`StorageError`] but never panic, and all read paths gracefully tolerate
//! missing files (returning an empty record set).

// We intentionally hold a write guard across the (insert + flush) pair for
// atomicity; clippy's `significant_drop_tightening` would push us into
// chained calls that lose readability without helping contention in practice.
#![allow(clippy::significant_drop_tightening)]

use std::collections::HashMap;
use std::fs::{self, OpenOptions};
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::sync::{Arc, RwLock};

use serde::{Deserialize, Serialize};

use crate::config::StoragePolicy;

mod associative;
#[allow(unsafe_code)]
mod file_mapped;
mod legacy;

/// One persisted token record.
///
/// `id` is the JWT `sub` (a UUID); the JWT itself is NOT stored — only the
/// metadata required to list/expire/revoke.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct TokenRecord {
    pub id: String,
    pub label: String,
    pub issued_at: i64,
    pub expires_at: i64,
    pub revoked: bool,
    /// Optional account identifier the token is bound to (multi-account mode).
    #[serde(default)]
    pub account: Option<String>,
    /// Optional cap on the number of upstream requests this token may make.
    /// `None` means unlimited. Used to bound how much a task can consume.
    #[serde(default)]
    pub max_requests: Option<u64>,
    /// Number of upstream requests already made with this token.
    #[serde(default)]
    pub used_requests: u64,
    /// Privilege scope carried by the token. Empty (the default) means an
    /// ordinary client token that may only proxy inference; `"admin"` marks a
    /// credential that also unlocks the administrative endpoints. See
    /// [`crate::token::ADMIN_SCOPE`].
    #[serde(default)]
    pub scope: String,
}

/// Errors a [`TokenStore`] can return.
#[derive(Debug)]
pub enum StorageError {
    Io(io::Error),
    Codec(String),
    LockPoisoned,
}

impl std::fmt::Display for StorageError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Io(e) => write!(f, "storage I/O error: {e}"),
            Self::Codec(msg) => write!(f, "storage codec error: {msg}"),
            Self::LockPoisoned => write!(f, "storage lock poisoned"),
        }
    }
}

impl std::error::Error for StorageError {}

impl From<io::Error> for StorageError {
    fn from(e: io::Error) -> Self {
        Self::Io(e)
    }
}

/// Persistent token store API.
///
/// Implementations must be cheap to clone (use `Arc` internally) — the
/// router shares them across handler tasks.
pub trait TokenStore: Send + Sync {
    fn list(&self) -> Result<Vec<TokenRecord>, StorageError>;
    fn get(&self, id: &str) -> Result<Option<TokenRecord>, StorageError>;
    fn put(&self, record: TokenRecord) -> Result<(), StorageError>;
    fn delete(&self, id: &str) -> Result<bool, StorageError>;
    fn revoke(&self, id: &str) -> Result<bool, StorageError> {
        if let Some(mut rec) = self.get(id)? {
            if rec.revoked {
                return Ok(false);
            }
            rec.revoked = true;
            self.put(rec)?;
            return Ok(true);
        }
        Ok(false)
    }
    fn revoked_ids(&self) -> Result<Vec<String>, StorageError> {
        Ok(self
            .list()?
            .into_iter()
            .filter(|r| r.revoked)
            .map(|r| r.id)
            .collect())
    }

    /// Atomically check the request budget for `id` and, when there is room,
    /// increment the used-request counter.
    ///
    /// Returns `Ok(true)` when the request is permitted (and was recorded),
    /// or `Ok(false)` when the token is already at or over its `max_requests`
    /// budget. Tokens that have no stored record, or whose record has no
    /// limit (`max_requests == None`), are always permitted.
    ///
    /// Note: the default implementation performs a `get` then a `put`, which
    /// is not strictly atomic across concurrent callers, so a small number of
    /// requests may slip over the limit under heavy parallelism. This is an
    /// acceptable trade-off for a per-task usage cap; backends that need exact
    /// enforcement can override this with a locked read-modify-write.
    fn try_consume_request(&self, id: &str) -> Result<bool, StorageError> {
        if let Some(mut rec) = self.get(id)? {
            if let Some(max) = rec.max_requests {
                if rec.used_requests >= max {
                    return Ok(false);
                }
            }
            rec.used_requests = rec.used_requests.saturating_add(1);
            self.put(rec)?;
        }
        Ok(true)
    }
}

/// Trivial in-memory store. No persistence. Useful for tests.
#[derive(Default, Clone)]
pub struct MemoryTokenStore {
    inner: Arc<RwLock<HashMap<String, TokenRecord>>>,
}

impl MemoryTokenStore {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }
}

impl TokenStore for MemoryTokenStore {
    fn list(&self) -> Result<Vec<TokenRecord>, StorageError> {
        let guard = self.inner.read().map_err(|_| StorageError::LockPoisoned)?;
        Ok(guard.values().cloned().collect())
    }

    fn get(&self, id: &str) -> Result<Option<TokenRecord>, StorageError> {
        let guard = self.inner.read().map_err(|_| StorageError::LockPoisoned)?;
        Ok(guard.get(id).cloned())
    }

    fn put(&self, record: TokenRecord) -> Result<(), StorageError> {
        let mut guard = self.inner.write().map_err(|_| StorageError::LockPoisoned)?;
        guard.insert(record.id.clone(), record);
        Ok(())
    }

    fn delete(&self, id: &str) -> Result<bool, StorageError> {
        let mut guard = self.inner.write().map_err(|_| StorageError::LockPoisoned)?;
        Ok(guard.remove(id).is_some())
    }

    fn try_consume_request(&self, id: &str) -> Result<bool, StorageError> {
        let mut guard = self.inner.write().map_err(|_| StorageError::LockPoisoned)?;
        Ok(consume_request(guard.get_mut(id)))
    }
}

/// Links Notation text token store.
///
/// Existing hand-built files are read once and atomically migrated to the
/// official `lino-objects-codec` representation.
#[derive(Clone)]
pub struct TextTokenStore {
    path: PathBuf,
    inner: Arc<RwLock<HashMap<String, TokenRecord>>>,
}

impl TextTokenStore {
    pub fn open(path: impl Into<PathBuf>) -> Result<Self, StorageError> {
        let path = path.into();
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)?;
        }
        let (records, migrated) = if path.exists() {
            let contents = fs::read_to_string(&path)?;
            match associative::decode_text(&contents) {
                Ok(records) => (records, false),
                Err(_) => (
                    legacy::decode_text(&contents).map_err(StorageError::Codec)?,
                    true,
                ),
            }
        } else {
            (Vec::new(), false)
        };
        let map: HashMap<_, _> = records.into_iter().map(|r| (r.id.clone(), r)).collect();
        let store = Self {
            path,
            inner: Arc::new(RwLock::new(map)),
        };
        if migrated {
            let guard = store.inner.read().map_err(|_| StorageError::LockPoisoned)?;
            store.flush(&guard)?;
        }
        Ok(store)
    }

    fn flush(&self, guard: &HashMap<String, TokenRecord>) -> Result<(), StorageError> {
        let mut sorted: Vec<&TokenRecord> = guard.values().collect();
        sorted.sort_by(|a, b| a.id.cmp(&b.id));
        let body = associative::encode_text(sorted.iter().copied());
        atomic_write(&self.path, body.as_bytes())
    }
}

impl TokenStore for TextTokenStore {
    fn list(&self) -> Result<Vec<TokenRecord>, StorageError> {
        let guard = self.inner.read().map_err(|_| StorageError::LockPoisoned)?;
        Ok(guard.values().cloned().collect())
    }

    fn get(&self, id: &str) -> Result<Option<TokenRecord>, StorageError> {
        let guard = self.inner.read().map_err(|_| StorageError::LockPoisoned)?;
        Ok(guard.get(id).cloned())
    }

    fn put(&self, record: TokenRecord) -> Result<(), StorageError> {
        let mut guard = self.inner.write().map_err(|_| StorageError::LockPoisoned)?;
        guard.insert(record.id.clone(), record);
        self.flush(&guard)
    }

    fn delete(&self, id: &str) -> Result<bool, StorageError> {
        let mut guard = self.inner.write().map_err(|_| StorageError::LockPoisoned)?;
        let removed = guard.remove(id).is_some();
        if removed {
            self.flush(&guard)?;
        }
        Ok(removed)
    }

    fn try_consume_request(&self, id: &str) -> Result<bool, StorageError> {
        let mut guard = self.inner.write().map_err(|_| StorageError::LockPoisoned)?;
        let permitted = consume_request(guard.get_mut(id));
        if permitted && guard.contains_key(id) {
            self.flush(&guard)?;
        }
        Ok(permitted)
    }
}

/// Native file-mapped doublets token store.
///
/// Existing `LARTOK01` length-prefixed JSON files are read once and
/// atomically migrated to the doublets representation.
#[derive(Clone)]
pub struct BinaryTokenStore {
    path: PathBuf,
    inner: Arc<RwLock<HashMap<String, TokenRecord>>>,
}

impl BinaryTokenStore {
    pub fn open(path: impl Into<PathBuf>) -> Result<Self, StorageError> {
        let path = path.into();
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)?;
        }
        let (records, migrated) = if path.exists() {
            if legacy::is_binary(&path)? {
                (legacy::decode_binary(&path)?, true)
            } else {
                (associative::read_binary(&path)?, false)
            }
        } else {
            (Vec::new(), false)
        };
        let map: HashMap<_, _> = records.into_iter().map(|r| (r.id.clone(), r)).collect();
        let store = Self {
            path,
            inner: Arc::new(RwLock::new(map)),
        };
        if migrated {
            let guard = store.inner.read().map_err(|_| StorageError::LockPoisoned)?;
            store.flush(&guard)?;
        }
        Ok(store)
    }

    fn flush(&self, guard: &HashMap<String, TokenRecord>) -> Result<(), StorageError> {
        let mut sorted: Vec<&TokenRecord> = guard.values().collect();
        sorted.sort_by(|a, b| a.id.cmp(&b.id));
        associative::write_binary(&self.path, sorted)
    }
}

impl TokenStore for BinaryTokenStore {
    fn list(&self) -> Result<Vec<TokenRecord>, StorageError> {
        let guard = self.inner.read().map_err(|_| StorageError::LockPoisoned)?;
        Ok(guard.values().cloned().collect())
    }

    fn get(&self, id: &str) -> Result<Option<TokenRecord>, StorageError> {
        let guard = self.inner.read().map_err(|_| StorageError::LockPoisoned)?;
        Ok(guard.get(id).cloned())
    }

    fn put(&self, record: TokenRecord) -> Result<(), StorageError> {
        let mut guard = self.inner.write().map_err(|_| StorageError::LockPoisoned)?;
        guard.insert(record.id.clone(), record);
        self.flush(&guard)
    }

    fn delete(&self, id: &str) -> Result<bool, StorageError> {
        let mut guard = self.inner.write().map_err(|_| StorageError::LockPoisoned)?;
        let removed = guard.remove(id).is_some();
        if removed {
            self.flush(&guard)?;
        }
        Ok(removed)
    }

    fn try_consume_request(&self, id: &str) -> Result<bool, StorageError> {
        let mut guard = self.inner.write().map_err(|_| StorageError::LockPoisoned)?;
        let permitted = consume_request(guard.get_mut(id));
        if permitted && guard.contains_key(id) {
            self.flush(&guard)?;
        }
        Ok(permitted)
    }
}

fn consume_request(record: Option<&mut TokenRecord>) -> bool {
    let Some(record) = record else {
        return true;
    };
    if record
        .max_requests
        .is_some_and(|max| record.used_requests >= max)
    {
        return false;
    }
    record.used_requests = record.used_requests.saturating_add(1);
    true
}

/// Dual-write store: every mutation goes to *both* primary and secondary;
/// reads consult the primary first, falling back to the secondary on miss.
///
/// Used for [`StoragePolicy::Both`] so the text and binary files stay in
/// sync. The configured order is `text` (primary) → `binary` (secondary).
pub struct DualTokenStore {
    pub primary: Arc<dyn TokenStore>,
    pub secondary: Arc<dyn TokenStore>,
}

impl TokenStore for DualTokenStore {
    fn list(&self) -> Result<Vec<TokenRecord>, StorageError> {
        let mut by_id: HashMap<String, TokenRecord> = HashMap::new();
        for rec in self.primary.list()? {
            by_id.insert(rec.id.clone(), rec);
        }
        for rec in self.secondary.list()? {
            by_id.entry(rec.id.clone()).or_insert(rec);
        }
        Ok(by_id.into_values().collect())
    }

    fn get(&self, id: &str) -> Result<Option<TokenRecord>, StorageError> {
        if let Some(rec) = self.primary.get(id)? {
            return Ok(Some(rec));
        }
        self.secondary.get(id)
    }

    fn put(&self, record: TokenRecord) -> Result<(), StorageError> {
        self.primary.put(record.clone())?;
        self.secondary.put(record)?;
        Ok(())
    }

    fn delete(&self, id: &str) -> Result<bool, StorageError> {
        let a = self.primary.delete(id)?;
        let b = self.secondary.delete(id)?;
        Ok(a || b)
    }

    fn try_consume_request(&self, id: &str) -> Result<bool, StorageError> {
        if !self.primary.try_consume_request(id)? {
            return Ok(false);
        }
        self.secondary.try_consume_request(id)
    }
}

/// Build a [`TokenStore`] following the configured [`StoragePolicy`].
pub fn build_token_store(
    policy: StoragePolicy,
    data_dir: &Path,
) -> Result<Arc<dyn TokenStore>, StorageError> {
    match policy {
        StoragePolicy::Memory => Ok(Arc::new(MemoryTokenStore::new())),
        StoragePolicy::Text => {
            let s = TextTokenStore::open(data_dir.join("tokens.lino"))?;
            Ok(Arc::new(s))
        }
        StoragePolicy::Binary => {
            let s = BinaryTokenStore::open(data_dir.join("tokens.bin"))?;
            Ok(Arc::new(s))
        }
        StoragePolicy::Both => {
            let text = Arc::new(TextTokenStore::open(data_dir.join("tokens.lino"))?);
            let binary = Arc::new(BinaryTokenStore::open(data_dir.join("tokens.bin"))?);
            Ok(Arc::new(DualTokenStore {
                primary: text,
                secondary: binary,
            }))
        }
    }
}

fn atomic_write(path: &Path, contents: &[u8]) -> Result<(), StorageError> {
    let parent = path
        .parent()
        .ok_or_else(|| io::Error::other("storage path has no parent directory"))?;
    fs::create_dir_all(parent)?;
    let file_name = path
        .file_name()
        .and_then(|name| name.to_str())
        .ok_or_else(|| io::Error::other("storage file name is not valid UTF-8"))?;
    let tmp = parent.join(format!(
        ".{file_name}.{}.{}.tmp",
        std::process::id(),
        uuid::Uuid::new_v4()
    ));

    let result = (|| -> Result<(), StorageError> {
        let mut options = OpenOptions::new();
        options.write(true).create_new(true);
        #[cfg(unix)]
        {
            use std::os::unix::fs::OpenOptionsExt;
            options.mode(0o600);
        }
        let mut file = options.open(&tmp)?;
        file.write_all(contents)?;
        file.sync_all()?;
        if let Ok(metadata) = fs::metadata(path) {
            fs::set_permissions(&tmp, metadata.permissions())?;
        }
        fs::rename(&tmp, path)?;
        Ok(())
    })();

    if result.is_err() {
        let _ = fs::remove_file(&tmp);
    }
    result
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Barrier;
    use std::thread;
    use tempfile::tempdir;

    fn sample_record(id: &str) -> TokenRecord {
        TokenRecord {
            id: id.into(),
            label: "test \"label\"".into(),
            issued_at: 1_700_000_000,
            expires_at: 1_700_001_000,
            revoked: false,
            account: Some("primary".into()),
            max_requests: None,
            used_requests: 0,
            scope: String::new(),
        }
    }

    #[test]
    fn memory_store_roundtrip() {
        let s = MemoryTokenStore::new();
        s.put(sample_record("a")).unwrap();
        assert_eq!(s.list().unwrap().len(), 1);
        assert!(s.get("a").unwrap().is_some());
        assert!(s.delete("a").unwrap());
        assert!(s.get("a").unwrap().is_none());
    }

    #[test]
    fn text_store_roundtrip() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("tokens.lino");
        let s = TextTokenStore::open(&path).unwrap();
        s.put(sample_record("a")).unwrap();
        s.put(sample_record("b")).unwrap();
        let s2 = TextTokenStore::open(&path).unwrap();
        let mut list = s2.list().unwrap();
        list.sort_by(|x, y| x.id.cmp(&y.id));
        assert_eq!(list.len(), 2);
        assert_eq!(list[0].id, "a");
        assert_eq!(list[0].label, "test \"label\"");
        assert_eq!(list[0].account.as_deref(), Some("primary"));
    }

    #[test]
    fn binary_store_roundtrip() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("tokens.bin");
        let s = BinaryTokenStore::open(&path).unwrap();
        s.put(sample_record("a")).unwrap();
        s.put(sample_record("b")).unwrap();
        let s2 = BinaryTokenStore::open(&path).unwrap();
        let mut list = s2.list().unwrap();
        list.sort_by(|x, y| x.id.cmp(&y.id));
        assert_eq!(list.len(), 2);
        assert_eq!(list[1].id, "b");
    }

    #[test]
    fn stores_persist_the_admin_scope() {
        let dir = tempdir().unwrap();
        let mut admin = sample_record("admin");
        admin.scope = crate::token::ADMIN_SCOPE.to_string();

        let text_path = dir.path().join("tokens.lino");
        let text = TextTokenStore::open(&text_path).unwrap();
        text.put(admin.clone()).unwrap();
        text.put(sample_record("client")).unwrap();
        let text = TextTokenStore::open(&text_path).unwrap();
        assert_eq!(
            text.get("admin").unwrap().unwrap().scope,
            crate::token::ADMIN_SCOPE
        );
        assert!(text.get("client").unwrap().unwrap().scope.is_empty());

        let bin_path = dir.path().join("tokens.bin");
        let bin = BinaryTokenStore::open(&bin_path).unwrap();
        bin.put(admin).unwrap();
        let bin = BinaryTokenStore::open(&bin_path).unwrap();
        assert_eq!(
            bin.get("admin").unwrap().unwrap().scope,
            crate::token::ADMIN_SCOPE
        );
    }

    #[test]
    fn dual_store_writes_both() {
        let dir = tempdir().unwrap();
        let text = Arc::new(TextTokenStore::open(dir.path().join("a.lino")).unwrap());
        let bin = Arc::new(BinaryTokenStore::open(dir.path().join("a.bin")).unwrap());
        let dual = DualTokenStore {
            primary: text.clone(),
            secondary: bin.clone(),
        };
        dual.put(sample_record("a")).unwrap();
        assert_eq!(text.list().unwrap().len(), 1);
        assert_eq!(bin.list().unwrap().len(), 1);
    }

    #[test]
    fn dual_store_concurrent_consumption_is_atomic_and_preserves_formats() {
        const REQUESTS: usize = 32;

        let dir = tempdir().unwrap();
        let store = build_token_store(StoragePolicy::Both, dir.path()).unwrap();
        store.put(sample_record("shared")).unwrap();

        let barrier = Arc::new(Barrier::new(REQUESTS));
        let handles: Vec<_> = (0..REQUESTS)
            .map(|_| {
                let store = Arc::clone(&store);
                let barrier = Arc::clone(&barrier);
                thread::spawn(move || {
                    barrier.wait();
                    store.try_consume_request("shared")
                })
            })
            .collect();

        for handle in handles {
            assert!(handle.join().unwrap().unwrap());
        }

        let text = TextTokenStore::open(dir.path().join("tokens.lino")).unwrap();
        let binary = BinaryTokenStore::open(dir.path().join("tokens.bin")).unwrap();
        assert_eq!(
            text.get("shared").unwrap().unwrap().used_requests,
            REQUESTS as u64
        );
        assert_eq!(
            binary.get("shared").unwrap().unwrap().used_requests,
            REQUESTS as u64
        );
    }

    #[test]
    fn revoke_marks_record() {
        let s = MemoryTokenStore::new();
        s.put(sample_record("a")).unwrap();
        assert!(s.revoke("a").unwrap());
        assert!(s.get("a").unwrap().unwrap().revoked);
        // second revoke is a no-op
        assert!(!s.revoke("a").unwrap());
        // unknown id returns false
        assert!(!s.revoke("missing").unwrap());
    }

    #[test]
    fn build_token_store_dispatches_correctly() {
        let dir = tempdir().unwrap();
        let mem = build_token_store(StoragePolicy::Memory, dir.path()).unwrap();
        mem.put(sample_record("m")).unwrap();
        assert!(mem.get("m").unwrap().is_some());

        let text = build_token_store(StoragePolicy::Text, dir.path()).unwrap();
        text.put(sample_record("t")).unwrap();
        assert!(dir.path().join("tokens.lino").exists());

        let bin = build_token_store(StoragePolicy::Binary, dir.path()).unwrap();
        bin.put(sample_record("b")).unwrap();
        assert!(dir.path().join("tokens.bin").exists());

        let dual = build_token_store(StoragePolicy::Both, dir.path()).unwrap();
        dual.put(sample_record("d")).unwrap();
        // both files updated
        let text_contents = std::fs::read_to_string(dir.path().join("tokens.lino")).unwrap();
        assert_eq!(associative::decode_text(&text_contents).unwrap()[0].id, "d");
    }

    #[test]
    fn lino_codec_handles_special_chars() {
        let rec = TokenRecord {
            id: "id1".into(),
            label: "with \"quote\" and \\ backslash and\nnewline".into(),
            issued_at: 1,
            expires_at: 2,
            revoked: true,
            account: None,
            max_requests: Some(100),
            used_requests: 7,
            scope: crate::token::ADMIN_SCOPE.to_string(),
        };
        let s = associative::encode_text(std::iter::once(&rec));
        let parsed = associative::decode_text(&s).unwrap();
        assert_eq!(parsed.len(), 1);
        assert_eq!(parsed[0], rec);
    }
}