llm_kernel/secrets/
vault.rs

1use std::collections::HashMap;
2use std::ops::{Deref, DerefMut};
3use std::path::Path;
4
5use crate::error::{KernelError, Result};
6
7use super::atomic::write_atomic;
8
9/// Credential store backed by a dotenv-style file.
10///
11/// Wraps a `HashMap<String, String>` with typed methods for load/save/normalize,
12/// keeping the ergonomics of a map via `Deref`/`DerefMut`.
13#[derive(Debug, Clone, Default)]
14pub struct SecretVault(HashMap<String, String>);
15
16impl SecretVault {
17    /// Create an empty vault with no credentials loaded.
18    pub fn empty() -> Self {
19        Self(HashMap::new())
20    }
21
22    /// Load a vault from a dotenv-style file at `path`.
23    ///
24    /// Returns an empty vault if the file does not exist.
25    /// Errors if the file is a symlink, has invalid UTF-8, or contains malformed lines.
26    pub fn load_from(path: impl AsRef<Path>) -> Result<Self> {
27        let path = path.as_ref();
28
29        // Symlink check BEFORE read to prevent TOCTOU race.
30        if path.exists() {
31            Self::guard_not_symlink(path)?;
32        }
33
34        let raw = match std::fs::read(path) {
35            Ok(d) => d,
36            Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Self::empty()),
37            Err(e) => return Err(e.into()),
38        };
39
40        raw.split(|&b| b == b'\n')
41            .enumerate()
42            .filter(|(_, line)| {
43                let text = std::str::from_utf8(line).unwrap_or("");
44                let trimmed = text.trim();
45                !trimmed.is_empty() && !trimmed.starts_with('#')
46            })
47            .try_fold(Self::empty(), |mut acc, (i, line)| {
48                let text = std::str::from_utf8(line)
49                    .map_err(|e| {
50                        KernelError::Vault(format!("invalid UTF-8 on line {}: {}", i + 1, e))
51                    })?
52                    .trim();
53                let (key, raw_val) = text.split_once('=').ok_or_else(|| {
54                    KernelError::Vault(format!("invalid secrets file line {}", i + 1))
55                })?;
56                if !is_valid_env_key(key) {
57                    return Err(KernelError::Vault(format!(
58                        "invalid secrets file line {}",
59                        i + 1
60                    )));
61                }
62                acc.0.insert(key.to_owned(), decode_shell_value(raw_val)?);
63                Ok(acc)
64            })
65    }
66
67    /// Persist the vault to a dotenv-style file at `path` using an atomic write.
68    pub fn persist_to(&self, path: impl AsRef<Path>) -> Result<()> {
69        let p = path.as_ref();
70        if let Some(parent) = p.parent() {
71            std::fs::create_dir_all(parent)?;
72        }
73
74        let body = self
75            .0
76            .keys()
77            .filter(|k| is_valid_env_key(k))
78            .collect::<std::collections::BTreeSet<_>>()
79            .iter()
80            .map(|k| format!("{}={}\n", k, encode_for_shell(&self.0[*k])))
81            .collect::<String>();
82
83        write_atomic(&p.to_string_lossy(), body.as_bytes(), 0o600)?;
84        #[cfg(unix)]
85        {
86            use std::os::unix::fs::PermissionsExt;
87            std::fs::set_permissions(p, std::fs::Permissions::from_mode(0o600))?;
88        }
89        Ok(())
90    }
91
92    fn guard_not_symlink(path: &Path) -> Result<()> {
93        let meta = std::fs::symlink_metadata(path)?;
94        if meta.file_type().is_symlink() {
95            return Err(KernelError::Vault(format!(
96                "secrets file is a symlink: {}",
97                path.display()
98            )));
99        }
100        Ok(())
101    }
102}
103
104// --- Deref/DerefMut so callers can use `.get()`, `.iter()`, etc. ---
105
106impl Deref for SecretVault {
107    type Target = HashMap<String, String>;
108    fn deref(&self) -> &Self::Target {
109        &self.0
110    }
111}
112
113impl DerefMut for SecretVault {
114    fn deref_mut(&mut self) -> &mut Self::Target {
115        &mut self.0
116    }
117}
118
119impl From<HashMap<String, String>> for SecretVault {
120    fn from(map: HashMap<String, String>) -> Self {
121        Self(map)
122    }
123}
124
125impl IntoIterator for SecretVault {
126    type Item = (String, String);
127    type IntoIter = std::collections::hash_map::IntoIter<String, String>;
128    fn into_iter(self) -> Self::IntoIter {
129        self.0.into_iter()
130    }
131}
132
133impl<'a> IntoIterator for &'a SecretVault {
134    type Item = (&'a String, &'a String);
135    type IntoIter = std::collections::hash_map::Iter<'a, String, String>;
136    fn into_iter(self) -> Self::IntoIter {
137        self.0.iter()
138    }
139}
140
141/// Mask a credential for display, showing only first/last 4 chars.
142pub fn redact_credential(value: &str) -> String {
143    match value.len() {
144        0 => String::new(),
145        1..=8 => "****".to_owned(),
146        _ => format!("{}****{}", &value[..4], &value[value.len() - 4..]),
147    }
148}
149
150// --- Internal helpers ---
151
152fn is_valid_env_key(key: &str) -> bool {
153    let first = key.as_bytes().first();
154    first.is_some_and(|&b| {
155        (b.is_ascii_uppercase() || b == b'_')
156            && key.as_bytes()[1..]
157                .iter()
158                .all(|b| b.is_ascii_uppercase() || b.is_ascii_digit() || *b == b'_')
159    })
160}
161
162fn decode_shell_value(value: &str) -> Result<String> {
163    let b = value.as_bytes();
164    match b.first() {
165        Some(b'\'') if b.last() == Some(&b'\'') && b.len() >= 2 => {
166            Ok(value[1..value.len() - 1].to_owned())
167        }
168        Some(b'$') if b.get(1) == Some(&b'\'') && b.last() == Some(&b'\'') => {
169            unescape_ansi(&value[2..value.len() - 1])
170        }
171        Some(b'"') if b.last() == Some(&b'"') && b.len() >= 2 => {
172            Ok(value[1..value.len() - 1].replace("\\\"", "\""))
173        }
174        _ => Ok(value.to_owned()),
175    }
176}
177
178fn unescape_ansi(s: &str) -> Result<String> {
179    let mut out = String::with_capacity(s.len());
180    let mut chars = s.as_bytes().iter().copied().peekable();
181    while let Some(b) = chars.next() {
182        if b != b'\\' {
183            out.push(b as char);
184            continue;
185        }
186        match chars.next() {
187            None => return Err(KernelError::Vault("unterminated escape".into())),
188            Some(b'n') => out.push('\n'),
189            Some(b't') => out.push('\t'),
190            Some(b'\\') => out.push('\\'),
191            Some(b'\'') => out.push('\''),
192            Some(other) => out.push(other as char),
193        }
194    }
195    Ok(out)
196}
197
198fn encode_for_shell(value: &str) -> String {
199    if value.is_empty() {
200        return "''".to_owned();
201    }
202    let needs_quoting = value
203        .as_bytes()
204        .iter()
205        .any(|b| matches!(b, b'\n' | b'\t' | b'\'' | b'\\' | b' '));
206    if !needs_quoting {
207        return value.to_owned();
208    }
209    let escaped = value
210        .replace('\\', "\\\\")
211        .replace('\n', "\\n")
212        .replace('\t', "\\t")
213        .replace('\'', "\\'");
214    format!("$'{}'", escaped)
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220
221    #[test]
222    fn test_redact_short() {
223        assert_eq!(redact_credential("ab"), "****");
224    }
225
226    #[test]
227    fn test_redact_empty() {
228        assert_eq!(redact_credential(""), "");
229    }
230
231    #[test]
232    fn test_redact_long() {
233        assert_eq!(redact_credential("abcdefghijklmnop"), "abcd****mnop");
234    }
235
236    #[test]
237    fn test_decode_single_quotes() {
238        assert_eq!(decode_shell_value("'hello world'").unwrap(), "hello world");
239    }
240
241    #[test]
242    fn test_decode_ansi_dollar_quotes() {
243        assert_eq!(
244            decode_shell_value("$'hello\\nworld'").unwrap(),
245            "hello\nworld"
246        );
247        assert_eq!(decode_shell_value("$'tab\\there'").unwrap(), "tab\there");
248        assert_eq!(
249            decode_shell_value("$'back\\\\slash'").unwrap(),
250            "back\\slash"
251        );
252        assert_eq!(decode_shell_value("$'quo\\'te'").unwrap(), "quo'te");
253    }
254
255    #[test]
256    fn test_decode_double_quotes() {
257        assert_eq!(
258            decode_shell_value("\"hello \\\"world\\\"\"").unwrap(),
259            "hello \"world\""
260        );
261    }
262
263    #[test]
264    fn test_decode_bare() {
265        assert_eq!(decode_shell_value("simple123").unwrap(), "simple123");
266    }
267
268    #[test]
269    fn test_encode_simple() {
270        assert_eq!(encode_for_shell("hello"), "hello");
271    }
272
273    #[test]
274    fn test_encode_empty() {
275        assert_eq!(encode_for_shell(""), "''");
276    }
277
278    #[test]
279    fn test_encode_special() {
280        let quoted = encode_for_shell("hello world");
281        assert!(
282            quoted.starts_with("$'"),
283            "expected $'...' for space, got {}",
284            quoted
285        );
286    }
287
288    #[test]
289    fn test_is_valid_env_key() {
290        assert!(is_valid_env_key("VALID_KEY"));
291        assert!(is_valid_env_key("_LEADING"));
292        assert!(!is_valid_env_key(""));
293        assert!(!is_valid_env_key("lowercase"));
294        assert!(!is_valid_env_key("1STARTS_NUM"));
295        assert!(is_valid_env_key("HAS_123"));
296        assert!(!is_valid_env_key("HAS-DASH"));
297    }
298
299    #[test]
300    fn test_roundtrip_via_impl_methods() {
301        let dir = tempfile::tempdir().expect("tempdir");
302        let path = dir.path().join("secrets.env");
303
304        let secrets = SecretVault::from(HashMap::from([
305            ("MY_KEY".to_string(), "my-value".to_string()),
306            ("OTHER_KEY".to_string(), "other".to_string()),
307        ]));
308
309        secrets.persist_to(&path).expect("persist");
310        let loaded = SecretVault::load_from(&path).expect("load");
311
312        assert_eq!(loaded.get("MY_KEY").map(|s| s.as_str()), Some("my-value"));
313        assert_eq!(loaded.get("OTHER_KEY").map(|s| s.as_str()), Some("other"));
314    }
315
316    #[test]
317    fn test_load_missing_returns_empty() {
318        let secrets =
319            SecretVault::load_from("/nonexistent/path/secrets.env").expect("load missing");
320        assert!(secrets.is_empty());
321    }
322
323    #[test]
324    fn test_roundtrip_with_special_chars() {
325        let dir = tempfile::tempdir().expect("tempdir");
326        let path = dir.path().join("secrets.env");
327
328        let secrets = SecretVault::from(HashMap::from([(
329            "MY_KEY".to_string(),
330            "value with spaces\nand newlines".to_string(),
331        )]));
332
333        secrets.persist_to(&path).expect("persist");
334        let loaded = SecretVault::load_from(&path).expect("load");
335
336        assert_eq!(
337            loaded.get("MY_KEY").map(|s| s.as_str()),
338            Some("value with spaces\nand newlines")
339        );
340    }
341}
llm_kernel/secrets/vault.rs

llm_kernel/secrets/
vault.rs
