From 0f8c1c7f152daf781e8e1f4a9af7f68ecb82af15 Mon Sep 17 00:00:00 2001 From: WoodenMaiden Date: Wed, 25 Oct 2023 10:34:32 +0200 Subject: [PATCH] fix(test): create trait to test file compression Signed-off-by: WoodenMaiden --- initramfs/Cargo.toml | 3 - initramfs/src/image.rs | 161 ++++++++++++++++++++++++++++++++--------- initramfs/src/main.rs | 4 +- 3 files changed, 131 insertions(+), 37 deletions(-) diff --git a/initramfs/Cargo.toml b/initramfs/Cargo.toml index 3b61015..e7008c2 100644 --- a/initramfs/Cargo.toml +++ b/initramfs/Cargo.toml @@ -21,6 +21,3 @@ bytes = "1.4.0" libflate = "0.1" cpio = "0.2.2" openssl = { version = "0.10", features = ["vendored"] } - -[dev-dependencies] -infer = "0.15.0" \ No newline at end of file diff --git a/initramfs/src/image.rs b/initramfs/src/image.rs index b89731d..9958185 100644 --- a/initramfs/src/image.rs +++ b/initramfs/src/image.rs @@ -1,7 +1,7 @@ use std::{ collections::HashMap, fs::File, - io::{Cursor, Read}, + io::{Cursor, Read, Write}, }; use anyhow::{anyhow, Result}; @@ -12,6 +12,26 @@ use log::debug; use serde::Deserialize; use tar::Archive; +/// Trait to abstract File reading +pub trait FileHandler: Write + Read { + fn create(path: &str) -> Result + where + Self: std::marker::Sized; + fn open(path: &str) -> Result + where + Self: std::marker::Sized; +} + +impl FileHandler for File { + fn create(path: &str) -> Result { + File::create(path).map_err(|e| anyhow!(e).context("Failed to create file")) + } + + fn open(path: &str) -> Result { + File::open(path).map_err(|e| anyhow!(e).context("Failed to open file")) + } +} + #[derive(Deserialize, Debug, Clone)] #[allow(non_snake_case)] pub struct LayerMetadata { @@ -61,7 +81,7 @@ impl Image { &self.tag } - pub fn export_to_initramfs( + pub fn export_to_initramfs( &self, init_path: &str, agent_path: &str, @@ -70,7 +90,7 @@ impl Image { // Write the cpio to disk let file_name = format!("initramfs-{}-{}.img", self.name.replace('/', "-"), self.tag); let archive = Encoder::new( - File::create(file_name).map_err(|e| anyhow!(e).context("Failed to create file"))?, + ::create(&file_name).map_err(|e| anyhow!(e).context("Failed to create file"))?, ) .map_err(|e| anyhow!(e).context("Failed to create gzip encoder"))?; @@ -127,10 +147,10 @@ impl Image { } let mut init_file = - File::open(init_path).map_err(|e| anyhow!(e).context("Failed to open init file"))?; + ::open(init_path).map_err(|e| anyhow!(e).context("Failed to open init file"))?; let mut agent_file = - File::open(agent_path).map_err(|e| anyhow!(e).context("Failed to open init file"))?; - let mut agent_config_file = File::open(agent_config_path) + ::open(agent_path).map_err(|e| anyhow!(e).context("Failed to open init file"))?; + let mut agent_config_file = ::open(agent_config_path) .map_err(|e| anyhow!(e).context("Failed to open agent config file"))?; let mut init_content = Vec::new(); @@ -178,23 +198,109 @@ impl Image { #[cfg(test)] mod test { + use super::{FileHandler, Image}; + use anyhow::{anyhow, Ok}; use std::env; - use std::io; - use std::io::Read; - use std::{error::Error, result::Result}; + use std::io::{Read, Write}; + use std::vec::Drain; + + const VALID_IMAGE_NAME: &str = "my_awesome_img:14md35"; + const VALID_IMAGE_NAME_FROM_HUB: &str = "bitnami/mongodb:latest"; + + #[derive(Debug, Clone)] + struct MockFileHandler { + contents: Vec, + path: String, + } + + // holds all the files created by the tests + static mut FILES: Vec = vec![]; + + impl Write for MockFileHandler { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.contents.extend_from_slice(buf); - use super::Image; - use std::fs::{File, OpenOptions}; + unsafe { + FILES.iter_mut().for_each(|f| { + if f.path == self.path { + f.contents.extend_from_slice(buf); + } + }); + } + + std::result::Result::Ok(buf.len()) + } - fn touch(path: &str) -> Result<&str, Box> { - match OpenOptions::new().create(true).write(true).open(path) { - Ok(_) => Ok(path), - Err(e) => Err(Box::new(e)), + fn flush(&mut self) -> std::io::Result<()> { + std::result::Result::Ok(()) } } - const VALID_IMAGE_NAME: &str = "my_awesome_img:14md35"; - const VALID_IMAGE_NAME_FROM_HUB: &str = "bitnami/mongodb:latest"; + impl Read for MockFileHandler { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + let length = if buf.len() > self.contents.len() { + self.contents.len() + } else { + buf.len() + }; + + let drained = &self.contents.drain(..length); + let read = <&Drain<'_, u8>>::clone(&drained).as_slice(); + + for i in 0..length { + buf[i] = if i < length { read[i] } else { 0 } + } + + std::result::Result::Ok(length) + } + + fn read_to_end(&mut self, buf: &mut Vec) -> std::io::Result { + loop { + let nb = self.read(buf)?; + if nb == 0 { + break; + } + buf.extend_from_slice(&self.contents); + } + std::result::Result::Ok(buf.len()) + } + } + + impl FileHandler for MockFileHandler { + fn create(path: &str) -> anyhow::Result + where + Self: std::marker::Sized, + { + unsafe { + FILES.push(Self { + contents: vec![], + path: path.to_string(), + }); + } + + Ok(Self { + contents: vec![], + path: path.to_string(), + }) + } + + fn open(path: &str) -> anyhow::Result + where + Self: std::marker::Sized, + { + unsafe { + let mut iter = FILES.iter_mut(); + + loop { + let file = iter.next().ok_or(anyhow!("No file found"))?; + + if file.path == path { + return Ok(file.clone()); + } + } + } + } + } #[test] pub fn valid_image_name() { @@ -223,11 +329,9 @@ mod test { let agent_filename = format!("{}/agent", env::temp_dir().display()); let agent_config_filename = format!("{}/agent_config", env::temp_dir().display()); - let files = [ - touch(&image_filename).unwrap(), - touch(&agent_filename).unwrap(), - touch(&agent_config_filename).unwrap(), - ]; + MockFileHandler::create(&image_filename).unwrap(); + MockFileHandler::create(&agent_filename).unwrap(); + MockFileHandler::create(&agent_config_filename).unwrap(); let image = image.unwrap(); let generated_filename = format!( @@ -238,17 +342,14 @@ mod test { // checks assert!(image - .export_to_initramfs( + .export_to_initramfs::( image_filename.as_str(), agent_filename.as_str(), agent_config_filename.as_str() ) .is_ok()); - let generated_file = File::open(&generated_filename); - - assert!(std::path::Path::new(&generated_filename).exists()); - assert!(generated_file.is_ok()); + let generated_file = MockFileHandler::open(&generated_filename); let magic_number = generated_file .unwrap() @@ -257,12 +358,6 @@ mod test { .map(|b| b.unwrap()) .collect::>(); - let kind = infer::get_from_path(&generated_filename).unwrap().unwrap(); - assert_eq!(kind.matcher_type(), infer::MatcherType::Archive); assert_eq!(magic_number, &[0x1F, 0x8b]); - - // clean - files.iter().for_each(|f| std::fs::remove_file(f).unwrap()); - std::fs::remove_file(&generated_filename).unwrap(); } } diff --git a/initramfs/src/main.rs b/initramfs/src/main.rs index a516e72..ff56603 100644 --- a/initramfs/src/main.rs +++ b/initramfs/src/main.rs @@ -1,3 +1,5 @@ +use std::fs::File; + use anyhow::anyhow; use clap::Parser; use env_logger::Env; @@ -58,7 +60,7 @@ async fn main() -> Result<(), Box> { info!("Writing to disk ..."); image - .export_to_initramfs(&args.init, &args.agent, &args.agent_config) + .export_to_initramfs::(&args.init, &args.agent, &args.agent_config) .map_err(|e| anyhow!(e).context("Failed to write filesystem to disk"))?; info!("Writing done!");