statespace_tool_runtime/
env_validation.rs

1use std::collections::HashMap;
2
3pub const MAX_ENV_VAR_COUNT: usize = 64;
4pub const MAX_ENV_VAR_KEY_BYTES: usize = 64;
5pub const MAX_ENV_VAR_VALUE_BYTES: usize = 4 * 1024;
6pub const MAX_ENV_TOTAL_BYTES: usize = 16 * 1024;
7
8const RESERVED_ENV_PREFIXES: &[&str] = &["AWS_", "LD_", "DYLD_", "_LAMBDA", "_HANDLER"];
9const RESERVED_ENV_KEYS: &[&str] = &[
10    "HOME",
11    "LANG",
12    "LC_ALL",
13    "PATH",
14    "STATESPACE_SCRATCH",
15    "STATESPACE_WORKSPACE",
16];
17
18#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
19pub enum EnvValidationError {
20    #[error("too many environment variables (max {max})")]
21    TooManyEntries { max: usize },
22
23    #[error("invalid environment variable name '{key}'")]
24    InvalidKeyName { key: String },
25
26    #[error("environment variable '{key}' value is too long (max {max} bytes)")]
27    ValueTooLong { key: String, max: usize },
28
29    #[error("environment variable '{key}' contains control characters")]
30    ValueContainsControlChars { key: String },
31
32    #[error("environment variables exceed total size limit (max {max} bytes)")]
33    TotalBytesExceeded { max: usize },
34}
35
36#[must_use]
37pub fn is_reserved_env_key(key: &str) -> bool {
38    RESERVED_ENV_KEYS.contains(&key) || RESERVED_ENV_PREFIXES.iter().any(|p| key.starts_with(p))
39}
40
41/// Validate user-provided environment variables before merging/execution.
42///
43/// This enforces structural safety only (name shape + size + control chars),
44/// not semantic typing for values.
45///
46/// # Errors
47///
48/// Returns [`EnvValidationError`] when key names, value bytes, or total map
49/// size exceed the runtime limits.
50pub fn validate_env_map<S: std::hash::BuildHasher>(
51    env: &HashMap<String, String, S>,
52) -> Result<(), EnvValidationError> {
53    if env.len() > MAX_ENV_VAR_COUNT {
54        return Err(EnvValidationError::TooManyEntries {
55            max: MAX_ENV_VAR_COUNT,
56        });
57    }
58
59    let mut total_bytes = 0usize;
60
61    for (key, value) in env {
62        if !is_valid_env_key(key) {
63            return Err(EnvValidationError::InvalidKeyName {
64                key: display_key(key),
65            });
66        }
67
68        if value.len() > MAX_ENV_VAR_VALUE_BYTES {
69            return Err(EnvValidationError::ValueTooLong {
70                key: key.clone(),
71                max: MAX_ENV_VAR_VALUE_BYTES,
72            });
73        }
74
75        if value.chars().any(|ch| ch == '\0' || ch.is_ascii_control()) {
76            return Err(EnvValidationError::ValueContainsControlChars { key: key.clone() });
77        }
78
79        total_bytes += key.len() + value.len();
80        if total_bytes > MAX_ENV_TOTAL_BYTES {
81            return Err(EnvValidationError::TotalBytesExceeded {
82                max: MAX_ENV_TOTAL_BYTES,
83            });
84        }
85    }
86
87    Ok(())
88}
89
90fn is_valid_env_key(key: &str) -> bool {
91    if key.is_empty() || key.len() > MAX_ENV_VAR_KEY_BYTES {
92        return false;
93    }
94
95    let mut chars = key.chars();
96    let Some(first) = chars.next() else {
97        return false;
98    };
99
100    if !(first.is_ascii_alphabetic() || first == '_') {
101        return false;
102    }
103
104    chars.all(|ch| ch.is_ascii_alphanumeric() || ch == '_')
105}
106
107fn display_key(key: &str) -> String {
108    if key.is_empty() {
109        "<empty>".to_string()
110    } else {
111        key.chars().flat_map(char::escape_default).collect()
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    #[test]
120    fn accepts_valid_env_map() {
121        let env = HashMap::from([
122            ("USER_ID".to_string(), "42".to_string()),
123            ("PAGE".to_string(), "stats".to_string()),
124        ]);
125
126        assert!(validate_env_map(&env).is_ok());
127    }
128
129    #[test]
130    fn rejects_invalid_key_name() {
131        let env = HashMap::from([("USER-ID".to_string(), "42".to_string())]);
132
133        assert!(matches!(
134            validate_env_map(&env),
135            Err(EnvValidationError::InvalidKeyName { .. })
136        ));
137    }
138
139    #[test]
140    fn rejects_control_characters_in_value() {
141        let env = HashMap::from([("USER_ID".to_string(), "abc\nxyz".to_string())]);
142
143        assert!(matches!(
144            validate_env_map(&env),
145            Err(EnvValidationError::ValueContainsControlChars { .. })
146        ));
147    }
148
149    #[test]
150    fn rejects_oversized_value() {
151        let env = HashMap::from([(
152            "USER_ID".to_string(),
153            "x".repeat(MAX_ENV_VAR_VALUE_BYTES + 1),
154        )]);
155
156        assert!(matches!(
157            validate_env_map(&env),
158            Err(EnvValidationError::ValueTooLong { .. })
159        ));
160    }
161
162    #[test]
163    fn recognizes_reserved_env_keys() {
164        assert!(is_reserved_env_key("HOME"));
165        assert!(is_reserved_env_key("LC_ALL"));
166        assert!(is_reserved_env_key("AWS_ACCESS_KEY_ID"));
167        assert!(!is_reserved_env_key("USER_ID"));
168    }
169
170    #[test]
171    fn display_key_escapes_control_characters() {
172        assert_eq!(display_key("BAD\nKEY"), "BAD\\nKEY");
173    }
174}
statespace_tool_runtime/env_validation.rs

statespace_tool_runtime/
env_validation.rs
