cargo_lambda_metadata/
env.rs

1use clap::{ArgAction, Args, ValueHint};
2use env_file_reader::read_file;
3use miette::Result;
4use serde::{Deserialize, Serialize, ser::SerializeStruct};
5use std::{collections::HashMap, env::VarError, path::PathBuf};
6
7use crate::{cargo::deserialize_vec_or_map, error::MetadataError};
8
9pub type Environment = HashMap<String, String>;
10
11#[derive(Args, Clone, Debug, Default, Deserialize, Serialize)]
12pub struct EnvOptions {
13    /// Option to add one or many environment variables,
14    /// allows multiple repetitions (--env-var KEY=VALUE --env-var OTHER=NEW-VALUE).
15    /// It also allows to set a list of environment variables separated by commas
16    /// (e.g. --env-var KEY=VALUE,OTHER=NEW-VALUE).
17    #[arg(long, value_delimiter = ',', action = ArgAction::Append, visible_alias = "env-vars")]
18    #[serde(default, alias = "env", deserialize_with = "deserialize_vec_or_map")]
19    pub env_var: Option<Vec<String>>,
20
21    /// Read environment variables from a file.
22    /// Variables are separated by new lines in KEY=VALUE format.
23    #[arg(long, value_hint = ValueHint::FilePath)]
24    #[serde(default)]
25    pub env_file: Option<PathBuf>,
26}
27
28impl EnvOptions {
29    pub fn lambda_environment(
30        &self,
31        base: &HashMap<String, String>,
32    ) -> Result<Environment, MetadataError> {
33        lambda_environment(Some(base), &self.env_file, self.env_var.as_ref())
34    }
35
36    pub fn count_fields(&self) -> usize {
37        self.env_var.is_some() as usize + self.env_file.is_some() as usize
38    }
39
40    pub fn serialize_fields<S>(
41        &self,
42        state: &mut <S as serde::Serializer>::SerializeStruct,
43    ) -> Result<(), S::Error>
44    where
45        S: serde::Serializer,
46    {
47        if let Some(env_var) = &self.env_var {
48            state.serialize_field("env_var", env_var)?;
49        }
50        if let Some(env_file) = &self.env_file {
51            state.serialize_field("env_file", env_file)?;
52        }
53        Ok(())
54    }
55}
56
57pub(crate) fn lambda_environment(
58    base: Option<&HashMap<String, String>>,
59    env_file: &Option<PathBuf>,
60    vars: Option<&Vec<String>>,
61) -> Result<Environment, MetadataError> {
62    let mut env = HashMap::new();
63
64    if let Some(base) = base.cloned() {
65        env.extend(base);
66    }
67
68    if let Some(path) = env_file {
69        if path.is_file() {
70            let env_variables =
71                read_file(path).map_err(|e| MetadataError::InvalidEnvFile(path.into(), e))?;
72            for (key, value) in env_variables {
73                env.insert(key, value);
74            }
75        } else {
76            return Err(MetadataError::InvalidEnvFile(
77                path.into(),
78                std::io::Error::new(std::io::ErrorKind::NotFound, "file not found"),
79            ));
80        }
81    }
82
83    if let Some(vars) = vars {
84        for var in vars {
85            let (key, value) = extract_var(var)?;
86            env.insert(key.to_string(), value.to_string());
87        }
88    }
89
90    Ok(env)
91}
92
93fn extract_var(line: &str) -> Result<(&str, &str), MetadataError> {
94    let mut iter = line.trim().splitn(2, '=');
95
96    let key = iter
97        .next()
98        .map(|s| s.trim())
99        .ok_or_else(|| MetadataError::InvalidEnvVar(line.into()))?;
100    if key.is_empty() {
101        Err(MetadataError::InvalidEnvVar(line.into()))?;
102    }
103
104    let value = iter
105        .next()
106        .map(|s| s.trim())
107        .ok_or_else(|| MetadataError::InvalidEnvVar(line.into()))?;
108    if value.is_empty() {
109        Err(MetadataError::InvalidEnvVar(line.into()))?;
110    }
111
112    Ok((key, value))
113}
114
115pub trait EnvVarExtractor {
116    fn var(&self, name: &str) -> Result<String, VarError>;
117}
118
119pub struct SystemEnvExtractor;
120
121impl EnvVarExtractor for SystemEnvExtractor {
122    fn var(&self, name: &str) -> Result<String, VarError> {
123        std::env::var(name)
124    }
125}
126
127pub struct HashMapEnvExtractor {
128    env: HashMap<String, String>,
129}
130
131impl From<Vec<(&str, &str)>> for HashMapEnvExtractor {
132    fn from(env: Vec<(&str, &str)>) -> Self {
133        Self {
134            env: env
135                .into_iter()
136                .map(|(k, v)| (k.to_string(), v.to_string()))
137                .collect(),
138        }
139    }
140}
141
142impl EnvVarExtractor for HashMapEnvExtractor {
143    fn var(&self, name: &str) -> Result<String, VarError> {
144        self.env.get(name).cloned().ok_or(VarError::NotPresent)
145    }
146}
147
148#[cfg(test)]
149mod test {
150    use std::env::temp_dir;
151
152    use super::*;
153
154    #[test]
155    fn test_extract_var() {
156        let (k, v) = extract_var("FOO=BAR").unwrap();
157        assert_eq!("FOO", k);
158        assert_eq!("BAR", v);
159
160        let (k, v) = extract_var(" FOO = BAR ").unwrap();
161        assert_eq!("FOO", k);
162        assert_eq!("BAR", v);
163
164        extract_var("=BAR").expect_err("missing key");
165        extract_var("FOO=").expect_err("missing value");
166        extract_var("  ").expect_err("missing variable");
167    }
168
169    #[test]
170    fn test_empty_environment() {
171        let env = lambda_environment(None, &None, None).unwrap();
172        assert!(env.is_empty());
173    }
174
175    #[test]
176    fn test_base_environment() {
177        let mut base = HashMap::new();
178        base.insert("FOO".into(), "BAR".into());
179        let env = lambda_environment(Some(&base), &None, None).unwrap();
180
181        assert_eq!("BAR".to_string(), env["FOO"]);
182    }
183
184    #[test]
185    fn test_environment_with_flags() {
186        let mut base = HashMap::new();
187        base.insert("FOO".into(), "BAR".into());
188
189        let flags = vec!["FOO=QUX".to_string(), "BAZ=QUUX".to_string()];
190        let env = lambda_environment(Some(&base), &None, Some(&flags)).unwrap();
191
192        assert_eq!("QUX".to_string(), env["FOO"]);
193        assert_eq!("QUUX".to_string(), env["BAZ"]);
194    }
195
196    #[test]
197    fn test_environment_with_file() {
198        let file = temp_dir().join(".env");
199        std::fs::write(&file, "BAR=BAZ\n\nexport QUUX = 'QUUUX'\n#IGNORE=ME").unwrap();
200
201        let mut base = HashMap::new();
202        base.insert("FOO".into(), "BAR".into());
203
204        let flags = vec!["FOO=QUX".to_string(), "BAZ=QUUX".to_string()];
205        let vars = lambda_environment(Some(&base), &Some(file), Some(&flags)).unwrap();
206
207        assert_eq!("QUX".to_string(), vars["FOO"]);
208        assert_eq!("QUUX".to_string(), vars["BAZ"]);
209        assert_eq!("BAZ".to_string(), vars["BAR"]);
210        assert_eq!("QUUUX".to_string(), vars["QUUX"]);
211        assert!(!vars.contains_key("IGNORE"));
212        assert!(!vars.contains_key(""));
213    }
214
215    #[test]
216    fn test_environment_with_missing_file() {
217        let file = temp_dir().join(".env.nonexistent");
218        // Ensure the file doesn't exist
219        let _ = std::fs::remove_file(&file);
220
221        let result = lambda_environment(None, &Some(file), None);
222        assert!(result.is_err());
223        assert!(matches!(result, Err(MetadataError::InvalidEnvFile(_, _))));
224    }
225}
cargo_lambda_metadata/env.rs

cargo_lambda_metadata/
env.rs
