riglr_core/util/
mod.rs

1//! Utility modules and functions for riglr-core
2
3pub mod rate_limit_strategy;
4pub mod rate_limiter;
5pub mod secure_keys;
6pub mod token_bucket;
7
8// Re-export main types for convenience
9pub use rate_limit_strategy::{FixedWindowStrategy, RateLimitStrategy};
10pub use rate_limiter::{RateLimitStrategyType, RateLimiter, RateLimiterBuilder};
11pub use secure_keys::{
12    ensure_key_directory, get_default_key_directory, load_private_key_from_file,
13    load_private_key_with_fallback,
14};
15pub use token_bucket::TokenBucketStrategy;
16
17use std::env;
18
19// Constants for doctest environment variables
20#[doc(hidden)]
21pub const DOCTEST_API_KEY: &str = "MY_API_KEY";
22#[doc(hidden)]
23pub const DOCTEST_OPTIONAL_SETTING: &str = "OPTIONAL_SETTING";
24#[doc(hidden)]
25pub const DOCTEST_API_KEY_VALIDATE: &str = "API_KEY";
26#[doc(hidden)]
27pub const DOCTEST_DATABASE_URL: &str = "DATABASE_URL";
28#[doc(hidden)]
29pub const DOCTEST_VAR1: &str = "VAR1";
30#[doc(hidden)]
31pub const DOCTEST_VAR2: &str = "VAR2";
32#[doc(hidden)]
33pub const DOCTEST_VAR3: &str = "VAR3";
34#[doc(hidden)]
35pub const DOCTEST_ENV_FILE_VAR: &str = "TEST_ENV_FILE_VAR";
36
37/// Error type for environment variable operations
38#[derive(Debug, thiserror::Error)]
39pub enum EnvError {
40    /// Required environment variable is not set
41    #[error("Environment variable '{0}' is required but not set")]
42    MissingRequired(String),
43
44    /// Environment variable contains invalid UTF-8
45    #[error("Environment variable '{0}' contains invalid UTF-8")]
46    InvalidUtf8(String),
47}
48
49/// Result type alias for environment operations
50pub type EnvResult<T> = Result<T, EnvError>;
51
52/// Gets a required environment variable, returning an error if not set.
53///
54/// This is the recommended approach for libraries, allowing the application
55/// to decide how to handle missing configuration.
56///
57/// # Examples
58///
59/// ```rust
60/// use riglr_core::util::{get_required_env, DOCTEST_API_KEY};
61///
62/// # std::env::set_var(DOCTEST_API_KEY, "secret123");
63/// let api_key = get_required_env(DOCTEST_API_KEY).expect("MY_API_KEY must be set");
64/// assert_eq!(api_key, "secret123");
65/// # std::env::remove_var(DOCTEST_API_KEY);
66/// ```
67///
68/// # Errors
69///
70/// Returns [`EnvError::MissingRequired`] if the environment variable is not set.
71pub fn get_required_env(key: &str) -> EnvResult<String> {
72    env::var(key).map_err(|_| EnvError::MissingRequired(key.to_string()))
73}
74
75/// Gets an optional environment variable with a default value.
76///
77/// # Examples
78///
79/// ```rust
80/// use riglr_core::util::{get_env_or_default, DOCTEST_OPTIONAL_SETTING};
81///
82/// # std::env::remove_var(DOCTEST_OPTIONAL_SETTING);
83/// let setting = get_env_or_default(DOCTEST_OPTIONAL_SETTING, "default_value");
84/// assert_eq!(setting, "default_value");
85/// ```
86pub fn get_env_or_default(key: &str, default: &str) -> String {
87    env::var(key).unwrap_or_else(|_| default.to_string())
88}
89
90/// Validates that all required environment variables are set.
91///
92/// This is useful during application initialization to fail fast if
93/// configuration is incomplete.
94///
95/// # Examples
96///
97/// ```rust
98/// use riglr_core::util::{validate_required_env, DOCTEST_API_KEY_VALIDATE, DOCTEST_DATABASE_URL};
99///
100/// # std::env::set_var(DOCTEST_API_KEY_VALIDATE, "value1");
101/// # std::env::set_var(DOCTEST_DATABASE_URL, "value2");
102/// let required = vec![DOCTEST_API_KEY_VALIDATE, DOCTEST_DATABASE_URL];
103/// validate_required_env(&required).expect("Missing required environment variables");
104/// # std::env::remove_var(DOCTEST_API_KEY_VALIDATE);
105/// # std::env::remove_var(DOCTEST_DATABASE_URL);
106/// ```
107///
108/// # Errors
109///
110/// Returns the first [`EnvError::MissingRequired`] encountered.
111pub fn validate_required_env(keys: &[&str]) -> EnvResult<()> {
112    for key in keys {
113        get_required_env(key)?;
114    }
115    Ok(())
116}
117
118/// Gets multiple environment variables at once, returning a map.
119///
120/// # Examples
121///
122/// ```rust
123/// use riglr_core::util::{get_env_vars, DOCTEST_VAR1, DOCTEST_VAR2, DOCTEST_VAR3};
124/// use std::collections::HashMap;
125///
126/// # std::env::set_var(DOCTEST_VAR1, "value1");
127/// # std::env::set_var(DOCTEST_VAR2, "value2");
128/// let vars = get_env_vars(&[DOCTEST_VAR1, DOCTEST_VAR2, DOCTEST_VAR3]);
129/// assert_eq!(vars.get(DOCTEST_VAR1), Some(&"value1".to_string()));
130/// assert_eq!(vars.get(DOCTEST_VAR3), None);
131/// # std::env::remove_var(DOCTEST_VAR1);
132/// # std::env::remove_var(DOCTEST_VAR2);
133/// ```
134pub fn get_env_vars(keys: &[&str]) -> std::collections::HashMap<String, String> {
135    keys.iter()
136        .filter_map(|&key| env::var(key).ok().map(|value| (key.to_string(), value)))
137        .collect()
138}
139
140/// Application-level helper that initializes environment from a `.env` file if present.
141///
142/// This is a convenience function for applications (not libraries) that want to
143/// support `.env` files for local development.
144///
145/// # Examples
146///
147/// ```rust
148/// use riglr_core::util::init_env_from_file;
149///
150/// // Load .env file if it exists (usually at application startup)
151/// init_env_from_file(".env").ok(); // Ignore if file doesn't exist
152/// ```
153pub fn init_env_from_file(path: &str) -> std::io::Result<()> {
154    if std::path::Path::new(path).exists() {
155        dotenv::from_filename(path)
156            .map_err(|e| std::io::Error::other(format!("Failed to load .env file: {}", e)))?;
157    }
158    Ok(())
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164    use std::env;
165
166    // Test environment variable constants
167    const TEST_VAR_EXISTS: &str = "TEST_VAR_EXISTS";
168    const TEST_VAR_MISSING: &str = "TEST_VAR_MISSING";
169    const TEST_REQUIRED_VAR: &str = "TEST_REQUIRED_VAR";
170    const TEST_MISSING_REQUIRED: &str = "TEST_MISSING_REQUIRED";
171    const TEST_VALIDATE_VAR1: &str = "TEST_VALIDATE_VAR1";
172    const TEST_VALIDATE_VAR2: &str = "TEST_VALIDATE_VAR2";
173    const TEST_VALIDATE_MISSING_VAR1: &str = "TEST_VALIDATE_MISSING_VAR1";
174    const TEST_VALIDATE_MISSING_VAR2: &str = "TEST_VALIDATE_MISSING_VAR2";
175    const TEST_MULTI_1: &str = "TEST_MULTI_1";
176    const TEST_MULTI_2: &str = "TEST_MULTI_2";
177    const TEST_MULTI_3: &str = "TEST_MULTI_3";
178    const TEST_NONEXISTENT_VAR: &str = "NONEXISTENT_VAR_EMPTY_DEFAULT";
179    const TEST_EMPTY_VALUE: &str = "TEST_EMPTY_VALUE";
180    const TEST_EMPTY_REQUIRED: &str = "TEST_EMPTY_REQUIRED";
181    const TEST_FIRST_MISSING: &str = "FIRST_MISSING";
182    const TEST_SECOND_EXISTS: &str = "SECOND_EXISTS";
183    const TEST_SPECIAL_CHARS: &str = "TEST_SPECIAL_CHARS";
184
185    #[test]
186    fn test_get_env_or_default_with_existing_var() {
187        env::set_var(TEST_VAR_EXISTS, "test_value");
188        let result = get_env_or_default(TEST_VAR_EXISTS, "default");
189        assert_eq!(result, "test_value");
190        env::remove_var(TEST_VAR_EXISTS);
191    }
192
193    #[test]
194    fn test_get_env_or_default_with_missing_var() {
195        env::remove_var(TEST_VAR_MISSING);
196        let result = get_env_or_default(TEST_VAR_MISSING, "default_value");
197        assert_eq!(result, "default_value");
198    }
199
200    #[test]
201    fn test_get_required_env_with_existing_var() {
202        env::set_var(TEST_REQUIRED_VAR, "required_value");
203        let result = get_required_env(TEST_REQUIRED_VAR).unwrap();
204        assert_eq!(result, "required_value");
205        env::remove_var(TEST_REQUIRED_VAR);
206    }
207
208    #[test]
209    fn test_get_required_env_with_missing_var() {
210        env::remove_var(TEST_MISSING_REQUIRED);
211        let result = get_required_env(TEST_MISSING_REQUIRED);
212        assert!(result.is_err());
213        match result {
214            Err(EnvError::MissingRequired(key)) => {
215                assert_eq!(key, TEST_MISSING_REQUIRED);
216            }
217            _ => panic!("Expected MissingRequired error"),
218        }
219    }
220
221    #[test]
222    fn test_validate_required_env_all_present() {
223        env::set_var(TEST_VALIDATE_VAR1, "value1");
224        env::set_var(TEST_VALIDATE_VAR2, "value2");
225
226        let result = validate_required_env(&[TEST_VALIDATE_VAR1, TEST_VALIDATE_VAR2]);
227        assert!(result.is_ok());
228
229        env::remove_var(TEST_VALIDATE_VAR1);
230        env::remove_var(TEST_VALIDATE_VAR2);
231    }
232
233    #[test]
234    fn test_validate_required_env_missing_one() {
235        env::set_var(TEST_VALIDATE_MISSING_VAR1, "value1");
236        env::remove_var(TEST_VALIDATE_MISSING_VAR2);
237
238        let result =
239            validate_required_env(&[TEST_VALIDATE_MISSING_VAR1, TEST_VALIDATE_MISSING_VAR2]);
240        assert!(result.is_err());
241
242        env::remove_var(TEST_VALIDATE_MISSING_VAR1);
243    }
244
245    #[test]
246    fn test_get_env_vars() {
247        env::set_var(TEST_MULTI_1, "value1");
248        env::set_var(TEST_MULTI_2, "value2");
249        env::remove_var(TEST_MULTI_3);
250
251        let vars = get_env_vars(&[TEST_MULTI_1, TEST_MULTI_2, TEST_MULTI_3]);
252
253        assert_eq!(vars.get(TEST_MULTI_1), Some(&"value1".to_string()));
254        assert_eq!(vars.get(TEST_MULTI_2), Some(&"value2".to_string()));
255        assert_eq!(vars.get(TEST_MULTI_3), None);
256
257        env::remove_var(TEST_MULTI_1);
258        env::remove_var(TEST_MULTI_2);
259    }
260
261    #[test]
262    fn test_get_env_vars_with_empty_array() {
263        let vars = get_env_vars(&[]);
264        assert!(vars.is_empty());
265    }
266
267    #[test]
268    fn test_validate_required_env_with_empty_array() {
269        let result = validate_required_env(&[]);
270        assert!(result.is_ok());
271    }
272
273    #[test]
274    fn test_env_error_display_missing_required() {
275        let error = EnvError::MissingRequired("TEST_KEY".to_string());
276        assert_eq!(
277            error.to_string(),
278            "Environment variable 'TEST_KEY' is required but not set"
279        );
280    }
281
282    #[test]
283    fn test_env_error_display_invalid_utf8() {
284        let error = EnvError::InvalidUtf8("TEST_KEY".to_string());
285        assert_eq!(
286            error.to_string(),
287            "Environment variable 'TEST_KEY' contains invalid UTF-8"
288        );
289    }
290
291    #[test]
292    fn test_env_error_debug_format() {
293        let error = EnvError::MissingRequired("TEST_KEY".to_string());
294        let debug_str = format!("{:?}", error);
295        assert!(debug_str.contains("MissingRequired"));
296        assert!(debug_str.contains("TEST_KEY"));
297    }
298
299    #[test]
300    fn test_init_env_from_file_with_nonexistent_file() {
301        let result = init_env_from_file("nonexistent_file.env");
302        assert!(result.is_ok());
303    }
304
305    #[test]
306    fn test_init_env_from_file_with_existing_file() {
307        use std::fs;
308        use std::io::Write;
309
310        // Create a temporary .env file
311        let temp_file = "test_temp.env";
312        let mut file = fs::File::create(temp_file).expect("Failed to create temp file");
313        writeln!(file, "{}=test_value", DOCTEST_ENV_FILE_VAR)
314            .expect("Failed to write to temp file");
315        drop(file);
316
317        // Test loading the file
318        let result = init_env_from_file(temp_file);
319        assert!(result.is_ok());
320
321        // Verify the environment variable was loaded
322        let loaded_value = env::var(DOCTEST_ENV_FILE_VAR).ok();
323        assert_eq!(loaded_value, Some("test_value".to_string()));
324
325        // Clean up
326        fs::remove_file(temp_file).ok();
327        env::remove_var(DOCTEST_ENV_FILE_VAR);
328    }
329
330    #[test]
331    fn test_init_env_from_file_with_invalid_file() {
332        use std::fs;
333        use std::io::Write;
334
335        // Create a temporary file with invalid content that will cause dotenv to fail
336        let temp_file = "test_invalid.env";
337        let mut file = fs::File::create(temp_file).expect("Failed to create temp file");
338        // Write invalid UTF-8 bytes
339        file.write_all(&[0xFF, 0xFE])
340            .expect("Failed to write invalid bytes");
341        drop(file);
342
343        // Test loading the invalid file
344        let result = init_env_from_file(temp_file);
345        assert!(result.is_err());
346
347        // Clean up
348        fs::remove_file(temp_file).ok();
349    }
350
351    #[test]
352    fn test_get_env_or_default_with_empty_string_default() {
353        env::remove_var(TEST_NONEXISTENT_VAR);
354        let result = get_env_or_default(TEST_NONEXISTENT_VAR, "");
355        assert_eq!(result, "");
356    }
357
358    #[test]
359    fn test_get_env_or_default_with_empty_string_value() {
360        env::set_var(TEST_EMPTY_VALUE, "");
361        let result = get_env_or_default(TEST_EMPTY_VALUE, "default");
362        assert_eq!(result, "");
363        env::remove_var(TEST_EMPTY_VALUE);
364    }
365
366    #[test]
367    fn test_get_required_env_with_empty_string_value() {
368        env::set_var(TEST_EMPTY_REQUIRED, "");
369        let result = get_required_env(TEST_EMPTY_REQUIRED).unwrap();
370        assert_eq!(result, "");
371        env::remove_var(TEST_EMPTY_REQUIRED);
372    }
373
374    #[test]
375    fn test_validate_required_env_fails_on_first_missing() {
376        env::remove_var(TEST_FIRST_MISSING);
377        env::set_var(TEST_SECOND_EXISTS, "value");
378
379        let result = validate_required_env(&[TEST_FIRST_MISSING, TEST_SECOND_EXISTS]);
380        assert!(result.is_err());
381
382        match result {
383            Err(EnvError::MissingRequired(key)) => {
384                assert_eq!(key, TEST_FIRST_MISSING);
385            }
386            _ => panic!("Expected MissingRequired error for first missing variable"),
387        }
388
389        env::remove_var(TEST_SECOND_EXISTS);
390    }
391
392    #[test]
393    fn test_get_env_vars_with_special_characters() {
394        env::set_var(TEST_SPECIAL_CHARS, "value with spaces & symbols!");
395
396        let vars = get_env_vars(&[TEST_SPECIAL_CHARS]);
397        assert_eq!(
398            vars.get(TEST_SPECIAL_CHARS),
399            Some(&"value with spaces & symbols!".to_string())
400        );
401
402        env::remove_var(TEST_SPECIAL_CHARS);
403    }
404}
riglr_core/util/mod.rs

riglr_core/util/
mod.rs
