Refactor update_images for multiple images/files

This will make it easier to deduplicate the images across multiple
files. It also lets us do all the registry/database/update comparison
logic in more closed kernel while pushing the update logic to the end of
that operation. May unlock some future optimization if we choose to
go down that rabbithole.

1 files changed, 164 insertions, 141 deletions

diff --git a/src/main.rs b/src/main.rs
index fff151e..ed3c9d5 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -16,8 +16,6 @@ use crate::registry::{Registry, HttpRegistry, AuthInfo, Credentials, basic_auth}
 
 use rand::distr::{Alphanumeric, SampleString};
 use std::ops::Range;
-use std::io::Read;
-use std::io::Seek;
 use std::io::Write;
 
 fn help(cmd: &str) {
@@ -48,112 +46,151 @@ struct FilePatch {
     content: String,
 }
 
-fn update_images(db: &dyn Db, reg: &dyn Registry, file: &str, img: DockerRef, edits: &mut Vec<FilePatch>) -> Result<(), String> {
-    let registry = img.registry.map(|x| &file[x]).unwrap_or("registry.hub.docker.com");
-    let mut tag = img.tag.as_ref().map(|x| file[x.clone()].to_string());
-    let image = &file[img.image.clone()];
-
-    metrics::get().images_checked.with_label_values(&[registry]).inc();
+struct FileInput {
+    path: std::path::PathBuf,
+    content: String,
+    images: Vec<Range<usize>>,
+}
 
-    let now = chrono::offset::Utc::now();
-    let cache_max_age = chrono::Duration::hours(24);
+fn update_images(db: &dyn Db, reg: &dyn Registry, files: &[FileInput], outcomes: &mut Vec<Result<Vec<FilePatch>, String>>) {
+    for file in files {
+        let mut patches = vec![];
+        let mut error: Option<String> = None;
 
-    let (image_id, last_checked) = match db.get_image(registry, image) {
-        Some(x) => x,
-        None => {
-            let id = db.insert_image(registry, image, now);
-            (id, now)
-        },
-    };
+        for image in &file.images {
+            println!("Checking image {}", &file.content[image.clone()]);
+            let img = DockerRef::parse(&file.content, image);
+            let registry = img.registry.as_ref().map(|x| &file.content[x.clone()]).unwrap_or("registry.hub.docker.com");
+            let mut tag = img.tag.as_ref().map(|x| file.content[x.clone()].to_string());
+            let image_name = &file.content[img.image.clone()];
 
-    let cache_stale = now - last_checked > cache_max_age;
+            metrics::get().images_checked.with_label_values(&[registry]).inc();
 
-    let cached_tags: Vec<String> = db.get_tags_sorted(image_id);
+            let now = chrono::offset::Utc::now();
+            let cache_max_age = chrono::Duration::hours(24);
 
-    let (tags, cached_digest): (Vec<String>, Option<String>) = if cache_stale {
-        println!("Invalid in cache");
-        let mut fetched_tags = reg.get_tags(registry, image)
-            .ok_or_else(|| format!("image not found: {}", image))?;
+            let (image_id, last_checked) = match db.get_image(registry, image_name) {
+                Some(x) => x,
+                None => {
+                    let id = db.insert_image(registry, image_name, now);
+                    (id, now)
+                },
+            };
+
+            let cache_stale = now - last_checked > cache_max_age;
+
+            let cached_tags: Vec<String> = db.get_tags_sorted(image_id);
+
+            let tags_result: Result<(Vec<String>, Option<String>), String> = if cache_stale {
+                println!("Invalid in cache");
+                match reg.get_tags(registry, image_name) {
+                    Some(mut fetched_tags) => {
+                        fetched_tags.sort();
+                        let mut existing = cached_tags.iter().peekable();
+                        for t in &fetched_tags {
+                            while existing.peek().is_some_and(|e| *e < t) {
+                                let gone = existing.next().unwrap();
+                                db.delete_tag(image_id, gone);
+                            }
+                            if existing.peek() == Some(&t) {
+                                existing.next();
+                            } else {
+                                db.insert_tag(image_id, t, now);
+                            }
+                        }
+                        for gone in existing {
+                            db.delete_tag(image_id, gone);
+                        }
+                        db.update_last_checked(image_id, now);
 
-        fetched_tags.sort();
-        let mut existing = cached_tags.iter().peekable();
-        for t in &fetched_tags {
-            while existing.peek().is_some_and(|e| *e < t) {
-                let gone = existing.next().unwrap();
-                db.delete_tag(image_id, gone);
-            }
-            if existing.peek() == Some(&t) {
-                existing.next();
+                        Ok((fetched_tags, None))
+                    }
+                    None => Err(format!("image not found: {}", image_name)),
+                }
             } else {
-                db.insert_tag(image_id, t, now);
-            }
-        }
-        for gone in existing {
-            db.delete_tag(image_id, gone);
-        }
-        db.update_last_checked(image_id, now);
-
-        (fetched_tags, None)
-    } else {
-        println!("Valid in cache");
-        let tag_for_digest = tag.as_deref().unwrap_or("latest");
-        let digest = db.get_tag_digest(image_id, tag_for_digest);
-
-        (cached_tags, digest)
-    };
-
-    if let Some(ref tag_str) = tag {
-        let mut current = VersionPattern::parse(&tag_str);
-        let mut new_tag = None;
-
-        for candidate_str in &tags {
-            let candidate = VersionPattern::parse(&candidate_str);
-            match current.compare(&candidate) {
-                CompareOutcome::Higher => {
-                    new_tag = Some(candidate_str.clone());
-                    current = candidate;
-                },
-                CompareOutcome::Lower => {},
-                CompareOutcome::Incompatible => {},
-                CompareOutcome::Identical => {},
-            }
-        }
+                println!("Valid in cache");
+                let tag_for_digest = tag.as_deref().unwrap_or("latest");
+                let digest = db.get_tag_digest(image_id, tag_for_digest);
+
+                Ok((cached_tags, digest))
+            };
+
+            let (tags, cached_digest) = match tags_result {
+                Ok(v) => v,
+                Err(msg) => {
+                    error = Some(msg);
+                    break;
+                }
+            };
+
+            if let Some(ref tag_str) = tag {
+                let mut current = VersionPattern::parse(&tag_str);
+                let mut new_tag = None;
+
+                for candidate_str in &tags {
+                    let candidate = VersionPattern::parse(&candidate_str);
+                    match current.compare(&candidate) {
+                        CompareOutcome::Higher => {
+                            new_tag = Some(candidate_str.clone());
+                            current = candidate;
+                        },
+                        CompareOutcome::Lower => {},
+                        CompareOutcome::Incompatible => {},
+                        CompareOutcome::Identical => {},
+                    }
+                }
 
-        if let Some(new_tag) = new_tag {
-            tag = Some(new_tag.clone());
-            edits.push(FilePatch {
-                position: img.tag.unwrap(),
-                content: new_tag,
-            });
-            metrics::get().images_updated.with_label_values(&[registry]).inc();
-        }
-    }
+                if let Some(new_tag) = new_tag {
+                    tag = Some(new_tag.clone());
+                    patches.push(FilePatch {
+                        position: img.tag.unwrap(),
+                        content: new_tag,
+                    });
+                    metrics::get().images_updated.with_label_values(&[registry]).inc();
+                }
+            }
 
-    if let Some(ref digest) = img.digest {
-        let tag_for_digest = tag.as_deref().unwrap_or("latest");
+            if let Some(ref digest) = img.digest {
+                let tag_for_digest = tag.as_deref().unwrap_or("latest");
 
-        let digest_string = match cached_digest {
-            Some(d) => d,
-            None => {
-                let fetched = reg.get_digest(registry, image, tag_for_digest)
-                    .ok_or_else(|| format!("digest not found for {}:{}", image, tag_for_digest))?;
+                let digest_result = match cached_digest {
+                    Some(d) => Ok(d),
+                    None => {
+                        match reg.get_digest(registry, image_name, tag_for_digest) {
+                            Some(fetched) => {
+                                db.update_tag_digest(image_id, tag_for_digest, &fetched, now);
+                                Ok(fetched)
+                            }
+                            None => Err(format!("digest not found for {}:{}", image_name, tag_for_digest)),
+                        }
+                    }
+                };
 
-                db.update_tag_digest(image_id, tag_for_digest, &fetched, now);
+                let digest_string = match digest_result {
+                    Ok(d) => d,
+                    Err(msg) => {
+                        error = Some(msg);
+                        break;
+                    }
+                };
 
-                fetched
+                if file.content[digest.clone()] != digest_string {
+                    patches.push(FilePatch {
+                        position: img.digest.clone().unwrap(),
+                        content: digest_string,
+                    });
+                    metrics::get().images_updated.with_label_values(&[registry]).inc();
+                }
             }
-        };
+        }
 
-        if file[digest.clone()] != digest_string {
-            edits.push(FilePatch {
-                position: img.digest.unwrap(),
-                content: digest_string,
-            });
-            metrics::get().images_updated.with_label_values(&[registry]).inc();
+        if let Some(msg) = error {
+            outcomes.push(Err(msg));
+        } else {
+            patches.sort_by_key(|e| e.position.start);
+            outcomes.push(Ok(patches));
         }
     }
-
-    return Ok(());
 }
 
 enum Output {
@@ -345,64 +382,50 @@ fn run_tool(db: &dyn Db, reg: &dyn Registry, repos: &Vec<Repository>, mut infile
         };
     }
 
-    for infile_path in inpaths {
-        let mut file = std::fs::File::open(&infile_path).unwrap();
-        let mut file_content = String::new();
-        file.read_to_string(&mut file_content).unwrap();
-
-        let images: Vec<std::ops::Range<usize>> = if is_dockerfile(&infile_path) {
-            DockerfileFile::parse(&file_content).image_refs
+    let mut file_inputs: Vec<FileInput> = vec![];
+    for path in inpaths {
+        let content = std::fs::read_to_string(&path).unwrap();
+        let images = if is_dockerfile(&path) {
+            DockerfileFile::parse(&content).image_refs
         } else {
-            ManifestFile::parse(&file_content).image_tags
+            ManifestFile::parse(&content).image_tags
         };
+        file_inputs.push(FileInput { path, content, images });
+    }
 
-        let mut failed = None;
+    let mut outcomes: Vec<Result<Vec<FilePatch>, String>> = Vec::with_capacity(file_inputs.len());
+    update_images(db, reg, &file_inputs, &mut outcomes);
 
-        let mut edits = vec![];
-        for ref image in images {
-            println!("Checking image {}", &file_content[image.clone()]);
-            let image_ref = DockerRef::parse(&file_content, image);
-            if let Err(msg) = update_images(db, reg, &file_content, image_ref, &mut edits) {
-                failed = Some(msg);
-                break;
-            }
-        }
+    for i in 0..file_inputs.len() {
+        let file = &file_inputs[i];
+        match &outcomes[i] {
+            Ok(patches) => {
+                if patches.is_empty() {
+                    continue;
+                }
 
-        if let Some(msg) = failed {
-            println!("{}: {} ", infile_path.display(), msg);
-            continue;
-        }
+                let mut out = if overwrite {
+                    Output::overlay_file(&file.path)
+                } else {
+                    Output::Stdout(std::io::stdout())
+                };
 
-        let mut out = if overwrite {
-            Output::overlay_file(&infile_path)
-        } else {
-            Output::Stdout(std::io::stdout())
-        };
+                let mut current_position = 0;
+                for patch in patches {
+                    if patch.position.start > current_position {
+                        out.write_all(file.content[current_position..patch.position.start].as_bytes()).unwrap();
+                    }
+                    out.write_all(patch.content.as_bytes()).unwrap();
+                    current_position = patch.position.end;
+                }
+                out.write_all(file.content[current_position..].as_bytes()).unwrap();
 
-        let mut current_position = 0;
-        file.seek(std::io::SeekFrom::Start(0)).unwrap();
-        let mut file_block = file.take(0);
-        for edit in edits {
-            if edit.position.start > current_position {
-                file_block.set_limit((edit.position.start - current_position) as u64);
-                std::io::copy(&mut file_block, &mut out).unwrap();
+                out.commit();
+            }
+            Err(msg) => {
+                println!("{}: {}", file.path.display(), msg);
             }
-
-            out.write_all(edit.content.as_bytes()).unwrap();
-
-            // We should have been able to just seek to the position.end here, but that doesn't work
-            // for whatever reason. What we can do is calculate the amount to skip ahead, and then do
-            // that.
-            let skip = (edit.position.end - edit.position.start) as i64;
-            file_block.set_limit(skip as u64);
-            file_block.seek(std::io::SeekFrom::Current(skip)).unwrap();
-            current_position = edit.position.end;
         }
-
-        file_block.set_limit(u64::MAX);
-        std::io::copy(&mut file_block, &mut out).unwrap();
-
-        out.commit();
     }
 
     for repo in repos {