rialo-build-lib 0.11.2

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! S3 storage backend for toolchain distribution
//!
//! This module provides functionality for uploading and downloading toolchain
//! binaries to/from Amazon S3. It supports:
//! - Uploading toolchain tarballs with SHA256 checksums
//! - Downloading toolchains with checksum verification
//! - Checking toolchain availability in S3
//! - Automatic credential detection

use std::path::Path;

use anyhow::{Context, Result};
use s3::{creds::Credentials, Bucket, Region};

/// S3 storage backend for toolchain distribution
///
/// Provides methods for uploading and downloading toolchain tarballs to/from S3.
/// Requires AWS credentials to be configured via environment variables.
///
/// # Examples
///
/// ```no_run
/// # use rialo_build_lib::toolchain::s3_backend::S3StorageBackend;
/// # use std::path::Path;
/// # async fn example() -> anyhow::Result<()> {
/// let backend = S3StorageBackend::new("rialo-toolchains".to_string()).await?;
///
/// if let Some(backend) = backend {
///     // Upload a toolchain
///     backend.upload_toolchain(
///         "gnu-riscv",
///         "13.2.0",
///         "x86_64-apple-darwin",
///         Path::new("/path/to/toolchain.tar.gz"),
///     ).await?;
///
///     // Download a toolchain
///     backend.download_toolchain(
///         "gnu-riscv",
///         "13.2.0",
///         "x86_64-apple-darwin",
///         Path::new("/tmp/toolchain.tar.gz"),
///     ).await?;
/// }
/// # Ok(())
/// # }
/// ```
pub struct S3StorageBackend {
    bucket: Box<Bucket>,
}

impl S3StorageBackend {
    /// Create a new S3 storage backend
    ///
    /// Returns `Ok(None)` if AWS credentials are not available.
    /// Returns `Ok(Some(backend))` if credentials are available and the client was created.
    ///
    /// # Arguments
    ///
    /// * `bucket_name` - The name of the S3 bucket to use for storing toolchains
    ///
    /// # Errors
    ///
    /// Returns an error if the bucket configuration cannot be created.
    pub async fn new(bucket_name: String) -> Result<Option<Self>> {
        // Check if credentials are available in environment
        let access_key = std::env::var("AWS_ACCESS_KEY_ID").ok();
        let secret_key = std::env::var("AWS_SECRET_ACCESS_KEY").ok();

        if access_key.is_none() || secret_key.is_none() {
            // No credentials available, S3 backend not usable
            return Ok(None);
        }

        // Get region from environment or default to us-east-1
        let region_str = std::env::var("AWS_DEFAULT_REGION")
            .or_else(|_| std::env::var("AWS_REGION"))
            .unwrap_or_else(|_| "us-east-1".to_string());

        // Check for custom endpoint (for LocalStack or other S3-compatible services)
        let endpoint_url = std::env::var("AWS_ENDPOINT_URL_S3")
            .or_else(|_| std::env::var("AWS_ENDPOINT_URL"))
            .ok();

        let region = if let Some(endpoint) = endpoint_url {
            log::debug!("Using custom S3 endpoint: {}", endpoint);
            Region::Custom {
                region: region_str,
                endpoint,
            }
        } else {
            region_str.parse().unwrap_or(Region::UsEast1)
        };

        // Create credentials
        let credentials = Credentials::new(
            access_key.as_deref(),
            secret_key.as_deref(),
            None, // security_token
            None, // session_token
            None, // expiration
        )
        .context("Failed to create AWS credentials")?;

        // Create bucket with path-style addressing for LocalStack compatibility
        let mut bucket =
            Bucket::new(&bucket_name, region, credentials).context("Failed to create S3 bucket")?;

        // Enable path-style addressing for LocalStack/custom endpoints
        if std::env::var("AWS_ENDPOINT_URL_S3").is_ok() || std::env::var("AWS_ENDPOINT_URL").is_ok()
        {
            log::debug!("Enabling path-style addressing for custom endpoint");
            bucket = bucket.with_path_style();
        }

        Ok(Some(Self { bucket }))
    }

    /// Upload a toolchain tarball to S3
    ///
    /// Uploads the toolchain tarball and a corresponding SHA256 checksum file.
    ///
    /// # Arguments
    ///
    /// * `toolchain_name` - Name of the toolchain (e.g., "gnu-riscv", "rialo-rust")
    /// * `version` - Version of the toolchain (e.g., "13.2.0", "latest")
    /// * `platform` - Target platform (e.g., "x86_64-apple-darwin")
    /// * `archive_path` - Path to the local tarball to upload
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The archive file cannot be read
    /// - The S3 upload fails
    /// - Checksum computation fails
    pub async fn upload_toolchain(
        &self,
        toolchain_name: &str,
        version: &str,
        platform: &str,
        archive_path: &Path,
    ) -> Result<()> {
        let archive_name = self.get_archive_name(toolchain_name, version, platform);
        let key = format!("toolchains/{toolchain_name}/{version}/{archive_name}.tar.gz");

        log::info!("Uploading to s3://{}/{}...", self.bucket.name(), key);

        // Read file contents
        let bytes = tokio::fs::read(archive_path)
            .await
            .with_context(|| format!("Failed to read archive file {}", archive_path.display()))?;

        // Upload to S3
        self.bucket
            .put_object(&key, &bytes)
            .await
            .context("Failed to upload toolchain to S3")?;

        // Upload SHA256 checksum
        let checksum = compute_sha256(archive_path)?;
        self.upload_checksum(&key, &checksum).await?;

        log::info!("Uploaded to S3: {}", key);
        Ok(())
    }

