argument parsing: add size to memory argument

This commit adds an optional size indicator to the memory argument, to
simplify declaration of memory. For example, 2G would be equivalent to
2048MiB. The reason to add this is that when using large enclave sizes,
such as 2TiB, writing down the actual number requires a calculator.

Signed-off-by: Petre Eftime <[email protected]>

src/common/commands_parser.rs

@@ -9,6 +9,7 @@ use libc::VMADDR_CID_HOST;
 use libc::VMADDR_CID_LOCAL;
 use serde::{Deserialize, Serialize};
 use std::fs::File;
+use std::str::FromStr;
 
 use crate::common::{NitroCliErrorEnum, NitroCliFailure, NitroCliResult, VMADDR_CID_PARENT};
 use crate::get_id_by_name;
@@ -290,21 +291,74 @@ fn parse_file_path(args: &ArgMatches, val_name: &str) -> NitroCliResult<String>
     Ok(path.to_string())
 }
 
+#[derive(Debug)]
+enum MemoryUnit {
+    Mebibytes,
+    Gibibytes,
+    Tebibytes,
+}
+
+#[derive(Debug)]
+struct UnknownMemoryUnitErr;
+
+impl MemoryUnit {
+    fn to_mebibytes(&self) -> u64 {
+        match self {
+            MemoryUnit::Mebibytes => 1,
+            MemoryUnit::Gibibytes => 1024,
+            MemoryUnit::Tebibytes => 1024 * 1024,
+        }
+    }
+}
+
+impl FromStr for MemoryUnit {
+    type Err = UnknownMemoryUnitErr;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "M" | "m" | "" => Ok(MemoryUnit::Mebibytes),
+            "G" | "g" => Ok(MemoryUnit::Gibibytes),
+            "T" | "t" => Ok(MemoryUnit::Tebibytes),
+            _ => Err(UnknownMemoryUnitErr),
+        }
+    }
+}
+
 /// Parse the requested amount of enclave memory from the command-line arguments.
-fn parse_memory(args: &ArgMatches) -> NitroCliResult<u64> {
+/// It can be just a number like 123, or it can end in a size indicator like 100M or 10G.
+/// If the size indicator is missing, it defaults to M.
+/// If the size indicator is not M, G or T, it returns an error.
+///
+/// # Arguments
+/// * `args` - The command-line arguments.
+pub fn parse_memory(args: &ArgMatches) -> NitroCliResult<u64> {
     let memory = args.value_of("memory").ok_or_else(|| {
         new_nitro_cli_failure!(
             "`memory` argument not found",
             NitroCliErrorEnum::MissingArgument
         )
     })?;
-    memory.parse().map_err(|_| {
+
+    let (num_str, size_str) = match memory.find(|c: char| !c.is_numeric()) {
+        Some(index) => memory.split_at(index),
+        None => (memory, ""),
+    };
+    let num = num_str.parse::<u64>().map_err(|_| {
+        new_nitro_cli_failure!(
+            "`memory` argument does not contain a number",
+            NitroCliErrorEnum::InvalidArgument
+        )
+        .add_info(vec!["memory", memory])
+    })?;
+
+    let unit = size_str.parse::<MemoryUnit>().map_err(|_| {
         new_nitro_cli_failure!(
-            "`memory` is not a number",
+            "`memory` argument does not contain a valid size indicator",
             NitroCliErrorEnum::InvalidArgument
         )
         .add_info(vec!["memory", memory])
-    })
+    })?;
+    Ok(num * unit.to_mebibytes())
 }
 
 /// Parse the Docker tag from the command-line arguments.
@@ -565,6 +619,102 @@ mod tests {
             let err_str = construct_error_message(&err_info);
             assert!(err_str.contains("Invalid argument provided"))
         }
+
+        let app = create_app!();
+        let args = vec![
+            "nitro-cli",
+            "run-enclave",
+            "--memory",
+            "256",
+            "--cpu-count",
+            "2",
+            "--eif-path",
+            "non_existing_eif.eif",
+        ];
+
+        let matches = app.get_matches_from_safe(args);
+        assert!(matches.is_ok());
+
+        let result = parse_memory(
+            matches
+                .as_ref()
+                .unwrap()
+                .subcommand_matches("run-enclave")
+                .unwrap(),
+        );
+        assert_eq!(result, Ok(256));
+
+        let app = create_app!();
+        let args = vec![
+            "nitro-cli",
+            "run-enclave",
+            "--memory",
+            "100M",
+            "--cpu-count",
+            "2",
+            "--eif-path",
+            "non_existing_eif.eif",
+        ];
+
+        let matches = app.get_matches_from_safe(args);
+        assert!(matches.is_ok());
+
+        let result = parse_memory(
+            matches
+                .as_ref()
+                .unwrap()
+                .subcommand_matches("run-enclave")
+                .unwrap(),
+        );
+        assert_eq!(result, Ok(100));
+
+        let app = create_app!();
+        let args = vec![
+            "nitro-cli",
+            "run-enclave",
+            "--memory",
+            "10G",
+            "--cpu-count",
+            "2",
+            "--eif-path",
+            "non_existing_eif.eif",
+        ];
+
+        let matches = app.get_matches_from_safe(args);
+        assert!(matches.is_ok());
+
+        let result = parse_memory(
+            matches
+                .as_ref()
+                .unwrap()
+                .subcommand_matches("run-enclave")
+                .unwrap(),
+        );
+        assert_eq!(result, Ok(10_240));
+
+        let app = create_app!();
+        let args = vec![
+            "nitro-cli",
+            "run-enclave",
+            "--memory",
+            "2T",
+            "--cpu-count",
+            "2",
+            "--eif-path",
+            "non_existing_eif.eif",
+        ];
+
+        let matches = app.get_matches_from_safe(args);
+        assert!(matches.is_ok());
+
+        let result = parse_memory(
+            matches
+                .as_ref()
+                .unwrap()
+                .subcommand_matches("run-enclave")
+                .unwrap(),
+        );
+        assert_eq!(result, Ok(2 * 1024 * 1024));
     }
 
     #[test]

src/common/mod.rs

@@ -225,7 +225,7 @@ pub enum EnclaveProcessReply {
 }
 
 /// Struct that is passed along the backtrace and accumulates error messages.
-#[derive(Debug, Default)]
+#[derive(Debug, Default, PartialEq)]
 pub struct NitroCliFailure {
     /// Main action which was attempted and failed.
     pub action: String,