From ec96b58244665969eb3a4d5ca34b07b022360749 Mon Sep 17 00:00:00 2001
From: Nupur Garg <37600866+gargnupur@users.noreply.github.com>
Date: Thu, 12 Mar 2020 15:07:36 -0700
Subject: [PATCH] [Release-1.5] Fix memory leak in TCP (#187)

* Fix memory leak in TCP

Signed-off-by: gargnupur <gargnupur@google.com>

* fix lint

Signed-off-by: gargnupur <gargnupur@google.com>
---
 source/extensions/common/wasm/context.cc | 21 ++++++++++++++++++++-
 source/extensions/common/wasm/context.h  |  9 ++++++++-
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/source/extensions/common/wasm/context.cc b/source/extensions/common/wasm/context.cc
index 02e040066e28..b460569a6d33 100644
--- a/source/extensions/common/wasm/context.cc
+++ b/source/extensions/common/wasm/context.cc
@@ -269,6 +269,16 @@ WasmResult Context::setTickPeriod(std::chrono::milliseconds tick_period) {
   return WasmResult::Ok;
 }
 
+void Context::onCloseTCP() {
+  if (tcp_connection_closed_) {
+    return;
+  }
+  tcp_connection_closed_ = true;
+  onDone();
+  onLog();
+  onDelete();
+}
+
 uint64_t Context::getCurrentTimeNanoseconds() {
   return std::chrono::duration_cast<std::chrono::nanoseconds>(
              wasm_->time_source_.systemTime().time_since_epoch())
@@ -431,7 +441,6 @@ Context::FindValue(absl::string_view name, Protobuf::Arena* arena) const {
     break;
   case PropertyToken::FILTER_STATE:
     if (info) {
-
       return CelValue::CreateMap(Protobuf::Arena::Create<WasmStateWrapper>(
           arena, info->filterState(), info->upstreamFilterState().get()));
     }
@@ -1204,6 +1213,12 @@ void Context::onDownstreamConnectionClose(PeerType peer_type) {
     DeferAfterCallActions actions(this);
     wasm_->on_downstream_connection_close_(this, id_, static_cast<uint32_t>(peer_type));
   }
+  downstream_closed_ = true;
+  // Call close on TCP connection, if upstream connection closed or there was a failure seen in this
+  // connection.
+  if (upstream_closed_ || getRequestStreamInfo()->hasAnyResponseFlag()) {
+    onCloseTCP();
+  }
 }
 
 void Context::onUpstreamConnectionClose(PeerType peer_type) {
@@ -1211,6 +1226,10 @@ void Context::onUpstreamConnectionClose(PeerType peer_type) {
     DeferAfterCallActions actions(this);
     wasm_->on_upstream_connection_close_(this, id_, static_cast<uint32_t>(peer_type));
   }
+  upstream_closed_ = true;
+  if (downstream_closed_) {
+    onCloseTCP();
+  }
 }
 
 // Empty headers/trailers have zero size.
diff --git a/source/extensions/common/wasm/context.h b/source/extensions/common/wasm/context.h
index 08a69431778f..0a1847ed0ae6 100644
--- a/source/extensions/common/wasm/context.h
+++ b/source/extensions/common/wasm/context.h
@@ -334,6 +334,8 @@ class Context : public Logger::Loggable<Logger::Id::wasm>,
 protected:
   friend class Wasm;
 
+  void onCloseTCP();
+
   struct AsyncClientHandler : public Http::AsyncClient::Callbacks {
     // Http::AsyncClient::Callbacks
     void onSuccess(Envoy::Http::MessagePtr&& response) override {
@@ -487,8 +489,13 @@ class Context : public Logger::Loggable<Logger::Id::wasm>,
   std::map<uint32_t, GrpcCallClientHandler> grpc_call_request_;
   std::map<uint32_t, GrpcStreamClientHandler> grpc_stream_;
 
-  // Opaque state
+  // Opaque state.
   absl::flat_hash_map<std::string, std::unique_ptr<StorageObject>> data_storage_;
+
+  // TCP State.
+  bool upstream_closed_ = false;
+  bool downstream_closed_ = false;
+  bool tcp_connection_closed_ = false;
 };
 
 using ContextSharedPtr = std::shared_ptr<Context>;