char/avril

git clone https://git.t4t.associates/char/avril

Charlotte Somadd `avril:kv` module78f47482026-10-21

use std::collections::BTreeMap;
use std::io;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use avril_core::{Extension, Op};
use avril_kv::{Access, Decision, Store};
use sha2::{Digest, Sha256};

pub const QUOTA_BYTES: u64 = 10 * 1024 * 1024;

// Single keyspace holding raw UTF-8 bytes. Files in the pre-unification
// layout (old "kv"/"meta" tables) are not read; first write creates the
// keyspace as usual.
const KEYSPACE: &str = "webstorage";
const USED_KEY: &str = "webstorage/bytes_used";

fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
    mutex.lock().unwrap_or_else(|error| error.into_inner())
}

fn field(hasher: &mut Sha256, value: &[u8]) {
    hasher.update((value.len() as u64).to_be_bytes());
    hasher.update(value);
}

/// Persistent localStorage for one (config, entry point) namespace.
///
/// Constructing this performs no filesystem work; the database file is only
/// opened, inside a single operation, when storage is actually read or written.
#[derive(Clone)]
pub struct Persistent {
    path: Arc<PathBuf>,
    // Serializes operations within this process so our own isolates queue
    // here instead of spinning on the cross-process file lock.
    gate: Arc<Mutex<()>>,
}

impl Persistent {
    pub fn new(cache_dir: impl AsRef<Path>, config: &Path, entry: &str) -> Self {
        let mut hasher = Sha256::new();
        hasher.update(b"avril-webstorage");
        field(&mut hasher, b"v1");
        field(&mut hasher, config.to_string_lossy().as_bytes());
        field(&mut hasher, entry.as_bytes());
        let digest = format!("{:x}", hasher.finalize());
        Self {
            path: Arc::new(
                cache_dir
                    .as_ref()
                    .join("webstorage")
                    .join("v1")
                    .join(format!("{digest}.redb")),
            ),
            gate: Arc::new(Mutex::new(())),
        }
    }

    pub fn path(&self) -> &Path {
        &self.path
    }

