fix grpc create initial metadata (envoyproxy#158)

Signed-off-by: Pengyuan Bian <bianpengyuan@google.com>
PiotrSikora · Feb 24, 2020 · 9bb0cc6 · 9bb0cc6
1 parent d2c83d4
commit 9bb0cc6
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 16 deletions.
diff --git a/api/wasm/cpp/proxy_wasm_api.h b/api/wasm/cpp/proxy_wasm_api.h
@@ -352,7 +352,9 @@ class RootContext : public ContextBase {
   const std::string root_id_;
   std::unordered_map<uint32_t, HttpCallCallback> http_calls_;
   std::unordered_map<uint32_t, GrpcSimpleCallCallback> simple_grpc_calls_;
+  std::unique_ptr<GrpcCallHandlerBase> cur_grpc_call_;
   std::unordered_map<uint32_t, std::unique_ptr<GrpcCallHandlerBase>> grpc_calls_;
+  std::unique_ptr<GrpcStreamHandlerBase> cur_grpc_stream_;
   std::unordered_map<uint32_t, std::unique_ptr<GrpcStreamHandlerBase>> grpc_streams_;
 };
 
@@ -1244,19 +1246,15 @@ inline void GrpcStreamHandlerBase::send(StringView message, bool end_of_stream)
   }
 }
 
-inline void RootContext::onGrpcCreateInitialMetadata(uint32_t token, uint32_t headers) {
+inline void RootContext::onGrpcCreateInitialMetadata(uint32_t, uint32_t headers) {
   {
-    auto it = grpc_calls_.find(token);
-    if (it != grpc_calls_.end()) {
-      it->second->onCreateInitialMetadata(headers);
-      return;
+    if (cur_grpc_call_ != nullptr) {
+      cur_grpc_call_->onCreateInitialMetadata(headers);
     }
   }
   {
-    auto it = grpc_streams_.find(token);
-    if (it != grpc_streams_.end()) {
-      it->second->onCreateInitialMetadata(headers);
-      return;
+    if (cur_grpc_stream_ != nullptr) {
+      cur_grpc_stream_->onCreateInitialMetadata(headers);
     }
   }
 }
@@ -1354,25 +1352,29 @@ inline WasmResult RootContext::grpcCallHandler(StringView service, StringView se
                                                uint32_t timeout_milliseconds,
                                                std::unique_ptr<GrpcCallHandlerBase> handler) {
   uint32_t token = 0;
+  cur_grpc_call_ = std::move(handler);
   auto result = grpcCall(service, service_name, method_name, request, timeout_milliseconds, &token);
   if (result == WasmResult::Ok) {
-    handler->token_ = token;
-    handler->context_ = this;
-    grpc_calls_[token] = std::move(handler);
+    cur_grpc_call_->token_ = token;
+    cur_grpc_call_->context_ = this;
+    grpc_calls_[token] = std::move(cur_grpc_call_);
   }
+  cur_grpc_call_ = nullptr;
   return result;
 }
 
 inline WasmResult RootContext::grpcStreamHandler(StringView service, StringView service_name,
                                                  StringView method_name,
                                                  std::unique_ptr<GrpcStreamHandlerBase> handler) {
   uint32_t token = 0;
+  cur_grpc_stream_ = std::move(handler);
   auto result = grpcStream(service, service_name, method_name, &token);
   if (result == WasmResult::Ok) {
-    handler->token_ = token;
-    handler->context_ = this;
-    grpc_streams_[token] = std::move(handler);
+    cur_grpc_stream_->token_ = token;
+    cur_grpc_stream_->context_ = this;
+    grpc_streams_[token] = std::move(cur_grpc_stream_);
   }
+  cur_grpc_stream_ = nullptr;
   return result;
 }
 

diff --git a/test/extensions/filters/http/wasm/test_data/grpc_call_cpp.cc b/test/extensions/filters/http/wasm/test_data/grpc_call_cpp.cc
@@ -18,7 +18,10 @@ static RegisterContextFactory register_ExampleContext(CONTEXT_FACTORY(ExampleCon
 class MyGrpcCallHandler : public GrpcCallHandler<google::protobuf::Value> {
 public:
   MyGrpcCallHandler() : GrpcCallHandler<google::protobuf::Value>() {}
-  void onCreateInitialMetadata(uint32_t) override {}
+  void onCreateInitialMetadata(uint32_t) override {
+    addHeaderMapValue(
+        HeaderMapType::GrpcCreateInitialMetadata, "key", "val");
+  }
   void onSuccess(size_t body_size) override {
     auto response = getBufferBytes(BufferType::GrpcReceiveBuffer, 0, body_size);
     logDebug(response->proto<google::protobuf::Value>().string_value());

diff --git a/test/extensions/filters/http/wasm/test_data/grpc_call_cpp.wasm b/test/extensions/filters/http/wasm/test_data/grpc_call_cpp.wasm
diff --git a/test/extensions/filters/http/wasm/wasm_filter_test.cc b/test/extensions/filters/http/wasm/wasm_filter_test.cc
@@ -381,6 +381,11 @@ TEST_P(WasmHttpFilterTest, GrpcCallAfterDestroyed) {
         EXPECT_TRUE(value.ParseFromArray(message->linearize(message->length()), message->length()));
         EXPECT_EQ(value.string_value(), "request");
         callbacks = &cb;
+        Http::HeaderMapImpl initial_metadata;
+        callbacks->onCreateInitialMetadata(initial_metadata);
+        auto* metadata_entry = initial_metadata.get(Envoy::Http::LowerCaseString("key"));
+        EXPECT_TRUE(metadata_entry != nullptr);
+        EXPECT_EQ(metadata_entry->value().getStringView(), "val");
         parent_span = &span;
         EXPECT_EQ(options.timeout->count(), 1000);
         return &request;