~ knotless Search/
↑↓ move↵ open⇧↵ all results

johnandersen777.bsky.social/atlibs

0
0
main
345 lines · 12.8 KB Raw History
//! The Node host's state dir (javascript/src/server createNodeContext): <state>/store.db rows
//! (ns, id, value JSON, revision) and secret settings sealed with <state>/secret.key (AES-256-GCM,
//! additional data "<plugin>:settings:<name>"), byte compatible, so the JS and Rust CLIs share
//! one saved session, attest keys, did:keys and managed_records ledger.

use std::fs::{self, OpenOptions};
use std::io::{ErrorKind, Write};
use std::os::unix::fs::{DirBuilderExt, OpenOptionsExt};
use std::path::Path;
use std::sync::{Arc, Mutex};

use aes_gcm::aead::{Aead, KeyInit, Payload};
use aes_gcm::{Aes256Gcm, Nonce};
use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use rand_core::{OsRng, RngCore};
use rusqlite::{Connection, OptionalExtension, params};
use serde_json::{Value, json};

use crate::BoxFut;
use crate::run::Ledger;

pub fn random_bytes<const N: usize>() -> [u8; N] {
    let mut out = [0; N];
    OsRng.fill_bytes(&mut out);
    out
}

/// crypto.randomUUID().
pub fn uuid() -> String {
    let mut b = random_bytes::<16>();
    b[6] = (b[6] & 0x0f) | 0x40;
    b[8] = (b[8] & 0x3f) | 0x80;
    let h: String = b.iter().map(|x| format!("{x:02x}")).collect();
    format!(
        "{}-{}-{}-{}-{}",
        &h[..8],
        &h[8..12],
        &h[12..16],
        &h[16..20],
        &h[20..]
    )
}

/// Base64 as atob / Uint8Array.fromBase64 read it: either alphabet, padding optional.
pub fn unbase64(text: &str) -> Result<Vec<u8>, String> {
    let clean: String = text
        .chars()
        .filter(|c| !c.is_whitespace() && *c != '=')
        .map(|c| match c {
            '-' => '+',
            '_' => '/',
            c => c,
        })
        .collect();
    base64::engine::general_purpose::STANDARD_NO_PAD
        .decode(clean)
        .map_err(|e| format!("invalid base64: {e}"))
}

#[derive(Clone)]
pub struct Store {
    db: Arc<Mutex<Connection>>,
    key: [u8; 32],
    pub id: String,
}

impl Store {
    pub fn open(state_dir: &Path, id: &str) -> Result<Self, String> {
        fs::DirBuilder::new()
            .recursive(true)
            .mode(0o700)
            .create(state_dir)
            .map_err(|e| format!("{}: {e}", state_dir.display()))?;
        let db = Connection::open(state_dir.join("store.db")).map_err(|e| e.to_string())?;
        db.busy_timeout(std::time::Duration::from_secs(5))
            .map_err(|e| e.to_string())?;
        db.execute("CREATE TABLE IF NOT EXISTS rows (ns TEXT NOT NULL, id TEXT NOT NULL, value TEXT NOT NULL, revision TEXT NOT NULL, PRIMARY KEY (ns, id))", []).map_err(|e| e.to_string())?;
        let key = secret_key(&state_dir.join("secret.key"))?;
        Ok(Self {
            db: Arc::new(Mutex::new(db)),
            key,
            id: id.into(),
        })
    }

    pub fn get(&self, ns: &str, id: &str) -> Result<Option<Value>, String> {
        let db = self.db.lock().unwrap();
        let found: Option<String> = db
            .query_row(
                "SELECT value FROM rows WHERE ns = ? AND id = ?",
                params![ns, id],
                |r| r.get(0),
            )
            .optional()
            .map_err(|e| e.to_string())?;
        found
            .map(|v| serde_json::from_str(&v).map_err(|e| e.to_string()))
            .transpose()
    }

    pub fn put(&self, ns: &str, id: &str, value: &Value) -> Result<(), String> {
        let db = self.db.lock().unwrap();
        db.execute(
			"INSERT INTO rows (ns, id, value, revision) VALUES (?, ?, ?, ?) ON CONFLICT DO UPDATE SET value = excluded.value, revision = excluded.revision",
			params![ns, id, value.to_string(), uuid()],
		)
		.map(drop)
		.map_err(|e| e.to_string())
    }

    pub fn delete(&self, ns: &str, id: &str) -> Result<bool, String> {
        let db = self.db.lock().unwrap();
        db.execute("DELETE FROM rows WHERE ns = ? AND id = ?", params![ns, id])
            .map(|n| n > 0)
            .map_err(|e| e.to_string())
    }

