Use fallible callbacks for Mapping placeholders (#1217)

microsoft · Aug 31, 2021 · 883518e · 883518e
1 parent a656c9d
commit 883518e
Showing 1 changed file with 26 additions and 15 deletions.
diff --git a/src/agent/onefuzz/src/expand.rs b/src/agent/onefuzz/src/expand.rs
@@ -2,7 +2,7 @@
 // Licensed under the MIT License.
 
 use crate::sha256::digest_file_blocking;
-use anyhow::{Context, Result};
+use anyhow::{format_err, Context, Result};
 use onefuzz_telemetry::{InstanceTelemetryKey, MicrosoftTelemetryKey};
 use std::path::{Path, PathBuf};
 use std::{collections::HashMap, hash::Hash};
@@ -14,7 +14,7 @@ pub enum ExpandedValue<'a> {
     Path(String),
     Scalar(String),
     List(&'a [String]),
-    Mapping(Box<dyn Fn(&Expand<'a>, &str) -> Option<ExpandedValue<'a>> + Send>),
+    Mapping(Box<dyn Fn(&Expand<'a>, &str) -> Result<Option<ExpandedValue<'a>>> + Send>),
 }
 
 #[derive(PartialEq, Eq, Hash, EnumIter)]
@@ -111,38 +111,49 @@ impl<'a> Expand<'a> {
         Self { values }
     }
 
-    fn input_file_sha256(&self, _format_str: &str) -> Option<ExpandedValue<'a>> {
-        match self.values.get(&PlaceHolder::Input.get_string()) {
+    fn input_file_sha256(&self, _format_str: &str) -> Result<Option<ExpandedValue<'a>>> {
+        let val = match self.values.get(&PlaceHolder::Input.get_string()) {
             Some(ExpandedValue::Path(fp)) => {
                 let file = PathBuf::from(fp);
-                digest_file_blocking(file).ok().map(ExpandedValue::Scalar)
+                let hash = digest_file_blocking(file)?;
+                Some(ExpandedValue::Scalar(hash))
             }
             _ => None,
-        }
+        };
+
+        Ok(val)
     }
 
-    fn extract_file_name_no_ext(&self, _format_str: &str) -> Option<ExpandedValue<'a>> {
-        match self.values.get(&PlaceHolder::Input.get_string()) {
+    fn extract_file_name_no_ext(&self, _format_str: &str) -> Result<Option<ExpandedValue<'a>>> {
+        let val = match self.values.get(&PlaceHolder::Input.get_string()) {
             Some(ExpandedValue::Path(fp)) => {
                 let file = PathBuf::from(fp);
-                let stem = file.file_stem()?;
+                let stem = file
+                    .file_stem()
+                    .ok_or_else(|| format_err!("missing file stem: {}", file.display()))?;
                 let name_as_str = stem.to_string_lossy().to_string();
                 Some(ExpandedValue::Scalar(name_as_str))
             }
             _ => None,
-        }
+        };
+
+        Ok(val)
     }
 
-    fn extract_file_name(&self, _format_str: &str) -> Option<ExpandedValue<'a>> {
-        match self.values.get(&PlaceHolder::Input.get_string()) {
+    fn extract_file_name(&self, _format_str: &str) -> Result<Option<ExpandedValue<'a>>> {
+        let val = match self.values.get(&PlaceHolder::Input.get_string()) {
             Some(ExpandedValue::Path(fp)) => {
                 let file = PathBuf::from(fp);
-                let name = file.file_name()?;
+                let name = file
+                    .file_name()
+                    .ok_or_else(|| format_err!("missing file name: {}", file.display()))?;
                 let name_as_str = name.to_string_lossy().to_string();
                 Some(ExpandedValue::Scalar(name_as_str))
             }
             _ => None,
-        }
+        };
+
+        Ok(val)
     }
 
     pub fn set_value(self, name: PlaceHolder, value: ExpandedValue<'a>) -> Self {
@@ -343,7 +354,7 @@ impl<'a> Expand<'a> {
                 Ok(arg)
             }
             ExpandedValue::Mapping(func) => {
-                if let Some(value) = func(self, fmtstr) {
+                if let Some(value) = func(self, fmtstr)? {
                     let arg = self.replace_value(fmtstr, arg, &value)?;
                     Ok(arg)
                 } else {