iris_chat_protocol/
storage.rs

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

iris_chat_protocol/
storage.rs
