diff --git a/README.md b/README.md index 3d1972f..8deb5fb 100644 --- a/README.md +++ b/README.md @@ -95,6 +95,57 @@ Write complete: 5120.00 MB in 281.23s (18.21 MB/s) Compression ratio: 9.80x ``` +### OCI Images + +`fls` can pull OCI images from registries and flash them either to a block device +or to fastboot partitions. + +#### Flash an OCI image to a block device + +Use `from-url` with an `oci://` prefix: + +```bash +fls from-url \ + -u "$REGISTRY_USER" \ + -p "$REGISTRY_PASS" \ + "oci://quay.io/org/image:latest" \ + /dev/mmcblk1 +``` + +#### Flash an OCI image via fastboot + +`fls fastboot` pulls the OCI image, extracts partition images, and flashes them +using the system `fastboot` CLI: + +```bash +fls fastboot oci://quay.io/org/image:latest +``` + +To avoid using `/tmp` for fastboot extraction, set `FLS_TMP_DIR` to a directory +on persistent storage (default is `/var/lib/fls`). + +If the OCI manifest includes +`automotive.sdv.cloud.redhat.com/default-partitions` (comma-separated), +`fls fastboot` flashes only those partitions by default. Otherwise it flashes +all annotated partitions. + +Provide explicit partition mappings when the OCI image contains multiple files: + +```bash +fls fastboot oci://quay.io/org/image:latest \ + -t boot_a:boot_a.simg \ + -t system_a:system_a.simg +``` + +When `-t` is provided, `fls` applies those mappings on top of the OCI layer +annotations and includes those partitions in the flash set (e.g., add +`-t vbmeta_a:vbmeta_a.simg` to flash vbmeta). If the image lacks annotations, +it falls back to looking for the specified files directly in the image. + +Registry credentials can be provided with `-u/--username` and `-p/--password` +(`FLS_REGISTRY_PASSWORD` env var is supported for the password). Both are +required for authenticated access. + ## Command Options ### `fls from-url` diff --git a/src/fls/automotive.rs b/src/fls/automotive.rs new file mode 100644 index 0000000..53246bb --- /dev/null +++ b/src/fls/automotive.rs @@ -0,0 +1,52 @@ +//! CentOS Automotive Suite OCI constant +//! + +/// OCI annotation keys used by automotive images for partition mapping +pub mod annotations { + pub const PARTITION_ANNOTATION: &str = "automotive.sdv.cloud.redhat.com/partition"; + + pub const DECOMPRESSED_SIZE: &str = "automotive.sdv.cloud.redhat.com/decompressed-size"; + + pub const TARGET: &str = "automotive.sdv.cloud.redhat.com/target"; + + pub const ARCH: &str = "automotive.sdv.cloud.redhat.com/arch"; + + pub const DISTRO: &str = "automotive.sdv.cloud.redhat.com/distro"; + + pub const MULTI_LAYER: &str = "automotive.sdv.cloud.redhat.com/multi-layer"; + + pub const PARTS: &str = "automotive.sdv.cloud.redhat.com/parts"; + + /// Comma-separated list of default partitions to flash + pub const DEFAULT_PARTITIONS: &str = "automotive.sdv.cloud.redhat.com/default-partitions"; +} + +pub mod targets { + pub const RIDESX4: &str = "ridesx4"; + pub const AUTOSD: &str = "autosd"; +} + +/// Extract partition name from layer annotations +pub fn extract_partition_name( + layer_annotations: &std::collections::HashMap, +) -> Option { + layer_annotations + .get(annotations::PARTITION_ANNOTATION) + .cloned() +} + +/// Extract decompressed size from layer annotations +pub fn extract_decompressed_size( + layer_annotations: &std::collections::HashMap, +) -> Option { + layer_annotations + .get(annotations::DECOMPRESSED_SIZE) + .and_then(|s| s.parse().ok()) +} + +/// Extract target platform from OCI annotations +pub fn extract_target_from_annotations( + manifest_annotations: &std::collections::HashMap, +) -> Option { + manifest_annotations.get(annotations::TARGET).cloned() +} diff --git a/src/fls/fastboot.rs b/src/fls/fastboot.rs new file mode 100644 index 0000000..899dc4f --- /dev/null +++ b/src/fls/fastboot.rs @@ -0,0 +1,651 @@ +//! Fastboot flashing implementation +//! +//! This module uses the system fastboot CLI for flashing partitions. +//! OCI images are downloaded and extracted, then written to a temporary +//! directory and flashed via `fastboot flash`. + +use std::collections::{HashMap, HashSet}; +use std::error::Error; +use std::fmt; +use std::path::{Path, PathBuf}; +use std::process::Stdio; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use tokio::process::Command; +use tokio::time::timeout; + +use super::options::FastbootOptions; + +#[derive(Debug)] +#[allow(dead_code)] +pub enum FastbootError { + DeviceNotFound(Option), + CommandError(String), + PartitionError(String), + TimeoutError, +} + +impl fmt::Display for FastbootError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + FastbootError::DeviceNotFound(serial) => { + if let Some(s) = serial { + write!(f, "Fastboot device with serial '{}' not found", s) + } else { + write!(f, "No fastboot devices found") + } + } + FastbootError::CommandError(msg) => write!(f, "Fastboot command error: {}", msg), + FastbootError::PartitionError(msg) => write!(f, "Partition error: {}", msg), + FastbootError::TimeoutError => write!(f, "Operation timed out"), + } + } +} + +impl Error for FastbootError {} + +fn temp_base_dir() -> Result { + let env_value = std::env::var("FLS_TMP_DIR").ok(); + let base = match env_value + .as_deref() + .map(str::trim) + .filter(|v| !v.is_empty()) + { + Some(value) => PathBuf::from(value), + None => PathBuf::from("/var/lib/fls"), + }; + + std::fs::create_dir_all(&base).map_err(|e| { + let hint = if env_value.is_some() { + format!( + "Failed to create temp base directory {}: {}", + base.display(), + e + ) + } else { + format!( + "Failed to create temp base directory {}: {}. Set FLS_TMP_DIR to override.", + base.display(), + e + ) + }; + FastbootError::CommandError(hint) + })?; + + Ok(base) +} + +fn build_oci_options(options: &FastbootOptions) -> crate::fls::options::OciOptions { + crate::fls::options::OciOptions { + common: crate::fls::options::FlashOptions { + insecure_tls: options.http.insecure_tls, + cacert: options.http.cacert.clone(), + debug: options.http.debug, + ..crate::fls::options::FlashOptions::default() + }, + username: options.username.clone(), + password: options.password.clone(), + file_pattern: None, + } +} + +/// Main entry point for fastboot flashing +/// +/// Downloads an OCI image and flashes it to a fastboot device via the system fastboot CLI. +/// The image_ref should be an OCI image reference without a scheme +/// (e.g., "registry.example.com/my-image:latest"). +pub async fn flash_from_fastboot( + image_ref: &str, + options: FastbootOptions, +) -> Result<(), Box> { + let temp_dir = TempDirGuard::new("fls-fastboot")?; + let partition_map = process_oci_image_to_dir(image_ref, &options, temp_dir.path()).await?; + + if partition_map.is_empty() { + return Err("No partition files found for flashing".into()); + } + + flash_partitions_with_fastboot_cli(&partition_map, &options).await?; + + println!("Fastboot flash completed successfully!"); + Ok(()) +} + +struct TempDirGuard { + path: PathBuf, +} + +impl TempDirGuard { + fn new(prefix: &str) -> Result { + let base = temp_base_dir()?; + let pid = std::process::id(); + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis(); + let dir = base.join(format!("{}-{}-{}", prefix, pid, timestamp)); + std::fs::create_dir(&dir).map_err(|e| { + FastbootError::CommandError(format!( + "Failed to create temp directory {}: {}", + dir.display(), + e + )) + })?; + Ok(Self { path: dir }) + } + + fn path(&self) -> &Path { + &self.path + } +} + +impl Drop for TempDirGuard { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.path); + } +} + +async fn run_fastboot_command( + fastboot_path: &str, + args: &[String], + timeout_duration: Duration, +) -> Result { + let mut cmd = Command::new(fastboot_path); + cmd.args(args) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true); + + let output = match timeout(timeout_duration, cmd.output()).await { + Ok(Ok(output)) => output, + Ok(Err(e)) => { + return Err(FastbootError::CommandError(format!( + "Failed to run fastboot: {}", + e + ))) + } + Err(_) => return Err(FastbootError::TimeoutError), + }; + + if !output.status.success() { + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + return Err(FastbootError::CommandError(format!( + "fastboot {:?} failed (code {:?}). stdout: '{}' stderr: '{}'", + args, + output.status.code(), + stdout.trim(), + stderr.trim() + ))); + } + + Ok(output) +} + +async fn detect_fastboot_device_id( + fastboot_path: &str, + timeout_duration: Duration, +) -> Result { + let output = run_fastboot_command( + fastboot_path, + &["devices".to_string(), "-l".to_string()], + timeout_duration, + ) + .await?; + + let stdout = String::from_utf8_lossy(&output.stdout); + let mut ids: Vec = stdout + .lines() + .filter_map(|line| line.split_whitespace().next().map(|s| s.to_string())) + .collect(); + + if ids.is_empty() { + return Err(FastbootError::DeviceNotFound(None)); + } + + if ids.len() > 1 { + println!( + "Warning: Multiple fastboot devices found, using first: {}", + ids[0] + ); + } + + Ok(ids.remove(0)) +} + +async fn flash_partitions_with_fastboot_cli( + partition_map: &HashMap, + options: &FastbootOptions, +) -> Result<(), Box> { + let fastboot_path = "fastboot"; + let timeout_duration = Duration::from_secs(options.timeout_secs as u64); + + let device_id = if let Some(serial) = options.device_serial.as_deref() { + serial.to_string() + } else { + detect_fastboot_device_id(fastboot_path, timeout_duration).await? + }; + + println!("Using fastboot CLI: {}", fastboot_path); + println!("Using device: {}", device_id); + + let mut partitions: Vec<(&String, &PathBuf)> = partition_map.iter().collect(); + partitions.sort_by(|a, b| a.0.cmp(b.0)); + + for (partition, path) in &partitions { + println!("Flashing partition '{}' from {}", partition, path.display()); + let mut args = vec!["-s".to_string(), device_id.clone()]; + args.push("flash".to_string()); + args.push((*partition).clone()); + args.push(path.display().to_string()); + run_fastboot_command(fastboot_path, &args, timeout_duration).await?; + } + + println!("Running fastboot continue..."); + let args = vec!["-s".to_string(), device_id.clone(), "continue".to_string()]; + run_fastboot_command(fastboot_path, &args, timeout_duration).await?; + + println!("Fastboot CLI flash completed successfully."); + Ok(()) +} + +/// Process OCI image and extract files for flashing +async fn process_oci_image_to_dir( + image_ref: &str, + options: &FastbootOptions, + output_dir: &Path, +) -> Result, Box> { + println!("Processing OCI image for fastboot: {}", image_ref); + + if options.partition_mappings.is_empty() { + println!("No explicit partition mappings provided, attempting auto-detection from OCI annotations"); + return extract_files_by_auto_detection_to_dir(image_ref, options, output_dir).await; + } + + println!("Partition mappings provided; applying overrides on top of OCI annotations"); + + let oci_options = build_oci_options(options); + + match super::oci::extract_files_by_annotations_with_overrides_to_dir( + image_ref, + &oci_options, + output_dir, + &options.partition_mappings, + ) + .await + { + Ok(Some(partition_files)) => return Ok(partition_files), + Ok(None) => { + println!("No OCI annotations found, falling back to file-pattern extraction"); + } + Err(e) => { + return Err(format!("Annotation-based extraction failed: {}", e).into()); + } + } + + extract_files_by_patterns_to_dir(image_ref, options, output_dir).await +} + +async fn extract_files_by_auto_detection_to_dir( + image_ref: &str, + options: &FastbootOptions, + output_dir: &Path, +) -> Result, Box> { + println!("Auto-detecting partitions from OCI layer annotations..."); + + // Create OCI options + let oci_options = build_oci_options(options); + + // Use annotation-aware extraction to get files from correct layers + let partition_files = + super::oci::extract_files_by_annotations_with_overrides_to_dir( + image_ref, + &oci_options, + output_dir, + &[], + ) + .await + .map_err(|e| format!("Annotation-based extraction failed: {}", e))? + .ok_or_else(|| { + "No partitions found in OCI annotations. Expected layers with 'automotive.sdv.cloud.redhat.com/partition' annotations".to_string() + })?; + + if partition_files.is_empty() { + return Err("No partitions found in OCI annotations. Expected layers with 'automotive.sdv.cloud.redhat.com/partition' annotations".into()); + } + + println!( + "Successfully extracted {} partitions:", + partition_files.len() + ); + for (partition, path) in &partition_files { + println!(" {} -> {}", partition, path.display()); + } + + Ok(partition_files) +} + +async fn extract_files_by_patterns_to_dir( + image_ref: &str, + options: &FastbootOptions, + output_dir: &Path, +) -> Result, Box> { + let mut partition_files = HashMap::new(); + let target_files: HashSet = options + .partition_mappings + .iter() + .map(|(_, filename)| filename.clone()) + .collect(); + + println!("Extracting files using partition mappings:"); + for (partition_name, file_pattern) in &options.partition_mappings { + println!(" {} = {}", partition_name, file_pattern); + println!(" Searching for file: {}", file_pattern); + } + + let file_map = + extract_files_from_oci_to_dir(image_ref, options, output_dir, &target_files).await?; + + for (partition_name, file_pattern) in &options.partition_mappings { + match file_map.get(file_pattern) { + Some(path) => { + println!( + " Found file for partition '{}' at {}", + partition_name, + path.display() + ); + partition_files.insert(partition_name.clone(), path.clone()); + } + None => { + let available_files: Vec<&String> = file_map.keys().collect(); + return Err(format!( + "Missing partition file: {} (pattern: {}): File '{}' not found in OCI image. Available files: {:?}", + partition_name, file_pattern, file_pattern, available_files + ) + .into()); + } + } + } + + if partition_files.is_empty() { + return Err("No partition files found matching the specified patterns".into()); + } + + Ok(partition_files) +} + +async fn extract_files_from_oci_to_dir( + image_ref: &str, + options: &FastbootOptions, + output_dir: &Path, + target_files: &HashSet, +) -> Result, Box> { + println!("Downloading OCI image: {}", image_ref); + + println!("Looking for files: {:?}", target_files); + + let oci_options = build_oci_options(options); + + let file_data = super::oci::extract_files_from_oci_image_to_dir( + image_ref, + target_files, + &oci_options, + output_dir, + ) + .await + .map_err(|e| -> Box { + format!("OCI extraction failed: {}", e).into() + })?; + + println!( + "Successfully extracted {} files from OCI image", + file_data.len() + ); + Ok(file_data) +} + +#[cfg(test)] +mod tests { + use super::*; + #[cfg(unix)] + use std::fs; + #[cfg(unix)] + use std::future::Future; + #[cfg(unix)] + use std::os::unix::fs::PermissionsExt; + #[cfg(unix)] + use std::sync::{Mutex, OnceLock}; + #[cfg(unix)] + use tempfile::tempdir; + + #[test] + fn test_fastboot_options_default() { + let options = FastbootOptions::default(); + assert_eq!(options.timeout_secs, 1200); + assert!(options.partition_mappings.is_empty()); + assert!(options.username.is_none()); + assert!(options.password.is_none()); + } + + #[test] + fn test_fastboot_error_display() { + let err = FastbootError::DeviceNotFound(Some("ABC123".to_string())); + assert_eq!( + err.to_string(), + "Fastboot device with serial 'ABC123' not found" + ); + } + + #[cfg(unix)] + fn create_fastboot_script(contents: &str) -> (tempfile::TempDir, PathBuf) { + let dir = tempdir().expect("create temp dir"); + let path = dir.path().join("fastboot"); + fs::write(&path, contents).expect("write fastboot script"); + let mut perms = fs::metadata(&path) + .expect("read script metadata") + .permissions(); + perms.set_mode(0o755); + fs::set_permissions(&path, perms).expect("set script permissions"); + (dir, path) + } + + #[cfg(unix)] + fn path_lock() -> &'static Mutex<()> { + static LOCK: OnceLock> = OnceLock::new(); + LOCK.get_or_init(|| Mutex::new(())) + } + + #[cfg(unix)] + struct EnvVarGuard { + key: &'static str, + old: Option, + } + + #[cfg(unix)] + impl EnvVarGuard { + fn set(key: &'static str, value: String) -> Self { + let old = std::env::var(key).ok(); + std::env::set_var(key, &value); + Self { key, old } + } + } + + #[cfg(unix)] + impl Drop for EnvVarGuard { + fn drop(&mut self) { + if let Some(old) = &self.old { + std::env::set_var(self.key, old); + } else { + std::env::remove_var(self.key); + } + } + } + + #[cfg(unix)] + fn prepend_path(dir: &Path) -> String { + let current = std::env::var("PATH").unwrap_or_default(); + format!("{}:{}", dir.display(), current) + } + + #[cfg(unix)] + fn block_on(future: F) -> F::Output { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("build tokio runtime") + .block_on(future) + } + + #[cfg(unix)] + #[tokio::test] + async fn test_detect_fastboot_device_id_parses_first_device() { + let script = r#"#!/bin/sh +if [ "$1" = "devices" ]; then + echo "ABC123 fastboot" + echo "DEF456 fastboot" +fi +"#; + let (_dir, fastboot_path) = create_fastboot_script(script); + let fastboot_path = fastboot_path.to_string_lossy().to_string(); + + let device_id = detect_fastboot_device_id(&fastboot_path, Duration::from_secs(2)).await; + + assert_eq!(device_id.unwrap(), "ABC123"); + } + + #[cfg(unix)] + #[tokio::test] + async fn test_detect_fastboot_device_id_no_devices() { + let script = r#"#!/bin/sh +if [ "$1" = "devices" ]; then + exit 0 +fi +"#; + let (_dir, fastboot_path) = create_fastboot_script(script); + let fastboot_path = fastboot_path.to_string_lossy().to_string(); + + let err = detect_fastboot_device_id(&fastboot_path, Duration::from_secs(2)).await; + + assert!(matches!(err, Err(FastbootError::DeviceNotFound(None)))); + } + + #[cfg(unix)] + #[tokio::test] + async fn test_run_fastboot_command_failure_includes_output() { + let script = r#"#!/bin/sh +if [ "$1" = "fail" ]; then + echo "stdout message" + echo "stderr message" >&2 + exit 7 +fi +"#; + let (_dir, fastboot_path) = create_fastboot_script(script); + let fastboot_path = fastboot_path.to_string_lossy().to_string(); + + let err = run_fastboot_command( + &fastboot_path, + &["fail".to_string()], + Duration::from_secs(2), + ) + .await + .unwrap_err(); + + match err { + FastbootError::CommandError(msg) => { + assert!(msg.contains("stdout message")); + assert!(msg.contains("stderr message")); + } + other => panic!("Unexpected error: {:?}", other), + } + } + + #[cfg(unix)] + #[tokio::test] + async fn test_run_fastboot_command_timeout() { + let script = r#"#!/bin/sh +if [ "$1" = "sleep" ]; then + sleep 2 +fi +"#; + let (_dir, fastboot_path) = create_fastboot_script(script); + let fastboot_path = fastboot_path.to_string_lossy().to_string(); + + let err = run_fastboot_command( + &fastboot_path, + &["sleep".to_string()], + Duration::from_millis(100), + ) + .await + .unwrap_err(); + + assert!(matches!(err, FastbootError::TimeoutError)); + } + + #[cfg(unix)] + #[test] + fn test_flash_partitions_with_fastboot_cli_runs_flash_and_continue() { + let _lock = path_lock().lock().expect("lock PATH"); + let script = r#"#!/bin/sh +if [ -n "$LOG_FILE" ]; then + echo "$@" >> "$LOG_FILE" +fi +if [ "$1" = "devices" ]; then + echo "ABC123 fastboot" +fi +exit 0 +"#; + let (dir, _fastboot_path) = create_fastboot_script(script); + let log_path = dir.path().join("fastboot.log"); + let _log_guard = EnvVarGuard::set("LOG_FILE", log_path.to_string_lossy().to_string()); + let _path_guard = EnvVarGuard::set("PATH", prepend_path(dir.path())); + + let images_dir = tempdir().expect("create images dir"); + let boot_path = images_dir.path().join("boot_a.simg"); + let system_path = images_dir.path().join("system_a.simg"); + fs::write(&boot_path, b"boot image").expect("write boot image"); + fs::write(&system_path, b"system image").expect("write system image"); + + let mut partition_map = HashMap::new(); + partition_map.insert("system_a".to_string(), system_path.clone()); + partition_map.insert("boot_a".to_string(), boot_path.clone()); + + let options = FastbootOptions { + device_serial: Some("SER123".to_string()), + timeout_secs: 2, + ..Default::default() + }; + + block_on(async { + flash_partitions_with_fastboot_cli(&partition_map, &options) + .await + .expect("fastboot flash should succeed"); + }); + + let log = fs::read_to_string(&log_path).expect("read fastboot log"); + let lines: Vec = log.lines().map(|line| line.to_string()).collect(); + + let expected = vec![ + format!("-s SER123 flash boot_a {}", boot_path.display()), + format!("-s SER123 flash system_a {}", system_path.display()), + "-s SER123 continue".to_string(), + ]; + + assert_eq!(lines, expected); + } + + #[cfg(unix)] + #[test] + fn test_temp_dir_guard_cleans_up() { + let base_dir = tempdir().expect("create temp base dir"); + let _env_guard = + EnvVarGuard::set("FLS_TMP_DIR", base_dir.path().to_string_lossy().to_string()); + let path = { + let guard = TempDirGuard::new("fls-fastboot-test").expect("create temp dir guard"); + let path = guard.path().to_path_buf(); + assert!(path.exists(), "temp dir should exist"); + path + }; + + assert!(!path.exists(), "temp dir should be removed on drop"); + } +} diff --git a/src/fls/mod.rs b/src/fls/mod.rs index 40e29f8..67736a5 100644 --- a/src/fls/mod.rs +++ b/src/fls/mod.rs @@ -1,9 +1,11 @@ // Module declarations +pub mod automotive; mod block_writer; pub(crate) mod compression; mod decompress; mod download_error; mod error_handling; +mod fastboot; mod format_detector; mod from_url; pub(crate) mod http; @@ -16,9 +18,12 @@ mod simg; mod stream_utils; // Public re-exports +pub use fastboot::flash_from_fastboot; pub use from_url::flash_from_url; pub use oci::flash_from_oci; -pub use options::{BlockFlashOptions, FlashOptions, HttpClientOptions, OciOptions}; +pub use options::{ + BlockFlashOptions, FastbootOptions, FlashOptions, HttpClientOptions, OciOptions, +}; #[cfg(test)] mod tests { diff --git a/src/fls/oci/from_oci.rs b/src/fls/oci/from_oci.rs index 8733bf9..56f86c6 100644 --- a/src/fls/oci/from_oci.rs +++ b/src/fls/oci/from_oci.rs @@ -2,10 +2,12 @@ /// /// Implements the streaming pipeline: /// Registry blob -> gzip decompress -> tar extract -> xzcat -> block device +use std::fs::File; use std::io::{Read, Write}; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::time::Duration; +use bytes::Bytes; use flate2::read::GzDecoder; use futures_util::StreamExt; use tokio::io::AsyncWriteExt; @@ -15,17 +17,22 @@ use xz2::read::XzDecoder; use super::manifest::{LayerCompression, Manifest}; use super::reference::ImageReference; use super::registry::RegistryClient; +use crate::fls::automotive::annotations as automotive_annotations; use crate::fls::block_writer::AsyncBlockWriter; use crate::fls::compression::Compression; use crate::fls::decompress::{spawn_stderr_reader, start_decompressor_process}; use crate::fls::error_handling::process_error_messages; use crate::fls::format_detector::{DetectionResult, FileFormat, FormatDetector}; -use crate::fls::magic_bytes::{detect_content_and_compression, ContentType}; +use crate::fls::magic_bytes::{detect_compression, detect_content_and_compression, ContentType}; use crate::fls::options::OciOptions; use crate::fls::progress::ProgressTracker; use crate::fls::simg::{SparseParser, WriteCommand}; use crate::fls::stream_utils::ChannelReader; +use std::collections::{HashMap, HashSet}; + +const OCI_TITLE_ANNOTATION: &str = "org.opencontainers.image.title"; + /// Parameters for download coordination functions struct DownloadCoordinationParams { http_tx: mpsc::Sender, @@ -79,6 +86,598 @@ struct TarPipelineComponents { decompressor_name: &'static str, } +/// Connect to an OCI registry and resolve the image manifest. +/// +/// Handles image reference parsing, client creation, authentication, and +/// manifest resolution (including multi-platform index negotiation). +async fn connect_and_resolve( + image: &str, + options: &OciOptions, +) -> Result<(RegistryClient, Manifest), Box> { + let image_ref = ImageReference::parse(image)?; + println!("Pulling OCI image: {}", image_ref); + + let mut client = RegistryClient::new(image_ref.clone(), options).await?; + println!("Connecting to registry: {}", image_ref.registry); + client.authenticate().await?; + + let manifest = resolve_manifest(&mut client).await?; + Ok((client, manifest)) +} + +async fn stream_blob_to_file( + mut stream: impl futures_util::Stream> + Unpin + Send + 'static, + output_path: PathBuf, +) -> Result<(), Box> { + let mut prefix = Vec::new(); + let mut initial_chunks: Vec = Vec::new(); + while prefix.len() < 6 { + match stream.next().await { + Some(Ok(chunk)) => { + if prefix.len() < 6 { + let needed = 6 - prefix.len(); + let take = needed.min(chunk.len()); + prefix.extend_from_slice(&chunk[..take]); + } + initial_chunks.push(chunk); + } + Some(Err(e)) => return Err(format!("Stream error: {}", e).into()), + None => break, + } + } + + if initial_chunks.is_empty() { + return Err("Empty OCI layer stream".into()); + } + + let blob_compression = detect_compression(&prefix); + + let (tx, rx) = mpsc::channel::(16); + for chunk in initial_chunks { + tx.send(chunk) + .await + .map_err(|_| "Failed to send initial chunk to reader")?; + } + + let forward_handle = tokio::spawn(async move { + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(|e| format!("Stream error: {}", e))?; + tx.send(chunk) + .await + .map_err(|_| "Reader channel closed".to_string())?; + } + Ok::<(), String>(()) + }); + + let writer_handle = tokio::task::spawn_blocking(move || -> Result<(), String> { + let reader = ChannelReader::new(rx); + let mut file = File::create(&output_path) + .map_err(|e| format!("Failed to create {}: {}", output_path.display(), e))?; + + match blob_compression { + Compression::Gzip => { + let mut decoder = GzDecoder::new(reader); + std::io::copy(&mut decoder, &mut file) + .map_err(|e| format!("Gzip decompression failed: {}", e))?; + } + Compression::Xz => { + let mut decoder = XzDecoder::new(reader); + std::io::copy(&mut decoder, &mut file) + .map_err(|e| format!("XZ decompression failed: {}", e))?; + } + _ => { + let mut reader = reader; + std::io::copy(&mut reader, &mut file) + .map_err(|e| format!("Stream copy failed: {}", e))?; + } + } + + file.flush() + .map_err(|e| format!("Failed to flush {}: {}", output_path.display(), e))?; + Ok(()) + }); + + let (forward_result, writer_result) = tokio::join!(forward_handle, writer_handle); + + // Prioritize writer errors — a writer failure (e.g. disk full, decompression + // error) drops rx, which causes the forwarder's tx.send to fail with a + // misleading "Reader channel closed" error. + writer_result + .map_err(|e| format!("Writer task failed: {}", e))? + .map_err(|e| format!("Writer task error: {}", e))?; + forward_result + .map_err(|e| format!("Stream task failed: {}", e))? + .map_err(|e| format!("Stream task error: {}", e))?; + + Ok(()) +} + +async fn stream_blob_to_tar_files( + mut stream: impl futures_util::Stream> + Unpin + Send + 'static, + target_files: std::collections::HashSet, + output_dir: PathBuf, +) -> Result, Box> { + let mut prefix = Vec::new(); + let mut initial_chunks: Vec = Vec::new(); + while prefix.len() < 6 { + match stream.next().await { + Some(Ok(chunk)) => { + if prefix.len() < 6 { + let needed = 6 - prefix.len(); + let take = needed.min(chunk.len()); + prefix.extend_from_slice(&chunk[..take]); + } + initial_chunks.push(chunk); + } + Some(Err(e)) => return Err(format!("Stream error: {}", e).into()), + None => break, + } + } + + if initial_chunks.is_empty() { + return Err("Empty OCI layer stream".into()); + } + + let blob_compression = detect_compression(&prefix); + + let (tx, rx) = mpsc::channel::(16); + for chunk in initial_chunks { + tx.send(chunk) + .await + .map_err(|_| "Failed to send initial chunk to reader")?; + } + + let forward_handle = tokio::spawn(async move { + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(|e| format!("Stream error: {}", e))?; + tx.send(chunk) + .await + .map_err(|_| "Reader channel closed".to_string())?; + } + Ok::<(), String>(()) + }); + + let writer_handle = + tokio::task::spawn_blocking(move || -> Result, String> { + let reader = ChannelReader::new(rx); + let reader: Box = match blob_compression { + Compression::Gzip => Box::new(GzDecoder::new(reader)), + Compression::Xz => Box::new(XzDecoder::new(reader)), + _ => Box::new(reader), + }; + + let mut found = HashMap::new(); + let mut archive = tar::Archive::new(reader); + for entry_result in archive.entries().map_err(|e| format!("Tar error: {}", e))? { + let mut entry = entry_result.map_err(|e| format!("Tar entry error: {}", e))?; + let path = entry.path().map_err(|e| format!("Invalid path: {}", e))?; + let file_name = match path.file_name() { + Some(name) => name.to_string_lossy().to_string(), + None => continue, + }; + + if target_files.contains(&file_name) { + let output_path = output_dir.join(&file_name); + let mut file = File::create(&output_path).map_err(|e| { + format!("Failed to create {}: {}", output_path.display(), e) + })?; + let mut prefix = [0u8; 6]; + let prefix_len = entry + .read(&mut prefix) + .map_err(|e| format!("Failed to read {}: {}", file_name, e))?; + let prefix_slice = &prefix[..prefix_len]; + let cursor = std::io::Cursor::new(prefix_slice.to_vec()); + let mut combined = cursor.chain(entry); + + match detect_compression(prefix_slice) { + Compression::Gzip => { + let mut decoder = GzDecoder::new(combined); + std::io::copy(&mut decoder, &mut file).map_err(|e| { + format!("Failed to write {}: {}", output_path.display(), e) + })?; + } + Compression::Xz => { + let mut decoder = XzDecoder::new(combined); + std::io::copy(&mut decoder, &mut file).map_err(|e| { + format!("Failed to write {}: {}", output_path.display(), e) + })?; + } + _ => { + std::io::copy(&mut combined, &mut file).map_err(|e| { + format!("Failed to write {}: {}", output_path.display(), e) + })?; + } + } + found.insert(file_name, output_path); + + if found.len() == target_files.len() { + break; + } + } + } + + Ok(found) + }); + + let (forward_result, writer_result) = tokio::join!(forward_handle, writer_handle); + + // Prioritize writer errors (see stream_blob_to_file for rationale). + let found = writer_result + .map_err(|e| format!("Writer task failed: {}", e))? + .map_err(|e| format!("Writer task error: {}", e))?; + forward_result + .map_err(|e| format!("Stream task failed: {}", e))? + .map_err(|e| format!("Stream task error: {}", e))?; + + Ok(found) +} + +/// Extract specific files from an OCI image and write them to output_dir +pub async fn extract_files_from_oci_image_to_dir( + image: &str, + target_files: &std::collections::HashSet, + options: &OciOptions, + output_dir: &Path, +) -> Result, Box> { + let (client, manifest) = connect_and_resolve(image, options).await?; + + // Get the layer to download + let layer = manifest.get_single_layer()?; + let layer_size = layer.size; + + println!("Layer digest: {}", layer.digest); + println!( + "Layer size: {} bytes ({:.2} MB)", + layer_size, + layer_size as f64 / (1024.0 * 1024.0) + ); + + ensure_supported_layer_compression(layer.compression(), &layer.media_type)?; + + // Start blob download + println!("Starting download..."); + let response = client.get_blob_stream(&layer.digest).await?; + let stream = response.bytes_stream(); + + let is_tar_layer = layer.media_type.contains("tar") || looks_like_tar_layer(layer); + let should_use_tar = is_tar_layer || target_files.len() > 1; + + if should_use_tar { + if is_tar_layer { + println!("Processing as tar archive based on layer media type..."); + } else { + println!("Processing as tar archive due to multiple target files..."); + } + let file_map = + stream_blob_to_tar_files(stream, target_files.clone(), output_dir.to_path_buf()) + .await?; + return Ok(file_map); + } + + println!("Processing as direct file based on layer content..."); + if target_files.len() != 1 { + return Err("Multiple target files specified but OCI layer contains a single direct file. Use separate layers for each file.".into()); + } + + let filename = target_files.iter().next().unwrap(); + let basename = Path::new(filename) + .file_name() + .ok_or_else(|| format!("Invalid filename '{}'", filename))?; + let output_path = output_dir.join(basename); + + stream_blob_to_file(stream, output_path.clone()).await?; + + let mut map = HashMap::new(); + map.insert(filename.clone(), output_path); + Ok(map) +} + +/// Extract files based on layer annotations for automotive images and write to output_dir +pub async fn extract_files_by_annotations_to_dir( + image: &str, + options: &OciOptions, + output_dir: &Path, +) -> Result, Box> { + let (client, manifest) = connect_and_resolve(image, options).await?; + + // Get layers and extract from each based on annotations + let layers = manifest.get_layers()?; + let mut partition_files = HashMap::new(); + + for layer in layers { + if let Some(ref annotations) = layer.annotations { + if let Some(partition) = annotations.get(automotive_annotations::PARTITION_ANNOTATION) { + ensure_supported_layer_compression(layer.compression(), &layer.media_type)?; + let sanitized_name = sanitize_partition_name(partition) + .map_err(|e| format!("Invalid partition annotation '{}': {}", partition, e))?; + let title = annotations + .get(OCI_TITLE_ANNOTATION) + .map(|s| s.as_str()) + .unwrap_or("layer"); + println!( + "Extracting {} from layer {} for partition {}", + title, + &layer.digest[0..12], + partition + ); + + // Download this specific layer + let response = client.get_blob_stream(&layer.digest).await?; + let stream = response.bytes_stream(); + + let output_path = output_dir.join(format!("{}.img", sanitized_name)); + stream_blob_to_file(stream, output_path.clone()).await?; + partition_files.insert(sanitized_name, output_path); + } + } + } + + println!( + "Extracted {} partitions by annotations", + partition_files.len() + ); + Ok(partition_files) +} + +fn parse_default_partitions(manifest: &Manifest) -> Result>, String> { + let Manifest::Image(image) = manifest else { + return Ok(None); + }; + + let Some(ref annotations) = image.annotations else { + return Ok(None); + }; + + let Some(raw) = annotations.get(automotive_annotations::DEFAULT_PARTITIONS) else { + return Ok(None); + }; + + let mut partitions = HashSet::new(); + for entry in raw.split(',') { + let trimmed = entry.trim(); + if trimmed.is_empty() { + continue; + } + let sanitized = sanitize_partition_name(trimmed) + .map_err(|e| format!("Invalid default partition '{}': {}", trimmed, e))?; + partitions.insert(sanitized); + } + + if partitions.is_empty() { + return Err(format!( + "Default partition annotation '{}' is empty", + automotive_annotations::DEFAULT_PARTITIONS + )); + } + + Ok(Some(partitions)) +} + +fn looks_like_tar_layer(layer: &super::manifest::Descriptor) -> bool { + if layer.media_type.contains("tar") { + return true; + } + + let Some(ref annotations) = layer.annotations else { + return false; + }; + + let Some(title) = annotations.get(OCI_TITLE_ANNOTATION) else { + return false; + }; + + let title = title.to_ascii_lowercase(); + title.ends_with(".tar") + || title.ends_with(".tar.gz") + || title.ends_with(".tgz") + || title.ends_with(".tar.xz") + || title.ends_with(".tar.lz4") +} + +/// Extract files based on layer annotations, applying optional overrides. +/// +/// When overrides are provided, each entry maps a target partition to a layer +/// title (org.opencontainers.image.title). Auto-detected partitions remain +/// unless explicitly overridden. If the manifest provides default partitions, +/// only those partitions are auto-extracted. Returns Ok(None) when no +/// annotations exist and overrides were provided, allowing callers to fall back +/// to pattern-based extraction. +pub async fn extract_files_by_annotations_with_overrides_to_dir( + image: &str, + options: &OciOptions, + output_dir: &Path, + overrides: &[(String, String)], +) -> Result>, Box> { + let (client, manifest) = connect_and_resolve(image, options).await?; + + let default_partitions = parse_default_partitions(&manifest) + .map_err(|e| format!("Invalid default partitions annotation: {}", e))?; + if let Some(ref partitions) = default_partitions { + println!("Using default partitions from manifest: {:?}", partitions); + } + + // Get layers and build lookup tables + let layers = manifest.get_layers()?; + let mut partition_files = HashMap::new(); + let mut title_to_layer: HashMap = HashMap::new(); + let mut title_to_path = HashMap::new(); + let mut has_partition_annotations = false; + let mut available_partitions = HashSet::new(); + let overridden_partitions: HashSet = overrides + .iter() + .map(|(partition, _)| partition.clone()) + .collect(); + let allowed_partitions = default_partitions.as_ref(); + let is_single_tar_layer = layers.len() == 1 && looks_like_tar_layer(&layers[0]); + + for layer in layers { + if let Some(ref annotations) = layer.annotations { + if let Some(title) = annotations.get(OCI_TITLE_ANNOTATION) { + if let Some(existing) = title_to_layer.get(title) { + eprintln!( + "Warning: duplicate OCI title '{}': keeping layer {}, ignoring layer {}", + title, + &existing.digest[..12.min(existing.digest.len())], + &layer.digest[..12.min(layer.digest.len())], + ); + } else { + title_to_layer.insert(title.clone(), layer); + } + } + + if let Some(partition) = annotations.get(automotive_annotations::PARTITION_ANNOTATION) { + has_partition_annotations = true; + available_partitions.insert(partition.clone()); + if overridden_partitions.contains(partition) { + continue; + } + if let Some(allowed) = allowed_partitions { + if !allowed.contains(partition) { + continue; + } + } + ensure_supported_layer_compression(layer.compression(), &layer.media_type)?; + let sanitized_name = sanitize_partition_name(partition) + .map_err(|e| format!("Invalid partition annotation '{}': {}", partition, e))?; + let title = annotations + .get(OCI_TITLE_ANNOTATION) + .map(|s| s.as_str()) + .unwrap_or("layer"); + println!( + "Extracting {} from layer {} for partition {}", + title, + &layer.digest[0..12], + partition + ); + + // Download this specific layer + let response = client.get_blob_stream(&layer.digest).await?; + let stream = response.bytes_stream(); + + let output_path = output_dir.join(format!("{}.img", sanitized_name)); + stream_blob_to_file(stream, output_path.clone()).await?; + partition_files.insert(sanitized_name, output_path.clone()); + if let Some(title) = annotations.get(OCI_TITLE_ANNOTATION) { + title_to_path.insert(title.clone(), output_path); + } + } + } + } + + if !has_partition_annotations && title_to_layer.is_empty() { + if overrides.is_empty() { + return Err( + "No partitions found in OCI annotations. Expected layers with 'automotive.sdv.cloud.redhat.com/partition' annotations" + .into(), + ); + } + return Ok(None); + } + + if !has_partition_annotations && !overrides.is_empty() && is_single_tar_layer { + let single_title = layers[0] + .annotations + .as_ref() + .and_then(|annotations| annotations.get(OCI_TITLE_ANNOTATION)); + let override_titles: HashSet<&str> = overrides + .iter() + .map(|(_, filename)| filename.as_str()) + .collect(); + + if single_title + .map(|title| !override_titles.contains(title.as_str())) + .unwrap_or(true) + { + return Ok(None); + } + } + + // Apply overrides by layer title + for (partition, filename) in overrides { + let sanitized_partition = sanitize_partition_name(partition) + .map_err(|e| format!("Invalid partition mapping '{}': {}", partition, e))?; + + if let Some(path) = title_to_path.get(filename) { + partition_files.insert(sanitized_partition, path.clone()); + continue; + } + + let layer = match title_to_layer.get(filename) { + Some(layer) => *layer, + None => { + let available: Vec<&String> = title_to_layer.keys().collect(); + return Err(format!( + "Override file '{}' not found in OCI layer titles. Available files: {:?}", + filename, available + ) + .into()); + } + }; + + ensure_supported_layer_compression(layer.compression(), &layer.media_type)?; + let output_path = output_dir.join(format!("{}.img", sanitized_partition)); + + let response = client.get_blob_stream(&layer.digest).await?; + let stream = response.bytes_stream(); + stream_blob_to_file(stream, output_path.clone()).await?; + + title_to_path.insert(filename.clone(), output_path.clone()); + partition_files.insert(sanitized_partition, output_path); + } + + if partition_files.is_empty() { + if let Some(allowed) = allowed_partitions { + return Err(format!( + "No partitions matched default partitions {:?}. Available partitions: {:?}", + allowed, available_partitions + ) + .into()); + } + return Err( + "No partitions found in OCI annotations. Expected layers with 'automotive.sdv.cloud.redhat.com/partition' annotations" + .into(), + ); + } + + Ok(Some(partition_files)) +} + +fn ensure_supported_layer_compression( + compression: LayerCompression, + media_type: &str, +) -> Result<(), Box> { + match compression { + LayerCompression::None | LayerCompression::Gzip => Ok(()), + other => Err(format!( + "Unsupported OCI layer compression {:?} (media type: {}). Supported: uncompressed, gzip", + other, media_type + ) + .into()), + } +} + +fn sanitize_partition_name(name: &str) -> Result { + if name.is_empty() { + return Err("partition name is empty".to_string()); + } + if std::path::Path::new(name).is_absolute() { + return Err("partition name is an absolute path".to_string()); + } + if name.contains("..") { + return Err("partition name contains '..'".to_string()); + } + if name.contains('/') || name.contains('\\') { + return Err("partition name contains path separators".to_string()); + } + if !name + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-') + { + return Err("partition name contains invalid characters".to_string()); + } + Ok(name.to_string()) +} + /// Execute a sequence of write commands on the block writer async fn execute_write_commands( commands: Vec, @@ -1019,10 +1618,7 @@ pub async fn flash_from_oci( println!("Layer media type: {}", layer.media_type); println!("Layer compression: {:?}", compression); - // Validate compression - we only support gzip for now - if compression == LayerCompression::Zstd { - return Err("Zstd-compressed layers are not yet supported".into()); - } + ensure_supported_layer_compression(compression, &layer.media_type)?; // Start blob download println!("\nStarting download..."); @@ -1797,3 +2393,52 @@ fn extract_tar_archive_from_stream( extract_tar_stream_impl(decompressed_reader, tar_tx, file_pattern, debug) } + +#[cfg(test)] +mod tests { + use super::sanitize_partition_name; + + #[test] + fn sanitize_partition_name_accepts_valid() { + let valid = ["boot_a", "system-a", "slot_1", "boot_a.1", "A1-b_c.2"]; + for name in valid { + assert_eq!(sanitize_partition_name(name).unwrap(), name); + } + } + + #[test] + fn sanitize_partition_name_rejects_empty() { + let err = sanitize_partition_name("").unwrap_err(); + assert!(err.contains("empty")); + } + + #[test] + fn sanitize_partition_name_rejects_path_separators() { + for name in ["boot/a", "boot\\a"] { + let err = sanitize_partition_name(name).unwrap_err(); + assert!(err.contains("path separators")); + } + } + + #[test] + fn sanitize_partition_name_rejects_parent_traversal() { + for name in ["..", "../boot", "boot/..", "boot.."] { + let err = sanitize_partition_name(name).unwrap_err(); + assert!(err.contains("..")); + } + } + + #[test] + fn sanitize_partition_name_rejects_absolute_path() { + let err = sanitize_partition_name("/tmp/boot").unwrap_err(); + assert!(err.contains("absolute path")); + } + + #[test] + fn sanitize_partition_name_rejects_invalid_chars() { + for name in ["boot:1", "boot a", "boot@a"] { + let err = sanitize_partition_name(name).unwrap_err(); + assert!(err.contains("invalid characters")); + } + } +} diff --git a/src/fls/oci/mod.rs b/src/fls/oci/mod.rs index 6f3090e..41116b7 100644 --- a/src/fls/oci/mod.rs +++ b/src/fls/oci/mod.rs @@ -9,5 +9,8 @@ mod reference; mod registry; // Public re-exports -pub use from_oci::flash_from_oci; +pub use from_oci::{ + extract_files_by_annotations_to_dir, extract_files_by_annotations_with_overrides_to_dir, + extract_files_from_oci_image_to_dir, flash_from_oci, +}; pub use reference::ImageReference; diff --git a/src/fls/options.rs b/src/fls/options.rs index 257bc1a..e48416e 100644 --- a/src/fls/options.rs +++ b/src/fls/options.rs @@ -61,8 +61,32 @@ pub struct OciOptions { pub file_pattern: Option, } -/// Options for HTTP client setup (subset of FlashOptions) +/// Options for fastboot flash operations #[derive(Debug, Clone)] +pub struct FastbootOptions { + pub http: HttpClientOptions, + pub device_serial: Option, + pub partition_mappings: Vec<(String, String)>, // (partition_name, file_pattern) - fallback for manual mapping + pub timeout_secs: u32, + pub username: Option, + pub password: Option, +} + +impl Default for FastbootOptions { + fn default() -> Self { + Self { + http: HttpClientOptions::default(), + device_serial: None, + partition_mappings: Vec::new(), + timeout_secs: 1200, + username: None, + password: None, + } + } +} + +/// Options for HTTP client setup (subset of FlashOptions) +#[derive(Debug, Clone, Default)] pub struct HttpClientOptions { pub insecure_tls: bool, pub cacert: Option, @@ -90,3 +114,9 @@ impl From<&OciOptions> for HttpClientOptions { Self::from(&opts.common) } } + +impl From<&FastbootOptions> for HttpClientOptions { + fn from(opts: &FastbootOptions) -> Self { + opts.http.clone() + } +} diff --git a/src/main.rs b/src/main.rs index 62d462e..119cf37 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,6 +4,22 @@ use std::path::PathBuf; // Use the library module use fls::fls; +/// Parse target argument in the format "partition:filename" +fn parse_target_mapping(s: &str) -> Result<(String, String), String> { + match s.split_once(':') { + Some((partition, filename)) => { + if partition.is_empty() { + Err("Partition name cannot be empty".to_string()) + } else if filename.is_empty() { + Err("Filename cannot be empty".to_string()) + } else { + Ok((partition.to_string(), filename.to_string())) + } + } + None => Err("Format should be 'partition:filename' (e.g., 'boot_a:aboot.img')".to_string()), + } +} + #[derive(Parser)] #[command(name = "fls")] #[command(about = "A small Rust utility for flashing devices")] @@ -57,7 +73,7 @@ enum Commands { #[arg(long)] show_memory: bool, /// Registry username for OCI authentication - #[arg(short = 'u', long)] + #[arg(short = 'u', long, env = "FLS_REGISTRY_USERNAME")] username: Option, /// Registry password for OCI authentication (or use FLS_REGISTRY_PASSWORD env) #[arg(short = 'p', long, env = "FLS_REGISTRY_PASSWORD")] @@ -66,6 +82,35 @@ enum Commands { #[arg(long)] file_pattern: Option, }, + /// Flash an OCI image to fastboot partitions via USB + Fastboot { + /// OCI image reference to download and flash (must be prefixed with "oci://") + image_ref: String, + /// Device serial number (optional, will use first device if not specified) + #[arg(short = 's', long)] + serial: Option, + /// Target partition and file override (e.g., "boot_a:boot_a.simg"), can be used multiple times + #[arg(short = 't', long = "target", value_parser = parse_target_mapping)] + targets: Vec<(String, String)>, + /// Fastboot operation timeout in seconds (default: 1200) + #[arg(long, default_value = "1200")] + timeout: u32, + /// Path to CA certificate PEM file for TLS validation + #[arg(long)] + cacert: Option, + /// Ignore SSL certificate verification + #[arg(short = 'k', long = "insecure-tls")] + insecure_tls: bool, + /// Enable debug output + #[arg(long)] + debug: bool, + /// Registry username for OCI authentication + #[arg(short = 'u', long, env = "FLS_REGISTRY_USERNAME")] + username: Option, + /// Registry password for OCI authentication (or use FLS_REGISTRY_PASSWORD env) + #[arg(short = 'p', long, env = "FLS_REGISTRY_PASSWORD")] + password: Option, + }, } #[tokio::main] @@ -109,9 +154,13 @@ async fn main() { println!(" Image: {}", image_ref); println!(" Device: {}", device); match (&username, &password) { + (Some(_), None) | (None, Some(_)) => { + eprintln!( + "Error: OCI authentication requires both --username and --password" + ); + std::process::exit(1); + } (Some(_), Some(_)) => println!(" Auth: Using provided credentials"), - (Some(_), None) => println!(" Auth: Username provided but password missing"), - (None, Some(_)) => println!(" Auth: Password provided but username missing"), (None, None) => println!(" Auth: Anonymous"), } if let Some(ref pattern) = file_pattern { @@ -233,5 +282,80 @@ async fn main() { } } } + Commands::Fastboot { + image_ref, + serial, + targets, + timeout, + cacert, + insecure_tls, + debug, + username, + password, + } => { + let image_ref_input = image_ref; + let image_ref = match image_ref_input.strip_prefix("oci://") { + Some(reference) => reference, + None => { + eprintln!( + "Error: fastboot expects an OCI image reference prefixed with 'oci://'" + ); + eprintln!(" Example: fls fastboot oci://quay.io/org/image:latest"); + std::process::exit(1); + } + }; + + println!("Fastboot flash command:"); + println!(" Image: {}", image_ref_input); + if let Some(ref serial) = serial { + println!(" Device serial: {}", serial); + } + if !targets.is_empty() { + println!(" Target partitions:"); + for (partition, filename) in &targets { + println!(" {} → {}", partition, filename); + } + } + if let Some(ref cert_path) = cacert { + println!(" CA Certificate: {}", cert_path.display()); + } + println!(" Ignore certificates: {}", insecure_tls); + println!(" Timeout: {} seconds", timeout); + println!(" Debug: {}", debug); + match (&username, &password) { + (Some(_), None) | (None, Some(_)) => { + eprintln!("Error: OCI authentication requires both --username and --password"); + std::process::exit(1); + } + (Some(_), Some(_)) => println!(" Auth: Using provided credentials"), + (None, None) => println!(" Auth: Anonymous"), + } + println!(); + + let options = fls::FastbootOptions { + http: fls::HttpClientOptions { + insecure_tls, + cacert, + debug, + }, + device_serial: serial, + partition_mappings: targets, + timeout_secs: timeout, + username, + password, + }; + + match fls::flash_from_fastboot(image_ref, options).await { + Ok(_) => { + println!("Result: FLASH_COMPLETED"); + std::process::exit(0); + } + Err(e) => { + eprintln!("Error: {}", e); + println!("Result: FLASH_FAILED"); + std::process::exit(1); + } + } + } } } diff --git a/tests/common/mod.rs b/tests/common/mod.rs index f316008..3790a71 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -18,6 +18,7 @@ pub fn create_test_data(size: usize) -> Vec { } /// Compress data using xz compression +#[allow(dead_code)] pub fn compress_xz(data: &[u8]) -> Vec { let mut encoder = XzEncoder::new(Vec::new(), 6); encoder diff --git a/tests/oci_extract.rs b/tests/oci_extract.rs new file mode 100644 index 0000000..cfb7f56 --- /dev/null +++ b/tests/oci_extract.rs @@ -0,0 +1,380 @@ +mod common; + +use fls::fls::oci::extract_files_from_oci_image_to_dir; +use fls::{FlashOptions, OciOptions}; +use http_body_util::Full; +use hyper::body::Bytes; +use hyper::server::conn::http1; +use hyper::service::service_fn; +use hyper::{Request as HyperRequest, Response, StatusCode}; +use hyper_util::rt::TokioIo; +use std::collections::{HashMap, HashSet}; +use std::fs; +use std::net::SocketAddr; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use tar::Header; +use tempfile::tempdir; +use tokio::net::TcpListener; +use tokio_rustls::TlsAcceptor; + +/// Spawn an HTTPS server that serves OCI registry endpoints: /v2/, manifest at manifest_path, blob at blob_path. +/// Returns (socket_addr, server_join_handle). Caller should abort the handle when done. +async fn spawn_oci_server( + manifest_path: String, + blob_path: String, + manifest_body: Vec, + blob_body: Vec, +) -> (SocketAddr, tokio::task::JoinHandle<()>) { + let cert_dir = PathBuf::from("tests/test_certs"); + let server_cert_pem = + fs::read_to_string(cert_dir.join("server-cert.pem")).expect("read server cert"); + let server_key_pem = + fs::read_to_string(cert_dir.join("server-key.pem")).expect("read server key"); + let server_cert = rustls_pemfile::certs(&mut server_cert_pem.as_bytes()) + .collect::, _>>() + .expect("parse server cert"); + let server_key = rustls_pemfile::private_key(&mut server_key_pem.as_bytes()) + .expect("parse server key") + .expect("missing server key"); + let mut server_config = rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(server_cert, server_key) + .expect("build server config"); + server_config.alpn_protocols = vec![b"http/1.1".to_vec()]; + let tls_acceptor = TlsAcceptor::from(Arc::new(server_config)); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let local_addr = listener.local_addr().unwrap(); + + let manifest_body = Arc::new(manifest_body); + let blob_body = Arc::new(blob_body); + let manifest_path = Arc::new(manifest_path); + let blob_path = Arc::new(blob_path); + + let handle = tokio::spawn(async move { + loop { + let (stream, _) = listener.accept().await.unwrap(); + let tls_acceptor = tls_acceptor.clone(); + let manifest_body = manifest_body.clone(); + let blob_body = blob_body.clone(); + let manifest_path = manifest_path.clone(); + let blob_path = blob_path.clone(); + + tokio::spawn(async move { + let tls_stream = tls_acceptor.accept(stream).await.unwrap(); + let io = TokioIo::new(tls_stream); + let service = service_fn(move |req: HyperRequest| { + let manifest_body = manifest_body.clone(); + let blob_body = blob_body.clone(); + let manifest_path = manifest_path.clone(); + let blob_path = blob_path.clone(); + async move { + let path = req.uri().path(); + let response = if path == "/v2/" { + Response::builder() + .status(StatusCode::OK) + .body(Full::new(Bytes::new())) + .unwrap() + } else if path == manifest_path.as_str() { + Response::builder() + .status(StatusCode::OK) + .header( + "Content-Type", + "application/vnd.oci.image.manifest.v1+json", + ) + .body(Full::new(Bytes::copy_from_slice(&manifest_body))) + .unwrap() + } else if path == blob_path.as_str() { + Response::builder() + .status(StatusCode::OK) + .body(Full::new(Bytes::copy_from_slice(&blob_body))) + .unwrap() + } else { + Response::builder() + .status(StatusCode::NOT_FOUND) + .body(Full::new(Bytes::new())) + .unwrap() + }; + Ok::<_, std::convert::Infallible>(response) + } + }); + + let _ = http1::Builder::new().serve_connection(io, service).await; + }); + } + }); + + (local_addr, handle) +} + +fn build_tar(entries: Vec<(&str, Vec)>) -> Vec { + let mut tar_buf = Vec::new(); + { + let mut builder = tar::Builder::new(&mut tar_buf); + for (name, data) in entries { + let mut header = Header::new_gnu(); + header.set_mode(0o644); + header.set_size(data.len() as u64); + header.set_cksum(); + builder + .append_data(&mut header, name, data.as_slice()) + .expect("append tar entry"); + } + builder.finish().expect("finish tar"); + } + tar_buf +} + +const REPO: &str = "test/repo"; +const TAG: &str = "latest"; + +fn default_options(cert_dir: &Path) -> OciOptions { + OciOptions { + common: FlashOptions { + insecure_tls: false, + cacert: Some(cert_dir.join("ca-cert.pem")), + ..Default::default() + }, + username: None, + password: None, + file_pattern: None, + } +} + +#[tokio::test] +async fn test_extract_files_from_oci_image_to_dir_tar_layer() { + let _ = rustls::crypto::ring::default_provider().install_default(); + + let boot_data = common::create_test_data(1024); + let system_data = common::create_test_data(2048); + let tar_bytes = build_tar(vec![ + ("boot_a.simg", boot_data.clone()), + ("system_a.simg", system_data.clone()), + ]); + let blob_bytes = common::compress_gz(&tar_bytes); + let blob_digest = "sha256:layer"; + + let manifest_json = serde_json::json!({ + "schemaVersion": 2, + "mediaType": "application/vnd.oci.image.manifest.v1+json", + "config": { + "mediaType": "application/vnd.oci.image.config.v1+json", + "digest": "sha256:config", + "size": 2 + }, + "layers": [ + { + "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip", + "digest": blob_digest, + "size": blob_bytes.len() + } + ] + }) + .to_string(); + + let (local_addr, server_handle) = spawn_oci_server( + format!("/v2/{}/manifests/{}", REPO, TAG), + format!("/v2/{}/blobs/{}", REPO, blob_digest), + manifest_json.into_bytes(), + blob_bytes, + ) + .await; + + let cert_dir = PathBuf::from("tests/test_certs"); + let host = format!("127.0.0.1.nip.io:{}", local_addr.port()); + let image_ref = format!("{}/{}:{}", host, REPO, TAG); + let options = default_options(&cert_dir); + + let out_dir = tempdir().expect("create temp dir"); + let target_files: HashSet = ["boot_a.simg".to_string(), "system_a.simg".to_string()] + .into_iter() + .collect(); + + let result: Result, Box> = + extract_files_from_oci_image_to_dir(&image_ref, &target_files, &options, out_dir.path()) + .await; + + server_handle.abort(); + + assert!(result.is_ok(), "Extraction failed: {:?}", result.err()); + + let boot_path = out_dir.path().join("boot_a.simg"); + let system_path = out_dir.path().join("system_a.simg"); + let boot_written = fs::read(boot_path).expect("read boot output"); + let system_written = fs::read(system_path).expect("read system output"); + + assert_eq!(boot_written, boot_data); + assert_eq!(system_written, system_data); +} + +/// Direct-file (non-tar) layer: media type without "tar", single target file, blob is raw gzip content. +#[tokio::test] +async fn test_extract_direct_file_layer() { + let _ = rustls::crypto::ring::default_provider().install_default(); + + let file_data = common::create_test_data(512); + let blob_bytes = common::compress_gz(&file_data); + let blob_digest = "sha256:direct"; + + // Media type without "tar" so the code uses stream_blob_to_file instead of tar extraction. + let manifest_json = serde_json::json!({ + "schemaVersion": 2, + "mediaType": "application/vnd.oci.image.manifest.v1+json", + "config": { + "mediaType": "application/vnd.oci.image.config.v1+json", + "digest": "sha256:config", + "size": 2 + }, + "layers": [ + { + "mediaType": "application/vnd.oci.image.layer.v1+gzip", + "digest": blob_digest, + "size": blob_bytes.len() + } + ] + }) + .to_string(); + + let (local_addr, server_handle) = spawn_oci_server( + format!("/v2/{}/manifests/{}", REPO, TAG), + format!("/v2/{}/blobs/{}", REPO, blob_digest), + manifest_json.into_bytes(), + blob_bytes, + ) + .await; + + let cert_dir = PathBuf::from("tests/test_certs"); + let host = format!("127.0.0.1.nip.io:{}", local_addr.port()); + let image_ref = format!("{}/{}:{}", host, REPO, TAG); + let options = default_options(&cert_dir); + + let out_dir = tempdir().expect("create temp dir"); + let target_files: HashSet = ["single.img".to_string()].into_iter().collect(); + + let result = + extract_files_from_oci_image_to_dir(&image_ref, &target_files, &options, out_dir.path()) + .await; + + server_handle.abort(); + + assert!(result.is_ok(), "Extraction failed: {:?}", result.err()); + let map = result.unwrap(); + assert_eq!(map.len(), 1); + assert!(map.contains_key("single.img")); + + let path = out_dir.path().join("single.img"); + let written = fs::read(&path).expect("read output"); + assert_eq!(written, file_data); +} + +/// Tar layer but request only one of two files in the layer. +#[tokio::test] +async fn test_extract_tar_single_requested_file() { + let _ = rustls::crypto::ring::default_provider().install_default(); + + let boot_data = common::create_test_data(256); + let system_data = common::create_test_data(512); + let tar_bytes = build_tar(vec![ + ("boot_a.simg", boot_data.clone()), + ("system_a.simg", system_data.clone()), + ]); + let blob_bytes = common::compress_gz(&tar_bytes); + let blob_digest = "sha256:layer"; + + let manifest_json = serde_json::json!({ + "schemaVersion": 2, + "mediaType": "application/vnd.oci.image.manifest.v1+json", + "config": { "mediaType": "application/vnd.oci.image.config.v1+json", "digest": "sha256:config", "size": 2 }, + "layers": [{ "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip", "digest": blob_digest, "size": blob_bytes.len() }] + }) + .to_string(); + + let (local_addr, server_handle) = spawn_oci_server( + format!("/v2/{}/manifests/{}", REPO, TAG), + format!("/v2/{}/blobs/{}", REPO, blob_digest), + manifest_json.into_bytes(), + blob_bytes, + ) + .await; + + let cert_dir = PathBuf::from("tests/test_certs"); + let options = default_options(&cert_dir); + let out_dir = tempdir().expect("create temp dir"); + let host = format!("127.0.0.1.nip.io:{}", local_addr.port()); + let image_ref = format!("{}/{}:{}", host, REPO, TAG); + + // Request only one file; layer contains two. + let target_files: HashSet = ["system_a.simg".to_string()].into_iter().collect(); + + let result = + extract_files_from_oci_image_to_dir(&image_ref, &target_files, &options, out_dir.path()) + .await; + + server_handle.abort(); + + assert!(result.is_ok(), "Extraction failed: {:?}", result.err()); + let map = result.unwrap(); + assert_eq!(map.len(), 1); + assert!(map.contains_key("system_a.simg")); + let written = fs::read(out_dir.path().join("system_a.simg")).expect("read"); + assert_eq!(written, system_data); + // boot_a.simg was not requested so should not be present + assert!(!out_dir.path().join("boot_a.simg").exists()); +} + +/// Tar layer contains one file; request two. Only the present file is returned (partial success). +#[tokio::test] +async fn test_extract_tar_missing_file() { + let _ = rustls::crypto::ring::default_provider().install_default(); + + let boot_data = common::create_test_data(256); + // Tar has only boot_a.simg + let tar_bytes = build_tar(vec![("boot_a.simg", boot_data.clone())]); + let blob_bytes = common::compress_gz(&tar_bytes); + let blob_digest = "sha256:layer"; + + let manifest_json = serde_json::json!({ + "schemaVersion": 2, + "mediaType": "application/vnd.oci.image.manifest.v1+json", + "config": { "mediaType": "application/vnd.oci.image.config.v1+json", "digest": "sha256:config", "size": 2 }, + "layers": [{ "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip", "digest": blob_digest, "size": blob_bytes.len() }] + }) + .to_string(); + + let (local_addr, server_handle) = spawn_oci_server( + format!("/v2/{}/manifests/{}", REPO, TAG), + format!("/v2/{}/blobs/{}", REPO, blob_digest), + manifest_json.into_bytes(), + blob_bytes, + ) + .await; + + let cert_dir = PathBuf::from("tests/test_certs"); + let options = default_options(&cert_dir); + let out_dir = tempdir().expect("create temp dir"); + let host = format!("127.0.0.1.nip.io:{}", local_addr.port()); + let image_ref = format!("{}/{}:{}", host, REPO, TAG); + + // Request both; only boot_a.simg exists in the layer. + let target_files: HashSet = ["boot_a.simg".to_string(), "system_a.simg".to_string()] + .into_iter() + .collect(); + + let result = + extract_files_from_oci_image_to_dir(&image_ref, &target_files, &options, out_dir.path()) + .await; + + server_handle.abort(); + + // Current behavior: returns Ok with only the files that were found. + assert!(result.is_ok(), "Extraction failed: {:?}", result.err()); + let map = result.unwrap(); + assert_eq!(map.len(), 1, "only boot_a.simg should be found"); + assert!(map.contains_key("boot_a.simg")); + assert!(!map.contains_key("system_a.simg")); + + let written = fs::read(out_dir.path().join("boot_a.simg")).expect("read"); + assert_eq!(written, boot_data); + assert!(!out_dir.path().join("system_a.simg").exists()); +}