    // Another process holds the file for the duration of one of its
    // operations. The lock is OS-held, so a crashed holder releases it and
    // waiting cannot deadlock on stale state.
    fn wait(&self, delay: &mut Duration, started: Instant, noticed: &mut bool) {
        if !*noticed && started.elapsed() >= Duration::from_secs(2) {
            eprintln!(
                "localStorage: waiting for {} (held by another process)",
                self.path.display()
            );
            *noticed = true;
        }
        let jitter = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|now| now.subsec_nanos() as u64 % (delay.as_millis() as u64 + 1))
            .unwrap_or(0);
        std::thread::sleep(*delay + Duration::from_millis(jitter));
        *delay = (*delay * 2).min(Duration::from_millis(50));
    }

    fn is_empty_file(&self) -> bool {
        std::fs::metadata(&*self.path).is_ok_and(|metadata| metadata.len() == 0)
    }

    // Runs one whole operation, waiting out cross-process file locks instead
    // of failing on them. Every operation is idempotent, so retrying after
    // a wait is always safe.
    fn run<R>(&self, mut op: impl FnMut() -> io::Result<R>) -> io::Result<R> {
        let _gate = lock(&self.gate);
        let mut delay = Duration::from_millis(1);
        let started = Instant::now();
        let mut noticed = false;
        loop {
            match op() {
                Err(error) if error.kind() == io::ErrorKind::AlreadyExists => {
                    self.wait(&mut delay, started, &mut noticed)
                }
                result => return result,
            }
        }
    }

    fn missing(&self) -> bool {
        !self.path.exists() || self.is_empty_file()
    }

    fn store(&self, create: bool) -> io::Result<Option<Store>> {
        match Store::open(&self.path, Access::Shared, create) {
            Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(None),
            store => store.map(Some),
        }
    }

    // Reads go through a read-only handle so merely reading storage never
    // modifies the database file (a writable open could repair or touch
    // file metadata even when it changes no entries).
    fn read_store(&self) -> io::Result<Option<Store>> {
        if self.missing() {
            return Ok(None);
        }
        Store::open_read_only_repairing(&self.path).map(Some)
    }

    fn get(&self, key: &str) -> io::Result<Option<String>> {
        self.run(|| {
            let Some(store) = self.read_store()? else {
                return Ok(None);
            };
            let value = store.read_tx(|tx| tx.get(KEYSPACE, key.as_bytes()))?;
            match value {
                Some((bytes, _)) => String::from_utf8(bytes).map(Some).map_err(io::Error::other),
                None => Ok(None),
            }
        })
    }

    fn keys(&self) -> io::Result<Vec<String>> {
        self.run(|| {
            let Some(store) = self.read_store()? else {
                return Ok(Vec::new());
            };
            let entries = store.read_tx(|tx| tx.entries(KEYSPACE))?;
            entries
                .into_iter()
                .map(|(key, _, _)| String::from_utf8(key).map_err(io::Error::other))
                .collect()
        })
    }

    fn key_at(&self, index: u64) -> io::Result<Option<String>> {
        self.run(|| {
            let Some(store) = self.read_store()? else {
                return Ok(None);
            };
            let entries = store.read_tx(|tx| tx.entries(KEYSPACE))?;
            match entries
                .into_iter()
                .map(|(key, _, _)| key)
                .nth(usize::try_from(index).unwrap_or(usize::MAX))
            {
                Some(key) => String::from_utf8(key).map(Some).map_err(io::Error::other),
                None => Ok(None),
            }
        })
    }

    fn length(&self) -> io::Result<u64> {
        self.run(|| {
            let Some(store) = self.read_store()? else {
                return Ok(0);
            };
            store.read_tx(|tx| tx.count(KEYSPACE))
        })
    }

    fn set(&self, key: &str, value: &str) -> io::Result<bool> {
        let size = (key.len() + value.len()) as u64;
        if size > QUOTA_BYTES {
            return Ok(false);
        }
        self.run(|| {
            if let Some(parent) = self.path.parent() {
                std::fs::create_dir_all(parent)?;
            }
            let store = Store::open(&self.path, Access::Shared, true)?;
            store.write_tx(|tx| {
                let used = tx.get_meta(USED_KEY)?.unwrap_or(0);
                let replaced = tx
                    .get(KEYSPACE, key.as_bytes())?
                    .map_or(0, |(old, _)| (key.len() + old.len()) as u64);
                // A drifted counter must not lock writes out forever: rebuild
                // it from the table itself rather than failing.
                let used = if used < replaced {
                    tx.entries(KEYSPACE)?
                        .iter()
                        .map(|(key, value, _)| (key.len() + value.len()) as u64)
                        .sum()
                } else {
                    used
                };
                let total = used
                    .checked_sub(replaced)
                    .and_then(|total| total.checked_add(size))
                    .ok_or_else(|| io::Error::other("Storage size overflow"))?;
                if total <= QUOTA_BYTES {
                    tx.put(KEYSPACE, key.as_bytes(), value.as_bytes())?;
                    tx.put_meta(USED_KEY, total)?;
                    Ok(Decision::Commit(true))
                } else {
                    Ok(Decision::Abort(false))
                }
            })
        })
    }

    fn remove(&self, key: &str) -> io::Result<()> {
        self.run(|| {
            let Some(store) = self.store(false)? else {
                return Ok(());
            };
            store.write_tx(|tx| match tx.get(KEYSPACE, key.as_bytes())? {
                Some((old, _)) => {
                    tx.delete(KEYSPACE, key.as_bytes())?;
                    let used = tx.get_meta(USED_KEY)?.unwrap_or(0);
                    tx.put_meta(
                        USED_KEY,
                        used.saturating_sub((key.len() + old.len()) as u64),
                    )?;
                    Ok(Decision::Commit(()))
                }
                None => Ok(Decision::Abort(())),
            })
        })
    }

    fn clear(&self) -> io::Result<()> {
        self.run(|| {
            let Some(store) = self.store(false)? else {
                return Ok(());
            };
            store.write_tx(|tx| {
                tx.clear(KEYSPACE)?;
                tx.put_meta(USED_KEY, 0)?;
                Ok(Decision::Commit(()))
            })
        })
    }
}

#[derive(Default)]
struct Session {
    entries: BTreeMap<String, String>,
    used: u64,
}

struct Stores {
    local: Persistent,
    session: Mutex<Session>,
}

impl Stores {
    fn context<T>(&self, result: io::Result<T>) -> io::Result<T> {
        result.map_err(|error| {
            let message = error.to_string();
            let lower = message.to_lowercase();
            // The delete-to-reset hint only helps when the file itself is
            // unusable, not for e.g. permission errors.
            let hint = if ["magic number", "corrupt", "file format", "upgrade"]
                .iter()
                .any(|marker| lower.contains(marker))
            {
                "; delete the file to reset storage"
            } else {
                ""
            };
            io::Error::other(format!(
                "localStorage ({}): {message}{hint}",
                self.local.path().display()
            ))
        })
    }

    fn get(&self, persistent: bool, key: &str) -> io::Result<Option<String>> {
        if persistent {
            self.context(self.local.get(key))
        } else {
            Ok(lock(&self.session).entries.get(key).cloned())
        }
    }

    fn keys(&self, persistent: bool) -> io::Result<Vec<String>> {
        if persistent {
            self.context(self.local.keys())
        } else {
            Ok(lock(&self.session).entries.keys().cloned().collect())
        }
    }