    /// Download a toolchain tarball from S3
    ///
    /// Downloads the toolchain tarball and verifies its SHA256 checksum.
    ///
    /// # Arguments
    ///
    /// * `toolchain_name` - Name of the toolchain (e.g., "gnu-riscv", "rialo-rust")
    /// * `version` - Version of the toolchain (e.g., "13.2.0", "latest")
    /// * `platform` - Target platform (e.g., "x86_64-apple-darwin")
    /// * `dest_path` - Path where the downloaded tarball should be saved
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The S3 object does not exist
    /// - The download fails
    /// - The checksum verification fails
    /// - The destination file cannot be written
    pub async fn download_toolchain(
        &self,
        toolchain_name: &str,
        version: &str,
        platform: &str,
        dest_path: &Path,
    ) -> Result<()> {
        let archive_name = self.get_archive_name(toolchain_name, version, platform);
        let key = format!("toolchains/{toolchain_name}/{version}/{archive_name}.tar.gz");

        log::info!("Downloading from s3://{}/{}", self.bucket.name(), key);

        // Download from S3
        let response = self.bucket.get_object(&key).await.with_context(|| {
            format!(
                "Failed to download from S3: s3://{}/{}",
                self.bucket.name(),
                key
            )
        })?;

        // Check if the request was successful
        if response.status_code() != 200 {
            return Err(anyhow::anyhow!(
                "Failed to download from S3: HTTP {}",
                response.status_code()
            ));
        }

        // Write to file
        tokio::fs::write(dest_path, response.bytes())
            .await
            .with_context(|| format!("Failed to write to {}", dest_path.display()))?;

        // Verify checksum
        let expected_checksum = self.download_checksum(&key).await?;
        if let Err(e) = verify_checksum(dest_path, &expected_checksum) {
            // Clean up corrupted file on checksum failure
            let _ = std::fs::remove_file(dest_path);
            return Err(e);
        }

        log::info!("Downloaded from S3 (checksum verified)");
        Ok(())
    }

    /// Check if a toolchain exists in S3
    ///
    /// Uses HEAD request to check if the toolchain tarball exists without downloading it.
    ///
    /// # Arguments
    ///
    /// * `toolchain_name` - Name of the toolchain (e.g., "gnu-riscv", "rialo-rust")
    /// * `version` - Version of the toolchain (e.g., "13.2.0", "latest")
    /// * `platform` - Target platform (e.g., "x86_64-apple-darwin")
    ///
    /// # Returns
    ///
    /// `true` if the toolchain exists in S3, `false` otherwise
    pub async fn check_availability(
        &self,
        toolchain_name: &str,
        version: &str,
        platform: &str,
    ) -> bool {
        let archive_name = self.get_archive_name(toolchain_name, version, platform);
        let key = format!("toolchains/{toolchain_name}/{version}/{archive_name}.tar.gz");

        self.bucket.head_object(&key).await.is_ok()
    }

    /// Get the archive name based on toolchain naming conventions
    ///
    /// Different toolchains use different naming conventions for compatibility:
    /// - **GNU RISC-V**: `riscv64-elf-{platform}-{version}` (version at end)
    ///   This matches the naming convention used by the official riscv-gnu-toolchain
    ///   releases, making it easier to identify and compare versions.
    /// - **Rialo Rust**: `rialo-rust-{version}-{platform}` (version before platform)
    ///   This follows Rust's standard toolchain naming where version comes first,
    ///   consistent with rustup conventions.
    /// - **Generic**: `{toolchain_name}-{version}-{platform}`
    ///   Default pattern for future toolchains.
    fn get_archive_name(&self, toolchain_name: &str, version: &str, platform: &str) -> String {
        match toolchain_name {
            "gnu-riscv" => format!("riscv64-elf-{platform}-{version}"),
            "rialo-rust" => format!("rialo-rust-{version}-{platform}"),
            _ => format!("{toolchain_name}-{version}-{platform}"),
        }
    }

    /// Upload a checksum file to S3
    ///
    /// Stores the SHA256 checksum as a text file alongside the tarball.
    async fn upload_checksum(&self, archive_key: &str, checksum: &str) -> Result<()> {
        let checksum_key = format!("{archive_key}.sha256");

        self.bucket
            .put_object(&checksum_key, checksum.as_bytes())
            .await
            .context("Failed to upload checksum to S3")?;

        Ok(())
    }

