use std::borrow::Cow;
pub fn parse_toml<T: serde::de::DeserializeOwned>(
text: &str,
source: &str,
) -> Result<T, toml::de::Error> {
if !text.contains('$') {
return toml::from_str(text);
}
let mut document: toml::Value = toml::from_str(text)?;
expand_document(&mut document, source);
document.try_into()
}
pub fn expand_document(value: &mut toml::Value, source: &str) {
match value {
toml::Value::String(text) => {
if let Cow::Owned(expanded) = expand(text, source) {
*text = expanded;
}
}
toml::Value::Array(items) => {
for item in items {
expand_document(item, source);
}
}
toml::Value::Table(table) => {
for (_key, item) in table.iter_mut() {
expand_document(item, source);
}
}
_ => {}
}
}
pub fn expand<'a>(text: &'a str, source: &str) -> Cow<'a, str> {
if !text.contains('$') {
return Cow::Borrowed(text);
}
let bytes = text.as_bytes();
let mut out = String::with_capacity(text.len());
let mut i = 0;
while i < bytes.len() {
if bytes[i] != b'$' {
let next = text[i..].find('$').map(|n| i + n).unwrap_or(bytes.len());
out.push_str(&text[i..next]);
i = next;
continue;
}
if bytes.get(i + 1) == Some(&b'$') {
out.push('$');
i += 2;
continue;
}
match parse_reference(&text[i..]) {
Some(reference) => {
out.push_str(&resolve(&reference, source));
i += reference.length;
}
None => {
out.push('$');
i += 1;
}
}
}
Cow::Owned(out)
}
struct Reference {
name: String,
default: Option<String>,
length: usize,
}
fn parse_reference(text: &str) -> Option<Reference> {
let rest = &text[1..];
if let Some(body) = rest.strip_prefix('{') {
let end = body.find('}')?;
let inner = &body[..end];
let (name, default) = match inner.split_once(":-") {
Some((name, default)) => (name, Some(default.to_string())),
None => (inner, None),
};
if !is_name(name) {
return None;
}
return Some(Reference {
name: name.to_string(),
default,
length: 2 + end + 1,
});
}
let name: String = rest
.chars()
.take_while(|c| c.is_ascii_alphanumeric() || *c == '_')
.collect();
if !is_name(&name) {
return None;
}
let length = 1 + name.len();
Some(Reference {
name,
default: None,
length,
})
}
fn is_name(name: &str) -> bool {
let mut chars = name.chars();
match chars.next() {
Some(first) if first.is_ascii_alphabetic() || first == '_' => {}
_ => return false,
}
chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
}
fn resolve(reference: &Reference, source: &str) -> String {
match std::env::var(&reference.name) {
Ok(value) if !value.is_empty() => value,
_ => match &reference.default {
Some(default) => default.clone(),
None => {
tracing::warn!(
variable = %reference.name,
file = source,
"environment variable is not set; using an empty string \
(write ${{{}:-…}} to give a default)",
reference.name
);
String::new()
}
},
}
}
#[cfg(test)]
mod tests {
use super::*;
struct Vars(&'static [(&'static str, &'static str)]);
impl Vars {
fn set(pairs: &'static [(&'static str, &'static str)]) -> Vars {
for (name, value) in pairs {
std::env::set_var(name, value);
}
Vars(pairs)
}
}
impl Drop for Vars {
fn drop(&mut self) {
for (name, _) in self.0 {
std::env::remove_var(name);
}
}
}
fn expanded(text: &str) -> String {
expand(text, "test.toml").into_owned()
}
#[test]
fn both_spellings_of_a_reference_expand() {
let _vars = Vars::set(&[("APIPLANT_T_HOST", "db.example.com")]);
assert_eq!(expanded("$APIPLANT_T_HOST"), "db.example.com");
assert_eq!(expanded("${APIPLANT_T_HOST}"), "db.example.com");
assert_eq!(
expanded("postgres://$APIPLANT_T_HOST:5432/db"),
"postgres://db.example.com:5432/db"
);
assert_eq!(expanded("${APIPLANT_T_HOST}_1"), "db.example.com_1");
}
#[test]
fn several_references_expand_in_one_string() {
let _vars = Vars::set(&[
("APIPLANT_T_USER", "user01"),
("APIPLANT_T_PASS", "veryToughPas$w0rd"),
("APIPLANT_T_HOST", "some-host.tld"),
("APIPLANT_T_NAME", "my_database"),
]);
assert_eq!(
expanded(
"mysql://$APIPLANT_T_USER:$APIPLANT_T_PASS@$APIPLANT_T_HOST:\
${APIPLANT_T_PORT:-3306}/$APIPLANT_T_NAME"
),
"mysql://user01:veryToughPas$w0rd@some-host.tld:3306/my_database"
);
}
#[test]
fn a_default_covers_an_unset_or_empty_variable() {
let _vars = Vars::set(&[("APIPLANT_T_EMPTY", "")]);
assert_eq!(expanded("${APIPLANT_T_UNSET:-us-east-1}"), "us-east-1");
assert_eq!(expanded("${APIPLANT_T_EMPTY:-us-east-1}"), "us-east-1");
assert_eq!(expanded("${APIPLANT_T_UNSET:-}"), "");
assert_eq!(
expanded("${APIPLANT_T_UNSET:-postgres://a:b@c/d}"),
"postgres://a:b@c/d"
);
let _set = Vars::set(&[("APIPLANT_T_REGION", "eu-west-1")]);
assert_eq!(expanded("${APIPLANT_T_REGION:-us-east-1}"), "eu-west-1");
}
#[test]
fn an_unset_variable_without_a_default_expands_to_nothing() {
assert_eq!(expanded("$APIPLANT_T_MISSING"), "");
assert_eq!(expanded("a${APIPLANT_T_MISSING}b"), "ab");
}
#[test]
fn dollars_that_are_not_references_survive() {
assert_eq!(expanded("$$19.99"), "$19.99");
assert_eq!(expanded("$$"), "$");
assert_eq!(expanded("$$$$"), "$$");
assert_eq!(expanded("$19.99"), "$19.99");
assert_eq!(expanded("100 US$"), "100 US$");
assert_eq!(expanded("a $ b"), "a $ b");
assert_eq!(expanded("${unterminated"), "${unterminated");
assert_eq!(expanded("${}"), "${}");
assert_eq!(expanded("${1BAD}"), "${1BAD}");
}
#[test]
fn text_without_a_dollar_is_returned_untouched() {
assert!(matches!(
expand("postgres://localhost/db", "test.toml"),
Cow::Borrowed(_)
));
assert_eq!(expanded(""), "");
}
#[test]
fn a_document_is_expanded_through_tables_and_arrays() {
let _vars = Vars::set(&[
("APIPLANT_T_DOC_URL", "postgres://db/app"),
("APIPLANT_T_DOC_ORIGIN", "https://example.com"),
]);
let mut document: toml::Value = toml::from_str(
r#"
title = "no references here"
port = 5432
[database]
url = "$APIPLANT_T_DOC_URL"
[server]
origins = ["$APIPLANT_T_DOC_ORIGIN", "http://localhost:3000"]
[[hooks]]
target = "${APIPLANT_T_DOC_ORIGIN}/hook"
"#,
)
.unwrap();
expand_document(&mut document, "main.toml");
assert_eq!(
document["database"]["url"].as_str(),
Some("postgres://db/app")
);
assert_eq!(
document["server"]["origins"][0].as_str(),
Some("https://example.com")
);
assert_eq!(
document["server"]["origins"][1].as_str(),
Some("http://localhost:3000")
);
assert_eq!(
document["hooks"][0]["target"].as_str(),
Some("https://example.com/hook")
);
assert_eq!(document["port"].as_integer(), Some(5432));
assert_eq!(document["title"].as_str(), Some("no references here"));
}
#[test]
fn an_expanded_value_cannot_inject_toml() {
let _vars = Vars::set(&[("APIPLANT_T_EVIL", "\"\nadmin = true\n[x]\ny = \"")]);
let mut document: toml::Value = toml::from_str(r#"password = "$APIPLANT_T_EVIL""#).unwrap();
expand_document(&mut document, "main.toml");
assert_eq!(
document["password"].as_str(),
Some("\"\nadmin = true\n[x]\ny = \"")
);
assert_eq!(document.as_table().unwrap().len(), 1);
}
#[test]
fn keys_are_not_expanded() {
let _vars = Vars::set(&[("APIPLANT_T_KEY", "surprise")]);
let mut document: toml::Value = toml::from_str(r#"'$APIPLANT_T_KEY' = "value""#).unwrap();
expand_document(&mut document, "main.toml");
assert!(document.get("$APIPLANT_T_KEY").is_some());
assert!(document.get("surprise").is_none());
}
}