docker_image_pusher/registry/

client.rs

//! Enhanced registry client with better configuration and error handling

use crate::error::{Result, PusherError};
use crate::error::handlers::NetworkErrorHandler;
use crate::output::OutputManager;
use crate::config::AuthConfig;
use crate::registry::auth::Auth;
use reqwest::Client;
use std::time::Duration;

#[derive(Clone)] // Add Clone derive
pub struct RegistryClient {
    client: Client,
    auth: Auth,
    address: String,
    output: OutputManager,
}

#[derive(Debug)]
pub struct RegistryClientBuilder {
    address: String,
    auth_config: Option<AuthConfig>,
    timeout: u64,
    skip_tls: bool,
    verbose: bool,
}

impl RegistryClientBuilder {
    pub fn new(address: String) -> Self {
        Self {
            address,
            auth_config: None,
            timeout: 7200, // 2 hours default
            skip_tls: false,
            verbose: false,
        }
    }

    pub fn with_auth(mut self, auth_config: Option<AuthConfig>) -> Self {
        self.auth_config = auth_config;
        self
    }

    pub fn with_timeout(mut self, timeout: u64) -> Self {
        self.timeout = timeout;
        self
    }

    pub fn with_skip_tls(mut self, skip_tls: bool) -> Self {
        self.skip_tls = skip_tls;
        self
    }

    pub fn with_verbose(mut self, verbose: bool) -> Self {
        self.verbose = verbose;
        self
    }

    pub fn build(self) -> Result<RegistryClient> {
        let output = OutputManager::new(self.verbose);
        output.verbose("Building HTTP client...");
        
        let client_builder = if self.skip_tls {
            output.verbose("TLS verification disabled");
            Client::builder()
                .danger_accept_invalid_certs(true)
                .danger_accept_invalid_hostnames(true)
        } else {
            output.verbose("TLS verification enabled");
            Client::builder()
        };
        
        let client = client_builder
            .timeout(Duration::from_secs(self.timeout))
            .connect_timeout(Duration::from_secs(60))
            .read_timeout(Duration::from_secs(3600))
            .pool_idle_timeout(Duration::from_secs(300))
            .pool_max_idle_per_host(10)
            .user_agent("docker-image-pusher/1.0")
            .build()
            .map_err(|e| {
                output.error(&format!("Failed to build HTTP client: {}", e));
                PusherError::Network(e.to_string())
            })?;
        
        output.verbose("HTTP client built successfully");
        
        let auth = Auth::new(&self.address, self.skip_tls)?;
        
        Ok(RegistryClient {
            client,
            auth,
            address: self.address,
            output,
        })
    }
}

impl RegistryClient {
    pub async fn test_connectivity(&self) -> Result<()> {
        self.output.verbose("Testing registry connectivity...");
        
        let url = format!("{}/v2/", self.address);
        let response = self.client.get(&url).send().await
            .map_err(|e| PusherError::Network(format!("Failed to connect to registry: {}", e)))?;
        
        self.output.verbose(&format!("Registry response status: {}", response.status()));
        
        if response.status().is_success() || response.status() == 401 {
            // 401 is expected for registries that require authentication
            self.output.verbose("Registry connectivity test passed");
            Ok(())
        } else {
            Err(PusherError::Registry(format!(
                "Registry connectivity test failed with status: {}", 
                response.status()
            )))
        }
    }

    pub async fn check_blob_exists(&self, digest: &str, repository: &str) -> Result<bool> {
        // Ensure digest has proper sha256: prefix
        let normalized_digest = if digest.starts_with("sha256:") {
            digest.to_string()
        } else {
            format!("sha256:{}", digest)
        };
        
        let url = format!("{}/v2/{}/blobs/{}", self.address, repository, normalized_digest);
        
        self.output.detail(&format!("Checking blob existence: {}", &normalized_digest[..23]));
        
        // Use HEAD request to check existence without downloading
        let request = self.client.head(&url);
        
        let response = request.send().await
            .map_err(|e| {
                self.output.warning(&format!("Failed to check blob existence: {}", e));
                NetworkErrorHandler::handle_network_error(&e, "blob existence check")
            })?;
        
        let status = response.status();
        
        match status.as_u16() {
            200 => {
                self.output.detail(&format!("Blob {} exists", &normalized_digest[..16]));
                Ok(true)
            },
            404 => {
                self.output.detail(&format!("Blob {} does not exist", &normalized_digest[..16]));
                Ok(false)
            },
            401 => {
                self.output.warning("Authentication required for blob check");
                // Assume blob doesn't exist if we can't authenticate to check
                Ok(false)
            },
            403 => {
                self.output.warning("Permission denied for blob check");
                // Assume blob doesn't exist if we can't check permissions
                Ok(false)
            },
            _ => {
                self.output.warning(&format!("Unexpected status {} when checking blob existence", status));
                // On other errors, assume blob doesn't exist to be safe
                Ok(false)
            }
        }
    }