    fn aad(&self, name: &str) -> String {
        format!("{}:settings:{name}", self.id)
    }

    /// ctx.settings.get for a secret setting (plain values pass through, as settingsAccess does).
    pub fn setting(&self, name: &str) -> Result<Option<Value>, String> {
        let Some(value) = self.get(&format!("{}:settings", self.id), name)? else {
            return Ok(None);
        };
        let Some(secret) = value.get("$secret") else {
            return Ok(Some(value));
        };
        let field = |k: &str| {
            secret
                .get(k)
                .and_then(Value::as_str)
                .ok_or_else(|| format!("setting {name}: bad $secret"))
                .and_then(unbase64)
        };
        let (iv, data) = (field("iv")?, field("data")?);
        let cipher = Aes256Gcm::new_from_slice(&self.key).expect("32-byte key");
        let iv: [u8; 12] = iv
            .try_into()
            .map_err(|_| format!("setting {name}: bad $secret iv"))?;
        let plain = cipher
            .decrypt(
                &Nonce::from(iv),
                Payload {
                    msg: &data,
                    aad: self.aad(name).as_bytes(),
                },
            )
            .map_err(|_| format!("setting {name}: cannot open (another secret.key?)"))?;
        serde_json::from_slice(&plain)
            .map(Some)
            .map_err(|e| e.to_string())
    }

    /// ctx.settings.set for a secret setting: sealed with secret.key.
    pub fn set_setting(&self, name: &str, value: &Value) -> Result<(), String> {
        let iv = random_bytes::<12>();
        let cipher = Aes256Gcm::new_from_slice(&self.key).expect("32-byte key");
        let data = cipher
            .encrypt(
                &Nonce::from(iv),
                Payload {
                    msg: value.to_string().as_bytes(),
                    aad: self.aad(name).as_bytes(),
                },
            )
            .map_err(|e| e.to_string())?;
        self.put(
            &format!("{}:settings", self.id),
            name,
            &json!({ "$secret": { "iv": STANDARD.encode(iv), "data": STANDARD.encode(data) } }),
        )
    }

    /// Take the lease `name` now if nobody holds it (or its holder's lease ran out), held until the guard
    /// drops: the lock rows every process sharing the state dir sees, as javascript/src/server
    /// createNodeContext's ctx.lock and ctx.lease take them ("<id>:locks" / name -> { owner, until }, in
    /// BEGIN IMMEDIATE). None while another task or process holds it; one that died loses it after `ttl`.
    pub fn try_lock(&self, name: &str, ttl: std::time::Duration) -> Result<Option<Lease>, String> {
        let ns = format!("{}:locks", self.id);
        let mut db = self.db.lock().unwrap();
        let tx = db
            .transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)
            .map_err(|e| e.to_string())?;
        let held: Option<String> = tx
            .query_row("SELECT value FROM rows WHERE ns = ? AND id = ?", params![ns, name], |r| r.get(0))
            .optional()
            .map_err(|e| e.to_string())?;
        let now = crate::now_ms();
        let until = held.and_then(|v| serde_json::from_str::<Value>(&v).ok()?.get("until")?.as_f64());
        if until.is_some_and(|until| until > now as f64) {
            return Ok(None);
        }
        let revision = uuid();
        tx.execute(
            "INSERT INTO rows (ns, id, value, revision) VALUES (?, ?, ?, ?) ON CONFLICT DO UPDATE SET value = excluded.value, revision = excluded.revision",
            params![ns, name, json!({ "owner": revision, "until": now + ttl.as_millis() as u64 }).to_string(), revision],
        )
        .map_err(|e| e.to_string())?;
        tx.commit().map_err(|e| e.to_string())?;
        Ok(Some(Lease { store: self.clone(), ns, name: name.into(), revision }))
    }

    /// Wait for the lease `name` (polling while another task or process holds it): a lock across every
    /// process sharing the state dir, held until the guard drops.
    pub async fn lock(&self, name: &str, ttl: std::time::Duration) -> Result<Lease, String> {
        loop {
            if let Some(lease) = self.try_lock(name, ttl)? {
                return Ok(lease);
            }
            tokio::time::sleep(std::time::Duration::from_millis(25)).await;
        }
    }

    /// The did:key sealed under `name` in the "didkeys" setting, made and saved on first use. The
    /// read-modify-write runs under the "didkeys" lease, so runs and processes that first use different
    /// modules at once keep every key (javascript/src/server/cli.mjs takes the same lease).
    pub async fn didkey(&self, name: &str) -> Result<crate::crypto::DidKey, String> {
        let _lease = self.lock("didkeys", std::time::Duration::from_secs(10)).await?;
        let mut keys = self.setting("didkeys")?.unwrap_or_else(|| json!({}));
        let key = crate::crypto::DidKey::new(keys.get(name))?;
        if keys.get(name).is_none() {
            keys[name] = key.jwk.clone();
            self.set_setting("didkeys", &keys)?;
        }
        Ok(key)
    }

    /// ctx.storage.<name> as the run-module ledger.
    pub fn ledger(&self, name: &str) -> Arc<dyn Ledger> {
        Arc::new(Collection {
            store: self.clone(),
            ns: format!("{}:storage:{name}", self.id),
        })
    }
}

