Fix Wasm gRPC initial metadata callback

api/wasm/cpp/proxy_wasm_api.h

@@ -303,7 +303,7 @@ class RootContext : public ContextBase {
   // Low level HTTP/gRPC interface.
   virtual void onHttpCallResponse(uint32_t token, uint32_t headers, size_t body_size,
                                   uint32_t trailers);
-  virtual void onGrpcCreateInitialMetadata(uint32_t token, uint32_t headers);
+  virtual void onGrpcCreateInitialMetadata(uint32_t headers);
   virtual void onGrpcReceiveInitialMetadata(uint32_t token, uint32_t headers);
   virtual void onGrpcReceiveTrailingMetadata(uint32_t token, uint32_t trailers);
   virtual void onGrpcReceive(uint32_t token, size_t body_size);
@@ -354,6 +354,11 @@ class RootContext : public ContextBase {
   std::unordered_map<uint32_t, GrpcSimpleCallCallback> simple_grpc_calls_;
   std::unordered_map<uint32_t, std::unique_ptr<GrpcCallHandlerBase>> grpc_calls_;
   std::unordered_map<uint32_t, std::unique_ptr<GrpcStreamHandlerBase>> grpc_streams_;
+
+  // Handles initial metadata creation callback, which is triggered inline when initializing the call, by which time token is not yet available in module and thus the handler cannot be added into grpc_calls_ and grpc_streams_;
+  // These two pointers should only be used in onCreateInitialMetadata.
+  std::unique_ptr<GrpcCallHandlerBase> cur_grpc_call_handler_;
+  std::unique_ptr<GrpcStreamHandlerBase> cur_grpc_stream_handler_;
 };
 
 RootContext* getRoot(StringView root_id);
@@ -1247,18 +1252,16 @@ 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 headers) {
   {
-    auto it = grpc_calls_.find(token);
-    if (it != grpc_calls_.end()) {
-      it->second->onCreateInitialMetadata(headers);
+    if (cur_grpc_call_handler_ != nullptr) {
+      cur_grpc_call_handler_->onCreateInitialMetadata(headers);
       return;
     }
   }
   {
-    auto it = grpc_streams_.find(token);
-    if (it != grpc_streams_.end()) {
-      it->second->onCreateInitialMetadata(headers);
+    if (cur_grpc_stream_handler_ != nullptr) {
+      cur_grpc_stream_handler_->onCreateInitialMetadata(headers);
       return;
     }
   }
@@ -1357,25 +1360,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_handler_ = 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_handler_->token_ = token;
+    cur_grpc_call_handler_->context_ = this;
+    grpc_calls_[token] = std::move(cur_grpc_call_handler_);
   }
+  cur_grpc_call_handler_ = 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_handler_ = 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_handler_->token_ = token;
+    cur_grpc_stream_handler_->context_ = this;
+    grpc_streams_[token] = std::move(cur_grpc_stream_handler_);
   }
+  cur_grpc_stream_handler_ = nullptr;
   return result;
 }
 

api/wasm/cpp/proxy_wasm_externs.h