    pub async fn authenticate(&self, auth_config: &AuthConfig) -> Result<Option<String>> {
        self.output.verbose("Authenticating with registry...");
        
        let token = self.auth.login(&auth_config.username, &auth_config.password, &self.output).await?;
        
        if token.is_some() {
            self.output.success("Authentication successful");
        } else {
            self.output.info("No authentication required");
        }
        
        Ok(token)
    }

    pub async fn authenticate_for_repository(&self, auth_config: &AuthConfig, repository: &str) -> Result<Option<String>> {
        self.output.verbose(&format!("Authenticating for repository access: {}", repository));
        
        let token = self.auth.get_repository_token(
            &auth_config.username, 
            &auth_config.password, 
            repository,
            &self.output
        ).await?;
        
        if token.is_some() {
            self.output.success(&format!("Repository authentication successful for: {}", repository));
        } else {
            self.output.info("No repository-specific authentication required");
        }
        
        Ok(token)
    }

    pub async fn upload_blob(&self, data: &[u8], digest: &str, repository: &str) -> Result<String> {
        self.output.info(&format!("Uploading blob {} ({}) to {}", 
            &digest[..16], self.output.format_size(data.len() as u64), repository));
        
        // Step 1: Start upload session
        let upload_url = self.start_upload_session(repository).await?;
        
        // Step 2: Upload data
        let upload_response = self.client
            .put(&format!("{}?digest={}", upload_url, digest))
            .header("Content-Type", "application/octet-stream")
            .header("Content-Length", data.len().to_string())
            .body(data.to_vec())
            .send()
            .await
            .map_err(|e| PusherError::Network(format!("Failed to upload blob: {}", e)))?;

        if upload_response.status().is_success() {
            self.output.success(&format!("Blob {} uploaded successfully", &digest[..16]));
            Ok(digest.to_string())
        } else {
            // Store status before consuming response
            let status = upload_response.status();
            let error_text = upload_response.text().await
                .unwrap_or_else(|_| "Failed to read error response".to_string());
            Err(PusherError::Upload(format!(
                "Blob upload failed (status {}): {}", 
                status, 
                error_text
            )))
        }
    }

    pub async fn start_upload_session(&self, repository: &str) -> Result<String> {
        self.start_upload_session_with_token(repository, &None).await
    }

    pub async fn start_upload_session_with_token(&self, repository: &str, token: &Option<String>) -> Result<String> {
        let url = format!("{}/v2/{}/blobs/uploads/", self.address, repository);
        
        self.output.detail(&format!("Starting upload session for {}", repository));
        
        let mut request = self.client.post(&url);
        
        if let Some(token) = token {
            request = request.bearer_auth(token);
            self.output.detail("Using authentication token for upload session");
        }
        
        let response = request.send().await
            .map_err(|e| PusherError::Network(format!("Failed to start upload session: {}", e)))?;

        if response.status() == 202 {
            // Extract upload URL from Location header
            let location = response.headers()
                .get("Location")
                .and_then(|h| h.to_str().ok())
                .ok_or_else(|| PusherError::Registry("No Location header in upload session response".to_string()))?;
            
            // Convert relative URL to absolute if needed
            let upload_url = if location.starts_with("http") {
                location.to_string()
            } else {
                format!("{}{}", self.address, location)
            };
            
            self.output.detail(&format!("Upload session started: {}", &upload_url[..50]));
            Ok(upload_url)
        } else {
            // Store status before consuming response
            let status = response.status();
            let error_text = response.text().await
                .unwrap_or_else(|_| "Failed to read error response".to_string());
            
            let error_msg = match status.as_u16() {
                401 => format!("Unauthorized to access repository: {} - {}", repository, error_text),
                403 => format!("Forbidden: insufficient permissions for repository: {} - {}", repository, error_text),
                404 => format!("Repository not found: {} - {}", repository, error_text),
                _ => format!("Failed to start upload session (status {}): {}", status, error_text)
            };
            
            Err(PusherError::Registry(error_msg))
        }
    }

    pub async fn upload_manifest(&self, manifest: &str, repository: &str, tag: &str) -> Result<()> {
        let url = format!("{}/v2/{}/manifests/{}", self.address, repository, tag);
        
        self.output.info(&format!("Uploading manifest for {}:{}", repository, tag));
        self.output.detail(&format!("Manifest size: {}", self.output.format_size(manifest.len() as u64)));
        
        let response = self.client
            .put(&url)
            .header("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
            .body(manifest.to_string())
            .send()
            .await
            .map_err(|e| PusherError::Network(format!("Failed to upload manifest: {}", e)))?;

        if response.status().is_success() {
            self.output.success(&format!("Manifest uploaded successfully for {}:{}", repository, tag));
            Ok(())
        } else {
            // Store status before consuming response
            let status = response.status();
            let error_text = response.text().await
                .unwrap_or_else(|_| "Failed to read error response".to_string());
            Err(PusherError::Registry(format!(
                "Manifest upload failed (status {}): {}", 
                status, 
                error_text
            )))
        }
    }

    // Add getter for address
    pub fn get_address(&self) -> &str {
        &self.address
    }

    // Add getter for HTTP client
    pub fn get_http_client(&self) -> &Client {
        &self.client
    }

    pub fn get_output_manager(&self) -> &OutputManager {
        &self.output
    }
}