restflow_core/storage/
secrets.rs

1use crate::models::Secret;
2use anyhow::Result;
3use base64::{Engine as _, engine::general_purpose::STANDARD};
4use redb::{Database, ReadableDatabase, ReadableTable, TableDefinition};
5use std::sync::Arc;
6
7const SECRETS_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("secrets");
8
9#[derive(Debug, Clone)]
10pub struct SecretStorage {
11    db: Arc<Database>,
12}
13
14impl SecretStorage {
15    pub fn new(db: Arc<Database>) -> Result<Self> {
16        let write_txn = db.begin_write()?;
17        write_txn.open_table(SECRETS_TABLE)?;
18        write_txn.commit()?;
19
20        Ok(Self { db })
21    }
22
23    pub fn set_secret(&self, key: &str, value: &str, description: Option<String>) -> Result<()> {
24        let existing = self.get_secret_model(key)?;
25
26        let write_txn = self.db.begin_write()?;
27        {
28            let mut table = write_txn.open_table(SECRETS_TABLE)?;
29
30            let secret = if let Some(mut existing_secret) = existing {
31                existing_secret.update(value.to_string(), description);
32                existing_secret
33            } else {
34                Secret::new(key.to_string(), value.to_string(), description)
35            };
36
37            let json = serde_json::to_string(&secret)?;
38            let encoded = STANDARD.encode(json.as_bytes());
39            table.insert(key, encoded.as_bytes())?;
40        }
41        write_txn.commit()?;
42        Ok(())
43    }
44
45    // Create a new secret (fails if already exists)
46    pub fn create_secret(&self, key: &str, value: &str, description: Option<String>) -> Result<()> {
47        if self.get_secret_model(key)?.is_some() {
48            return Err(anyhow::anyhow!("Secret {} already exists", key));
49        }
50        self.set_secret(key, value, description)
51    }
52
53    // Update an existing secret (fails if not exists)
54    pub fn update_secret(&self, key: &str, value: &str, description: Option<String>) -> Result<()> {
55        if self.get_secret_model(key)?.is_none() {
56            return Err(anyhow::anyhow!("Secret {} not found", key));
57        }
58        self.set_secret(key, value, description)
59    }
60
61    // Internal use only
62    fn get_secret_model(&self, key: &str) -> Result<Option<Secret>> {
63        let read_txn = self.db.begin_read()?;
64        let table = read_txn.open_table(SECRETS_TABLE)?;
65
66        if let Some(data) = table.get(key)? {
67            let encoded = std::str::from_utf8(data.value())?;
68            let decoded = STANDARD.decode(encoded)?;
69            let json = String::from_utf8(decoded)?;
70            Ok(Some(serde_json::from_str(&json)?))
71        } else {
72            Ok(None)
73        }
74    }
75
76    pub fn get_secret(&self, key: &str) -> Result<Option<String>> {
77        if let Some(secret) = self.get_secret_model(key)? {
78            Ok(Some(secret.value))
79        } else {
80            // Fallback to environment variable (e.g., OPENAI_API_KEY)
81            Ok(std::env::var(key.to_uppercase().replace('-', "_")).ok())
82        }
83    }
84
85    pub fn delete_secret(&self, key: &str) -> Result<()> {
86        let write_txn = self.db.begin_write()?;
87        {
88            let mut table = write_txn.open_table(SECRETS_TABLE)?;
89            table.remove(key)?;
90        }
91        write_txn.commit()?;
92        Ok(())
93    }
94
95    // Returns all secrets without values for security
96    pub fn list_secrets(&self) -> Result<Vec<Secret>> {
97        let read_txn = self.db.begin_read()?;
98        let table = read_txn.open_table(SECRETS_TABLE)?;
99
100        let mut secrets = Vec::new();
101        for item in table.iter()? {
102            let (_, value) = item?;
103            let encoded = std::str::from_utf8(value.value())?;
104            let decoded = STANDARD.decode(encoded)?;
105            let json = String::from_utf8(decoded)?;
106            let mut secret: Secret = serde_json::from_str(&json)?;
107            // Clear the value for security
108            secret.value = String::new();
109            secrets.push(secret);
110        }
111
112        Ok(secrets)
113    }
114
115    pub fn has_secret(&self, key: &str) -> Result<bool> {
116        let read_txn = self.db.begin_read()?;
117        let table = read_txn.open_table(SECRETS_TABLE)?;
118        Ok(table.get(key)?.is_some())
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125    use tempfile::tempdir;
126
127    fn setup() -> (SecretStorage, tempfile::TempDir) {
128        let temp_dir = tempdir().unwrap();
129        let db_path = temp_dir.path().join("test.db");
130        let db = Arc::new(Database::create(db_path).unwrap());
131        let storage = SecretStorage::new(db).unwrap();
132        (storage, temp_dir)
133    }
134
135    #[test]
136    fn test_set_and_get_secret() {
137        let (storage, _temp_dir) = setup();
138
139        storage
140            .set_secret(
141                "OPENAI_API_KEY",
142                "sk-test123",
143                Some("OpenAI API key".to_string()),
144            )
145            .unwrap();
146
147        let value = storage.get_secret("OPENAI_API_KEY").unwrap();
148        assert_eq!(value, Some("sk-test123".to_string()));
149    }
150
151    #[test]
152    fn test_list_secrets_with_metadata() {
153        let (storage, _temp_dir) = setup();
154
155        storage
156            .set_secret("API_KEY_1", "value1", Some("First key".to_string()))
157            .unwrap();
158        storage.set_secret("API_KEY_2", "value2", None).unwrap();
159        storage
160            .set_secret("API_KEY_3", "value3", Some("Third key".to_string()))
161            .unwrap();
162
163        let secrets = storage.list_secrets().unwrap();
164        assert_eq!(secrets.len(), 3);
165
166        let key1 = secrets.iter().find(|s| s.key == "API_KEY_1").unwrap();
167        assert_eq!(key1.description, Some("First key".to_string()));
168        assert_eq!(key1.value, ""); // Value should be cleared
169
170        let key2 = secrets.iter().find(|s| s.key == "API_KEY_2").unwrap();
171        assert_eq!(key2.description, None);
172    }
173
174    #[test]
175    fn test_update_preserves_created_at() {
176        let (storage, _temp_dir) = setup();
177
178        storage
179            .set_secret("KEY", "initial", Some("Test key".to_string()))
180            .unwrap();
181
182        let secrets = storage.list_secrets().unwrap();
183        let initial = secrets.iter().find(|s| s.key == "KEY").unwrap();
184        let created_at = initial.created_at;
185        let initial_updated_at = initial.updated_at;
186
187        // Wait to ensure time difference
188        std::thread::sleep(std::time::Duration::from_millis(10));
189
190        storage
191            .set_secret("KEY", "updated", Some("Updated description".to_string()))
192            .unwrap();
193
194        let secrets = storage.list_secrets().unwrap();
195        let updated = secrets.iter().find(|s| s.key == "KEY").unwrap();
196
197        println!(
198            "created_at: {}, initial_updated_at: {}, new_updated_at: {}",
199            created_at, initial_updated_at, updated.updated_at
200        );
201
202        assert_eq!(updated.created_at, created_at); // created_at preserved
203        assert!(
204            updated.updated_at > initial_updated_at,
205            "updated_at should be greater: {} > {}",
206            updated.updated_at,
207            initial_updated_at
208        );
209        assert_eq!(updated.description, Some("Updated description".to_string()));
210    }
211
212    #[test]
213    fn test_delete_secret() {
214        let (storage, _temp_dir) = setup();
215
216        storage.set_secret("TEST_KEY", "test_value", None).unwrap();
217        storage.delete_secret("TEST_KEY").unwrap();
218
219        let value = storage.get_secret("TEST_KEY").unwrap();
220        assert_eq!(value, None);
221    }
222
223    #[test]
224    fn test_has_secret() {
225        let (storage, _temp_dir) = setup();
226
227        storage.set_secret("EXISTS", "value", None).unwrap();
228
229        assert!(storage.has_secret("EXISTS").unwrap());
230        assert!(!storage.has_secret("NOT_EXISTS").unwrap());
231    }
232}
restflow_core/storage/secrets.rs

restflow_core/storage/
secrets.rs