/// A held Store::lock or try_lock; dropping it gives the lease back.
pub struct Lease {
    store: Store,
    ns: String,
    name: String,
    revision: String,
}

impl Lease {
    /// Extend it by ttl from now (a long holder renews before it runs out); false once it is no longer ours.
    pub fn renew(&self, ttl: std::time::Duration) -> Result<bool, String> {
        let value = json!({ "owner": self.revision, "until": crate::now_ms() + ttl.as_millis() as u64 }).to_string();
        let db = self.store.db.lock().unwrap();
        db.execute("UPDATE rows SET value = ? WHERE ns = ? AND id = ? AND revision = ?", params![value, self.ns, self.name, self.revision])
            .map(|n| n > 0)
            .map_err(|e| e.to_string())
    }
}

impl Drop for Lease {
    fn drop(&mut self) {
        let db = self.store.db.lock().unwrap();
        let _ = db.execute("DELETE FROM rows WHERE ns = ? AND id = ? AND revision = ?", params![self.ns, self.name, self.revision]);
    }
}

// 32 bytes in <state>/secret.key (0600), created on first use.
fn secret_key(path: &Path) -> Result<[u8; 32], String> {
    let raw = match fs::read(path) {
        Ok(raw) => raw,
        Err(e) if e.kind() == ErrorKind::NotFound => {
            let fresh = random_bytes::<32>();
            match OpenOptions::new()
                .write(true)
                .create_new(true)
                .mode(0o600)
                .open(path)
            {
                Ok(mut f) => f
                    .write_all(&fresh)
                    .map(|_| fresh.to_vec())
                    .map_err(|e| e.to_string())?,
                Err(e) if e.kind() == ErrorKind::AlreadyExists => {
                    fs::read(path).map_err(|e| e.to_string())?
                }
                Err(e) => return Err(e.to_string()),
            }
        }
        Err(e) => return Err(e.to_string()),
    };
    raw.try_into()
        .map_err(|_| format!("{} must hold 32 bytes", path.display()))
}

struct Collection {
    store: Store,
    ns: String,
}

impl Ledger for Collection {
    fn get(&self, uri: &str) -> BoxFut<'_, Result<Option<Value>, String>> {
        let found = self
            .store
            .get(&self.ns, uri)
            .map(|v| v.filter(|v| !v.is_null()));
        Box::pin(async move { found })
    }
    fn put(&self, uri: String, value: Value) -> BoxFut<'_, Result<(), String>> {
        let done = self.store.put(&self.ns, &uri, &value);
        Box::pin(async move { done })
    }
    fn delete(&self, uri: String) -> BoxFut<'_, Result<(), String>> {
        let done = self.store.delete(&self.ns, &uri).map(drop);
        Box::pin(async move { done })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn sealed_settings_round_trip_and_stay_sealed_on_disk() {
        let dir = std::env::temp_dir().join(format!("atlibs-host-store-{}", uuid()));
        let store = Store::open(&dir, "atlibs-server").unwrap();
        store
            .set_setting("session", &json!({ "accessJwt": "secret-token" }))
            .unwrap();
        assert_eq!(
            store.setting("session").unwrap(),
            Some(json!({ "accessJwt": "secret-token" }))
        );
        let raw = store
            .get("atlibs-server:settings", "session")
            .unwrap()
            .unwrap();
        assert!(raw["$secret"]["iv"].is_string() && !raw.to_string().contains("secret-token"));
        let again = Store::open(&dir, "atlibs-server").unwrap();
        assert_eq!(
            again.setting("session").unwrap(),
            Some(json!({ "accessJwt": "secret-token" }))
        );
        assert_eq!(fs::metadata(dir.join("secret.key")).unwrap().len(), 32);
        fs::remove_dir_all(dir).unwrap();
    }
}