Something went wrong. Try again.
a tiny pet with an even tinier brain, living in your browser
Something went wrong. Try again.
9.8 kB · 277 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278//! Private saves and atomic, epoch-fenced operations. Access keys stay hashed.use crate::simulation::HostedPet;use anyhow::{Context, Result, ensure};use rusqlite::{Connection, OptionalExtension, TransactionBehavior, params};use sha2::{Digest, Sha256};use std::path::Path;use subtle::ConstantTimeEq;
#[derive(Debug)]pub struct AccessDenied;impl std::fmt::Display for AccessDenied { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "invalid pet access key") }}impl std::error::Error for AccessDenied {}#[derive(Debug)]pub struct Conflict(pub &'static str);impl std::fmt::Display for Conflict { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}", self.0) }}impl std::error::Error for Conflict {}
pub struct Store { conn: Connection,}pub fn hash(text: &str) -> Vec<u8> { Sha256::digest(text.as_bytes()).to_vec()}fn authorized(expected: &[u8], key: &str) -> bool { key.len() == 64 && key.bytes().all(|b| b.is_ascii_hexdigit()) && bool::from(expected.ct_eq(&hash(key)))}fn load(conn: &Connection, id: &str) -> Result<Option<(Vec<u8>, HostedPet)>> { let raw: Option<(Vec<u8>, String)> = conn .query_row("SELECT key_hash, state FROM pets WHERE id=?1", [id], |r| { Ok((r.get(0)?, r.get(1)?)) }) .optional()?; raw.map(|(key, state)| { let row: HostedPet = serde_json::from_str(&state) .context("stored pet is unreadable; it has not been replaced")?; row.state.validate()?; ensure!( matches!(row.authority.as_str(), "local" | "remote"), "stored execution owner is invalid" ); Ok((key, row)) }) .transpose()}fn update(conn: &Connection, id: &str, row: &HostedPet) -> Result<()> { conn.execute( "UPDATE pets SET state=?2 WHERE id=?1", params![id, serde_json::to_string(row)?], )?; Ok(())}impl Store { pub fn open(path: &Path) -> Result<Self> { if let Some(parent) = path.parent().filter(|p| !p.as_os_str().is_empty()) { std::fs::create_dir_all(parent)?; } let conn = Connection::open(path)?; let version: u32 = conn.query_row("PRAGMA user_version", [], |r| r.get(0))?; ensure!( version <= 1, "database was created by a newer host; it has not been changed" ); if version == 0 { let tables: u32 = conn.query_row( "SELECT count(*) FROM sqlite_master WHERE type='table'", [], |r| r.get(0), )?; ensure!( tables == 0, "database has an unrecognized schema; choose a separate database" ); } conn.pragma_update(None, "journal_mode", "WAL")?; conn.pragma_update(None, "foreign_keys", true)?; conn.busy_timeout(std::time::Duration::from_secs(5))?; conn.execute_batch("CREATE TABLE IF NOT EXISTS pets(id TEXT PRIMARY KEY, key_hash BLOB NOT NULL, state TEXT NOT NULL); CREATE TABLE IF NOT EXISTS operations(pet_id TEXT NOT NULL REFERENCES pets(id) ON DELETE CASCADE, id TEXT NOT NULL, body_hash BLOB NOT NULL, created INTEGER NOT NULL, PRIMARY KEY(pet_id,id)); PRAGMA user_version=1;")?; #[cfg(unix)] { use std::os::unix::fs::PermissionsExt; for candidate in [ path.to_path_buf(), std::path::PathBuf::from(format!("{}-wal", path.display())), std::path::PathBuf::from(format!("{}-shm", path.display())), ] { if candidate.exists() { std::fs::set_permissions(candidate, std::fs::Permissions::from_mode(0o600))?; } } } Ok(Self { conn }) } pub fn create( &mut self, id: &str, key: &str, row: HostedPet, max_pets: usize, ) -> Result<HostedPet> { ensure!( uuid::Uuid::parse_str(id)?.to_string() == id, "invalid pet identifier" ); ensure!( key.len() == 64 && key.bytes().all(|b| b.is_ascii_hexdigit()), "invalid access key" ); row.state.validate()?; let tx = self .conn .transaction_with_behavior(TransactionBehavior::Immediate)?; if let Some((expected, existing)) = load(&tx, id)? { if !authorized(&expected, key) { return Err(AccessDenied.into()); } // Retrying creation must not replace a running pet or its save. tx.commit()?; return Ok(existing); } let count: i64 = tx.query_row("SELECT count(*) FROM pets", [], |r| r.get(0))?; ensure!( count < max_pets as i64, "this host has reached its pet limit" ); tx.execute( "INSERT INTO pets(id,key_hash,state) VALUES (?1,?2,?3)", params![id, hash(key), serde_json::to_string(&row)?], )?; tx.commit()?; Ok(row) } pub fn get(&self, id: &str, key: &str) -> Result<HostedPet> { let Some((expected, row)) = load(&self.conn, id)? else { return Err(AccessDenied.into()); }; if !authorized(&expected, key) { return Err(AccessDenied.into()); } Ok(row) } pub fn operation( &mut self, id: &str, key: &str, epoch: u64, operation: &str, body: &str, f: impl FnOnce(&mut HostedPet) -> Result<()>, ) -> Result<HostedPet> { ensure!( uuid::Uuid::parse_str(operation)?.to_string() == operation, "invalid operation identifier" ); let tx = self .conn .transaction_with_behavior(TransactionBehavior::Immediate)?; let Some((expected, mut row)) = load(&tx, id)? else { return Err(AccessDenied.into()); }; if !authorized(&expected, key) { return Err(AccessDenied.into()); } let old: Option<Vec<u8>> = tx .query_row( "SELECT body_hash FROM operations WHERE pet_id=?1 AND id=?2", params![id, operation], |r| r.get(0), ) .optional()?; if let Some(old) = old { if old != hash(body) { return Err(Conflict("operation ID was reused for different data").into()); } tx.commit()?; return Ok(row); } if row.epoch != epoch { return Err(Conflict("execution changed; refresh this pet before acting").into()); } f(&mut row)?; update(&tx, id, &row)?; tx.execute( "INSERT INTO operations(pet_id,id,body_hash,created) VALUES(?1,?2,?3,?4)", params![ id, operation, hash(body), chrono::Utc::now().timestamp_millis() ], )?; tx.commit()?; Ok(row) } pub fn delete(&mut self, id: &str, key: &str) -> Result<()> { self.get(id, key)?; self.conn.execute("DELETE FROM pets WHERE id=?1", [id])?; Ok(()) } /// Counts pets this host executes whose save selects another brain. The /// scheduler runs one model for every pet, so starting with any would be /// an implicit brain change. pub fn foreign_brains(&self, model_id: &str) -> Result<usize> { let mut s = self.conn.prepare("SELECT state FROM pets")?; let raw: Vec<String> = s .query_map([], |r| r.get(0))? .collect::<rusqlite::Result<_>>()?; let mut count = 0; for raw in raw { let row: HostedPet = serde_json::from_str(&raw).context("stored pet is unreadable")?; if row.authority == "remote" && row .state .hosting .brain_id .as_deref() .is_some_and(|id| id != model_id) { count += 1; } } Ok(count) } pub fn advance(&mut self, t: f64, seed: &str) -> Result<Vec<(String, HostedPet)>> { let tx = self .conn .transaction_with_behavior(TransactionBehavior::Immediate)?; let raw: Vec<(String, String)> = { let mut s = tx.prepare("SELECT id,state FROM pets")?; s.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))? .collect::<rusqlite::Result<_>>()? }; let mut rows = Vec::new(); for (id, raw) in raw { let mut row: HostedPet = serde_json::from_str(&raw).context("stored pet is unreadable")?; row.state.validate()?; if row.authority == "remote" && t - row.state.last_tick >= 1000.0 { row.advance(t, seed); update(&tx, &id, &row)?; } rows.push((id, row)); } tx.commit()?; Ok(rows) } pub fn commit( &mut self, id: &str, epoch: u64, context: u64, f: impl FnOnce(&mut HostedPet), ) -> Result<bool> { let tx = self .conn .transaction_with_behavior(TransactionBehavior::Immediate)?; let Some((_, mut row)) = load(&tx, id)? else { return Ok(false); }; if row.epoch != epoch || row.context_revision != context || row.authority != "remote" { return Ok(false); } f(&mut row); update(&tx, id, &row)?; tx.commit()?; Ok(true) }}