docker_image_pusher/cli/

runner.rs

//! Enhanced runner with parallel upload support
//!
//! This module implements the main workflow for pushing Docker images, including parsing,
//! validation, registry authentication, and orchestrating the upload process. It supports
//! dry-run, verbose/quiet output, and robust error handling for all major steps.

use crate::cli::args::Args;
use crate::error::{PusherError, Result};
use crate::image::parser::ImageParser;
use crate::output::OutputManager;
use crate::registry::{AuthConfig, RegistryClient, RegistryClientBuilder};
use crate::upload::ParallelUploader;
use futures::FutureExt;
use serde_json::json;
use std::io::Read;
use std::path::Path;
use std::sync::Arc;

pub struct Runner {
    args: Args,
    output: OutputManager,
}

impl Runner {
    pub fn new(args: Args) -> Result<Self> {
        // Create output manager based on args
        let output = if args.quiet {
            OutputManager::new_quiet()
        } else {
            OutputManager::new(args.verbose)
        };

        Ok(Self { args, output })
    }

    pub async fn run(&self) -> Result<()> {
        self.output.section("Docker Image Pusher");

        // Validate arguments first
        self.validate_arguments()?;

        // Parse the image
        let image_info = self.parse_image().await?;

        // Don't proceed if dry run
        if self.args.dry_run {
            self.output
                .success("Dry run completed successfully - no data was uploaded");
            return Ok(());
        }

        // Create registry client and push
        let client = self.create_registry_client().await?;
        self.push_image(&client, &image_info).await?;

        self.output.success("Image push completed successfully!");
        Ok(())
    }

    fn validate_arguments(&self) -> Result<()> {
        // Validate file exists
        if !Path::new(&self.args.file).exists() {
            return Err(PusherError::Validation(format!(
                "Image file '{}' does not exist",
                self.args.file
            )));
        }

        // Validate URL format
        let _parsed_url = url::Url::parse(&self.args.repository_url)
            .map_err(|e| PusherError::Validation(format!("Invalid repository URL: {}", e)))?;

        // Validate concurrent settings
        if self.args.max_concurrent == 0 || self.args.max_concurrent > 10 {
            return Err(PusherError::Validation(
                "max_concurrent must be between 1 and 10".to_string(),
            ));
        }

        Ok(())
    }

    async fn parse_image(&self) -> Result<crate::image::parser::ImageInfo> {
        let mut parser = ImageParser::new(self.output.clone());
        parser.set_large_layer_threshold(self.args.large_layer_threshold);

        let tar_path = Path::new(&self.args.file);
        parser.parse_tar_file(tar_path).await
    }

    async fn create_registry_client(&self) -> Result<RegistryClient> {
        let parsed_url = url::Url::parse(&self.args.repository_url)?;
        let registry_address = format!(
            "{}://{}",
            parsed_url.scheme(),
            parsed_url.host_str().unwrap_or("")
        );

        let auth_config =
            if let (Some(username), Some(password)) = (&self.args.username, &self.args.password) {
                Some(AuthConfig::new(username.clone(), password.clone()))
            } else {
                None
            };

        RegistryClientBuilder::new(registry_address)
            .with_auth(auth_config)
            .with_timeout(self.args.timeout)
            .with_skip_tls(self.args.skip_tls)
            .with_verbose(self.args.verbose)
            .build()
    }

