From 92f339372c236d3b84188489417df2d65550b1c4 Mon Sep 17 00:00:00 2001
From: Tao Bao <tbao@google.com>
Date: Mon, 25 Jun 2018 12:11:53 -0700
Subject: updater: Check the number of args in Command::Parse.

Additionally checks for excess args when parsing ERASE, FREE, NEW, STASH
and ZERO. Note that the check for MOVE, BSDIFF, IMGDIFF has been covered
in Command::ParseTargetInfoAndSourceInfo.

Test: Run recovery_unit_test on marlin.
Change-Id: Ic8bc9b7a8dcf98f1f8db2e259607564508726857
---
 tests/unit/commands_test.cpp | 25 +++++++++++++++++++++++++
 updater/commands.cpp         | 25 ++++++++++++++++---------
 2 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/tests/unit/commands_test.cpp b/tests/unit/commands_test.cpp
index bbc83b9a8..cb2be9176 100644
--- a/tests/unit/commands_test.cpp
+++ b/tests/unit/commands_test.cpp
@@ -310,3 +310,28 @@ TEST(CommandsTest, Parse_ZERO) {
   ASSERT_EQ(StashInfo(), command.stash());
   ASSERT_EQ(PatchInfo(), command.patch());
 }
+
+TEST(CommandsTest, Parse_InvalidNumberOfArgs) {
+  // Note that the case of having excess args in BSDIFF, IMGDIFF and MOVE is covered by
+  // ParseTargetInfoAndSourceInfo_InvalidInput.
+  std::vector<std::string> inputs{
+    "bsdiff",
+    "erase",
+    "erase 4,3,5,10,12 hash1",
+    "free",
+    "free id1 id2",
+    "imgdiff",
+    "move",
+    "new",
+    "new 4,3,5,10,12 hash1",
+    "stash",
+    "stash id1",
+    "stash id1 4,3,5,10,12 id2",
+    "zero",
+    "zero 4,3,5,10,12 hash2",
+  };
+  for (const auto& input : inputs) {
+    std::string err;
+    ASSERT_FALSE(Command::Parse(input, 0, &err));
+  }
+}
diff --git a/updater/commands.cpp b/updater/commands.cpp
index 6d4b5310b..fb19ebc9a 100644
--- a/updater/commands.cpp
+++ b/updater/commands.cpp
@@ -54,7 +54,7 @@ bool Command::ParseTargetInfoAndSourceInfo(const std::vector<std::string>& token
                                            const std::string& tgt_hash, TargetInfo* target,
                                            const std::string& src_hash, SourceInfo* source,
                                            std::string* err) {
-  // We expect the given tokens parameter in one of the following formats.
+  // We expect the given args (in 'tokens' vector) in one of the following formats.
   //
   //    <tgt_ranges> <src_block_count> - <[stash_id:location] ...>
   //        (loads data from stashes only)
@@ -65,10 +65,9 @@ bool Command::ParseTargetInfoAndSourceInfo(const std::vector<std::string>& token
   //    <tgt_ranges> <src_block_count> <src_ranges> <src_ranges_location> <[stash_id:location] ...>
   //        (loads data from both of source image and stashes)
 
-  // At least it needs to provide three parameters: <tgt_ranges>, <src_block_count> and
-  // "-"/<src_ranges>.
+  // At least it needs to provide three args: <tgt_ranges>, <src_block_count> and "-"/<src_ranges>.
   if (tokens.size() < 3) {
-    *err = "invalid number of parameters";
+    *err = "invalid number of args";
     return false;
   }
 
@@ -170,6 +169,11 @@ Command Command::Parse(const std::string& line, size_t index, std::string* err)
 
   if (op == Type::ZERO || op == Type::NEW || op == Type::ERASE) {
     // zero/new/erase <rangeset>
+    if (pos + 1 != tokens.size()) {
+      *err = android::base::StringPrintf("invalid number of args: %zu (expected 1)",
+                                         tokens.size() - pos);
+      return {};
+    }
     RangeSet tgt_ranges = RangeSet::Parse(tokens[pos++]);
     if (!tgt_ranges) {
       return {};
@@ -178,8 +182,9 @@ Command Command::Parse(const std::string& line, size_t index, std::string* err)
     target_info = TargetInfo(kUnknownHash, tgt_ranges);
   } else if (op == Type::STASH) {
     // stash <stash_id> <src_ranges>
-    if (pos + 2 > tokens.size()) {
-      *err = "missing stash id and/or source ranges";
+    if (pos + 2 != tokens.size()) {
+      *err = android::base::StringPrintf("invalid number of args: %zu (expected 2)",
+                                         tokens.size() - pos);
       return {};
     }
     const std::string& id = tokens[pos++];
@@ -191,8 +196,9 @@ Command Command::Parse(const std::string& line, size_t index, std::string* err)
     stash_info = StashInfo(id, src_ranges);
   } else if (op == Type::FREE) {
     // free <stash_id>
-    if (pos + 1 > tokens.size()) {
-      *err = "missing stash id in free command";
+    if (pos + 1 != tokens.size()) {
+      *err = android::base::StringPrintf("invalid number of args: %zu (expected 1)",
+                                         tokens.size() - pos);
       return {};
     }
     stash_info = StashInfo(tokens[pos++], {});
@@ -211,7 +217,8 @@ Command Command::Parse(const std::string& line, size_t index, std::string* err)
   } else if (op == Type::BSDIFF || op == Type::IMGDIFF) {
     // <offset> <length> <srchash> <dsthash>
     if (pos + 4 > tokens.size()) {
-      *err = "invalid number of tokens";
+      *err = android::base::StringPrintf("invalid number of args: %zu (expected 4+)",
+                                         tokens.size() - pos);
       return {};
     }
     size_t offset;
-- 
cgit v1.2.3