dk_engine/
storage.rs

1//! Object storage abstraction.
2
3use async_trait::async_trait;
4
5/// Trait for object storage backends.
6#[async_trait]
7pub trait ObjectStore: Send + Sync + 'static {
8    async fn get(&self, key: &str) -> anyhow::Result<Vec<u8>>;
9    async fn put(&self, key: &str, data: Vec<u8>) -> anyhow::Result<()>;
10    async fn delete(&self, key: &str) -> anyhow::Result<()>;
11    async fn list(&self, prefix: &str) -> anyhow::Result<Vec<String>>;
12    async fn exists(&self, key: &str) -> anyhow::Result<bool>;
13}
14
15/// Local filesystem object store.
16pub struct LocalStore {
17    root: std::path::PathBuf,
18}
19
20impl LocalStore {
21    pub fn new(root: std::path::PathBuf) -> Self {
22        Self { root }
23    }
24}
25
26/// Validate that a storage key does not escape the root directory.
27///
28/// Rejects keys containing path traversal (`..`), null bytes, or absolute paths.
29fn validate_key(key: &str) -> anyhow::Result<()> {
30    if key.is_empty() {
31        anyhow::bail!("storage key cannot be empty");
32    }
33    if key.starts_with('/') || key.starts_with('\\') {
34        anyhow::bail!("storage key must be relative");
35    }
36    if key.contains('\0') {
37        anyhow::bail!("storage key contains null byte");
38    }
39    for component in key.split(&['/', '\\'] as &[char]) {
40        if component == ".." {
41            anyhow::bail!("storage key contains '..' traversal");
42        }
43    }
44    Ok(())
45}
46
47#[async_trait]
48impl ObjectStore for LocalStore {
49    async fn get(&self, key: &str) -> anyhow::Result<Vec<u8>> {
50        validate_key(key)?;
51        let path = self.root.join(key);
52        Ok(tokio::fs::read(path).await?)
53    }
54
55    async fn put(&self, key: &str, data: Vec<u8>) -> anyhow::Result<()> {
56        validate_key(key)?;
57        let path = self.root.join(key);
58        if let Some(parent) = path.parent() {
59            tokio::fs::create_dir_all(parent).await?;
60        }
61        Ok(tokio::fs::write(path, data).await?)
62    }
63
64    async fn delete(&self, key: &str) -> anyhow::Result<()> {
65        validate_key(key)?;
66        let path = self.root.join(key);
67        match tokio::fs::remove_file(path).await {
68            Ok(()) => Ok(()),
69            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
70            Err(e) => Err(e.into()),
71        }
72    }
73
74    async fn list(&self, prefix: &str) -> anyhow::Result<Vec<String>> {
75        // Empty prefix means "list root" — skip validation since there's
76        // nothing to traverse. Non-empty prefixes are validated normally.
77        if !prefix.is_empty() {
78            validate_key(prefix)?;
79        }
80        let dir = self.root.join(prefix);
81        let mut entries = Vec::new();
82        if dir.exists() {
83            let mut read_dir = tokio::fs::read_dir(&dir).await?;
84            while let Some(entry) = read_dir.next_entry().await? {
85                if let Some(name) = entry.file_name().to_str() {
86                    let key = if prefix.is_empty() {
87                        name.to_string()
88                    } else {
89                        format!("{prefix}/{name}")
90                    };
91                    entries.push(key);
92                }
93            }
94        }
95        Ok(entries)
96    }
97
98    async fn exists(&self, key: &str) -> anyhow::Result<bool> {
99        validate_key(key)?;
100        let path = self.root.join(key);
101        Ok(path.exists())
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    #[test]
110    fn validate_key_rejects_traversal() {
111        assert!(validate_key("../etc/passwd").is_err());
112        assert!(validate_key("foo/../../bar").is_err());
113    }
114
115    #[test]
116    fn validate_key_rejects_absolute() {
117        assert!(validate_key("/etc/passwd").is_err());
118    }
119
120    #[test]
121    fn validate_key_rejects_null() {
122        assert!(validate_key("foo\0bar").is_err());
123    }
124
125    #[test]
126    fn validate_key_accepts_valid() {
127        assert!(validate_key("repos/abc/data.json").is_ok());
128        assert!(validate_key("sessions/123").is_ok());
129        assert!(validate_key("file.txt").is_ok());
130    }
131
132    #[tokio::test]
133    async fn local_store_roundtrip() {
134        let dir = tempfile::tempdir().unwrap();
135        let store = LocalStore::new(dir.path().to_path_buf());
136
137        store.put("test/file.txt", b"hello".to_vec()).await.unwrap();
138        assert!(store.exists("test/file.txt").await.unwrap());
139
140        let data = store.get("test/file.txt").await.unwrap();
141        assert_eq!(data, b"hello");
142
143        let keys = store.list("test").await.unwrap();
144        assert_eq!(keys, vec!["test/file.txt"]);
145
146        store.delete("test/file.txt").await.unwrap();
147        assert!(!store.exists("test/file.txt").await.unwrap());
148    }
149
150    #[tokio::test]
151    async fn local_store_get_not_found() {
152        let dir = tempfile::tempdir().unwrap();
153        let store = LocalStore::new(dir.path().to_path_buf());
154        let result = store.get("nonexistent").await;
155        assert!(result.is_err());
156    }
157
158    #[tokio::test]
159    async fn local_store_delete_idempotent() {
160        let dir = tempfile::tempdir().unwrap();
161        let store = LocalStore::new(dir.path().to_path_buf());
162        store.delete("nonexistent").await.unwrap();
163    }
164}
dk_engine/storage.rs

dk_engine/
storage.rs