    async fn push_image(
        &self,
        client: &RegistryClient,
        image_info: &crate::image::parser::ImageInfo,
    ) -> Result<()> {
        self.output.section("Pushing image to registry");

        // Extract repository info from URL
        let parsed_url = url::Url::parse(&self.args.repository_url)?;
        let path = parsed_url.path().trim_start_matches('/');
        let (repository, tag) = if let Some(colon_pos) = path.rfind(':') {
            let (repo, tag_part) = path.split_at(colon_pos);
            (repo, &tag_part[1..])
        } else {
            (path, "latest")
        };

        // Authenticate for repository access if credentials provided
        let token =
            if let (Some(username), Some(password)) = (&self.args.username, &self.args.password) {
                let auth_config = AuthConfig::new(username.clone(), password.clone());
                client
                    .authenticate_for_repository(&auth_config, repository)
                    .await?
            } else {
                None
            };

        self.output.info(&format!(
            "Pushing {} layers to {}",
            image_info.layers.len(),
            repository
        ));
        self.output.info(&format!(
            "Total size: {}",
            self.output.format_size(image_info.total_size)
        ));

        // Step 1: Check which blobs already exist
        let mut missing_blobs = Vec::new();
        let mut existing_blobs = Vec::new();
        let mut upload_size = 0u64;
        let mut existing_size = 0u64;

        self.output.subsection("Checking existing blobs");
        for (i, layer) in image_info.layers.iter().enumerate() {
            self.output.detail(&format!(
                "Checking layer {}/{}: {}...",
                i + 1,
                image_info.layers.len(),
                &layer.digest[..16]
            ));

            let exists = client.check_blob_exists(&layer.digest, repository).await?;

            if !exists {
                missing_blobs.push(layer.clone());
                upload_size += layer.size;
                self.output.detail(&format!("Layer {} needs upload", i + 1));
            } else {
                existing_blobs.push(layer.clone());
                existing_size += layer.size;
                self.output
                    .success(&format!("Layer {} already exists", i + 1));
            }
        }

        // Report summary
        if existing_blobs.is_empty() {
            self.output
                .info("No existing layers found - full upload required");
        } else {
            self.output.success(&format!(
                "Found {} existing layers ({} total)",
                existing_blobs.len(),
                self.output.format_size(existing_size)
            ));
        }

        if missing_blobs.is_empty() {
            self.output.success("All layers already exist in registry");
        } else {
            self.output.info(&format!(
                "Need to upload {} layers ({} total)",
                missing_blobs.len(),
                self.output.format_size(upload_size)
            ));

            // Check if user wants to skip existing layers
            if self.args.skip_existing && !existing_blobs.is_empty() {
                self.output.warning(
                    "--skip-existing flag specified, but there are missing layers that need upload",
                );
                self.output
                    .info("Proceeding with upload of missing layers only");
            }

            if self.args.force_upload {
                self.output.warning(
                    "--force-upload specified, uploading all layers regardless of existence",
                );
                missing_blobs = image_info.layers.clone();
                upload_size = image_info.total_size;
                self.output.info(&format!(
                    "Force uploading {} layers ({} total)",
                    missing_blobs.len(),
                    self.output.format_size(upload_size)
                ));
            }

            // Step 2: Upload missing blobs in parallel/sequential
            if self.args.max_concurrent > 1 && missing_blobs.len() > 1 {
                self.upload_layers_parallel(client, missing_blobs, repository, &token)
                    .await?;
            } else {
                self.upload_layers_sequential(client, missing_blobs, repository, &token)
                    .await?;
            }
        }

        // Step 3: Upload config blob
        self.output.subsection("Uploading config");
        let config_exists = client
            .check_blob_exists(&image_info.config_digest, repository)
            .await?;

        if !config_exists {
            self.output.step("Uploading image config");
            self.upload_config_blob(client, image_info, repository, &token)
                .await?;
            self.output.success("Config uploaded successfully");
        } else {
            self.output.info("Config already exists in registry");
        }

        // Step 4: Create and upload manifest
        self.output.subsection("Creating manifest");
        let manifest = self.create_image_manifest(image_info)?;

        self.output
            .step(&format!("Uploading manifest for {}:{}", repository, tag));
        self.upload_manifest_with_token(client, &manifest, repository, tag, &token)
            .await?;

        self.output.success(&format!(
            "Image {}:{} pushed successfully!",
            repository, tag
        ));

        Ok(())
    }

    async fn upload_layers_parallel(
        &self,
        client: &RegistryClient,
        layers: Vec<crate::image::parser::LayerInfo>,
        repository: &str,
        token: &Option<String>, // Add token parameter
    ) -> Result<()> {
        self.output.subsection("Parallel Layer Upload");

        // Fix: Create Arc from owned RegistryClient, not reference
        let client_owned = Arc::new(client.clone()); // Assume RegistryClient implements Clone

        let parallel_uploader = ParallelUploader::new(
            client_owned,
            self.args.max_concurrent,
            self.args.large_layer_threshold,
            self.args.timeout,
            self.output.clone(),
        );

        let tar_path = Path::new(&self.args.file);

        parallel_uploader
            .upload_layers_parallel(
                layers, repository, tar_path, token, // Pass the token
            )
            .await
    }

