// SPDX-License-Identifier: AGPL-3.0-or-later use crate::model::{ CapacityBudget, MAX_ACCOUNT_BYTES, MAX_ACCOUNTS, MAX_INSTANCE_BYTES, MAX_KNOWN_INSTANCE_BYTES, MAX_KNOWN_INSTANCES, StoreError, WriteRefusal, check_lookup_key, }; use rusqlite::{Connection, OptionalExtension, Transaction, TransactionBehavior, params}; use serde::{Deserialize, Serialize}; use serde_json::Value; use serde_json::value::RawValue; #[derive(Deserialize, Serialize)] pub struct AccountEntry { #[serde(rename = "storageKey")] pub storage_key: String, pub record: Box, } #[derive(Deserialize)] pub struct AccountCompareAndSwap { pub expected: AccountEntry, pub replacement: AccountEntry, } #[derive(Deserialize, Serialize)] pub struct KnownInstanceEntry { pub key: String, pub record: Box, } pub enum AccountWrite { Written { instance_usable: bool }, Refused(WriteRefusal), } struct AccountFacts { user_id: String, last_active: f64, instance_usable: bool, } fn account_record_fields( record: &RawValue, ) -> Result, WriteRefusal> { match serde_json::from_str::(record.get()) { Ok(Value::Object(fields)) => Ok(fields), _ => Err(WriteRefusal::Invalid( "account record must be a JSON object", )), } } fn account_facts(record: &RawValue) -> Result { let fields = account_record_fields(record)?; Ok(AccountFacts { user_id: fields .get("userId") .and_then(Value::as_str) .unwrap_or_default() .to_owned(), last_active: numeric_field(&fields, "lastActive"), instance_usable: instance_is_usable(&fields), }) } fn numeric_field(fields: &serde_json::Map, name: &str) -> f64 { fields .get(name) .and_then(Value::as_f64) .filter(|value| value.is_finite()) .unwrap_or_default() } fn instance_is_usable(fields: &serde_json::Map) -> bool { let Some(Value::Object(instance)) = fields.get("instance") else { return false; }; let endpoint_is_present = instance .get("apiEndpoint") .and_then(Value::as_str) .is_some_and(|endpoint| !endpoint.is_empty()); endpoint_is_present && serde_json::to_string(instance) .is_ok_and(|serialized| serialized.len() <= MAX_INSTANCE_BYTES) } pub fn account_budget(conn: &Connection) -> Result { CapacityBudget::read( conn, "accounts", "SELECT COUNT(*), COALESCE(SUM(length(CAST(storage_key AS BLOB)) + length(CAST(record_json AS BLOB))), 0) FROM accounts", [], MAX_ACCOUNTS, MAX_ACCOUNT_BYTES, ) } pub fn write_account( conn: &Connection, budget: &mut CapacityBudget, entry: &AccountEntry, ) -> Result { if let Some(refusal) = check_lookup_key(&entry.storage_key, "account storageKey is unusable") { return Ok(AccountWrite::Refused(refusal)); } let facts = match account_facts(&entry.record) { Ok(facts) => facts, Err(refusal) => return Ok(AccountWrite::Refused(refusal)), }; let incoming = (entry.storage_key.len() + entry.record.get().len()) as i64; let replaces = stored_bytes( conn, "SELECT length(CAST(storage_key AS BLOB)) + length(CAST(record_json AS BLOB)) FROM accounts WHERE storage_key = ?1", params![entry.storage_key], )?; if let Err(refusal) = budget.check(replaces, incoming) { return Ok(AccountWrite::Refused(refusal.into())); } budget.commit(replaces, incoming); conn.execute( "INSERT INTO accounts(storage_key, user_id, last_active, record_json) VALUES (?1, ?2, ?3, ?4) ON CONFLICT(storage_key) DO UPDATE SET user_id = excluded.user_id, last_active = excluded.last_active, record_json = excluded.record_json", params![ entry.storage_key, facts.user_id, facts.last_active, entry.record.get() ], )?; Ok(AccountWrite::Written { instance_usable: facts.instance_usable, }) } pub fn upsert_account(conn: &Connection, entry: &AccountEntry) -> Result<(), StoreError> { let mut budget = account_budget(conn)?; match write_account(conn, &mut budget, entry)? { AccountWrite::Written { .. } => Ok(()), AccountWrite::Refused(refusal) => Err(refusal.into()), } } pub fn compare_and_swap_account( conn: &mut Connection, request: &AccountCompareAndSwap, ) -> Result { if request.expected.storage_key != request.replacement.storage_key { return Err(StoreError::Refused( "account compare-and-swap records must use the same storageKey", )); } if let Some(refusal) = check_lookup_key( &request.expected.storage_key, "account storageKey is unusable", ) { return Err(refusal.into()); } let expected = account_record_fields(&request.expected.record)?; let transaction = conn.transaction_with_behavior(TransactionBehavior::Immediate)?; let stored: Option = transaction .query_row( "SELECT record_json FROM accounts WHERE storage_key = ?1", params![request.expected.storage_key], |row| row.get(0), ) .optional()?; let Some(stored) = stored else { transaction.commit()?; return Ok(false); }; let stored = RawValue::from_string(stored)?; let stored = account_record_fields(&stored)?; if stored != expected { transaction.commit()?; return Ok(false); } let mut budget = account_budget(&transaction)?; match write_account(&transaction, &mut budget, &request.replacement)? { AccountWrite::Written { .. } => {} AccountWrite::Refused(refusal) => return Err(refusal.into()), } transaction.commit()?; Ok(true) } pub fn all_accounts(conn: &Connection) -> Result, StoreError> { let mut statement = conn.prepare( "SELECT storage_key, record_json FROM accounts ORDER BY last_active DESC, storage_key ASC", )?; let mut rows = statement.query([])?; let mut accounts = Vec::new(); while let Some(row) = rows.next()? { accounts.push(AccountEntry { storage_key: row.get(0)?, record: RawValue::from_string(row.get(1)?)?, }); } Ok(accounts) } pub fn account(conn: &Connection, storage_key: &str) -> Result, StoreError> { let record: Option = conn .query_row( "SELECT record_json FROM accounts WHERE storage_key = ?1", params![storage_key], |row| row.get(0), ) .optional()?; record .map(|record| { Ok(AccountEntry { storage_key: storage_key.to_owned(), record: RawValue::from_string(record)?, }) }) .transpose() } pub fn delete_account_rows(conn: &Connection, storage_key: &str) -> Result<(), StoreError> { conn.execute( "DELETE FROM accounts WHERE storage_key = ?1", params![storage_key], )?; conn.execute( "DELETE FROM scoped_storage WHERE scope = ?1", params![storage_key], )?; Ok(()) } pub fn delete_account(conn: &Connection, storage_key: &str) -> Result<(), StoreError> { let transaction = Transaction::new_unchecked(conn, TransactionBehavior::Immediate)?; delete_account_rows(&transaction, storage_key)?; transaction.commit()?; Ok(()) } pub fn known_instance_budget(conn: &Connection) -> Result { CapacityBudget::read( conn, "known instances", "SELECT COUNT(*), COALESCE(SUM(length(CAST(key AS BLOB)) + length(CAST(record_json AS BLOB))), 0) FROM known_instances", [], MAX_KNOWN_INSTANCES, MAX_KNOWN_INSTANCE_BYTES, ) } pub fn upsert_known_instance( conn: &Connection, entry: &KnownInstanceEntry, ) -> Result<(), StoreError> { if let Some(refusal) = check_lookup_key(&entry.key, "known instance key is unusable") { return Err(refusal.into()); } let Ok(Value::Object(fields)) = serde_json::from_str::(entry.record.get()) else { return Err(WriteRefusal::Invalid("known instance record must be a JSON object").into()); }; let budget = known_instance_budget(conn)?; let incoming = (entry.key.len() + entry.record.get().len()) as i64; let replaces = stored_bytes( conn, "SELECT length(CAST(key AS BLOB)) + length(CAST(record_json AS BLOB)) FROM known_instances WHERE key = ?1", params![entry.key], )?; budget.check(replaces, incoming)?; conn.execute( "INSERT INTO known_instances(key, last_used, record_json) VALUES (?1, ?2, ?3) ON CONFLICT(key) DO UPDATE SET last_used = excluded.last_used, record_json = excluded.record_json", params![ entry.key, numeric_field(&fields, "lastUsed"), entry.record.get() ], )?; Ok(()) } pub fn all_known_instances(conn: &Connection) -> Result, StoreError> { let mut statement = conn .prepare("SELECT key, record_json FROM known_instances ORDER BY last_used DESC, key ASC")?; let mut rows = statement.query([])?; let mut instances = Vec::new(); while let Some(row) = rows.next()? { instances.push(KnownInstanceEntry { key: row.get(0)?, record: RawValue::from_string(row.get(1)?)?, }); } Ok(instances) } pub fn delete_known_instance(conn: &Connection, key: &str) -> Result<(), StoreError> { conn.execute("DELETE FROM known_instances WHERE key = ?1", params![key])?; Ok(()) } fn stored_bytes( conn: &Connection, sql: &str, args: impl rusqlite::Params, ) -> Result, StoreError> { Ok(conn.query_row(sql, args, |row| row.get(0)).optional()?) } #[cfg(test)] mod tests { use super::*; use crate::model::{Store, temporary_store_path}; use crate::scoped_storage::{ScopedEntry, entry as scoped_entry, scope_entries, set_entry}; fn entry(storage_key: &str, record: &str) -> AccountEntry { AccountEntry { storage_key: storage_key.to_owned(), record: RawValue::from_string(record.to_owned()).expect("valid JSON"), } } #[test] fn a_record_with_unknown_keys_round_trips_byte_exactly() { let (_dir, path) = temporary_store_path(); let store = Store::open(path, None).expect("a fresh store opens"); let written = r#"{"userId":"1","instance":{"apiEndpoint":"https://a.example/api","futureKey":[1,2]},"unknownTopLevel":{"nested":true},"lastActive":7}"#; upsert_account(store.connection(), &entry("a.example::1", written)) .expect("the account is stored"); let read = account(store.connection(), "a.example::1") .expect("the account is readable") .expect("the account is present"); assert_eq!( read.record.get(), written, "a record written by a newer renderer must survive an older shell byte-exactly" ); assert_eq!(read.storage_key, "a.example::1"); } #[test] fn accounts_are_ordered_by_last_active_then_storage_key() { let (_dir, path) = temporary_store_path(); let store = Store::open(path, None).expect("a fresh store opens"); for (key, last_active) in [("b::2", 5.0), ("a::1", 9.0), ("a::3", 5.0)] { upsert_account( store.connection(), &entry( key, &format!(r#"{{"userId":"u","lastActive":{last_active}}}"#), ), ) .expect("the account is stored"); } let keys = all_accounts(store.connection()) .expect("the accounts are readable") .into_iter() .map(|account| account.storage_key) .collect::>(); assert_eq!(keys, vec!["a::1", "a::3", "b::2"]); } #[test] fn a_record_the_addon_cannot_use_is_still_stored() { let (_dir, path) = temporary_store_path(); let store = Store::open(path, None).expect("a fresh store opens"); let deployed_shape = r#"{"userId":"1","instance":{"gifProvider":"tenor","apiCodeVersion":9}}"#; upsert_account(store.connection(), &entry("legacy::1", deployed_shape)) .expect("a record with an unusable instance is still stored"); let mut budget = account_budget(store.connection()); let write = write_account( store.connection(), budget.as_mut().expect("the budget is readable"), &entry("legacy::1", deployed_shape), ) .expect("the write succeeds"); assert!(matches!( write, AccountWrite::Written { instance_usable: false } )); assert_eq!( account(store.connection(), "legacy::1") .expect("the account is readable") .expect("the account is present") .record .get(), deployed_shape ); } #[test] fn a_storage_key_crosses_the_boundary_opaquely() { let (_dir, path) = temporary_store_path(); let store = Store::open(path, None).expect("a fresh store opens"); let opaque = "HTTPS://Mixed.Case.Example:8443/api::user::12345"; upsert_account(store.connection(), &entry(opaque, r#"{"userId":"12345"}"#)) .expect("the account is stored"); assert!( account(store.connection(), opaque) .expect("the account is readable") .is_some(), "the addon must not normalise the storage key it was given" ); } #[test] fn a_record_that_is_not_a_json_object_is_refused_rather_than_stored() { let (_dir, path) = temporary_store_path(); let store = Store::open(path, None).expect("a fresh store opens"); assert!(upsert_account(store.connection(), &entry("a::1", "[1,2,3]")).is_err()); assert!( all_accounts(store.connection()) .expect("the accounts are readable") .is_empty() ); } #[test] fn deleting_an_account_takes_its_scoped_storage_with_it() { let (_dir, path) = temporary_store_path(); let store = Store::open(path, None).expect("a fresh store opens"); for key in ["a::1", "a::2"] { upsert_account(store.connection(), &entry(key, r#"{"userId":"1"}"#)) .expect("the account is stored"); } for scope in ["a::1", "a::2", "global"] { set_entry( store.connection(), &ScopedEntry { store: "app".to_owned(), scope: scope.to_owned(), key: "theme".to_owned(), value: "\"dark\"".to_owned(), updated_at: 1.0, }, ) .expect("the entry is stored"); } delete_account(store.connection(), "a::1").expect("the account is deleted"); assert!( account(store.connection(), "a::1") .expect("the account is readable") .is_none() ); assert!( scope_entries(store.connection(), "app", "a::1") .expect("entries are readable") .is_empty(), "a deleted account must not leave its scoped storage behind" ); assert!( scoped_entry(store.connection(), "app", "a::2", "theme") .expect("the entry is readable") .is_some() ); assert!( scoped_entry(store.connection(), "app", "global", "theme") .expect("the entry is readable") .is_some() ); } #[test] fn a_compare_and_swap_lands_once_and_refuses_the_stale_retry() { let (_dir, path) = temporary_store_path(); let mut store = Store::open(path, None).expect("a fresh store opens"); let first = r#"{"userId":"1","token":"first"}"#; let second = r#"{"userId":"1","token":"second"}"#; let third = r#"{"userId":"1","token":"third"}"#; upsert_account(store.connection(), &entry("a.example::1", first)) .expect("the account is stored"); assert!( compare_and_swap_account( store.connection_mut(), &AccountCompareAndSwap { expected: entry("a.example::1", first), replacement: entry("a.example::1", second), }, ) .expect("the swap runs") ); assert_eq!( account(store.connection(), "a.example::1") .expect("the account is readable") .expect("the account is present") .record .get(), second ); assert!( !compare_and_swap_account( store.connection_mut(), &AccountCompareAndSwap { expected: entry("a.example::1", first), replacement: entry("a.example::1", third), }, ) .expect("the swap runs"), "a swap holding a record the store has already moved past must report that it did not land" ); assert_eq!( account(store.connection(), "a.example::1") .expect("the account is readable") .expect("the account is present") .record .get(), second, "a refused swap must leave the newer record untouched" ); } #[test] fn a_compare_and_swap_across_two_storage_keys_is_refused() { let (_dir, path) = temporary_store_path(); let mut store = Store::open(path, None).expect("a fresh store opens"); let record = r#"{"userId":"1"}"#; upsert_account(store.connection(), &entry("a.example::1", record)) .expect("the account is stored"); let outcome = compare_and_swap_account( store.connection_mut(), &AccountCompareAndSwap { expected: entry("a.example::1", record), replacement: entry("a.example::2", record), }, ); assert!(matches!(outcome, Err(StoreError::Refused(_)))); assert!( account(store.connection(), "a.example::2") .expect("the accounts are readable") .is_none() ); } #[test] fn known_instances_are_stored_opaquely_and_ordered_by_last_used() { let (_dir, path) = temporary_store_path(); let store = Store::open(path, None).expect("a fresh store opens"); for (key, last_used) in [("beta.example", 1.0), ("alpha.example", 4.0)] { upsert_known_instance( store.connection(), &KnownInstanceEntry { key: key.to_owned(), record: RawValue::from_string(format!( r#"{{"lastUsed":{last_used},"displayName":"{key}"}}"# )) .expect("valid JSON"), }, ) .expect("the known instance is stored"); } let stored = all_known_instances(store.connection()).expect("the instances are readable"); assert_eq!( stored .iter() .map(|row| row.key.as_str()) .collect::>(), vec!["alpha.example", "beta.example"] ); delete_known_instance(store.connection(), "alpha.example") .expect("the instance is deleted"); assert_eq!( all_known_instances(store.connection()) .expect("the instances are readable") .len(), 1 ); } }