apollo_configuration/

expansion.rs

//! Traits and types to customize `${}` expansion functionality.

use crate::YamlLocationData;
use crate::errors::ExpandError;
use crate::errors::ExpandErrorKind;
use jsonschema::paths::LazyLocation;
use jsonschema::paths::Location;
use miette::SourceSpan;
use std::borrow::Cow;
use std::collections::HashMap;
use std::env::VarError;
use std::sync::Arc;

/// Trait for resolving `${env.VAR_NAME}` style expansions.
pub trait VariableProvider {
    /// Return the value to substitute for a `${env.VAR_NAME}` style expansion.
    fn get<'a>(&'a self, name: &str) -> Result<Cow<'a, str>, VarError>;
}

impl VariableProvider for () {
    fn get<'a>(&'a self, _name: &str) -> Result<Cow<'a, str>, VarError> {
        Err(VarError::NotPresent)
    }
}

/// A variable provider that returns values from environment variables using [`std::env::var`].
#[derive(Debug, Clone, Default)]
pub struct EnvVariables;
impl VariableProvider for EnvVariables {
    fn get<'a>(&'a self, name: &str) -> Result<Cow<'a, str>, VarError> {
        std::env::var(name).map(Cow::Owned)
    }
}

/// A variable provider that returns values from a fixed map of data, for example for mocking.
#[derive(Debug, Clone, Default)]
pub struct MapVariables(pub HashMap<String, String>);
impl VariableProvider for MapVariables {
    fn get<'a>(&'a self, name: &str) -> Result<Cow<'a, str>, VarError> {
        self.0
            .get(name)
            .map(|value| Cow::Borrowed(value.as_str()))
            .ok_or(VarError::NotPresent)
    }
}

fn expand_variable<'v>(
    variable: &str,
    span: SourceSpan,
    vars: &'v dyn VariableProvider,
) -> Result<Cow<'v, str>, ExpandErrorKind> {
    let Some((kind, argument)) = variable.split_once('.') else {
        return Err(ExpandErrorKind::UnprefixedExpansion {
            span,
            name: variable.to_string(),
        });
    };
    // XXX(@goto-bus-stop): We could instead give the whole value to the variable provider,
    // so users could add more than just `env.` (and `file.` which will probably be default in
    // future).
    match kind {
        "env" => match vars.get(argument) {
            Ok(value) => Ok(value),
            Err(VarError::NotPresent) => Err(ExpandErrorKind::NotPresent { span }),
            Err(VarError::NotUnicode(_)) => Err(ExpandErrorKind::NotUnicode { span }),
        },
        _ => Err(ExpandErrorKind::InvalidExpansion { span }),
    }
}

fn expand_inner(
    value: &mut serde_json::Value,
    vars: &dyn VariableProvider,
    location_data: &YamlLocationData,
    path: LazyLocation,
    errors: &mut Vec<ExpandErrorKind>,
) {
    match value {
        serde_json::Value::Object(map) => {
            for (key, value) in map.iter_mut() {
                expand_inner(value, vars, location_data, path.push(key.as_str()), errors);
            }
        }
        serde_json::Value::Array(vec) => {
            for (key, value) in vec.iter_mut().enumerate() {
                expand_inner(value, vars, location_data, path.push(key), errors);
            }
        }
        serde_json::Value::String(string) => {
            let location = Location::from(&path);
            // XXX(@goto-bus-stop): ideally this would point to the actual variable, not the whole
            // string
            let span = location_data
                .resolve_instance_span(&location)
                .unwrap_or_else(|| SourceSpan::new(0.into(), 0));
            match shellexpand::env_with_context(string, |name| {
                // XXX(@goto-bus-stop): we could allow using variable expansion on numeric fields
                // if we used schema type information here
                expand_variable(name, span, vars).map(Some)
            }) {
                Ok(Cow::Borrowed(_)) => {
                    // Unchanged
                }
                Ok(Cow::Owned(expanded)) => {
                    *string = expanded;
                }
                Err(err) => errors.push(err.cause),
            }
        }
        _ => {}
    }
}