    async fn upload_layers_sequential(
        &self,
        client: &RegistryClient,
        layers: Vec<crate::image::parser::LayerInfo>,
        repository: &str,
        token: &Option<String>,
    ) -> Result<()> {
        self.output.subsection("Sequential Layer Upload");

        let tar_path = Path::new(&self.args.file);

        for (i, layer) in layers.iter().enumerate() {
            self.output.info(&format!(
                "Uploading layer {}/{}: {} ({})",
                i + 1,
                layers.len(),
                &layer.digest[..16],
                self.output.format_size(layer.size)
            ));

            if layer.size == 0 {
                self.upload_empty_layer(client, layer, repository, token)
                    .await?;
            } else if layer.size > self.args.large_layer_threshold {
                self.upload_large_layer_streaming(client, layer, repository, tar_path, token)
                    .await?;
            } else {
                self.upload_regular_layer(client, layer, repository, tar_path, token)
                    .await?;
            }
        }

        Ok(())
    }

    async fn upload_empty_layer(
        &self,
        client: &RegistryClient,
        layer: &crate::image::parser::LayerInfo,
        repository: &str,
        token: &Option<String>,
    ) -> Result<()> {
        self.output.detail("Uploading empty layer");

        let upload_url = client
            .start_upload_session_with_token(repository, token)
            .await?;
        let empty_data = Vec::new();

        // Fix URL construction to match other uploaders
        let url = if upload_url.contains('?') {
            format!("{}&digest={}", upload_url, layer.digest)
        } else {
            format!("{}?digest={}", upload_url, layer.digest)
        };

        let mut request = client
            .get_http_client()
            .put(&url)
            .header("Content-Type", "application/octet-stream")
            .header("Content-Length", "0")
            .body(empty_data);

        if let Some(token) = token {
            request = request.bearer_auth(token);
        }

        let response = request
            .send()
            .await
            .map_err(|e| PusherError::Network(format!("Failed to upload empty layer: {}", e)))?;

        if response.status().is_success() {
            Ok(())
        } else {
            let status = response.status();
            let error_text = response
                .text()
                .await
                .unwrap_or_else(|_| "Failed to read error response".to_string());
            Err(PusherError::Upload(format!(
                "Empty layer upload failed (status {}): {}",
                status, error_text
            )))
        }
    }

