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

iris_chat_protocol/
storage.rs
