diff options
| -rw-r--r-- | src/main.rs | 305 |
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 { |