dragonfly-client-util 1.5.3

/*
 *     Copyright 2024 The Dragonfly Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use dragonfly_client_core::{Error as ClientError, Result as ClientResult};
use regex::Regex;
use sha2::Digest as Sha2Digest;
use std::fmt;
use std::io::{self, Read};
use std::path::Path;
use std::str::FromStr;
use std::sync::LazyLock;
use tracing::instrument;

/// The separator character for digest formatting.
pub const SEPARATOR: &str = ":";

/// Regex pattern for OCI blob URLs, e.g. http(s)://<registry>/v2/<repository>/blobs/<digest>.
static BLOB_URL_REGEX: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"^(.*)://(.*)/v2/(.*)/blobs/([^?]+)(?:\?.*)?$").unwrap());

/// Checks if the URL is an OCI blob URL.
pub fn is_blob_url(url: &str) -> bool {
    BLOB_URL_REGEX.is_match(url)
}

/// Regex pattern for OCI manifest URLs, e.g. http(s)://<registry>/v2/<repository>/manifests/<reference>.
static MANIFEST_URL_REGEX: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"^(.*)://(.*)/v2/(.*)/manifests/([^?]+)(?:\?.*)?$").unwrap());

/// Checks if the URL is an OCI manifest URL whose reference is a digest with a
/// supported algorithm. A manifest URL referenced by a tag returns false, so it
/// falls back to the url based task id.
pub fn is_manifest_digest_url(url: &str) -> bool {
    Digest::extract_from_manifest_url(url).is_some()
}

/// Algorithm for generating digests.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Algorithm {
    /// CRC32 algorithm for generating digests.
    Crc32,

    /// SHA-256 algorithm for generating digests.
    Sha256,

    /// SHA-512 algorithm for generating digests.
    Sha512,
}

/// Implements the Display.
impl fmt::Display for Algorithm {
    /// Formats the value using the given formatter.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Algorithm::Crc32 => write!(f, "crc32"),
            Algorithm::Sha256 => write!(f, "sha256"),
            Algorithm::Sha512 => write!(f, "sha512"),
        }
    }
}

/// Implements the FromStr.
impl FromStr for Algorithm {
    type Err = String;

    /// Parses an algorithm string.
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "crc32" => Ok(Algorithm::Crc32),
            "sha256" => Ok(Algorithm::Sha256),
            "sha512" => Ok(Algorithm::Sha512),
            _ => Err(format!("invalid digest algorithm: {s}")),
        }
    }
}

/// A digest value with its associated algorithm.
pub struct Digest {
    /// The algorithm used to generate the digest.
    algorithm: Algorithm,

    /// The encoded digest value.
    encoded: String,
}

/// Implements the Digest.
impl Digest {
    /// Creates a new digest with the specified algorithm and encoded value.
    pub fn new(algorithm: Algorithm, encoded: String) -> Self {
        Self { algorithm, encoded }
    }

    /// Extracts the digest from an OCI blob URL, e.g. http(s)://<registry>/v2/<repository>/blobs/<digest>.
    pub fn extract_from_blob_url(url: &str) -> Option<Self> {
        BLOB_URL_REGEX
            .captures(url)
            .and_then(|caps| caps.get(4))
            .map(|m| m.as_str())?
            .parse()
            .ok()
    }

    /// Extracts the digest from an OCI manifest URL, e.g. http(s)://<registry>/v2/<repository>/manifests/<digest>.
    pub fn extract_from_manifest_url(url: &str) -> Option<Self> {
        MANIFEST_URL_REGEX
            .captures(url)
            .and_then(|caps| caps.get(4))
            .map(|m| m.as_str())?
            .parse()
            .ok()
    }

    /// Returns the algorithm of the digest.
    pub fn algorithm(&self) -> Algorithm {
        self.algorithm
    }

    /// Returns the encoded digest value.
    pub fn encoded(&self) -> &str {
        &self.encoded
    }
}

/// Implements the Display.
impl fmt::Display for Digest {
    /// Formats the value using the given formatter.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}{}{}", self.algorithm, SEPARATOR, self.encoded)
    }
}

/// Implements the FromStr.
impl FromStr for Digest {
    type Err = String;

    /// Parses a digest string.
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let parts: Vec<&str> = s.splitn(2, SEPARATOR).collect();
        if parts.len() != 2 {
            return Err(format!("invalid digest: {s}"));
        }

        let algorithm = match parts[0] {
            "crc32" => {
                if parts[1].is_empty() {
                    return Err(format!("invalid crc32 digest: {s}"));
                }

                Algorithm::Crc32
            }
            "sha256" => {
                if parts[1].len() != 64 {
                    return Err(format!(
                        "invalid sha256 digest length: {}, expected 64",
                        parts[1].len()
                    ));
                }

                Algorithm::Sha256
            }
            "sha512" => {
                if parts[1].len() != 128 {
                    return Err(format!(
                        "invalid sha512 digest length: {}, expected 128",
                        parts[1].len()
                    ));
                }

                Algorithm::Sha512
            }
            _ => return Err(format!("invalid digest algorithm: {}", parts[0])),
        };

        Ok(Digest::new(algorithm, parts[1].to_string()))
    }
}

/// Calculates the digest of a file.
#[instrument(level = "debug", skip_all)]
pub fn calculate_file_digest(algorithm: Algorithm, path: &Path) -> ClientResult<Digest> {
    let f = std::fs::File::open(path)?;
    let mut reader = io::BufReader::new(f);
    match algorithm {
        Algorithm::Crc32 => {
            let mut buffer = [0; 4096];
            let mut hasher = crc32fast::Hasher::new();
            loop {
                match reader.read(&mut buffer) {
                    Ok(0) => break,
                    Ok(n) => hasher.update(&buffer[..n]),
                    Err(ref err) if err.kind() == io::ErrorKind::Interrupted => continue,
                    Err(err) => return Err(err.into()),
                };
            }

            Ok(Digest::new(algorithm, hasher.finalize().to_string()))
        }
        Algorithm::Sha256 => {
            let mut hasher = sha2::Sha256::new();
            io::copy(&mut reader, &mut hasher)?;
            Ok(Digest::new(algorithm, hex::encode(hasher.finalize())))
        }
        Algorithm::Sha512 => {
            let mut hasher = sha2::Sha512::new();
            io::copy(&mut reader, &mut hasher)?;
            Ok(Digest::new(algorithm, hex::encode(hasher.finalize())))
        }
    }
}

/// Verifies the digest of a file against an expected digest.
pub fn verify_file_digest(expected_digest: Digest, file_path: &Path) -> ClientResult<()> {
    let digest = match calculate_file_digest(expected_digest.algorithm(), file_path) {
        Ok(digest) => digest,
        Err(err) => {
            return Err(err);
        }
    };

    if digest.to_string() != expected_digest.to_string() {
        return Err(ClientError::DigestMismatch(
            expected_digest.to_string(),
            digest.to_string(),
        ));
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs::File;
    use std::io::Write;

    #[test]
    fn test_extract_from_blob_url() {
        let test_cases = vec![
            (
                "http://registry.example.com/v2/library/ubuntu/blobs/sha256:b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e",
                Some((Algorithm::Sha256, "b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e")),
            ),
            (
                "https://registry.example.com/v2/myorg/myrepo/blobs/sha512:94381a28e8c039fedfa78de025158a068226c3ccd041b22c2c8e73fc993584e9b167d9ae32bc8b372c66701c808ab134e0768c8f16b9a3e61eec1ccf8faa9db8",
                Some((Algorithm::Sha512, "94381a28e8c039fedfa78de025158a068226c3ccd041b22c2c8e73fc993584e9b167d9ae32bc8b372c66701c808ab134e0768c8f16b9a3e61eec1ccf8faa9db8")),
            ),
            (
                "https://registry.io/v2/org/team/project/blobs/sha256:b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e",
                Some((Algorithm::Sha256, "b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e")),
            ),
            (
                "http://localhost:5000/v2/myrepo/blobs/sha256:b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e",
                Some((Algorithm::Sha256, "b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e")),
            ),
            (
                "https://index.docker.io/v2/library/alpine/blobs/sha256:b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e?ns=docker.io",
                Some((Algorithm::Sha256, "b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e")),
            ),
            (
                "http://localhost:5000/v2/myrepo/blobs/sha256:b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e?id=12345",
                Some((Algorithm::Sha256, "b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e")),
            ),
            ("http://registry.example.com/blobs/sha256:abc", None),
            ("http://registry.example.com/v2/repo/manifests/sha256:abc", None),
            ("registry.example.com/v2/repo/blobs/sha256:abc", None),
            ("http://registry.example.com/v2/blobs/sha256:abc", None),
            ("", None),
            ("not-a-url", None),
            ("http://registry.example.com/v2/repo/blobs/invalid-digest", None),
        ];

        for (url, expected) in test_cases {
            match expected {
                Some((algorithm, encoded)) => {
                    let digest = Digest::extract_from_blob_url(url).unwrap();
                    assert_eq!(digest.algorithm(), algorithm);
                    assert_eq!(digest.encoded(), encoded);
                }
                None => assert!(Digest::extract_from_blob_url(url).is_none()),
            }
        }
    }

    #[test]
    fn test_is_manifest_digest_url() {
        let test_cases = vec![
            (
                "http://registry.example.com/v2/library/ubuntu/manifests/sha256:b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e",
                true,
            ),
            (
                "http://localhost:5000/v2/myrepo/manifests/sha256:b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e?ns=docker.io",
                true,
            ),
            (
                "http://registry.example.com/v2/library/ubuntu/manifests/latest",
                false,
            ),
            (
                "http://registry.example.com/v2/library/ubuntu/manifests/md5:8a04994a666b4e4b20a2fd9e5a44f44c",
                false,
            ),
            (
                "http://registry.example.com/v2/library/ubuntu/manifests/sha256:abc",
                false,
            ),
            (
                "http://registry.example.com/v2/library/ubuntu/blobs/sha256:b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e",
                false,
            ),
            ("https://example.com/file.txt", false),
        ];

        for (url, expected) in test_cases {
            assert_eq!(is_manifest_digest_url(url), expected);
        }
    }

    #[test]
    fn test_extract_from_manifest_url() {
        let test_cases = vec![
            (
                "http://registry.example.com/v2/library/ubuntu/manifests/sha256:b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e",
                Some((Algorithm::Sha256, "b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e")),
            ),
            (
                "http://localhost:5000/v2/myrepo/manifests/sha256:b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e?ns=docker.io",
                Some((Algorithm::Sha256, "b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e")),
            ),
            (
                "http://registry.example.com/v2/library/ubuntu/manifests/latest",
                None,
            ),
            (
                "http://registry.example.com/v2/library/ubuntu/manifests/sha256:abc",
                None,
            ),
            (
                "http://registry.example.com/v2/library/ubuntu/manifests/md5:8a04994a666b4e4b20a2fd9e5a44f44c",
                None,
            ),
            (
                "http://registry.example.com/v2/library/ubuntu/blobs/sha256:b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e",
                None,
            ),
            ("https://example.com/file.txt", None),
        ];

        for (url, expected) in test_cases {
            match expected {
                Some((algorithm, encoded)) => {
                    let digest = Digest::extract_from_manifest_url(url).unwrap();
                    assert_eq!(digest.algorithm(), algorithm);
                    assert_eq!(digest.encoded(), encoded);
                }
                None => assert!(Digest::extract_from_manifest_url(url).is_none()),
            }
        }
    }

    #[test]
    fn test_algorithm_display() {
        let test_cases = vec![
            (Algorithm::Crc32, "crc32"),
            (Algorithm::Sha256, "sha256"),
            (Algorithm::Sha512, "sha512"),
        ];

        for (algorithm, expected) in test_cases {
            assert_eq!(algorithm.to_string(), expected);
        }
    }

    #[test]
    fn test_algorithm_from_str() {
        let test_cases = vec![
            ("crc32", Some(Algorithm::Crc32)),
            ("sha256", Some(Algorithm::Sha256)),
            ("sha512", Some(Algorithm::Sha512)),
            ("invalid", None),
        ];

        for (s, expected) in test_cases {
            assert_eq!(s.parse::<Algorithm>().ok(), expected);
        }
    }

    #[test]
    fn test_digest_display() {
        let test_cases = vec![
            (Algorithm::Crc32, "1475635037", "crc32:1475635037"),
            (Algorithm::Sha256, "encoded_hash", "sha256:encoded_hash"),
            (Algorithm::Sha512, "encoded_hash", "sha512:encoded_hash"),
        ];

        for (algorithm, encoded, expected) in test_cases {
            assert_eq!(
                Digest::new(algorithm, encoded.to_string()).to_string(),
                expected
            );
        }
    }

    #[test]
    fn test_calculate_file_digest() {
        let content = b"test content";
        let temp_file = tempfile::NamedTempFile::new().expect("failed to create temp file");
        let path = temp_file.path();
        let mut file = File::create(path).expect("failed to create file");
        file.write_all(content).expect("failed to write to file");

        let test_cases = vec![
            (Algorithm::Crc32, "1475635037"),
            (
                Algorithm::Sha256,
                "6ae8a75555209fd6c44157c0aed8016e763ff435a19cf186f76863140143ff72",
            ),
            (
                Algorithm::Sha512,
                "0cbf4caef38047bba9a24e621a961484e5d2a92176a859e7eb27df343dd34eb98d538a6c5f4da1ce302ec250b821cc001e46cc97a704988297185a4df7e99602",
            ),
        ];

        for (algorithm, expected) in test_cases {
            let digest =
                calculate_file_digest(algorithm, path).expect("failed to calculate digest");
            assert_eq!(digest.encoded(), expected);
        }
    }

    #[test]
    fn test_verify_file_digest() {
        let content = b"test content";
        let temp_file = tempfile::NamedTempFile::new().expect("failed to create temp file");
        let path = temp_file.path();
        let mut file = File::create(path).expect("failed to create file");
        file.write_all(content).expect("failed to write to file");

        let test_cases = vec![
            (Algorithm::Crc32, "1475635037", true),
            (
                Algorithm::Sha256,
                "6ae8a75555209fd6c44157c0aed8016e763ff435a19cf186f76863140143ff72",
                true,
            ),
            (
                Algorithm::Sha512,
                "0cbf4caef38047bba9a24e621a961484e5d2a92176a859e7eb27df343dd34eb98d538a6c5f4da1ce302ec250b821cc001e46cc97a704988297185a4df7e99602",
                true,
            ),
            (
                Algorithm::Sha256,
                "b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e",
                false,
            ),
        ];

        for (algorithm, encoded, expected) in test_cases {
            assert_eq!(
                verify_file_digest(Digest::new(algorithm, encoded.to_string()), path).is_ok(),
                expected
            );
        }
    }
}