From 36b3af044cb863e5fa693280fd31758533b59330 Mon Sep 17 00:00:00 2001 From: Quang Date: Mon, 31 Mar 2025 00:26:49 -0700 Subject: [PATCH] mcookie: add support for human-readable sizes with -m option (#268) * Create size.rs * parse max size from human-readable strings * tests for human-readable strings * fmt * fix failed tests and edge cases * fix fmt * cpr & license * use usimpleerror * fmt --- src/uu/mcookie/src/mcookie.rs | 23 ++++++--- src/uu/mcookie/src/size.rs | 88 +++++++++++++++++++++++++++++++++++ tests/by-util/test_mcookie.rs | 37 ++++++++++++++- 3 files changed, 141 insertions(+), 7 deletions(-) create mode 100644 src/uu/mcookie/src/size.rs diff --git a/src/uu/mcookie/src/mcookie.rs b/src/uu/mcookie/src/mcookie.rs index ab28250..abbb857 100644 --- a/src/uu/mcookie/src/mcookie.rs +++ b/src/uu/mcookie/src/mcookie.rs @@ -8,7 +8,12 @@ use std::{fs::File, io::Read}; use clap::{crate_version, Arg, ArgAction, Command}; use md5::{Digest, Md5}; use rand::RngCore; -use uucore::{error::UResult, format_usage, help_about, help_usage}; +use uucore::{ + error::{UResult, USimpleError}, + format_usage, help_about, help_usage, +}; +mod size; +use size::Size; mod options { pub const FILE: &str = "file"; @@ -31,10 +36,16 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> { .map(|v| v.as_str()) .collect(); - // TODO: Parse max size from human-readable strings (KiB, MiB, GiB etc.) - let max_size = matches - .get_one::(options::MAX_SIZE) - .map(|v| v.parse::().expect("Failed to parse max-size value")); + let max_size = if let Some(size_str) = matches.get_one::(options::MAX_SIZE) { + match Size::parse(size_str) { + Ok(size) => Some(size.size_bytes()), + Err(_) => { + return Err(USimpleError::new(1, "Failed to parse max-size value")); + } + } + } else { + None + }; let mut hasher = Md5::new(); @@ -99,7 +110,7 @@ pub fn uu_app() -> Command { .long("max-size") .value_name("num") .action(ArgAction::Set) - .help("limit how much is read from seed files"), + .help("limit how much is read from seed files (supports B suffix or binary units: KiB, MiB, GiB, TiB)"), ) .arg( Arg::new(options::VERBOSE) diff --git a/src/uu/mcookie/src/size.rs b/src/uu/mcookie/src/size.rs new file mode 100644 index 0000000..b870978 --- /dev/null +++ b/src/uu/mcookie/src/size.rs @@ -0,0 +1,88 @@ +// This file is part of the uutils util-linux package. +// +// For the full copyright and license information, please view the LICENSE +// file that was distributed with this source code. + +use std::error::Error; +use std::fmt; + +#[derive(Debug)] +pub struct ParseSizeError(String); + +impl fmt::Display for ParseSizeError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "Invalid size format: {}", self.0) + } +} + +impl Error for ParseSizeError {} + +pub struct Size(u64); + +impl Size { + pub fn parse(s: &str) -> Result { + let s = s.trim(); + + // Handle bytes with "B" suffix + if s.ends_with('B') && !s.ends_with("iB") { + if let Some(nums) = s.strip_suffix('B') { + return nums + .trim() + .parse::() + .map(Self) + .map_err(|_| ParseSizeError(s.to_string())); + } + } + + // Handle binary units (KiB, MiB, GiB, TiB) + for (suffix, exponent) in [("KiB", 1), ("MiB", 2), ("GiB", 3), ("TiB", 4)] { + if let Some(nums) = s.strip_suffix(suffix) { + return nums + .trim() + .parse::() + .map(|n| Self(n * 1024_u64.pow(exponent))) + .map_err(|_| ParseSizeError(s.to_string())); + } + } + + // If no suffix, treat as bytes + s.parse::() + .map(Self) + .map_err(|_| ParseSizeError(s.to_string())) + } + + pub fn size_bytes(&self) -> u64 { + self.0 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_numeric() { + assert_eq!(Size::parse("1234").unwrap().size_bytes(), 1234); + } + + #[test] + fn test_parse_with_suffix() { + assert_eq!(Size::parse("1024B").unwrap().size_bytes(), 1024); + assert_eq!(Size::parse("1KiB").unwrap().size_bytes(), 1024); + assert_eq!(Size::parse("1MiB").unwrap().size_bytes(), 1024 * 1024); + assert_eq!( + Size::parse("1GiB").unwrap().size_bytes(), + 1024 * 1024 * 1024 + ); + assert_eq!( + Size::parse("1TiB").unwrap().size_bytes(), + 1024 * 1024 * 1024 * 1024 + ); + } + + #[test] + fn test_invalid_input() { + // Invalid format + assert!(Size::parse("invalid").is_err()); + } +} diff --git a/tests/by-util/test_mcookie.rs b/tests/by-util/test_mcookie.rs index 3fd2096..4766a95 100644 --- a/tests/by-util/test_mcookie.rs +++ b/tests/by-util/test_mcookie.rs @@ -32,7 +32,7 @@ fn test_verbose() { } #[test] -fn test_seed_files_and_max_size() { +fn test_seed_files_and_max_size_raw() { let mut file1 = NamedTempFile::new().unwrap(); const CONTENT1: &str = "Some seed data"; file1.write_all(CONTENT1.as_bytes()).unwrap(); @@ -63,3 +63,38 @@ fn test_seed_files_and_max_size() { file2.path().to_str().unwrap() )); } + +#[test] +fn test_seed_files_and_max_size_human_readable() { + let mut file = NamedTempFile::new().unwrap(); + const CONTENT: [u8; 4096] = [1; 4096]; + file.write_all(&CONTENT).unwrap(); + + let res = new_ucmd!() + .arg("--verbose") + .arg("-f") + .arg(file.path()) + .arg("-m") + .arg("2KiB") + .succeeds(); + + // Ensure we only read up to 2KiB (2048 bytes) + res.stderr_contains(format!( + "Got 2048 bytes from {}", + file.path().to_str().unwrap() + )); +} + +#[test] +fn test_invalid_size_format() { + let file = NamedTempFile::new().unwrap(); + + let res = new_ucmd!() + .arg("-f") + .arg(file.path()) + .arg("-m") + .arg("invalid") + .fails(); + + res.stderr_contains("Failed to parse max-size value"); +}