    fn key_at(&self, persistent: bool, index: u64) -> io::Result<Option<String>> {
        if persistent {
            self.context(self.local.key_at(index))
        } else {
            Ok(lock(&self.session)
                .entries
                .keys()
                .nth(usize::try_from(index).unwrap_or(usize::MAX))
                .cloned())
        }
    }

    fn length(&self, persistent: bool) -> io::Result<u64> {
        if persistent {
            self.context(self.local.length())
        } else {
            Ok(lock(&self.session).entries.len() as u64)
        }
    }

    fn set(&self, persistent: bool, key: &str, value: &str) -> io::Result<bool> {
        if persistent {
            return self.context(self.local.set(key, value));
        }
        let mut session = lock(&self.session);
        let size = (key.len() + value.len()) as u64;
        let replaced = session
            .entries
            .get(key)
            .map_or(0, |old| (key.len() + old.len()) as u64);
        let Some(total) = session
            .used
            .checked_sub(replaced)
            .and_then(|total| total.checked_add(size))
        else {
            return Err(io::Error::other("Storage size overflow"));
        };
        if total > QUOTA_BYTES {
            return Ok(false);
        }
        session.entries.insert(key.to_owned(), value.to_owned());
        session.used = total;
        Ok(true)
    }

    fn remove(&self, persistent: bool, key: &str) -> io::Result<()> {
        if persistent {
            self.context(self.local.remove(key))
        } else {
            let mut session = lock(&self.session);
            if let Some(old) = session.entries.remove(key) {
                session.used = session.used.saturating_sub((key.len() + old.len()) as u64);
            }
            Ok(())
        }
    }

    fn clear(&self, persistent: bool) -> io::Result<()> {
        if persistent {
            self.context(self.local.clear())
        } else {
            let mut session = lock(&self.session);
            session.entries.clear();
            session.used = 0;
            Ok(())
        }
    }
}

