add support for vendor override files

twoliter/src/lock.rs

@@ -1,7 +1,9 @@
 use crate::common::fs::{create_dir_all, read, remove_dir_all, write};
+use crate::lock_image::UniqueLockedImage;
 use crate::project::{Image, Project, ValidIdentifier, Vendor};
 use crate::schema_version::SchemaVersion;
 use anyhow::{bail, ensure, Context, Result};
+use async_walkdir::WalkDir;
 use base64::Engine;
 use futures::pin_mut;
 use futures::stream::{self, StreamExt, TryStreamExt};
@@ -72,6 +74,10 @@ impl LockedImage {
         })
     }
 
+    pub fn update_source(&mut self, registry: &str) {
+        self.source = format!("{}/{}:v{}", registry, self.name, self.version)
+    }
+
     pub fn digest_uri(&self, digest: &str) -> String {
         self.source.replace(
             format!(":v{}", self.version).as_str(),
@@ -347,8 +353,12 @@ impl OCIArchive {
     }
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct VendorOverrideFile {
+    vendor: HashMap<String, Vendor>,
+}
 /// Represents the structure of a `Twoliter.lock` lock file.
-#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd, Serialize, Deserialize)]
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "kebab-case")]
 pub(crate) struct Lock {
     /// The version of the Twoliter.toml this was generated from
@@ -357,16 +367,19 @@ pub(crate) struct Lock {
     pub sdk: LockedImage,
     /// Resolved kit dependencies
     pub kit: Vec<LockedImage>,
+    /// Overrides for vendors loaded from override files
+    #[serde(skip)]
+    pub vendor_override: HashMap<String, Vendor>,
 }
 
 #[allow(dead_code)]
 impl Lock {
     #[instrument(level = "trace", skip(project))]
     pub(crate) async fn create(project: &Project) -> Result<Self> {
         let lock_file_path = project.project_dir().join(TWOLITER_LOCK);
-
+        // On create we should always pass None for overrides as the lockfile should never be created off any override
         info!("Resolving project references to create lock file");
-        let lock_state = Self::resolve(project).await?;
+        let lock_state = Self::resolve(project, None).await?;
         let lock_str = toml::to_string(&lock_state).context("failed to serialize lock file")?;
 
         debug!("Writing new lock file to '{}'", lock_file_path.display());
@@ -383,17 +396,17 @@ impl Lock {
             lock_file_path.exists(),
             "Twoliter.lock does not exist, please run `twoliter update` first"
         );
+        let overrides = Self::resolve_overrides(project).await?;
+        info!("Resolving project references to check against lock file");
+        let lock_state = Self::resolve(project, Some(overrides.clone())).await?;
         debug!("Loading existing lockfile '{}'", lock_file_path.display());
         let lock_str = read_to_string(&lock_file_path)
             .await
             .context("failed to read lockfile")?;
-        let lock: Self =
+        let mut lock: Self =
             toml::from_str(lock_str.as_str()).context("failed to deserialize lockfile")?;
-
-        info!("Resolving project references to check against lock file");
-        let lock_state = Self::resolve(project).await?;
-
-        ensure!(lock_state == lock, "changes have occured to Twoliter.toml or the remote kit images that require an update to Twoliter.lock");
+        lock.vendor_override = overrides.clone();
+        ensure!(lock.is_match(&lock_state), "changes have occured to Twoliter.toml or the remote kit images that require an update to Twoliter.lock");
         Ok(lock)
     }
 
@@ -450,6 +463,23 @@ impl Lock {
         Ok(())
     }
 
+    fn is_match(&self, right: &Lock) -> bool {
+        // Build an image map for verification for both sides
+        let mut left_images: HashSet<UniqueLockedImage> = self
+            .kit
+            .iter()
+            .map(|x| UniqueLockedImage::new(x, &self.vendor_override))
+            .collect();
+        left_images.insert(UniqueLockedImage::new(&self.sdk, &self.vendor_override));
+        let mut right_images: HashSet<UniqueLockedImage> = right
+            .kit
+            .iter()
+            .map(|x| UniqueLockedImage::new(x, &self.vendor_override))
+            .collect();
+        right_images.insert(UniqueLockedImage::new(&self.sdk, &self.vendor_override));
+        left_images == right_images
+    }
+
     #[instrument(level = "trace", skip(image), fields(image = %image))]
     async fn get_manifest(
         &self,
@@ -492,16 +522,20 @@ impl Lock {
             image,
             path.as_ref().display()
         );
+        let mut image = image.clone();
         let vendor = image.vendor.clone();
+        if let Some(new_vendor) = self.vendor_override.get(&vendor) {
+            image.update_source(new_vendor.registry.as_str());
+        }
         let name = image.name.clone();
         let target_path = path.as_ref().join(format!("{vendor}/{name}/{arch}"));
         let cache_path = path.as_ref().join("cache");
         create_dir_all(&target_path).await?;
         create_dir_all(&cache_path).await?;
 
         // First get the manifest for the specific requested architecture
-        let manifest = self.get_manifest(image_tool, image, arch).await?;
-        let oci_archive = OCIArchive::new(image, manifest.digest.as_str(), &cache_path)?;
+        let manifest = self.get_manifest(image_tool, &image, arch).await?;
+        let oci_archive = OCIArchive::new(&image, manifest.digest.as_str(), &cache_path)?;
 
         // Checks for the saved image locally, or else pulls and saves it
         oci_archive.pull_image(image_tool).await?;
@@ -514,8 +548,16 @@ impl Lock {
     }
 
     #[instrument(level = "trace", skip(project))]
-    async fn resolve(project: &Project) -> Result<Self> {
-        let vendor_table = project.vendor();
+    async fn resolve(
+        project: &Project,
+        vendor_overrides: Option<HashMap<String, Vendor>>,
+    ) -> Result<Self> {
+        let mut vendor_table = project.vendor().clone();
+        if let Some(overrides) = vendor_overrides.as_ref() {
+            for (key, value) in overrides {
+                vendor_table.insert(ValidIdentifier(key.clone()), value.clone());
+            }
+        }
         let mut known: HashMap<(ValidIdentifier, ValidIdentifier), Version> = HashMap::new();
         let mut locked: Vec<LockedImage> = Vec::new();
         let image_tool = ImageTool::from_environment()?;
@@ -584,9 +626,48 @@ impl Lock {
             schema_version: project.schema_version(),
             sdk: LockedImage::new(&image_tool, vendor, sdk).await?,
             kit: locked,
+            vendor_override: Self::resolve_overrides(project).await?,
         })
     }
 
+    async fn resolve_overrides(project: &Project) -> Result<HashMap<String, Vendor>> {
+        // Vendor override should either be a vendor.toml file next to where Twoliter.toml is or multiple files defined in
+        // vendor.d/ folder
+        let dir_path = project.project_dir().join("vendor.d");
+        if dir_path.exists() {
+            let mut overrides = HashMap::new();
+            let mut walker = WalkDir::new(&dir_path);
+            let mut files = Vec::new();
+            while let Some(entry) = walker.next().await {
+                let entry = entry.context("failed to check file in vendor.d")?;
+                if let Some(extension) = entry.path().extension() {
+                    if extension == "toml" {
+                        files.push(entry.path());
+                    }
+                }
+            }
+            // Sort the files for override ordering (WalkDir/read_dir does not gaurantee any ordering)
+            files.sort();
+            for file in files.iter() {
+                let next_body = read_to_string(file).await?;
+                let next: VendorOverrideFile = toml::from_str(next_body.as_str())
+                    .context("failed to deserialize vendor override file")?;
+                for (key, value) in next.vendor.iter() {
+                    overrides.insert(key.clone(), value.clone());
+                }
+            }
+            return Ok(overrides);
+        }
+        let file_path = project.project_dir().join("vendor.toml");
+        if !file_path.exists() {
+            return Ok(HashMap::new());
+        }
+        let content = read_to_string(&file_path).await?;
+        let overrides: VendorOverrideFile = toml::from_str(content.as_str())
+            .context("failed to deserialize vendor override file")?;
+        Ok(overrides.vendor)
+    }
+
     #[instrument(level = "trace", skip(image), fields(image = %image))]
     async fn find_kit(
         image_tool: &ImageTool,
@@ -641,7 +722,7 @@ mod test {
             d3ZmxjTTdzaisrMmszSk5RWkovb3ZDUVRpUlkrRFpvaGdrNlk9IiwibmFtZSI6InRoYXItYmUtYmV0YS1zZGsiL\
             CJzb3VyY2UiOiJwdWJsaWMuZWNyLmF3cy91MWczYzh6NC90aGFyLWJlLWJldGEtc2RrOnYwLjQzLjAiLCJ2ZW5k\
             b3IiOiJib3R0bGVyb2NrZXQtbmV3IiwidmVyc2lvbiI6IjAuNDMuMCJ9LCJ2ZXJzaW9uIjoiMi4wLjAifQo="
-            .to_string()
+                .to_string()
         );
         assert!(encoded.debug_image_metadata().is_some());
     }

twoliter/src/lock_image.rs

@@ -0,0 +1,97 @@
+use crate::{lock::LockedImage, project::Vendor};
+use std::{collections::HashMap, hash::Hash};
+
+#[derive(Debug, Clone, Eq)]
+pub struct UniqueLockedImage<'lock> {
+    image: &'lock LockedImage,
+    vendor_overrides: &'lock HashMap<String, Vendor>,
+}
+
+impl<'lock> UniqueLockedImage<'lock> {
+    pub(crate) fn new(
+        image: &'lock LockedImage,
+        overrides: &'lock HashMap<String, Vendor>,
+    ) -> Self {
+        Self {
+            image,
+            vendor_overrides: overrides,
+        }
+    }
+}
+// 247242824936.dkr.ecr.us-west-2.amazonaws.com
+impl<'lock> PartialEq for UniqueLockedImage<'lock> {
+    fn eq(&self, other: &Self) -> bool {
+        if self.image.digest != other.image.digest || self.image.vendor != other.image.vendor {
+            // Always fail if the digest has changed or if the vendors are different
+            return false;
+        }
+        if !self.vendor_overrides.contains_key(&other.image.vendor) {
+            // If there is no vendor override the resolved source uri's should match
+            return self.image.source == other.image.source;
+        }
+        true
+    }
+}
+
+impl<'lock> Hash for UniqueLockedImage<'lock> {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        state.write(self.image.name.as_bytes());
+        state.write(self.image.vendor.as_bytes());
+        state.write(self.image.digest.as_bytes());
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use std::collections::HashMap;
+
+    use semver::Version;
+
+    use crate::{lock::LockedImage, project::Vendor};
+
+    use super::UniqueLockedImage;
+
+    #[test]
+    fn test_equivalency() {
+        let left = UniqueLockedImage {
+            image: &LockedImage {
+                name: "foo".to_string(),
+                version: Version::new(1, 2, 3),
+                vendor: "bar".to_string(),
+                source: "fake.ecr.aws/foo:v1.2.3".to_string(),
+                digest: "abcdef".to_string(),
+                manifest: Vec::new(),
+            },
+            vendor_overrides: &HashMap::new(),
+        };
+        let right = UniqueLockedImage {
+            image: &LockedImage {
+                name: "foo".to_string(),
+                version: Version::new(1, 2, 3),
+                vendor: "bar".to_string(),
+                source: "fake.ecr.aws/foo:v1.2.3".to_string(),
+                digest: "abcdef".to_string(),
+                manifest: Vec::new(),
+            },
+            vendor_overrides: &HashMap::new(),
+        };
+        assert_eq!(left, right, "basic equivalency");
+        let left = UniqueLockedImage {
+            image: &LockedImage {
+                name: "foo".to_string(),
+                version: Version::new(1, 2, 3),
+                vendor: "bar".to_string(),
+                source: "fake2.ecr.aws/foo:v1.2.3".to_string(),
+                digest: "abcdef".to_string(),
+                manifest: Vec::new(),
+            },
+            vendor_overrides: &HashMap::from([(
+                "bar".to_string(),
+                Vendor {
+                    registry: "fake.ecr.aws".to_string(),
+                },
+            )]),
+        };
+        assert_eq!(left, right, "vendor override equivalency");
+    }
+}

twoliter/src/main.rs

@@ -7,6 +7,7 @@ mod cmd;
 mod common;
 mod docker;
 mod lock;
+mod lock_image;
 mod project;
 mod schema_version;
 /// Test code that should only be compiled when running tests.

twoliter/src/test/cargo_make.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use semver::Version;
 use serde::Deserialize;
 
@@ -23,6 +25,7 @@ async fn test_cargo_make() {
             digest: "abc".to_string(),
             manifest: Vec::new(),
         },
+        vendor_override: HashMap::new(),
     };
     let cargo_make = CargoMake::new(&lock.sdk.source)
         .unwrap()