Merge pull request #9 from HewlettPackard/image-var-resolver

Add a function to resolve variables in ImageRefs

Cargo.toml

@@ -22,6 +22,8 @@ pest = "2.1"
 pest_derive = "2.1"
 snafu = "0.6"
 enquote = "1.0"
+regex = "1.4"
+lazy_static = "1.4"
 
 [dev-dependencies]
 indoc = "0.3"

src/dockerfile_parser.rs

@@ -162,7 +162,7 @@ fn parse_dockerfile(input: &str) -> Result<Dockerfile> {
         from_found = true;
       },
       Instruction::Arg(ref arg) => {
-        // args preceeding the first FROM instruction may be substituted into
+        // args preceding the first FROM instruction may be substituted into
         // all subsequent FROM image refs
         if !from_found {
           global_args.push(arg.clone());

src/image.rs

@@ -1,6 +1,13 @@
 // (C) Copyright 2019-2020 Hewlett Packard Enterprise Development LP
 
+use std::collections::HashMap;
 use std::fmt;
+use std::iter::FromIterator;
+
+use lazy_static::lazy_static;
+use regex::Regex;
+
+use crate::{Dockerfile, Span, Splicer};
 
 /// A parsed docker image reference
 ///
@@ -38,6 +45,44 @@ fn is_registry(token: &str) -> bool {
   token == "localhost" || token.contains('.') || token.contains(':')
 }
 
+/// Given a map of key/value pairs, perform variable substitution on a given
+/// input string. `max_recursion_depth` controls the maximum allowed recursion
+/// depth if variables refer to other strings themselves containing variable
+/// references. A small number but reasonable is recommended by default, e.g.
+/// 16.
+/// If None is returned, substitution was impossible, either because a
+/// referenced variable did not exist, or recursion depth was exceeded.
+fn substitute(
+  s: &str, vars: &HashMap<&str, &str>, max_recursion_depth: u8
+) -> Option<String> {
+  lazy_static! {
+    static ref VAR: Regex = Regex::new(r"\$(?:([A-Za-z0-9_]+)|\{([A-Za-z0-9_]+)\})").unwrap();
+  }
+
+  let mut splicer = Splicer::from_str(s);
+
+  for caps in VAR.captures_iter(s) {
+    if max_recursion_depth == 0 {
+      // can't substitute, so give up
+      return None;
+    }
+
+    let full_range = caps.get(0)?.range();
+    let var_name = caps.get(1).or(caps.get(2))?;
+    let var_content = vars.get(var_name.as_str())?;
+    let substituted_content = substitute(
+      var_content,
+      vars,
+      max_recursion_depth.saturating_sub(1)
+    )?;
+
+    // splice the substituted content back into the output string
+    splicer.splice(&Span::new(full_range.start, full_range.end), &substituted_content);
+  }
+
+  Some(splicer.content)
+}
+
 impl ImageRef {
   /// Parses an `ImageRef` from a string.
   ///
@@ -87,6 +132,26 @@ impl ImageRef {
       ImageRef { registry, image, tag, hash: None }
     }
   }
+
+  /// Given a Dockerfile (and its global `ARG`s), perform any necessary
+  /// variable substitution to resolve any variable references in this
+  /// `ImageRef`.
+  ///
+  /// If this `ImageRef` contains any unknown variables or if any references are
+  /// excessively recursive, returns None; otherwise, returns the
+  /// fully-substituted string.
+  pub fn resolve_vars(&self, dockerfile: &Dockerfile) -> Option<ImageRef> {
+    let vars: HashMap<&str, &str> = HashMap::from_iter(
+      dockerfile.global_args
+        .iter()
+        .filter_map(|a| match a.value.as_deref() {
+          Some(v) => Some((a.name.as_str(), v)),
+          None => None
+        })
+    );
+
+    substitute(&self.to_string(), &vars, 16).map(|s| ImageRef::parse(&s))
+  }
 }
 
 impl fmt::Display for ImageRef {
@@ -111,6 +176,10 @@ impl fmt::Display for ImageRef {
 mod tests {
   use super::*;
 
+  use std::convert::TryInto;
+  use indoc::indoc;
+  use crate::instructions::*;
+
   #[test]
   fn test_image_parse_dockerhub() {
     assert_eq!(
@@ -345,4 +414,155 @@ mod tests {
       }
     );
   }
+
+  #[test]
+  fn test_substitute() {
+    let mut vars = HashMap::new();
+    vars.insert("foo", "bar");
+    vars.insert("baz", "qux");
+    vars.insert("lorem", "$foo");
+    vars.insert("ipsum", "${lorem}");
+    vars.insert("recursion1", "$recursion2");
+    vars.insert("recursion2", "$recursion1");
+
+    assert_eq!(
+      substitute("hello world", &vars, 16).as_deref(),
+      Some("hello world")
+    );
+
+    assert_eq!(
+      substitute("hello $foo", &vars, 16).as_deref(),
+      Some("hello bar")
+    );
+
+    assert_eq!(
+      substitute("hello $foo", &vars, 0).as_deref(),
+      None
+    );
+
+    assert_eq!(
+      substitute("hello ${foo}", &vars, 16).as_deref(),
+      Some("hello bar")
+    );
+
+    assert_eq!(
+      substitute("$baz $foo", &vars, 16).as_deref(),
+      Some("qux bar")
+    );
+
+    assert_eq!(
+      substitute("hello $lorem", &vars, 16).as_deref(),
+      Some("hello bar")
+    );
+
+    assert_eq!(
+      substitute("hello $lorem", &vars, 1).as_deref(),
+      None
+    );
+
+    assert_eq!(
+      substitute("hello $ipsum", &vars, 16).as_deref(),
+      Some("hello bar")
+    );
+
+    assert_eq!(
+      substitute("hello $ipsum", &vars, 2).as_deref(),
+      None
+    );
+
+    assert_eq!(
+      substitute("hello $recursion1", &vars, 16).as_deref(),
+      None
+    );
+  }
+
+  #[test]
+  fn test_resolve_vars() {
+    let d = Dockerfile::parse(indoc!(r#"
+      ARG image=alpine:3.12
+      FROM $image
+    "#)).unwrap();
+
+    let from: &FromInstruction = d.instructions
+      .get(1).unwrap()
+      .try_into().unwrap();
+
+    assert_eq!(
+      from.image_parsed.resolve_vars(&d),
+      Some(ImageRef::parse("alpine:3.12"))
+    );
+  }
+
+  #[test]
+  fn test_resolve_vars_nested() {
+    let d = Dockerfile::parse(indoc!(r#"
+      ARG image=alpine
+      ARG unnecessarily_nested=${image}
+      ARG tag=3.12
+      FROM ${unnecessarily_nested}:${tag}
+    "#)).unwrap();
+
+    let from: &FromInstruction = d.instructions
+      .get(3).unwrap()
+      .try_into().unwrap();
+
+    assert_eq!(
+      from.image_parsed.resolve_vars(&d),
+      Some(ImageRef::parse("alpine:3.12"))
+    );
+  }
+
+  #[test]
+  fn test_resolve_vars_technically_invalid() {
+    // docker allows this, but we can't give an answer
+    let d = Dockerfile::parse(indoc!(r#"
+      ARG image
+      FROM $image
+    "#)).unwrap();
+
+    let from: &FromInstruction = d.instructions
+      .get(1).unwrap()
+      .try_into().unwrap();
+
+    assert_eq!(
+      from.image_parsed.resolve_vars(&d),
+      None
+    );
+  }
+
+  #[test]
+  fn test_resolve_vars_typo() {
+    // docker allows this, but we can't give an answer
+    let d = Dockerfile::parse(indoc!(r#"
+      ARG image="alpine:3.12"
+      FROM $foo
+    "#)).unwrap();
+
+    let from: &FromInstruction = d.instructions
+      .get(1).unwrap()
+      .try_into().unwrap();
+
+    assert_eq!(
+      from.image_parsed.resolve_vars(&d),
+      None
+    );
+  }
+
+  #[test]
+  fn test_resolve_vars_out_of_order() {
+    // docker allows this, but we can't give an answer
+    let d = Dockerfile::parse(indoc!(r#"
+      FROM $image
+      ARG image="alpine:3.12"
+    "#)).unwrap();
+
+    let from: &FromInstruction = d.instructions
+      .get(0).unwrap()
+      .try_into().unwrap();
+
+    assert_eq!(
+      from.image_parsed.resolve_vars(&d),
+      None
+    );
+  }
 }

src/splicer.rs

@@ -139,6 +139,13 @@ impl Splicer {
     }
   }
 
+  pub(crate) fn from_str(s: &str) -> Splicer {
+    Splicer {
+      content: s.to_string(),
+      splice_offsets: Vec::new()
+    }
+  }
+
   /// Replaces a Span with the given replacement string, mutating the `content`
   /// string.
   ///