@@ -155,8 +155,7 @@ extern "C" FilterMetadataStatus proxy_on_response_metadata(uint32_t context_id,
 // HTTP/gRPC.
 extern "C" void proxy_on_http_call_response(uint32_t context_id, uint32_t token, uint32_t headers,
                                             uint32_t body_size, uint32_t trailers);
-extern "C" void proxy_on_grpc_create_initial_metadata(uint32_t context_id, uint32_t token,
-                                                      uint32_t headers);
+extern "C" void proxy_on_grpc_create_initial_metadata(uint32_t context_id, uint32_t headers);
 extern "C" void proxy_on_grpc_receive_initial_metadata(uint32_t context_id, uint32_t token,
                                                        uint32_t headers);
 extern "C" void proxy_on_grpc_trailing_metadata(uint32_t context_id, uint32_t token,

api/wasm/cpp/proxy_wasm_intrinsics.cc

@@ -226,8 +226,8 @@ extern "C" PROXY_WASM_KEEPALIVE void proxy_on_http_call_response(uint32_t contex
 }
 
 extern "C" PROXY_WASM_KEEPALIVE void
-proxy_on_grpc_create_initial_metadata(uint32_t context_id, uint32_t token, uint32_t headers) {
-  getRootContext(context_id)->onGrpcCreateInitialMetadata(token, headers);
+proxy_on_grpc_create_initial_metadata(uint32_t context_id, uint32_t headers) {
+  getRootContext(context_id)->onGrpcCreateInitialMetadata(headers);
 }
 
 extern "C" PROXY_WASM_KEEPALIVE void

source/extensions/common/wasm/context.cc

@@ -1278,13 +1278,12 @@ void Context::onQueueReady(uint32_t token) {
   }
 }
 
-void Context::onGrpcCreateInitialMetadata(uint32_t token, Http::HeaderMap& metadata) {
+void Context::onGrpcCreateInitialMetadata(Http::HeaderMap& metadata) {
   if (!wasm_->on_grpc_create_initial_metadata_) {
     return;
   }
   grpc_create_initial_metadata_ = &metadata;
-  wasm_->on_grpc_create_initial_metadata_(this, id_, token,
-                                          headerSize(grpc_create_initial_metadata_));
+  wasm_->on_grpc_create_initial_metadata_(this, id_, headerSize(grpc_create_initial_metadata_));
   grpc_create_initial_metadata_ = nullptr;
 }
 

source/extensions/common/wasm/context.h

@@ -348,7 +348,7 @@ class Context : public Logger::Loggable<Logger::Id::wasm>,
   struct GrpcCallClientHandler : public Grpc::RawAsyncRequestCallbacks {
     // Grpc::AsyncRequestCallbacks
     void onCreateInitialMetadata(Http::HeaderMap& metadata) override {
-      context_->onGrpcCreateInitialMetadata(token_, metadata);
+      context_->onGrpcCreateInitialMetadata(metadata);
     }
     void onSuccessRaw(Buffer::InstancePtr&& response, Tracing::Span& /* span */) override {
       context_->onGrpcReceive(token_, std::move(response));
@@ -367,7 +367,7 @@ class Context : public Logger::Loggable<Logger::Id::wasm>,
   struct GrpcStreamClientHandler : public Grpc::RawAsyncStreamCallbacks {
     // Grpc::AsyncStreamCallbacks
     void onCreateInitialMetadata(Http::HeaderMap& metadata) override {
-      context_->onGrpcCreateInitialMetadata(token_, metadata);
+      context_->onGrpcCreateInitialMetadata(metadata);
     }
     void onReceiveInitialMetadata(Http::HeaderMapPtr&& metadata) override {
       context_->onGrpcReceiveInitialMetadata(token_, std::move(metadata));
@@ -395,8 +395,7 @@ class Context : public Logger::Loggable<Logger::Id::wasm>,
   void onHttpCallSuccess(uint32_t token, Envoy::Http::MessagePtr& response);
   void onHttpCallFailure(uint32_t token, Http::AsyncClient::FailureReason reason);
 
-  virtual void onGrpcCreateInitialMetadata(uint32_t token,
-                                           Http::HeaderMap& metadata); // For both Call and Stream.
+  virtual void onGrpcCreateInitialMetadata(Http::HeaderMap& metadata); // For both Call and Stream.
   virtual void onGrpcReceive(uint32_t token,
                              Buffer::InstancePtr response); // Call (implies OK close) and Stream.
   virtual void onGrpcClose(uint32_t token, const Grpc::Status::GrpcStatus& status,

source/extensions/common/wasm/null/null_plugin.cc

@@ -83,6 +83,11 @@ void NullPlugin::getFunction(absl::string_view function_name, WasmCallVoid<2>* f
       SaveRestoreContext saved_context(context);
       plugin->onQueueReady(context_id.u64_, token.u64_);
     };
+  } else if (function_name == "proxy_on_grpc_create_initial_metadata") {
+    *f = [plugin](Common::Wasm::Context* context, Word context_id, Word headers) {
+      SaveRestoreContext saved_context(context);
+      plugin->onGrpcCreateInitialMetadata(context_id.u64_, headers.u64_);
+    };
   } else {
     throw WasmVmException(fmt::format("Missing getFunction for: {}", function_name));
   }
@@ -100,11 +105,6 @@ void NullPlugin::getFunction(absl::string_view function_name, WasmCallVoid<3>* f
       SaveRestoreContext saved_context(context);
       plugin->onGrpcReceive(context_id.u64_, token.u64_, body_size.u64_);
     };
-  } else if (function_name == "proxy_on_grpc_create_initial_metadata") {
-    *f = [plugin](Common::Wasm::Context* context, Word context_id, Word token, Word headers) {
-      SaveRestoreContext saved_context(context);
-      plugin->onGrpcCreateInitialMetadata(context_id.u64_, token.u64_, headers.u64_);
-    };
   } else if (function_name == "proxy_on_grpc_receive_initial_metadata") {
     *f = [plugin](Common::Wasm::Context* context, Word context_id, Word token, Word headers) {
       SaveRestoreContext saved_context(context);
@@ -429,9 +429,8 @@ void NullPlugin::onGrpcClose(uint64_t context_id, uint64_t token, uint64_t statu
   getRootContext(context_id)->onGrpcClose(token, static_cast<Plugin::GrpcStatus>(status_code));
 }
 
-void NullPlugin::onGrpcCreateInitialMetadata(uint64_t context_id, uint64_t token,
-                                             uint64_t headers) {
-  getRootContext(context_id)->onGrpcCreateInitialMetadata(token, headers);
+void NullPlugin::onGrpcCreateInitialMetadata(uint64_t context_id, uint64_t headers) {
+  getRootContext(context_id)->onGrpcCreateInitialMetadata(headers);
 }
 
 void NullPlugin::onGrpcReceiveInitialMetadata(uint64_t context_id, uint64_t token,

source/extensions/common/wasm/null/null_plugin.h

@@ -103,7 +103,7 @@ class NullPlugin : public NullVmPlugin {
 
   void onGrpcReceive(uint64_t context_id, uint64_t token, size_t body_size);
   void onGrpcClose(uint64_t context_id, uint64_t token, uint64_t status_code);
-  void onGrpcCreateInitialMetadata(uint64_t context_id, uint64_t token, uint64_t headers);
+  void onGrpcCreateInitialMetadata(uint64_t context_id, uint64_t headers);
   void onGrpcReceiveInitialMetadata(uint64_t context_id, uint64_t token, uint64_t headers);
   void onGrpcReceiveTrailingMetadata(uint64_t context_id, uint64_t token, uint64_t trailers);
 

source/extensions/common/wasm/wasm.h

@@ -230,7 +230,7 @@ class Wasm : public Logger::Loggable<Logger::Id::wasm>, public std::enable_share
 
   WasmCallVoid<3> on_grpc_receive_;
   WasmCallVoid<3> on_grpc_close_;
-  WasmCallVoid<3> on_grpc_create_initial_metadata_;
+  WasmCallVoid<2> on_grpc_create_initial_metadata_;
   WasmCallVoid<3> on_grpc_receive_initial_metadata_;
   WasmCallVoid<3> on_grpc_receive_trailing_metadata_;
 

test/extensions/filters/http/wasm/test_data/Makefile.docker_cpp_builder

@@ -1,6 +1,6 @@
 DOCKER_SDK=/external_sdk
 
-all: headers_cpp.wasm async_call_cpp.wasm metadata_cpp.wasm grpc_call_cpp.wasm shared_cpp.wasm queue_cpp.wasm http_callout_cpp.wasm grpc_callout_cpp.wasm
+all: headers_cpp.wasm async_call_cpp.wasm metadata_cpp.wasm grpc_call_cpp.wasm shared_cpp.wasm queue_cpp.wasm http_callout_cpp.wasm grpc_callout_cpp.wasm root_id_cpp.wasm
 	chown ${uid}.${gid} *.wasm
 
 include ${DOCKER_SDK}/Makefile.base_lite