diff options
Diffstat (limited to 'src/main.rs')
| -rw-r--r-- | src/main.rs | 255 |
1 files changed, 11 insertions, 244 deletions
diff --git a/src/main.rs b/src/main.rs index 946a99d..4f42d45 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,17 +5,16 @@ mod manifest; mod dockerfile; mod metrics; mod db; +mod registry; -use crate::parser::*; use crate::version::{VersionPattern, CompareOutcome}; use crate::docker::DockerRef; use crate::manifest::ManifestFile; use crate::dockerfile::DockerfileFile; use crate::db::{Db, SqliteDb}; +use crate::registry::{Registry, HttpRegistry, AuthInfo, basic_auth}; -use base64::prelude::*; use rand::distr::{Alphanumeric, SampleString}; -use std::collections::HashMap; use std::ops::Range; use std::io::Read; use std::io::Seek; @@ -35,199 +34,6 @@ Options: ); } -#[derive(Debug)] -enum AuthMethod { - Basic, - Bearer { realm: String, service: String, scope: String }, -} - -impl AuthMethod { - fn from_header(header: &ureq::http::HeaderValue) -> Option<Self> { - let header_str = header.to_str().unwrap(); - let parse = parse_authenticate_header(header_str).unwrap(); - - if &header_str[parse.scheme.clone()] == "Basic" { - return Some(AuthMethod::Basic); - } - - if &header_str[parse.scheme.clone()] == "Bearer" { - return Some(AuthMethod::Bearer { - realm: header_str[parse.realm.unwrap()].to_string(), - scope: header_str[parse.scope.unwrap()].to_string(), - service: header_str[parse.service.unwrap()].to_string(), - }); - } - - return None; - } -} - -fn basic_auth(user: &str, pass: &str) -> String { - return BASE64_STANDARD.encode(format!("{}:{}", user, pass)); -} - -struct AuthInfo { - host: String, - auth: String, -} - -enum AuthStage { - Idle, - Authorized(String), -} - -struct AuthState { - info: AuthInfo, - stage: AuthStage, -} - -struct Auth { - states: HashMap<String, AuthState> -} - -impl Auth { - fn new(infos: Vec<AuthInfo>) -> Self { - let mut states = HashMap::new(); - for info in infos { - states.insert(info.host.clone(), AuthState { - info: info, - stage: AuthStage::Idle, - }); - } - - return Auth { - states - } - } -} - -impl AuthState{ - fn add_to_request<T>(&self, req: ureq::RequestBuilder<T>) -> ureq::RequestBuilder<T> { - if let AuthStage::Authorized(x) = &self.stage { - return req.header("Authorization", x); - } - - return req - } - - fn authenticate(&mut self, response: &ureq::http::Response<ureq::Body>) -> Result<(), String> { - match self.stage { - AuthStage::Authorized(_) => - // The token must have expired - self.stage = AuthStage::Idle, - AuthStage::Idle => {}, - } - - let auth_header = response.headers().get("www-authenticate").unwrap(); - match AuthMethod::from_header(auth_header) { - Some(AuthMethod::Basic) => { - let basic_auth = format!("Basic {}", self.info.auth); - self.stage = AuthStage::Authorized(basic_auth); - - return Ok(()); - }, - Some(AuthMethod::Bearer{realm, scope, service}) => { - let basic_auth = format!("Basic {}", self.info.auth); - - let url = format!("{}?service={}&scope={}", realm, service, scope); - - let body = ureq::get(url) - .config().http_status_as_error(false).build() - .header("Authorization", basic_auth) - .call(); - - let body = body.map_err(|x| format!("Server {} authentication request failed: {}", self.info.host, x.to_string()))? - .body_mut().read_to_string().expect(format!("Server {} responded with something non-string like", self.info.host).as_str()); - - let body = body.parse::<tinyjson::JsonValue>() - .unwrap(); - - let token = body["token"].get::<String>().unwrap(); - self.stage = AuthStage::Authorized(format!("Bearer {}", token)); - return Ok(()); - - }, - None => return Err(format!("Server {} provided us with a challenge, but we didn't understand it", self.info.host)), - } - } -} - -#[derive(Debug)] -enum RegistryError { - NotFound, - Other(String), -} - -fn perform_registry_request(registry: &str, url: &str, accept: &'static str, auth: &mut Auth) -> Result<ureq::http::Response<ureq::Body>, RegistryError> { - let mut state = auth.states.get_mut(registry); - - let url = format!("https://{}{}", registry, &url); - - let mut authentication_retry = false; - - loop { - let mut request = ureq::get(&url) - .header("Accept", accept) - .config().http_status_as_error(false).build(); - - if let Some(ref state) = state { - request = state.add_to_request(request); - } - - let start_time = std::time::Instant::now(); - let response = request.call().unwrap(); - let duration = start_time.elapsed().as_secs_f64(); - metrics::get().http_request_duration.with_label_values(&[registry]).observe(duration); - - if response.status() == 401 { - if !authentication_retry { - if let Some(ref mut state) = state { - match state.authenticate(&response) { - Ok(()) => { - metrics::get().auth_attempts.with_label_values(&[registry, "success"]).inc(); - authentication_retry = true; - continue; - } - Err(e) => { - metrics::get().auth_attempts.with_label_values(&[registry, "failure"]).inc(); - return Err(RegistryError::Other(e)); - } - } - } else { - metrics::get().http_request.with_label_values(&[registry, "error"]).inc(); - return Err(RegistryError::Other(format!("Server {} returned 401 but we have no credentials", registry))); - } - } else { - metrics::get().auth_attempts.with_label_values(&[registry, "failure"]).inc(); - metrics::get().http_request.with_label_values(&[registry, "error"]).inc(); - return Err(RegistryError::Other(format!("Authentication failed"))); - } - } - - if response.status() == 429 { - metrics::get().rate_limits.with_label_values(&[registry]).inc(); - metrics::get().state.set(metrics::STATE_THROTTLED); - println!("Too many requests"); - std::thread::sleep(std::time::Duration::from_secs(8)); - metrics::get().state.set(metrics::STATE_ACTIVE); - continue; - } - - if response.status() == 404 { - metrics::get().http_request.with_label_values(&[registry, "not_found"]).inc(); - return Err(RegistryError::NotFound); - } - - if response.status() != 200 { - metrics::get().http_request.with_label_values(&[registry, "error"]).inc(); - return Err(RegistryError::Other(format!("Unexpected status code: {}", response.status()))); - } - - metrics::get().http_request.with_label_values(&[registry, "success"]).inc(); - return Ok(response); - } -} - struct Repository { url: String, key: Option<std::path::PathBuf>, @@ -244,7 +50,7 @@ struct FilePatch { } -fn update_images(db: &dyn Db, file: &str, auth: &mut Auth, img: DockerRef, edits: &mut Vec<FilePatch>) -> Result<(), 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()]; @@ -268,38 +74,8 @@ fn update_images(db: &dyn Db, file: &str, auth: &mut Auth, img: DockerRef, edits let (tags, cached_digest): (Vec<String>, Option<String>) = if cache_stale { println!("Invalid in cache"); - let mut url = format!("/v2/{}/tags/list", &image); - let mut fetched_tags = vec![]; - loop { - let mut response = match perform_registry_request(registry, &url, "application/vnd.oci.image.index.v1+json", auth) { - Ok(r) => r, - Err(RegistryError::NotFound) => { - println!("Warning: image not found, skipping: {}", image); - return Ok(()); - }, - Err(RegistryError::Other(msg)) => return Err(msg), - }; - - let content_type = response.headers()["Content-Type"].to_str().unwrap(); - if !content_type.starts_with("application/json") { - return Err(format!("Unexpected Content-Type: {}", content_type)); - } - - let body = response.body_mut().read_to_string().unwrap(); - let body = body.parse::<tinyjson::JsonValue>().unwrap(); - - for it in body["tags"].get::<Vec<tinyjson::JsonValue>>().unwrap().iter() { - fetched_tags.push(it.get::<String>().unwrap().clone()); - } - - if let Some(link_header) = response.headers().get("link") { - let link_str = link_header.to_str().unwrap(); - let link = extract_next_page(&link_str).unwrap(); - url = link_str[link.next_uri.unwrap().clone()].to_string(); - } else { - break; - } - } + let mut fetched_tags = reg.get_tags(registry, image) + .ok_or_else(|| format!("image not found: {}", image))?; fetched_tags.sort(); let mut existing = cached_tags.iter().peekable(); @@ -361,17 +137,8 @@ fn update_images(db: &dyn Db, file: &str, auth: &mut Auth, img: DockerRef, edits let digest_string = match cached_digest { Some(d) => d, None => { - let url = format!("/v2/{}/manifests/{}", &file[img.image.clone()], tag_for_digest); - let response = match perform_registry_request(registry, &url, "application/vnd.oci.image.manifest.v1+json,application/vnd.oci.image.index.v1+json", auth) { - Ok(r) => r, - Err(RegistryError::NotFound) => { - println!("Warning: image not found, skipping: {}", image); - return Ok(()); - }, - Err(RegistryError::Other(msg)) => return Err(msg), - }; - - let fetched = response.headers()["docker-content-digest"].to_str().unwrap().to_string(); + let fetched = reg.get_digest(registry, image, tag_for_digest) + .ok_or_else(|| format!("digest not found for {}:{}", image, tag_for_digest))?; db.update_tag_digest(image_id, tag_for_digest, &fetched, now); @@ -481,7 +248,7 @@ fn is_yaml(path: &std::path::Path) -> bool { return false; } -fn run_tool(db: &dyn Db, auth: &mut Auth, repos: &Vec<Repository>, mut infile_paths: Vec<std::path::PathBuf>, overwrite: bool) { +fn run_tool(db: &dyn Db, reg: &dyn Registry, repos: &Vec<Repository>, mut infile_paths: Vec<std::path::PathBuf>, overwrite: bool) { for repo in repos { if repo.dest.exists() { let mut gitcmd = std::process::Command::new("git"); @@ -597,7 +364,7 @@ fn run_tool(db: &dyn Db, auth: &mut Auth, repos: &Vec<Repository>, mut infile_pa 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, &file_content, auth, image_ref, &mut edits) { + if let Err(msg) = update_images(db, reg, &file_content, image_ref, &mut edits) { failed = Some(msg); break; } @@ -952,7 +719,7 @@ fn main() { } } - let mut auth = Auth::new(auths); + let registry = HttpRegistry::new(auths); metrics::init(); metrics::get().state.set(metrics::STATE_IDLE); @@ -964,7 +731,7 @@ fn main() { metrics::get().state.set(metrics::STATE_ACTIVE); let run_start = std::time::Instant::now(); - run_tool(&sqlite_db, &mut auth, &repos, infile_paths.clone(), overwrite); + run_tool(&sqlite_db, ®istry, &repos, infile_paths.clone(), overwrite); let run_duration = run_start.elapsed().as_secs_f64(); metrics::get().run_duration.set(run_duration); |