Merge pull request #13721 from y-zeng/cli_metadata

Add escape sequences for the grpc_cli metadata flag

2 files changed, 108 insertions, 7 deletions

diff --git a/test/cpp/util/grpc_tool.cc b/test/cpp/util/grpc_tool.cc
index a6d08cd83c..30c43b206f 100644
--- a/test/cpp/util/grpc_tool.cc
+++ b/test/cpp/util/grpc_tool.cc
@@ -124,13 +124,32 @@ void ParseMetadataFlag(
     return;
   }
   std::vector<grpc::string> fields;
-  const char* delim = ":";
-  size_t cur, next = -1;
-  do {
-    cur = next + 1;
-    next = FLAGS_metadata.find_first_of(delim, cur);
-    fields.push_back(FLAGS_metadata.substr(cur, next - cur));
-  } while (next != grpc::string::npos);
+  const char delim = ':';
+  const char escape = '\\';
+  size_t cur = -1;
+  std::stringstream ss;
+  while (++cur < FLAGS_metadata.length()) {
+    switch (FLAGS_metadata.at(cur)) {
+      case escape:
+        if (cur < FLAGS_metadata.length() - 1) {
+          char c = FLAGS_metadata.at(++cur);
+          if (c == delim || c == escape) {
+            ss << c;
+            continue;
+          }
+        }
+        fprintf(stderr, "Failed to parse metadata flag.\n");
+        exit(1);
+      case delim:
+        fields.push_back(ss.str());
+        ss.str("");
+        ss.clear();
+        break;
+      default:
+        ss << FLAGS_metadata.at(cur);
+    }
+  }
+  fields.push_back(ss.str());
   if (fields.size() % 2) {
     fprintf(stderr, "Failed to parse metadata flag.\n");
     exit(1);
diff --git a/test/cpp/util/grpc_tool_test.cc b/test/cpp/util/grpc_tool_test.cc
index 1c07b2a8ef..0b599f4eeb 100644
--- a/test/cpp/util/grpc_tool_test.cc
+++ b/test/cpp/util/grpc_tool_test.cc
@@ -85,6 +85,8 @@ DECLARE_bool(binary_input);
 DECLARE_bool(binary_output);
 DECLARE_bool(l);
 DECLARE_bool(batch);
+DECLARE_string(metadata);
+DECLARE_string(protofiles);
 
 namespace {
 
@@ -618,6 +620,8 @@ TEST_F(GrpcToolTest, ParseCommand) {
   // Expected output: ECHO_RESPONSE_MESSAGE
   EXPECT_TRUE(0 == strcmp(output_stream.str().c_str(), ECHO_RESPONSE_MESSAGE));
 
+  FLAGS_binary_input = false;
+  FLAGS_binary_output = false;
   ShutdownServer();
 }
 
@@ -652,6 +656,84 @@ TEST_F(GrpcToolTest, TooManyArguments) {
   EXPECT_TRUE(0 == output_stream.tellp());
 }
 
+TEST_F(GrpcToolTest, CallCommandWithMetadata) {
+  // Test input "grpc_cli call localhost:<port> Echo "message: 'Hello'"
+  const grpc::string server_address = SetUpServer();
+  const char* argv[] = {"grpc_cli", "call", server_address.c_str(), "Echo",
+                        "message: 'Hello'"};
+
+  {
+    std::stringstream output_stream;
+    FLAGS_metadata = "key0:val0:key1:valq:key2:val2";
+    EXPECT_TRUE(0 == GrpcToolMainLib(ArraySize(argv), argv,
+                                     TestCliCredentials(),
+                                     std::bind(PrintStream, &output_stream,
+                                               std::placeholders::_1)));
+    // Expected output: "message: \"Hello\""
+    EXPECT_TRUE(nullptr !=
+                strstr(output_stream.str().c_str(), "message: \"Hello\""));
+  }
+
+  {
+    std::stringstream output_stream;
+    FLAGS_metadata = "key:val\\:val";
+    EXPECT_TRUE(0 == GrpcToolMainLib(ArraySize(argv), argv,
+                                     TestCliCredentials(),
+                                     std::bind(PrintStream, &output_stream,
+                                               std::placeholders::_1)));
+    // Expected output: "message: \"Hello\""
+    EXPECT_TRUE(nullptr !=
+                strstr(output_stream.str().c_str(), "message: \"Hello\""));
+  }
+
+  {
+    std::stringstream output_stream;
+    FLAGS_metadata = "key:val\\\\val";
+    EXPECT_TRUE(0 == GrpcToolMainLib(ArraySize(argv), argv,
+                                     TestCliCredentials(),
+                                     std::bind(PrintStream, &output_stream,
+                                               std::placeholders::_1)));
+    // Expected output: "message: \"Hello\""
+    EXPECT_TRUE(nullptr !=
+                strstr(output_stream.str().c_str(), "message: \"Hello\""));
+  }
+
+  FLAGS_metadata = "";
+  ShutdownServer();
+}
+
+TEST_F(GrpcToolTest, CallCommandWithBadMetadata) {
+  // Test input "grpc_cli call localhost:10000 Echo "message: 'Hello'"
+  const char* argv[] = {"grpc_cli", "call", "localhost:10000", "Echo",
+                        "message: 'Hello'"};
+  FLAGS_protofiles = "src/proto/grpc/testing/echo.proto";
+
+  {
+    std::stringstream output_stream;
+    FLAGS_metadata = "key0:val0:key1";
+    // Exit with 1
+    EXPECT_EXIT(
+        GrpcToolMainLib(
+            ArraySize(argv), argv, TestCliCredentials(),
+            std::bind(PrintStream, &output_stream, std::placeholders::_1)),
+        ::testing::ExitedWithCode(1), ".*Failed to parse metadata flag.*");
+  }
+
+  {
+    std::stringstream output_stream;
+    FLAGS_metadata = "key:val\\val";
+    // Exit with 1
+    EXPECT_EXIT(
+        GrpcToolMainLib(
+            ArraySize(argv), argv, TestCliCredentials(),
+            std::bind(PrintStream, &output_stream, std::placeholders::_1)),
+        ::testing::ExitedWithCode(1), ".*Failed to parse metadata flag.*");
+  }
+
+  FLAGS_metadata = "";
+  FLAGS_protofiles = "";
+}
+
 }  // namespace testing
 }  // namespace grpc