base: CommandLine::RemoveSwitch should also remove from argv

CommandLine::RemoveSwitch should also remove the switch from argv_.

Bug: 964068
Test: base_unittests --gtest_filter=*CommandLine*
Change-Id: If44106642c8089b83eb897e09c98ab8cf88fa24b
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1653929
Commit-Queue: Pavol Marko <pmarko@chromium.org>
Reviewed-by: Daniel Cheng <dcheng@chromium.org>
Cr-Commit-Position: refs/heads/master@{#669073}

base/command_line.cc

@@ -45,10 +45,10 @@ const CommandLine::CharType* const kSwitchPrefixes[] = {"--", "-"};
 #endif
 size_t switch_prefix_count = size(kSwitchPrefixes);
 
-size_t GetSwitchPrefixLength(const CommandLine::StringType& string) {
+size_t GetSwitchPrefixLength(CommandLine::StringPieceType string) {
   for (size_t i = 0; i < switch_prefix_count; ++i) {
     CommandLine::StringType prefix(kSwitchPrefixes[i]);
-    if (string.compare(0, prefix.length(), prefix) == 0)
+    if (string.substr(0, prefix.length()) == prefix)
       return prefix.length();
   }
   return 0;
@@ -72,6 +72,19 @@ bool IsSwitch(const CommandLine::StringType& string,
   return true;
 }
 
+// Returns true iff |string| represents a switch with key
+// |switch_key_without_prefix|, regardless of value.
+bool IsSwitchWithKey(CommandLine::StringPieceType string,
+                     CommandLine::StringPieceType switch_key_without_prefix) {
+  size_t prefix_length = GetSwitchPrefixLength(string);
+  if (prefix_length == 0 || prefix_length == string.length())
+    return false;
+
+  const size_t equals_position = string.find(kSwitchValueSeparator);
+  return string.substr(prefix_length, equals_position - prefix_length) ==
+         switch_key_without_prefix;
+}
+
 // Append switches and arguments, keeping switches before arguments.
 void AppendSwitchesAndArguments(CommandLine* command_line,
                                 const CommandLine::StringVector& argv) {
@@ -362,9 +375,38 @@ void CommandLine::AppendSwitchASCII(const std::string& switch_string,
 #endif
 }
 
-void CommandLine::RemoveSwitch(const StringPiece& switch_string) {
-  DCHECK_EQ(ToLowerASCII(switch_string), switch_string);
-  switches_.erase(switch_string.as_string());
+void CommandLine::RemoveSwitch(base::StringPiece switch_key_without_prefix) {
+#if defined(OS_WIN)
+  StringType switch_key_native = base::ASCIIToUTF16(switch_key_without_prefix);
+#elif defined(OS_POSIX) || defined(OS_FUCHSIA)
+  StringType switch_key_native = switch_key_without_prefix.as_string();
+#endif
+
+  DCHECK_EQ(ToLowerASCII(switch_key_without_prefix), switch_key_without_prefix);
+  DCHECK_EQ(0u, GetSwitchPrefixLength(switch_key_native));
+  size_t erased_from_switches =
+      switches_.erase(switch_key_without_prefix.as_string());
+  DCHECK(erased_from_switches <= 1);
+  if (!erased_from_switches)
+    return;
+
+  // Also erase from the switches section of |argv_| and update |begin_args_|
+  // accordingly.
+  // Switches in |argv_| have indices [1, begin_args_).
+  auto argv_switches_begin = argv_.begin() + 1;
+  auto argv_switches_end = argv_.begin() + begin_args_;
+  DCHECK(argv_switches_begin <= argv_switches_end);
+  DCHECK(argv_switches_end <= argv_.end());
+  auto arg_iter = std::find_if(argv_switches_begin, argv_switches_end,
+                               [&switch_key_native](const StringType& arg) {
+                                 return IsSwitchWithKey(arg, switch_key_native);
+                               });
+  if (arg_iter == argv_switches_end) {
+    NOTREACHED();
+    return;
+  }
+  argv_.erase(arg_iter);
+  --begin_args_;
 }
 
 void CommandLine::CopySwitchesFrom(const CommandLine& source,

base/command_line.h

@@ -35,8 +35,10 @@ class BASE_EXPORT CommandLine {
 #if defined(OS_WIN)
   // The native command line string type.
   using StringType = string16;
+  using StringPieceType = base::StringPiece16;
 #elif defined(OS_POSIX) || defined(OS_FUCHSIA)
   using StringType = std::string;
+  using StringPieceType = base::StringPiece;
 #endif
 
   using CharType = StringType::value_type;
@@ -183,8 +185,9 @@ class BASE_EXPORT CommandLine {
   void AppendSwitchASCII(const std::string& switch_string,
                          const std::string& value);
 
-  // Removes a switch.
-  void RemoveSwitch(const StringPiece& switch_string);
+  // Removes the switch that matches |switch_key_without_prefix|, regardless of
+  // prefix and value. If no such switch is present, this has no effect.
+  void RemoveSwitch(const base::StringPiece switch_key_without_prefix);
 
   // Copy a set of switches (and any values) from another command line.
   // Commonly used when launching a subprocess.

base/command_line_unittest.cc

@@ -13,6 +13,7 @@
 #include "base/strings/strcat.h"
 #include "base/strings/utf_string_conversions.h"
 #include "build/build_config.h"
+#include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace base {
@@ -441,9 +442,9 @@ TEST(CommandLineTest, PrependComplexWrapper) {
 }
 
 TEST(CommandLineTest, RemoveSwitch) {
-  std::string switch1 = "switch1";
-  std::string switch2 = "switch2";
-  std::string value2 = "value";
+  const std::string switch1 = "switch1";
+  const std::string switch2 = "switch2";
+  const std::string value2 = "value";
 
   CommandLine cl(FilePath(FILE_PATH_LITERAL("Program")));
 
@@ -453,12 +454,90 @@ TEST(CommandLineTest, RemoveSwitch) {
   EXPECT_TRUE(cl.HasSwitch(switch1));
   EXPECT_TRUE(cl.HasSwitch(switch2));
   EXPECT_EQ(value2, cl.GetSwitchValueASCII(switch2));
+  EXPECT_THAT(cl.argv(),
+              testing::ElementsAre(FILE_PATH_LITERAL("Program"),
+                                   FILE_PATH_LITERAL("--switch1"),
+                                   FILE_PATH_LITERAL("--switch2=value")));
 
   cl.RemoveSwitch(switch1);
 
   EXPECT_FALSE(cl.HasSwitch(switch1));
   EXPECT_TRUE(cl.HasSwitch(switch2));
   EXPECT_EQ(value2, cl.GetSwitchValueASCII(switch2));
+  EXPECT_THAT(cl.argv(),
+              testing::ElementsAre(FILE_PATH_LITERAL("Program"),
+                                   FILE_PATH_LITERAL("--switch2=value")));
+}
+
+TEST(CommandLineTest, RemoveSwitchWithValue) {
+  const std::string switch1 = "switch1";
+  const std::string switch2 = "switch2";
+  const std::string value2 = "value";
+
+  CommandLine cl(FilePath(FILE_PATH_LITERAL("Program")));
+
+  cl.AppendSwitch(switch1);
+  cl.AppendSwitchASCII(switch2, value2);
+
+  EXPECT_TRUE(cl.HasSwitch(switch1));
+  EXPECT_TRUE(cl.HasSwitch(switch2));
+  EXPECT_EQ(value2, cl.GetSwitchValueASCII(switch2));
+  EXPECT_THAT(cl.argv(),
+              testing::ElementsAre(FILE_PATH_LITERAL("Program"),
+                                   FILE_PATH_LITERAL("--switch1"),
+                                   FILE_PATH_LITERAL("--switch2=value")));
+
+  cl.RemoveSwitch(switch2);
+
+  EXPECT_TRUE(cl.HasSwitch(switch1));
+  EXPECT_FALSE(cl.HasSwitch(switch2));
+  EXPECT_THAT(cl.argv(), testing::ElementsAre(FILE_PATH_LITERAL("Program"),
+                                              FILE_PATH_LITERAL("--switch1")));
+}
+
+TEST(CommandLineTest, AppendAndRemoveSwitchWithDefaultPrefix) {
+  CommandLine cl(FilePath(FILE_PATH_LITERAL("Program")));
+
+  cl.AppendSwitch("foo");
+  EXPECT_THAT(cl.argv(), testing::ElementsAre(FILE_PATH_LITERAL("Program"),
+                                              FILE_PATH_LITERAL("--foo")));
+  EXPECT_EQ(0u, cl.GetArgs().size());
+
+  cl.RemoveSwitch("foo");
+  EXPECT_THAT(cl.argv(), testing::ElementsAre(FILE_PATH_LITERAL("Program")));
+  EXPECT_EQ(0u, cl.GetArgs().size());
+}
+
+TEST(CommandLineTest, AppendAndRemoveSwitchWithAlternativePrefix) {
+  CommandLine cl(FilePath(FILE_PATH_LITERAL("Program")));
+
+  cl.AppendSwitch("-foo");
+  EXPECT_THAT(cl.argv(), testing::ElementsAre(FILE_PATH_LITERAL("Program"),
+                                              FILE_PATH_LITERAL("-foo")));
+  EXPECT_EQ(0u, cl.GetArgs().size());
+
+  cl.RemoveSwitch("foo");
+  EXPECT_THAT(cl.argv(), testing::ElementsAre(FILE_PATH_LITERAL("Program")));
+  EXPECT_EQ(0u, cl.GetArgs().size());
+}
+
+TEST(CommandLineTest, AppendAndRemoveSwitchPreservesOtherSwitchesAndArgs) {
+  CommandLine cl(FilePath(FILE_PATH_LITERAL("Program")));
+
+  cl.AppendSwitch("foo");
+  cl.AppendSwitch("bar");
+  cl.AppendArg("arg");
+  EXPECT_THAT(cl.argv(), testing::ElementsAre(FILE_PATH_LITERAL("Program"),
+                                              FILE_PATH_LITERAL("--foo"),
+                                              FILE_PATH_LITERAL("--bar"),
+                                              FILE_PATH_LITERAL("arg")));
+  EXPECT_THAT(cl.GetArgs(), testing::ElementsAre(FILE_PATH_LITERAL("arg")));
+
+  cl.RemoveSwitch("foo");
+  EXPECT_THAT(cl.argv(), testing::ElementsAre(FILE_PATH_LITERAL("Program"),
+                                              FILE_PATH_LITERAL("--bar"),
+                                              FILE_PATH_LITERAL("arg")));
+  EXPECT_THAT(cl.GetArgs(), testing::ElementsAre(FILE_PATH_LITERAL("arg")));
 }
 
 TEST(CommandLineTest, MultipleSameSwitch) {