From a86a2711f8d23af047b05912ef1d9604e16a5a16 Mon Sep 17 00:00:00 2001
From: Rob Bradford <robert.bradford@intel.com>
Date: Tue, 8 Sep 2020 15:51:00 +0100
Subject: [PATCH] ch-remote: Support using unit suffices for "resize"

Remove the requirement for the user to calculate the size they want in
bytes.

Fixes: #1596

Signed-off-by: Rob Bradford <robert.bradford@intel.com>
---
 Cargo.lock               |  1 +
 Cargo.toml               |  1 +
 option_parser/src/lib.rs |  1 +
 src/bin/ch-remote.rs     | 27 +++++++++++++++++++--------
 4 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 08737cad6..d7bc8d8d4 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -206,6 +206,7 @@ dependencies = [
  "libc",
  "log 0.4.11",
  "net_util",
+ "option_parser",
  "seccomp",
  "serde_json",
  "ssh2",
diff --git a/Cargo.toml b/Cargo.toml
index 709136124..d8627c644 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -17,6 +17,7 @@ clap = { version = "2.33.3", features = ["wrap_help"] }
 hypervisor = { path = "hypervisor" }
 libc = "0.2.76"
 log = { version = "0.4.11", features = ["std"] }
+option_parser = { path = "option_parser" }
 seccomp = { git = "https://github.com/firecracker-microvm/firecracker", tag = "v0.22.0" }
 serde_json = "1.0.57"
 vhost_user_block = { path = "vhost_user_block"}
diff --git a/option_parser/src/lib.rs b/option_parser/src/lib.rs
index beac16297..ad71ca62b 100644
--- a/option_parser/src/lib.rs
+++ b/option_parser/src/lib.rs
@@ -143,6 +143,7 @@ impl FromStr for Toggle {
 
 pub struct ByteSized(pub u64);
 
+#[derive(Debug)]
 pub enum ByteSizedParseError {
     InvalidValue(String),
 }
diff --git a/src/bin/ch-remote.rs b/src/bin/ch-remote.rs
index e8e482bf9..39868715a 100644
--- a/src/bin/ch-remote.rs
+++ b/src/bin/ch-remote.rs
@@ -9,6 +9,7 @@ extern crate serde_json;
 extern crate vmm;
 
 use clap::{App, AppSettings, Arg, ArgMatches, SubCommand};
+use option_parser::{ByteSized, ByteSizedParseError};
 use std::fmt;
 use std::io::{Read, Write};
 use std::os::unix::net::UnixStream;
@@ -22,8 +23,8 @@ enum Error {
     ContentLengthParsing(std::num::ParseIntError),
     ServerResponse(StatusCode, Option<String>),
     InvalidCPUCount(std::num::ParseIntError),
-    InvalidMemorySize(std::num::ParseIntError),
-    InvalidBalloonSize(std::num::ParseIntError),
+    InvalidMemorySize(ByteSizedParseError),
+    InvalidBalloonSize(ByteSizedParseError),
     AddDeviceConfig(vmm::config::Error),
     AddDiskConfig(vmm::config::Error),
     AddFsConfig(vmm::config::Error),
@@ -49,8 +50,8 @@ impl fmt::Display for Error {
                 }
             }
             InvalidCPUCount(e) => write!(f, "Error parsing CPU count: {}", e),
-            InvalidMemorySize(e) => write!(f, "Error parsing memory size: {}", e),
-            InvalidBalloonSize(e) => write!(f, "Error parsing balloon size: {}", e),
+            InvalidMemorySize(e) => write!(f, "Error parsing memory size: {:?}", e),
+            InvalidBalloonSize(e) => write!(f, "Error parsing balloon size: {:?}", e),
             AddDeviceConfig(e) => write!(f, "Error parsing device syntax: {}", e),
             AddDiskConfig(e) => write!(f, "Error parsing disk syntax: {}", e),
             AddFsConfig(e) => write!(f, "Error parsing filesystem syntax: {}", e),
@@ -215,13 +216,23 @@ fn resize_api_command(
     };
 
     let desired_ram: Option<u64> = if let Some(memory) = memory {
-        Some(memory.parse().map_err(Error::InvalidMemorySize)?)
+        Some(
+            memory
+                .parse::<ByteSized>()
+                .map_err(Error::InvalidMemorySize)?
+                .0,
+        )
     } else {
         None
     };
 
     let desired_ram_w_balloon: Option<u64> = if let Some(balloon) = balloon {
-        Some(balloon.parse().map_err(Error::InvalidBalloonSize)?)
+        Some(
+            balloon
+                .parse::<ByteSized>()
+                .map_err(Error::InvalidBalloonSize)?
+                .0,
+        )
     } else {
         None
     };
@@ -529,14 +540,14 @@ fn main() {
                 .arg(
                     Arg::with_name("memory")
                         .long("memory")
-                        .help("New memory size (in MiB)")
+                        .help("New memory size in bytes (supports K/M/G suffix)")
                         .takes_value(true)
                         .number_of_values(1),
                 )
                 .arg(
                     Arg::with_name("balloon")
                         .long("balloon")
-                        .help("New memory with balloon size")
+                        .help("New memory with balloon size in bytes (supports K/M/G suffix)")
                         .takes_value(true)
                         .number_of_values(1),
                 ),