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

Author: quangcito

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::<String>(options::MAX_SIZE)
-        .map(|v| v.parse::<u64>().expect("Failed to parse max-size value"));
+    let max_size = if let Some(size_str) = matches.get_one::<String>(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)

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<Self, ParseSizeError> {
+        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::<u64>()
+                    .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::<u64>()
+                    .map(|n| Self(n * 1024_u64.pow(exponent)))
+                    .map_err(|_| ParseSizeError(s.to_string()));
+            }
+        }
+
+        // If no suffix, treat as bytes
+        s.parse::<u64>()
+            .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());
+    }
+}

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");
+}