pub fn init(local: Persistent) -> Extension {
    let stores = Arc::new(Stores {
        local,
        session: Mutex::new(Session::default()),
    });
    Extension {
        name: "ext:webstorage",
        dependencies: &["ext:web"],
        ops: vec![
            Op::new("op_webstorage_get", {
                let stores = stores.clone();
                move |(persistent, key): (bool, String)| stores.get(persistent, &key)
            }),
            Op::new("op_webstorage_keys", {
                let stores = stores.clone();
                move |persistent: bool| stores.keys(persistent)
            }),
            Op::new("op_webstorage_key", {
                let stores = stores.clone();
                move |(persistent, index): (bool, u64)| stores.key_at(persistent, index)
            }),
            Op::new("op_webstorage_length", {
                let stores = stores.clone();
                move |persistent: bool| stores.length(persistent)
            }),
            Op::new("op_webstorage_set", {
                let stores = stores.clone();
                move |(persistent, key, value): (bool, String, String)| {
                    stores.set(persistent, &key, &value)
                }
            }),
            Op::new("op_webstorage_remove", {
                let stores = stores.clone();
                move |(persistent, key): (bool, String)| stores.remove(persistent, &key)
            }),
            Op::new("op_webstorage_clear", {
                let stores = stores.clone();
                move |persistent: bool| stores.clear(persistent)
            }),
        ],
        modules: &[
            (
                "ext:webstorage/main.js",
                include_str!(concat!(env!("OUT_DIR"), "/main.js")),
            ),
            (
                "ext:webstorage/storage.js",
                include_str!(concat!(env!("OUT_DIR"), "/storage.js")),
            ),
        ],
        entry_point: Some("ext:webstorage/main.js"),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use redb::{Database, TableDefinition};

    fn stores(cache: &Path) -> Stores {
        Stores {
            local: Persistent::new(cache, Path::new("/project/avril.conf"), "/project/main.js"),
            session: Mutex::new(Session::default()),
        }
    }

    #[test]
    fn construction_and_reads_create_nothing() {
        let dir = tempfile::tempdir().unwrap();
        let stores = stores(dir.path());
        assert!(!stores.local.path().exists());
        assert_eq!(stores.get(true, "missing").unwrap(), None);
        assert_eq!(stores.keys(true).unwrap(), Vec::<String>::new());
        stores.remove(true, "missing").unwrap();
        stores.clear(true).unwrap();
        assert!(!stores.local.path().exists());
    }

    #[test]
    fn local_persists_across_instances_but_session_does_not() {
        let dir = tempfile::tempdir().unwrap();
        let first = stores(dir.path());
        assert!(first.set(true, "key", "value").unwrap());
        assert!(first.set(false, "key", "session").unwrap());
        let second = stores(dir.path());
        assert_eq!(second.get(true, "key").unwrap(), Some("value".to_owned()));
        assert_eq!(second.get(false, "key").unwrap(), None);
        assert_ne!(
            Persistent::new(
                dir.path(),
                Path::new("/other/avril.conf"),
                "/project/main.js"
            )
            .path(),
            second.local.path()
        );
        assert_ne!(
            Persistent::new(
                dir.path(),
                Path::new("/project/avril.conf"),
                "/other/main.js"
            )
            .path(),
            second.local.path()
        );
    }

    #[test]
    fn quota_counts_utf8_and_credits_replacements() {
        let dir = tempfile::tempdir().unwrap();
        let stores = stores(dir.path());
        // "é" is two UTF-8 bytes; key + value exactly fills the quota.
        let value = "x".repeat((QUOTA_BYTES - 2) as usize);
        assert!(stores.set(true, "é", &value).unwrap());
        assert!(!stores.set(true, "a", "").unwrap());
        assert_eq!(stores.get(true, "a").unwrap(), None);
        // Shrinking an existing entry succeeds and frees space.
        assert!(stores.set(true, "é", "").unwrap());
        assert!(stores.set(true, "a", "b").unwrap());
        stores.remove(true, "é").unwrap();
        stores.clear(true).unwrap();
        assert_eq!(stores.keys(true).unwrap(), Vec::<String>::new());
    }

    #[test]
    fn operations_wait_for_another_holders_handle() {
        let dir = tempfile::tempdir().unwrap();
        let stores = stores(dir.path());
        assert!(stores.set(true, "key", "before").unwrap());
        std::fs::create_dir_all(stores.local.path().parent().unwrap()).unwrap();
        let holder = Database::create(stores.local.path()).unwrap();
        let path = stores.local.path().to_owned();
        let cache = dir.path().to_owned();
        let (attempted, attempted_rx) = std::sync::mpsc::channel::<()>();
        let writer = std::thread::spawn(move || {
            let stores = Stores {
                local: Persistent::new(cache, Path::new("/project/avril.conf"), "/project/main.js"),
                session: Mutex::new(Session::default()),
            };
            assert_eq!(stores.local.path(), path);
            attempted.send(()).unwrap();
            stores.set(true, "key", "after")
        });
        attempted_rx.recv().unwrap();
        std::thread::sleep(Duration::from_millis(100));
        assert!(!writer.is_finished(), "writer did not wait for the holder");
        drop(holder);
        assert!(writer.join().unwrap().unwrap());
        assert_eq!(stores.get(true, "key").unwrap(), Some("after".to_owned()));
    }

    #[test]
    fn reads_repair_a_database_left_unclean() {
        // Child mode: commit a value through a held-open Store, then exit
        // without dropping it, leaving the file marked as needing repair.
        if let Ok(cache) = std::env::var("AVRIL_WEBSTORAGE_UNCLEAN_CACHE") {
            let stores = stores(Path::new(&cache));
            std::fs::create_dir_all(stores.local.path().parent().unwrap()).unwrap();
            let store = Store::open(stores.local.path(), Access::Exclusive, true).unwrap();
            store
                .write_tx(|tx| {
                    tx.put("webstorage", b"key", b"unclean")
                        .map(Decision::Commit)
                })
                .unwrap();
            std::process::exit(0);
        }
        let dir = tempfile::tempdir().unwrap();
        let status = std::process::Command::new(std::env::current_exe().unwrap())
            .args(["--exact", "tests::reads_repair_a_database_left_unclean"])
            .env("AVRIL_WEBSTORAGE_UNCLEAN_CACHE", dir.path())
            .stdout(std::process::Stdio::null())
            .status()
            .unwrap();
        assert!(status.success());
        let stores = stores(dir.path());
        assert_eq!(stores.get(true, "key").unwrap(), Some("unclean".to_owned()));
    }

    #[test]
    fn legacy_format_files_read_as_empty() {
        let dir = tempfile::tempdir().unwrap();
        let stores = stores(dir.path());
        std::fs::create_dir_all(stores.local.path().parent().unwrap()).unwrap();
        {
            let database = Database::create(stores.local.path()).unwrap();
            let transaction = database.begin_write().unwrap();
            {
                let mut table = transaction
                    .open_table(TableDefinition::<&str, &str>::new("kv"))
                    .unwrap();
                table.insert("old", "data").unwrap();
            }
            transaction.commit().unwrap();
        }
        // The pre-unification layout is not read: the file looks empty...
        assert_eq!(stores.get(true, "old").unwrap(), None);
        assert_eq!(stores.length(true).unwrap(), 0);
        assert_eq!(stores.keys(true).unwrap(), Vec::<String>::new());
        // ...and the database is fully usable from the new format onward.
        assert!(stores.set(true, "new", "value").unwrap());
        assert_eq!(stores.get(true, "new").unwrap(), Some("value".to_owned()));
        assert_eq!(stores.get(true, "old").unwrap(), None);
    }
}