Extract registry code

1 files changed, 481 insertions, 0 deletions

diff --git a/src/registry.rs b/src/registry.rs
new file mode 100644
index 0000000..afd2794
--- /dev/null
+++ b/src/registry.rs
@@ -0,0 +1,481 @@
+use crate::parser::{parse_authenticate_header, extract_next_page};
+use base64::prelude::*;
+use std::cell::RefCell;
+use std::collections::HashMap;
+
+pub trait Registry {
+    fn get_tags(&self, registry: &str, image: &str) -> Option<Vec<String>>;
+    fn get_digest(&self, registry: &str, image: &str, tag: &str) -> Option<String>;
+}
+
+#[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;
+    }
+}
+
+pub fn basic_auth(user: &str, pass: &str) -> String {
+    return BASE64_STANDARD.encode(format!("{}:{}", user, pass));
+}
+
+pub struct AuthInfo {
+    pub host: String,
+    pub 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),
+}
+
+pub struct HttpRegistry {
+    auth: RefCell<Auth>,
+}
+
+impl HttpRegistry {
+    pub fn new(infos: Vec<AuthInfo>) -> Self {
+        return HttpRegistry {
+            auth: RefCell::new(Auth::new(infos)),
+        };
+    }
+
+    fn perform_request(
+        &self,
+        registry: &str,
+        url: &str,
+        accept: &'static str,
+    ) -> Result<ureq::http::Response<ureq::Body>, RegistryError> {
+        let mut auth = self.auth.borrow_mut();
+        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 response = {
+                let _timer = crate::metrics::get()
+                    .http_request_duration
+                    .with_label_values(&[registry])
+                    .start_timer();
+                request.call().unwrap()
+            };
+            crate::metrics::get()
+                .http_request
+                .with_label_values(&[registry])
+                .inc();
+
+            if response.status() == 401 {
+                if !authentication_retry {
+                    if let Some(ref mut state) = state {
+                        match state.authenticate(&response) {
+                            Ok(()) => {
+                                authentication_retry = true;
+                                continue;
+                            }
+                            Err(e) => {
+                                return Err(RegistryError::Other(e));
+                            }
+                        }
+                    } else {
+                        return Err(RegistryError::Other(format!(
+                            "Server {} returned 401 but we have no credentials",
+                            registry
+                        )));
+                    }
+                } else {
+                    return Err(RegistryError::Other(format!("Authentication failed")));
+                }
+            }
+
+            if response.status() == 429 {
+                crate::metrics::get().state.set(crate::metrics::STATE_THROTTLED);
+                println!("Too many requests");
+                std::thread::sleep(std::time::Duration::from_secs(8));
+                crate::metrics::get().state.set(crate::metrics::STATE_ACTIVE);
+                continue;
+            }
+
+            if response.status() == 404 {
+                return Err(RegistryError::NotFound);
+            }
+
+            if response.status() != 200 {
+                return Err(RegistryError::Other(format!(
+                    "Unexpected status code: {}",
+                    response.status()
+                )));
+            }
+
+            return Ok(response);
+        }
+    }
+}
+
+impl Registry for HttpRegistry {
+    fn get_tags(&self, registry: &str, image: &str) -> Option<Vec<String>> {
+        let mut url = format!("/v2/{}/tags/list", image);
+        let mut tags = vec![];
+
+        loop {
+            let mut response = match self.perform_request(
+                registry,
+                &url,
+                "application/vnd.oci.image.index.v1+json",
+            ) {
+                Ok(r) => r,
+                Err(RegistryError::NotFound) => return None,
+                Err(RegistryError::Other(msg)) => {
+                    println!("Error fetching tags: {}", msg);
+                    return None;
+                }
+            };
+
+            let content_type = response.headers()["Content-Type"].to_str().unwrap();
+            if !content_type.starts_with("application/json") {
+                println!("Unexpected Content-Type: {}", content_type);
+                return None;
+            }
+
+            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() {
+                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;
+            }
+        }
+
+        return Some(tags);
+    }
+
+    fn get_digest(&self, registry: &str, image: &str, tag: &str) -> Option<String> {
+        let url = format!("/v2/{}/manifests/{}", image, tag);
+        let response = match self.perform_request(
+            registry,
+            &url,
+            "application/vnd.oci.image.manifest.v1+json,application/vnd.oci.image.index.v1+json",
+        ) {
+            Ok(r) => r,
+            Err(RegistryError::NotFound) => return None,
+            Err(RegistryError::Other(msg)) => {
+                println!("Error fetching digest: {}", msg);
+                return None;
+            }
+        };
+
+        let digest = response.headers()["docker-content-digest"]
+            .to_str()
+            .unwrap()
+            .to_string();
+        return Some(digest);
+    }
+}
+
+#[cfg(test)]
+struct StubTag {
+    registry: String,
+    image: String,
+    tag: String,
+    digest: String,
+}
+
+#[cfg(test)]
+pub struct StubRegistry {
+    entries: RefCell<Vec<StubTag>>,
+}
+
+#[cfg(test)]
+impl Default for StubRegistry {
+    fn default() -> Self {
+        return StubRegistry {
+            entries: RefCell::new(vec![]),
+        };
+    }
+}
+
+#[cfg(test)]
+impl StubRegistry {
+    pub fn add_tag(&self, registry: &str, image: &str, tag: &str, digest: &str) {
+        self.entries.borrow_mut().push(StubTag {
+            registry: registry.to_string(),
+            image: image.to_string(),
+            tag: tag.to_string(),
+            digest: digest.to_string(),
+        });
+    }
+
+    pub fn set_digest(&self, registry: &str, image: &str, tag: &str, digest: &str) {
+        for entry in self.entries.borrow_mut().iter_mut() {
+            if entry.registry == registry && entry.image == image && entry.tag == tag {
+                entry.digest = digest.to_string();
+                return;
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+impl Registry for StubRegistry {
+    fn get_tags(&self, registry: &str, image: &str) -> Option<Vec<String>> {
+        let mut tags = vec![];
+        let mut found = false;
+
+        for entry in self.entries.borrow().iter() {
+            if entry.registry == registry && entry.image == image {
+                tags.push(entry.tag.clone());
+                found = true;
+            }
+        }
+
+        if !found {
+            return None;
+        }
+
+        return Some(tags);
+    }
+
+    fn get_digest(&self, registry: &str, image: &str, tag: &str) -> Option<String> {
+        for entry in self.entries.borrow().iter() {
+            if entry.registry == registry && entry.image == image && entry.tag == tag {
+                return Some(entry.digest.clone());
+            }
+        }
+        return None;
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn test_get_tags_returns_none_for_unknown(reg: &dyn Registry) {
+        assert_eq!(reg.get_tags("docker.io", "unknown"), None);
+    }
+
+    fn test_get_tags_returns_all(reg: &dyn Registry, stub: &StubRegistry) {
+        stub.add_tag("docker.io", "nginx", "1.21", "sha256:abc");
+        stub.add_tag("docker.io", "nginx", "1.20", "sha256:def");
+        stub.add_tag("docker.io", "nginx", "latest", "sha256:ghi");
+        let mut tags = reg.get_tags("docker.io", "nginx").unwrap();
+
+        tags.sort();
+        assert_eq!(tags, vec!["1.20", "1.21", "latest"]);
+    }
+
+    fn test_get_digest_returns_none_for_unknown(reg: &dyn Registry) {
+        assert_eq!(reg.get_digest("docker.io", "nginx", "1.21"), None);
+    }
+
+    fn test_get_digest_returns_digest(reg: &dyn Registry, stub: &StubRegistry) {
+        stub.add_tag("docker.io", "nginx", "1.21", "sha256:abc123");
+
+        assert_eq!(reg.get_digest("docker.io", "nginx", "1.21"), Some("sha256:abc123".to_string()));
+    }
+
+    fn test_set_digest_updates_existing(reg: &dyn Registry, stub: &StubRegistry) {
+        stub.add_tag("docker.io", "nginx", "1.21", "sha256:old");
+        stub.set_digest("docker.io", "nginx", "1.21", "sha256:new");
+
+        assert_eq!(reg.get_digest("docker.io", "nginx", "1.21"), Some("sha256:new".to_string()));
+    }
+
+    #[test]
+    fn conformance_stub_get_tags_unknown() {
+        let stub = StubRegistry::default();
+        test_get_tags_returns_none_for_unknown(&stub);
+    }
+
+    #[test]
+    fn conformance_stub_get_tags_all() {
+        let stub = StubRegistry::default();
+        test_get_tags_returns_all(&stub, &stub);
+    }
+
+    #[test]
+    fn conformance_stub_get_digest_unknown() {
+        let stub = StubRegistry::default();
+        test_get_digest_returns_none_for_unknown(&stub);
+    }
+
+    #[test]
+    fn conformance_stub_get_digest() {
+        let stub = StubRegistry::default();
+        test_get_digest_returns_digest(&stub, &stub);
+    }
+
+    #[test]
+    fn conformance_stub_set_digest() {
+        let stub = StubRegistry::default();
+        test_set_digest_updates_existing(&stub, &stub);
+    }
+
+    #[cfg(feature = "integration")]
+    fn test_http_get_tags(reg: &dyn Registry) {
+        let tags = reg.get_tags("registry.hub.docker.com", "library/nginx");
+        assert!(tags.is_some());
+        let tags = tags.unwrap();
+        assert!(tags.contains(&"1.21".to_string()));
+    }
+
+    #[cfg(feature = "integration")]
+    fn test_http_get_digest(reg: &dyn Registry) {
+        let digest = reg.get_digest("registry.hub.docker.com", "library/nginx", "1.21");
+        assert!(digest.is_some());
+        assert!(digest.unwrap().starts_with("sha256:"));
+    }
+
+    #[test]
+    #[cfg(feature = "integration")]
+    fn conformance_http_get_tags() {
+        crate::metrics::init();
+        let reg = HttpRegistry::new(vec![]);
+        test_http_get_tags(&reg);
+    }
+
+    #[test]
+    #[cfg(feature = "integration")]
+    fn conformance_http_get_digest() {
+        crate::metrics::init();
+        let reg = HttpRegistry::new(vec![]);
+        test_http_get_digest(&reg);
+    }
+}