Extract registry code

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, &registry, &repos, infile_paths.clone(), overwrite);
 
         let run_duration = run_start.elapsed().as_secs_f64();
         metrics::get().run_duration.set(run_duration);