iris_chat_protocol/
storage.rs

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

iris_chat_protocol/
storage.rs