    /// Download a checksum file from S3
    ///
    /// Retrieves the SHA256 checksum file associated with a tarball.
    async fn download_checksum(&self, archive_key: &str) -> Result<String> {
        let checksum_key = format!("{archive_key}.sha256");

        let response = self
            .bucket
            .get_object(&checksum_key)
            .await
            .with_context(|| {
                format!(
                    "Failed to download checksum from S3: s3://{}/{}",
                    self.bucket.name(),
                    checksum_key
                )
            })?;

        if response.status_code() != 200 {
            return Err(anyhow::anyhow!(
                "Failed to download checksum: HTTP {}",
                response.status_code()
            ));
        }

        let checksum = String::from_utf8(response.bytes().to_vec())
            .context("Checksum file contains invalid UTF-8")?
            .trim()
            .to_string();

        Ok(checksum)
    }
}

/// Compute SHA256 checksum of a file
///
/// # Arguments
///
/// * `file_path` - Path to the file to hash
///
/// # Returns
///
/// Hex-encoded SHA256 hash of the file contents
///
/// # Errors
///
/// Returns an error if the file cannot be read
fn compute_sha256(file_path: &Path) -> Result<String> {
    use std::{fs::File, io};

    use sha2::{Digest, Sha256};

    let mut file =
        File::open(file_path).with_context(|| format!("Failed to open {}", file_path.display()))?;

    let mut hasher = Sha256::new();
    io::copy(&mut file, &mut hasher)
        .with_context(|| format!("Failed to read {}", file_path.display()))?;

    Ok(hex::encode(hasher.finalize()))
}

/// Verify that a file matches an expected SHA256 checksum
///
/// # Arguments
///
/// * `file_path` - Path to the file to verify
/// * `expected_checksum` - Expected SHA256 hash (hex-encoded)
///
/// # Errors
///
/// Returns an error if:
/// - The file cannot be read
/// - The checksum does not match
fn verify_checksum(file_path: &Path, expected_checksum: &str) -> Result<()> {
    let actual_checksum = compute_sha256(file_path)?;

    if actual_checksum != expected_checksum {
        return Err(anyhow::anyhow!(
            "Checksum mismatch for {}: expected {}, got {}",
            file_path.display(),
            expected_checksum,
            actual_checksum
        ));
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_archive_naming_gnu_riscv() {
        // Create a mock backend for testing
        let bucket = Bucket::new(
            "test-bucket",
            Region::UsEast1,
            Credentials::new(Some("test"), Some("test"), None, None, None).unwrap(),
        )
        .unwrap();
        let backend = S3StorageBackend { bucket };

        let name = backend.get_archive_name("gnu-riscv", "13.2.0", "x86_64-apple-darwin");
        assert_eq!(name, "riscv64-elf-x86_64-apple-darwin-13.2.0");
    }

    #[test]
    fn test_archive_naming_rialo_rust() {
        let bucket = Bucket::new(
            "test-bucket",
            Region::UsEast1,
            Credentials::new(Some("test"), Some("test"), None, None, None).unwrap(),
        )
        .unwrap();
        let backend = S3StorageBackend { bucket };

        let name = backend.get_archive_name("rialo-rust", "latest", "aarch64-apple-darwin");
        assert_eq!(name, "rialo-rust-latest-aarch64-apple-darwin");
    }

    #[test]
    fn test_archive_naming_generic() {
        let bucket = Bucket::new(
            "test-bucket",
            Region::UsEast1,
            Credentials::new(Some("test"), Some("test"), None, None, None).unwrap(),
        )
        .unwrap();
        let backend = S3StorageBackend { bucket };

        let name =
            backend.get_archive_name("custom-toolchain", "1.0.0", "x86_64-unknown-linux-gnu");
        assert_eq!(name, "custom-toolchain-1.0.0-x86_64-unknown-linux-gnu");
    }

    #[test]
    fn test_compute_sha256() {
        use std::io::Write;

        use tempfile::NamedTempFile;

        let mut temp_file = NamedTempFile::new().unwrap();
        temp_file.write_all(b"Hello, World!").unwrap();
        temp_file.flush().unwrap();

        let checksum = compute_sha256(temp_file.path()).unwrap();

        // SHA256 of "Hello, World!"
        assert_eq!(
            checksum,
            "dffd6021bb2bd5b0af676290809ec3a53191dd81c7f70a4b28688a362182986f"
        );
    }

    #[test]
    fn test_verify_checksum_success() {
        use std::io::Write;

        use tempfile::NamedTempFile;

        let mut temp_file = NamedTempFile::new().unwrap();
        temp_file.write_all(b"Hello, World!").unwrap();
        temp_file.flush().unwrap();

        let result = verify_checksum(
            temp_file.path(),
            "dffd6021bb2bd5b0af676290809ec3a53191dd81c7f70a4b28688a362182986f",
        );

        assert!(result.is_ok());
    }

    #[test]
    fn test_verify_checksum_failure() {
        use std::io::Write;

        use tempfile::NamedTempFile;

        let mut temp_file = NamedTempFile::new().unwrap();
        temp_file.write_all(b"Hello, World!").unwrap();
        temp_file.flush().unwrap();

        let result = verify_checksum(temp_file.path(), "wrong_checksum");

        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("Checksum mismatch"));
    }
}