pub(crate) fn expand(
    value: &mut serde_json::Value,
    vars: &dyn VariableProvider,
    location_data: Option<&YamlLocationData>,
) -> Result<(), ExpandError> {
    let no_yaml = YamlLocationData::empty();
    let mut errors = vec![];
    expand_inner(
        value,
        vars,
        location_data.unwrap_or(&no_yaml),
        LazyLocation::new(),
        &mut errors,
    );

    if errors.is_empty() {
        Ok(())
    } else {
        Err(ExpandError {
            source_code: location_data
                .map(|data| data.source_code.clone())
                .unwrap_or_else(|| Arc::from("")),
            errors,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::MapVariables;
    use super::VariableProvider;
    use super::expand;
    use miette::EyreContext as _;
    use serde_json::json;
    use std::collections::HashMap;
    use std::ffi::OsString;

    #[test]
    fn expand_env_variable() {
        let mut json = json!("${env.REPLACE}");
        let vars = MapVariables(HashMap::from([(
            "REPLACE".to_string(),
            "value".to_string(),
        )]));
        expand(&mut json, &vars, None).expect("should succeed");
        assert_eq!(json.as_str(), Some("value"));
    }

    #[test]
    fn expand_env_variable_inside_string() {
        let mut json = json!("this will be ${env.REPLACE}");
        let vars = MapVariables(HashMap::from([(
            "REPLACE".to_string(),
            "value".to_string(),
        )]));
        expand(&mut json, &vars, None).expect("should succeed");
        assert_eq!(json.as_str(), Some("this will be value"));
    }

    #[test]
    fn expand_env_variable_inside_nesting() {
        let mut json = json!({
            "object": {
                "array": ["${env.REPLACE}"],
            }
        });
        let vars = MapVariables(HashMap::from([(
            "REPLACE".to_string(),
            "value".to_string(),
        )]));
        expand(&mut json, &vars, None).expect("should succeed");
        assert_eq!(json["object"]["array"][0].as_str(), Some("value"));
    }

    #[test]
    fn missing_variable_error() {
        let mut json = json!({
            "object": {
                "array": ["${env.REPLACE}"],
            }
        });
        let vars = MapVariables(HashMap::new());
        let err = expand(&mut json, &vars, None).expect_err("should fail");
        println!("{:?}", err);
    }

    #[test]
    fn missing_defaulted_variable() {
        let mut json = json!({
            "object": {
                "array": ["${env.REPLACE:-default}"],
            }
        });
        let vars = MapVariables(HashMap::new());
        expand(&mut json, &vars, None).expect("should succeed");
        assert_eq!(json["object"]["array"][0].as_str(), Some("default"));
    }

    #[test]
    fn unprefixed_replacement() {
        let mut json = json!({
            "object": {
                "array": ["$REPLACE"],
            }
        });
        let vars = MapVariables(HashMap::new());
        let err = expand(&mut json, &vars, None).expect_err("should fail");
        println!("{:?}", miette::Report::new(err));
    }

    #[test]
    fn unsupported_prefix() {
        let mut json = json!({
            "object": {
                "array": ["${other.REPLACE}"],
            }
        });
        let vars = MapVariables(HashMap::new());
        let err = expand(&mut json, &vars, None).expect_err("should fail");
        println!("{:?}", miette::Report::new(err));
    }

    #[test]
    fn collects_all_errors() {
        let mut json = json!([
            "$REPLACE ${other.REPLACE}",
            { "works": "${env.REPLACE}" },
            { "nested": "${env.DOESNOTEXIST}" },
        ]);
        let vars = MapVariables(HashMap::from([(
            "REPLACE".to_string(),
            "value".to_string(),
        )]));
        let err = expand(&mut json, &vars, None).expect_err("should fail");
        println!("{:?}", miette::Report::new(err));
    }

    #[test]
    fn diagnostics() {
        let mut json = json!([
            "$REPLACE ${other.REPLACE}",
            { "works": "${env.REPLACE}" },
            { "nested": "${env.DOESNOTEXIST}" },
        ]);
        let source_code = r#"
          - '$REPLACE ${other.REPLACE}'
          - works: ${env.REPLACE}
          - { nested: "${env.DOESNOTEXIST}" }
        "#;
        let marked_yaml = crate::YamlLocationData::parse(source_code).expect("valid yaml");
        let vars = MapVariables(HashMap::from([(
            "REPLACE".to_string(),
            "value".to_string(),
        )]));
        let err = expand(&mut json, &vars, Some(&marked_yaml)).expect_err("should fail");
        insta::assert_snapshot!(print_miette(&err), @r#"
        apollo::configuration::expansion

          × could not expand configuration values

        Error: apollo::configuration::env_invalid

          × invalid bare expansion
           ╭─[2:13]
         1 │ 
         2 │           - '$REPLACE ${other.REPLACE}'
           ·             ─────────────┬─────────────
           ·                          ╰── change this to `${env.REPLACE}`
         3 │           - works: ${env.REPLACE}
           ╰────

        Error: apollo::configuration::env_missing

          × environment variable not present
           ╭─[4:23]
         3 │           - works: ${env.REPLACE}
         4 │           - { nested: "${env.DOESNOTEXIST}" }
           ·                       ───────────┬──────────
           ·                                  ╰── environment variable not present
         5 │         
           ╰────
        "#);
    }

    #[test]
    fn not_unicode() {
        // Like `MapVariables`, but with non-utf8 strings allowed
        struct MockNonUtf8Variables;
        impl VariableProvider for MockNonUtf8Variables {
            fn get<'a>(
                &'a self,
                name: &str,
            ) -> Result<std::borrow::Cow<'a, str>, std::env::VarError> {
                match name {
                    "NON_UTF8" => Err(std::env::VarError::NotUnicode(OsString::new())),
                    _ => Ok(std::borrow::Cow::Owned(name.to_lowercase())),
                }
            }
        }

        let mut json = json!({
            "UTF8": "${env.UTF8}",
            "NON_UTF8": "${env.NON_UTF8}",
        });
        let vars = MockNonUtf8Variables;
        let err = expand(&mut json, &vars, None).expect_err("should fail");
        assert!(matches!(
            err.errors.as_slice(),
            [super::ExpandErrorKind::NotUnicode { .. }]
        ));
    }

    fn print_miette(diagnostic: &dyn miette::Diagnostic) -> String {
        struct F<'a>(&'a dyn miette::Diagnostic);
        impl std::fmt::Display for F<'_> {
            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                miette::MietteHandlerOpts::new()
                    .color(false)
                    .build()
                    .debug(self.0, f)
            }
        }

        F(diagnostic).to_string()
    }
}