    async fn upload_manifest_with_token(
        &self,
        client: &RegistryClient,
        manifest: &str,
        repository: &str,
        tag: &str,
        token: &Option<String>,
    ) -> Result<()> {
        let url = format!(
            "{}/v2/{}/manifests/{}",
            client.get_address(),
            repository,
            tag
        );

        self.output
            .info(&format!("Uploading manifest for {}:{}", repository, tag));

        let mut request = client
            .get_http_client()
            .put(&url)
            .header(
                "Content-Type",
                "application/vnd.docker.distribution.manifest.v2+json",
            )
            .body(manifest.to_string());

        if let Some(token) = token {
            request = request.bearer_auth(token);
        }

        let response = request
            .send()
            .await
            .map_err(|e| PusherError::Network(format!("Failed to upload manifest: {}", e)))?;

        if response.status().is_success() {
            Ok(())
        } else {
            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
            )))
        }
    }

    async fn upload_config_blob(
        &self,
        client: &RegistryClient,
        image_info: &crate::image::parser::ImageInfo,
        repository: &str,
        token: &Option<String>, // Add token parameter
    ) -> Result<()> {
        let config_data = self.extract_config_data_from_tar(image_info).await?;
        let upload_url = client
            .start_upload_session_with_token(repository, token)
            .await?;

        let uploader = crate::upload::ChunkedUploader::new(self.args.timeout, self.output.clone());
        uploader
            .upload_large_blob(&upload_url, &config_data, &image_info.config_digest, token)
            .await
    }

    async fn upload_large_layer_streaming(
        &self,
        client: &RegistryClient,
        layer: &crate::image::parser::LayerInfo,
        repository: &str,
        tar_path: &Path,
        token: &Option<String>,
    ) -> Result<()> {
        let upload_url = client
            .start_upload_session_with_token(repository, token)
            .await?;
        let offset = 0; // Simplified - would need proper offset calculation

        let streaming_uploader = crate::upload::StreamingUploader::new(
            client.get_http_client().clone(),
            self.args.retry_attempts,
            self.args.timeout,
            self.output.clone(),
        );

        streaming_uploader
            .upload_from_tar_entry(
                tar_path,
                &layer.tar_path,
                offset,
                layer.size,
                &upload_url,
                &layer.digest,
                token,
                |_uploaded, _total| {
                    // Progress callback
                },
            )
            .await
    }
    async fn upload_regular_layer(
        &self,
        client: &RegistryClient,
        layer: &crate::image::parser::LayerInfo,
        repository: &str,
        tar_path: &Path,
        token: &Option<String>,
    ) -> Result<()> {
        let layer_data = self
            .extract_layer_data_from_tar(tar_path, &layer.tar_path)
            .await?;
        let upload_url = client
            .start_upload_session_with_token(repository, token)
            .await?;

        let uploader = crate::upload::ChunkedUploader::new(self.args.timeout, self.output.clone());
        uploader
            .upload_large_blob(&upload_url, &layer_data, &layer.digest, token)
            .await
    }

    /// Extract layer data from tar archive for upload
    ///
    /// 注意：返回的数据始终为 gzip 压缩的 tar 字节流，便于后续 digest 校验和上传。
    async fn extract_layer_data_from_tar(
        &self,
        tar_path: &Path,
        layer_path: &str,
    ) -> Result<Vec<u8>> {
        // 直接调用 TarUtils，保留原始格式
        let data = crate::tar_utils::TarUtils::extract_layer_data(tar_path, layer_path)?;
        self.output.detail(&format!(
            "Extracted {} bytes for layer {}",
            data.len(),
            layer_path
        ));

        // 检查并记录数据是否为 gzip 格式
        let is_gzipped = crate::tar_utils::TarUtils::is_gzipped(&data);
        self.output
            .detail(&format!("Extracted data is gzipped: {}", is_gzipped));

        // 计算并记录 SHA256 以便调试
        let computed = crate::digest::DigestUtils::compute_sha256(&data);
        self.output
            .detail(&format!("Extracted data SHA256: {}...", &computed[..16]));

        Ok(data)
    }

    async fn extract_config_data_from_tar(
        &self,
        image_info: &crate::image::parser::ImageInfo,
    ) -> Result<Vec<u8>> {
        use std::fs::File;
        use tar::Archive;

        let tar_path = Path::new(&self.args.file);
        let file = File::open(tar_path)
            .map_err(|e| PusherError::Io(format!("Failed to open tar file: {}", e)))?;
        let mut archive = Archive::new(file);

        // Look for config file (usually named like sha256:xxxxx.json)
        let config_filename = format!("{}.json", image_info.config_digest.replace("sha256:", ""));

        for entry_result in archive
            .entries()
            .map_err(|e| PusherError::ImageParsing(format!("Failed to read tar entries: {}", e)))?
        {
            let mut entry = entry_result.map_err(|e| {
                PusherError::ImageParsing(format!("Failed to read tar entry: {}", e))
            })?;

            let path = entry
                .path()
                .map_err(|e| {
                    PusherError::ImageParsing(format!("Failed to read entry path: {}", e))
                })?
                .to_string_lossy()
                .to_string();

            if path == config_filename {
                let mut config_string = String::new();
                entry.read_to_string(&mut config_string).map_err(|e| {
                    PusherError::ImageParsing(format!("Failed to read config data: {}", e))
                })?;
                return Ok(config_string.into_bytes());
            }
        }

        Err(PusherError::ImageParsing(format!(
            "Config file {} not found in tar archive",
            config_filename
        )))
    }

    fn create_image_manifest(
        &self,
        image_info: &crate::image::parser::ImageInfo,
    ) -> Result<String> {
        let layers: Vec<serde_json::Value> = image_info
            .layers
            .iter()
            .map(|layer| {
                json!({
                    "mediaType": layer.media_type,
                    "size": layer.size,
                    "digest": layer.digest
                })
            })
            .collect();

        // Validate config digest
        if !image_info.config_digest.starts_with("sha256:") || image_info.config_digest.len() != 71
        {
            return Err(PusherError::Parse(format!(
                "Invalid config digest format: {}",
                image_info.config_digest
            )));
        }

        let config_size = self.calculate_config_size(image_info)?;
        let manifest = json!({
            "schemaVersion": 2,
            "mediaType": "application/vnd.docker.distribution.manifest.v2+json",
            "config": {
                "mediaType": "application/vnd.docker.container.image.v1+json",
                "size": config_size,
                "digest": image_info.config_digest
            },
            "layers": layers
        });

        self.output
            .detail("✅ Created manifest with validated SHA256 digests");

        serde_json::to_string_pretty(&manifest)
            .map_err(|e| PusherError::Parse(format!("Failed to serialize manifest: {}", e)))
    }

    fn calculate_config_size(&self, image_info: &crate::image::parser::ImageInfo) -> Result<u64> {
        // Calculate config size by extracting the config data
        // or use a reasonable default
        match self.extract_config_data_from_tar(image_info).now_or_never() {
            Some(Ok(data)) => Ok(data.len() as u64),
            _ => Ok(1024), // Default size if we can't get the actual size
        }
    }
}