iris_chat_protocol/
storage.rs

1mod checkpoint;
2mod checkpoint_sqlite;
3#[cfg(test)]
4mod checkpoint_tests;
5mod segmented;
6#[cfg(test)]
7mod segmented_tests;
8
9pub use checkpoint::{StorageCheckpoint, StorageCheckpointBody};
10
11#[cfg(test)]
12pub(crate) fn checkpoint_io_bytes() -> (usize, usize) {
13    segmented::io_bytes()
14}
15
16use super::SharedConnection;
17use std::collections::HashMap;
18use std::fs;
19use std::io::Write;
20use std::path::PathBuf;
21use std::sync::{Arc, Mutex};
22use std::time::{Duration, Instant};
23
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct StorageError {
26    message: String,
27    interrupted: bool,
28}
29
30impl StorageError {
31    pub fn new(message: impl Into<String>) -> Self {
32        Self {
33            message: message.into(),
34            interrupted: false,
35        }
36    }
37
38    /// True only for a SQLite operation interrupted by its caller, never for
39    /// malformed records or an arbitrary error message with the same wording.
40    pub fn is_interrupted(&self) -> bool {
41        self.interrupted
42    }
43}
44
45impl std::fmt::Display for StorageError {
46    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47        self.message.fmt(f)
48    }
49}
50
51impl std::error::Error for StorageError {}
52
53pub type StorageResult<T> = Result<T, StorageError>;
54
55pub trait StorageAdapter: Send + Sync {
56    fn get(&self, key: &str) -> StorageResult<Option<String>>;
57    fn put(&self, key: &str, value: String) -> StorageResult<()>;
58    fn put_checkpoint(&self, key: &str, value: &StorageCheckpoint) -> StorageResult<()> {
59        self.put(key, value.materialize()?)
60    }
61    fn del(&self, key: &str) -> StorageResult<()>;
62    fn list(&self, prefix: &str) -> StorageResult<Vec<String>>;
63}
64
65#[derive(Clone)]
66pub struct InMemoryStorage {
67    store: Arc<Mutex<HashMap<String, String>>>,
68}
69
70impl InMemoryStorage {
71    pub fn new() -> Self {
72        Self {
73            store: Arc::new(Mutex::new(HashMap::new())),
74        }
75    }
76}
77
78impl Default for InMemoryStorage {
79    fn default() -> Self {
80        Self::new()
81    }
82}
83
84impl StorageAdapter for InMemoryStorage {
85    fn get(&self, key: &str) -> StorageResult<Option<String>> {
86        let store = self
87            .store
88            .lock()
89            .map_err(|_| StorageError::new("storage mutex poisoned"))?;
90        Ok(store.get(key).cloned())
91    }
92
93    fn put(&self, key: &str, value: String) -> StorageResult<()> {
94        let mut store = self
95            .store
96            .lock()
97            .map_err(|_| StorageError::new("storage mutex poisoned"))?;
98        store.insert(key.to_string(), value);
99        Ok(())
100    }
101
102    fn del(&self, key: &str) -> StorageResult<()> {
103        let mut store = self
104            .store
105            .lock()
106            .map_err(|_| StorageError::new("storage mutex poisoned"))?;
107        store.remove(key);
108        Ok(())
109    }
110
111    fn list(&self, prefix: &str) -> StorageResult<Vec<String>> {
112        let store = self
113            .store
114            .lock()
115            .map_err(|_| StorageError::new("storage mutex poisoned"))?;
116        Ok(store
117            .keys()
118            .filter(|key| key.starts_with(prefix))
119            .cloned()
120            .collect())
121    }
122}
123
124pub struct FileStorageAdapter {
125    base_path: PathBuf,
126}
127
128impl FileStorageAdapter {
129    /// Opens a dedicated directory for secret session state. On Unix, access
130    /// to this directory is restricted to the current user, including on reopen.
131    pub fn new(base_path: PathBuf) -> StorageResult<Self> {
132        let mut builder = fs::DirBuilder::new();
133        builder.recursive(true);
134        #[cfg(unix)]
135        {
136            use std::os::unix::fs::DirBuilderExt;
137            builder.mode(0o700);
138        }
139        builder
140            .create(&base_path)
141            .map_err(|err| storage_io_error("failed to create storage directory", err))?;
142        // Existing installations may contain session keys in readable files.
143        // Restrict their enclosing directory before loading any state.
144        #[cfg(unix)]
145        {
146            use std::os::unix::fs::PermissionsExt;
147            fs::set_permissions(&base_path, fs::Permissions::from_mode(0o700))
148                .map_err(|err| storage_io_error("failed to protect storage directory", err))?;
149        }
150        Ok(Self { base_path })
151    }
152
153    fn sanitize_key(key: &str) -> String {
154        key.replace(['/', '\\', ':'], "_")
155    }
156
157    fn key_to_path(&self, key: &str) -> PathBuf {
158        let sanitized = Self::sanitize_key(key);
159        self.base_path.join(format!("{}.json", sanitized))
160    }
161}
162
163impl StorageAdapter for FileStorageAdapter {
164    fn get(&self, key: &str) -> StorageResult<Option<String>> {
165        let path = self.key_to_path(key);
166        match fs::read_to_string(&path) {
167            Ok(contents) => Ok(Some(contents)),
168            Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(None),
169            Err(err) => Err(storage_io_error("failed to read storage file", err)),
170        }
171    }
172
173    fn put(&self, key: &str, value: String) -> StorageResult<()> {
174        let path = self.key_to_path(key);
175
176        let tmp_path = path.with_extension(format!("json.{}.tmp", rand::random::<u128>()));
177        let mut options = fs::OpenOptions::new();
178        options.write(true).create_new(true);
179        #[cfg(unix)]
180        {
181            use std::os::unix::fs::OpenOptionsExt;
182            options.mode(0o600);
183        }
184        // Set the mode at creation: chmod after writing would expose secrets
185        // briefly, and exclusive creation also refuses existing symlinks.
186        let mut file = options
187            .open(&tmp_path)
188            .map_err(|err| storage_io_error("failed to create storage temp file", err))?;
189        if let Err(err) = file.write_all(value.as_bytes()) {
190            drop(file);
191            let _ = fs::remove_file(&tmp_path);
192            return Err(storage_io_error("failed to write storage temp file", err));
193        }
194        drop(file);
195
196        #[cfg(windows)]
197        {
198            if path.exists() {
199                fs::remove_file(&path).map_err(|err| {
200                    storage_io_error("failed to replace existing storage file", err)
201                })?;
202            }
203        }
204
205        fs::rename(&tmp_path, &path)
206            .map_err(|err| storage_io_error("failed to commit storage file", err))?;
207
208        Ok(())
209    }
210
211    fn del(&self, key: &str) -> StorageResult<()> {
212        let path = self.key_to_path(key);
213        match fs::remove_file(&path) {
214            Ok(()) => Ok(()),
215            Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
216            Err(err) => Err(storage_io_error("failed to delete storage file", err)),
217        }
218    }
219
220    fn list(&self, prefix: &str) -> StorageResult<Vec<String>> {
221        let mut keys = Vec::new();
222        let sanitized_prefix = Self::sanitize_key(prefix);
223        let entries = fs::read_dir(&self.base_path)
224            .map_err(|err| storage_io_error("failed to read storage directory", err))?;
225
226        for entry in entries {
227            let entry =
228                entry.map_err(|err| storage_io_error("failed to read storage entry", err))?;
229            let file_name = entry.file_name();
230            let file_name_str = file_name.to_string_lossy();
231
232            if !file_name_str.ends_with(".json") {
233                continue;
234            }
235
236            let key = file_name_str
237                .strip_suffix(".json")
238                .unwrap_or(&file_name_str)
239                .to_string();
240
241            if prefix.is_empty() {
242                keys.push(key);
243                continue;
244            }
245
246            if key.starts_with(&sanitized_prefix) {
247                let remainder = key.strip_prefix(&sanitized_prefix).unwrap_or("");
248                keys.push(format!("{}{}", prefix, remainder));
249            }
250        }
251
252        Ok(keys)
253    }
254}
255
256pub struct DebouncedFileStorage {
257    adapter: FileStorageAdapter,
258    pending_writes: Mutex<HashMap<String, String>>,
259    last_flush: Mutex<Instant>,
260    flush_interval: Duration,
261}
262
263impl DebouncedFileStorage {
264    pub fn new(base_path: PathBuf, flush_interval_ms: u64) -> StorageResult<Self> {
265        Ok(Self {
266            adapter: FileStorageAdapter::new(base_path)?,
267            pending_writes: Mutex::new(HashMap::new()),
268            last_flush: Mutex::new(Instant::now()),
269            flush_interval: Duration::from_millis(flush_interval_ms),
270        })
271    }
272
273    pub fn flush(&self) -> StorageResult<()> {
274        let mut pending = self
275            .pending_writes
276            .lock()
277            .map_err(|_| StorageError::new("pending file storage mutex poisoned"))?;
278        for (key, value) in pending.drain() {
279            self.adapter.put(&key, value)?;
280        }
281        *self
282            .last_flush
283            .lock()
284            .map_err(|_| StorageError::new("file storage flush mutex poisoned"))? = Instant::now();
285        Ok(())
286    }
287
288    fn maybe_flush(&self) -> StorageResult<()> {
289        let last_flush = *self
290            .last_flush
291            .lock()
292            .map_err(|_| StorageError::new("file storage flush mutex poisoned"))?;
293        let pending_count = self
294            .pending_writes
295            .lock()
296            .map_err(|_| StorageError::new("pending file storage mutex poisoned"))?
297            .len();
298
299        if last_flush.elapsed() >= self.flush_interval && pending_count > 0 {
300            self.flush()?;
301        }
302        Ok(())
303    }
304}
305
306impl StorageAdapter for DebouncedFileStorage {
307    fn get(&self, key: &str) -> StorageResult<Option<String>> {
308        let pending = self
309            .pending_writes
310            .lock()
311            .map_err(|_| StorageError::new("pending file storage mutex poisoned"))?;
312        if let Some(value) = pending.get(key) {
313            return Ok(Some(value.clone()));
314        }
315        drop(pending);
316        self.adapter.get(key)
317    }
318
319    fn put(&self, key: &str, value: String) -> StorageResult<()> {
320        self.pending_writes
321            .lock()
322            .map_err(|_| StorageError::new("pending file storage mutex poisoned"))?
323            .insert(key.to_string(), value);
324        self.maybe_flush()
325    }
326
327    fn del(&self, key: &str) -> StorageResult<()> {
328        self.pending_writes
329            .lock()
330            .map_err(|_| StorageError::new("pending file storage mutex poisoned"))?
331            .remove(key);
332        self.adapter.del(key)
333    }
334
335    fn list(&self, prefix: &str) -> StorageResult<Vec<String>> {
336        let mut keys = self.adapter.list(prefix)?;
337        let pending = self
338            .pending_writes
339            .lock()
340            .map_err(|_| StorageError::new("pending file storage mutex poisoned"))?;
341
342        for key in pending.keys() {
343            if key.starts_with(prefix) && !keys.contains(key) {
344                keys.push(key.clone());
345            }
346        }
347
348        Ok(keys)
349    }
350}
351
352/// SQLite-backed implementation of `iris_chat_protocol::StorageAdapter`.
353/// Keys are namespaced by (owner_pubkey_hex, device_pubkey_hex) so a
354/// single database serves multiple owner accounts and devices without
355/// keyspace collisions, matching the per-(owner, device) directory
356/// scoping the previous file-backed adapter used.
357pub struct SqliteStorageAdapter {
358    conn: SharedConnection,
359    owner_pubkey_hex: String,
360    device_pubkey_hex: String,
361}
362
363impl SqliteStorageAdapter {
364    pub fn new(
365        conn: SharedConnection,
366        owner_pubkey_hex: String,
367        device_pubkey_hex: String,
368    ) -> Self {
369        Self {
370            conn,
371            owner_pubkey_hex,
372            device_pubkey_hex,
373        }
374    }
375
376    fn map_err<E: std::fmt::Display + std::any::Any>(error: E) -> StorageError {
377        let interrupted = matches!(
378            (&error as &dyn std::any::Any).downcast_ref::<rusqlite::Error>(),
379            Some(rusqlite::Error::SqliteFailure(inner, _))
380                if inner.code == rusqlite::ErrorCode::OperationInterrupted
381        );
382        StorageError {
383            message: error.to_string(),
384            interrupted,
385        }
386    }
387}
388
389impl StorageAdapter for SqliteStorageAdapter {
390    fn get(&self, key: &str) -> StorageResult<Option<String>> {
391        let mut conn = self
392            .conn
393            .lock()
394            .map_err(|_| StorageError::new("ndr_kv connection mutex poisoned"))?;
395        segmented::get(
396            &mut conn,
397            (&self.owner_pubkey_hex, &self.device_pubkey_hex, key),
398        )
399    }
400
401    fn put(&self, key: &str, value: String) -> StorageResult<()> {
402        let mut conn = self
403            .conn
404            .lock()
405            .map_err(|_| StorageError::new("ndr_kv connection mutex poisoned"))?;
406        segmented::put(
407            &mut conn,
408            (&self.owner_pubkey_hex, &self.device_pubkey_hex, key),
409            &value,
410        )
411    }
412
413    fn put_checkpoint(&self, key: &str, value: &StorageCheckpoint) -> StorageResult<()> {
414        // Keep small values on the existing simple path. Large checkpoints
415        // use independently addressed records, with no full-value mirror.
416        if value.len() < 64 * 1024 {
417            return self.put(key, value.materialize()?);
418        }
419        let mut conn = self
420            .conn
421            .lock()
422            .map_err(|_| StorageError::new("ndr_kv connection mutex poisoned"))?;
423        checkpoint_sqlite::put(
424            &mut conn,
425            (&self.owner_pubkey_hex, &self.device_pubkey_hex, key),
426            value,
427        )
428    }
429
430    fn del(&self, key: &str) -> StorageResult<()> {
431        let mut conn = self
432            .conn
433            .lock()
434            .map_err(|_| StorageError::new("ndr_kv connection mutex poisoned"))?;
435        segmented::delete(
436            &mut conn,
437            (&self.owner_pubkey_hex, &self.device_pubkey_hex, key),
438        )
439    }
440
441    fn list(&self, prefix: &str) -> StorageResult<Vec<String>> {
442        let conn = self
443            .conn
444            .lock()
445            .map_err(|_| StorageError::new("ndr_kv connection mutex poisoned"))?;
446        let mut stmt = conn
447            .prepare(
448                "SELECT key FROM ndr_kv
449                 WHERE owner_pubkey_hex = ?1 AND device_pubkey_hex = ?2 AND key LIKE ?3 ESCAPE '\\'",
450            )
451            .map_err(Self::map_err)?;
452        let pattern = format!("{}%", escape_like(prefix));
453        let rows = stmt
454            .query_map(
455                (&self.owner_pubkey_hex, &self.device_pubkey_hex, &pattern),
456                |row| row.get::<_, String>(0),
457            )
458            .map_err(Self::map_err)?;
459        let mut keys = Vec::new();
460        for row in rows {
461            keys.push(row.map_err(Self::map_err)?);
462        }
463        Ok(keys)
464    }
465}
466
467fn escape_like(input: &str) -> String {
468    let mut out = String::with_capacity(input.len());
469    for ch in input.chars() {
470        match ch {
471            '\\' | '%' | '_' => {
472                out.push('\\');
473                out.push(ch);
474            }
475            other => out.push(other),
476        }
477    }
478    out
479}
480
481fn storage_io_error(context: &str, error: std::io::Error) -> StorageError {
482    StorageError::new(format!("{}: {}", context, error))
483}
484
485#[cfg(test)]
486mod tests {
487    use super::*;
488    use std::sync::{Arc, Mutex};
489    use tempfile::TempDir;
490
491    fn fresh_connection() -> SharedConnection {
492        let conn = rusqlite::Connection::open_in_memory().unwrap();
493        conn.execute_batch(
494            "CREATE TABLE ndr_kv (
495                owner_pubkey_hex TEXT NOT NULL,
496                device_pubkey_hex TEXT NOT NULL,
497                key TEXT NOT NULL,
498                value TEXT NOT NULL,
499                PRIMARY KEY (owner_pubkey_hex, device_pubkey_hex, key)
500            );",
501        )
502        .unwrap();
503        Arc::new(Mutex::new(conn))
504    }
505
506    fn fresh_adapter() -> SqliteStorageAdapter {
507        SqliteStorageAdapter::new(
508            fresh_connection(),
509            "owner".to_string(),
510            "device".to_string(),
511        )
512    }
513
514    #[test]
515    fn put_get_del_round_trip() {
516        let adapter = fresh_adapter();
517        assert!(adapter.get("k").unwrap().is_none());
518        adapter.put("k", "v".to_string()).unwrap();
519        assert_eq!(adapter.get("k").unwrap(), Some("v".to_string()));
520        adapter.put("k", "v2".to_string()).unwrap();
521        assert_eq!(adapter.get("k").unwrap(), Some("v2".to_string()));
522        adapter.del("k").unwrap();
523        assert!(adapter.get("k").unwrap().is_none());
524    }
525
526    #[test]
527    fn list_returns_only_matching_prefix() {
528        let adapter = fresh_adapter();
529        adapter.put("user/alice", "1".to_string()).unwrap();
530        adapter.put("user/bob", "2".to_string()).unwrap();
531        adapter.put("invite/charlie", "3".to_string()).unwrap();
532        let mut keys = adapter.list("user/").unwrap();
533        keys.sort();
534        assert_eq!(keys, vec!["user/alice".to_string(), "user/bob".to_string()]);
535    }
536
537    #[test]
538    fn keys_are_isolated_per_owner_device() {
539        let conn = fresh_connection();
540        let alice = SqliteStorageAdapter::new(conn.clone(), "owner_a".into(), "device_a".into());
541        let bob = SqliteStorageAdapter::new(conn, "owner_b".into(), "device_b".into());
542        alice.put("shared-key", "alice".to_string()).unwrap();
543        bob.put("shared-key", "bob".to_string()).unwrap();
544        assert_eq!(alice.get("shared-key").unwrap(), Some("alice".to_string()));
545        assert_eq!(bob.get("shared-key").unwrap(), Some("bob".to_string()));
546    }
547
548    #[test]
549    fn file_storage_round_trips_values() {
550        let temp_dir = TempDir::new().unwrap();
551        let adapter = FileStorageAdapter::new(temp_dir.path().to_path_buf()).unwrap();
552
553        assert!(adapter.get("test-key").unwrap().is_none());
554
555        adapter.put("test-key", "test-value".to_string()).unwrap();
556        assert_eq!(
557            adapter.get("test-key").unwrap(),
558            Some("test-value".to_string())
559        );
560
561        adapter.del("test-key").unwrap();
562        assert!(adapter.get("test-key").unwrap().is_none());
563    }
564
565    #[cfg(unix)]
566    #[test]
567    fn file_storage_keeps_session_secrets_private_when_created_and_replaced() {
568        use std::os::unix::fs::PermissionsExt;
569
570        let temp_dir = TempDir::new().unwrap();
571        let storage_dir = temp_dir.path().join("protocol");
572        let adapter = FileStorageAdapter::new(storage_dir.clone()).unwrap();
573        assert_eq!(
574            fs::metadata(&storage_dir).unwrap().permissions().mode() & 0o777,
575            0o700
576        );
577
578        adapter.put("session", "secret".to_string()).unwrap();
579        let session_path = storage_dir.join("session.json");
580        assert_eq!(
581            fs::metadata(&session_path).unwrap().permissions().mode() & 0o777,
582            0o600
583        );
584
585        // Upgrading an old file must not preserve its permissive mode.
586        fs::set_permissions(&session_path, fs::Permissions::from_mode(0o644)).unwrap();
587        adapter.put("session", "new secret".to_string()).unwrap();
588        assert_eq!(
589            fs::metadata(&session_path).unwrap().permissions().mode() & 0o777,
590            0o600
591        );
592        assert_eq!(
593            adapter.get("session").unwrap().as_deref(),
594            Some("new secret")
595        );
596    }
597
598    #[cfg(unix)]
599    #[test]
600    fn file_storage_protects_existing_session_directory() {
601        use std::os::unix::fs::PermissionsExt;
602
603        let temp_dir = TempDir::new().unwrap();
604        let storage_dir = temp_dir.path().join("protocol");
605        fs::create_dir(&storage_dir).unwrap();
606        fs::set_permissions(&storage_dir, fs::Permissions::from_mode(0o755)).unwrap();
607        fs::write(storage_dir.join("session.json"), "legacy secret").unwrap();
608
609        let adapter = FileStorageAdapter::new(storage_dir.clone()).unwrap();
610
611        assert_eq!(
612            fs::metadata(&storage_dir).unwrap().permissions().mode() & 0o777,
613            0o700
614        );
615        assert_eq!(
616            adapter.get("session").unwrap().as_deref(),
617            Some("legacy secret")
618        );
619    }
620
621    #[test]
622    fn file_storage_lists_sanitized_runtime_keys() {
623        let temp_dir = TempDir::new().unwrap();
624        let adapter = FileStorageAdapter::new(temp_dir.path().to_path_buf()).unwrap();
625
626        adapter.put("user/alice", "1".to_string()).unwrap();
627        adapter.put("user/bob", "2".to_string()).unwrap();
628        adapter.put("invite/charlie", "3".to_string()).unwrap();
629
630        let mut user_keys = adapter.list("user/").unwrap();
631        user_keys.sort();
632        assert_eq!(
633            user_keys,
634            vec!["user/alice".to_string(), "user/bob".to_string()]
635        );
636
637        let mut all_keys = adapter.list("").unwrap();
638        all_keys.sort();
639        assert_eq!(
640            all_keys,
641            vec![
642                "invite_charlie".to_string(),
643                "user_alice".to_string(),
644                "user_bob".to_string()
645            ]
646        );
647    }
648
649    #[test]
650    fn debounced_file_storage_reads_pending_writes_and_flushes() {
651        let temp_dir = TempDir::new().unwrap();
652        let storage = DebouncedFileStorage::new(temp_dir.path().to_path_buf(), 1000).unwrap();
653
654        storage.put("key1", "value1".to_string()).unwrap();
655
656        assert_eq!(storage.get("key1").unwrap(), Some("value1".to_string()));
657        assert!(storage.pending_writes.lock().unwrap().contains_key("key1"));
658
659        storage.flush().unwrap();
660
661        assert!(storage.pending_writes.lock().unwrap().is_empty());
662        assert_eq!(
663            storage.adapter.get("key1").unwrap(),
664            Some("value1".to_string())
665        );
666    }
667}
iris_chat_protocol/storage.rs

iris_chat_protocol/
storage.rs
