From 793809cecd4894b31aa2da112a0d96df8e0f45a9 Mon Sep 17 00:00:00 2001
From: Kuat <kyessenov@users.noreply.github.com>
Date: Wed, 29 Jan 2020 15:18:35 -0800
Subject: [PATCH] Cherry-pick of https://github.com/envoyproxy/envoy/pull/9862

Signed-off-by: Kuat Yessenov <kuat@google.com>
---
 .../extensions/filters/common/expr/context.cc | 10 ++++---
 .../filters/common/expr/context_test.cc       | 26 +++++++++++++++++++
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/source/extensions/filters/common/expr/context.cc b/source/extensions/filters/common/expr/context.cc
index 8477662815e07..befd2de5af239 100644
--- a/source/extensions/filters/common/expr/context.cc
+++ b/source/extensions/filters/common/expr/context.cc
@@ -77,6 +77,9 @@ absl::optional<CelValue> RequestWrapper::operator[](CelValue key) const {
     } else {
       return CelValue::CreateInt64(info_.bytesReceived());
     }
+  } else if (value == TotalSize) {
+    return CelValue::CreateInt64(info_.bytesReceived() +
+                                 (headers_.value_ ? headers_.value_->byteSize().value() : 0));
   } else if (value == Duration) {
     auto duration = info_.requestComplete();
     if (duration.has_value()) {
@@ -106,8 +109,6 @@ absl::optional<CelValue> RequestWrapper::operator[](CelValue key) const {
       return convertHeaderEntry(headers_.value_->RequestId());
     } else if (value == UserAgent) {
       return convertHeaderEntry(headers_.value_->UserAgent());
-    } else if (value == TotalSize) {
-      return CelValue::CreateInt64(info_.bytesReceived() + headers_.value_->byteSize().value());
     }
   }
   return {};
@@ -132,8 +133,9 @@ absl::optional<CelValue> ResponseWrapper::operator[](CelValue key) const {
   } else if (value == Flags) {
     return CelValue::CreateInt64(info_.responseFlags());
   } else if (value == TotalSize) {
-    return CelValue::CreateInt64(info_.bytesSent() + headers_.value_->byteSize().value() +
-                                 trailers_.value_->byteSize().value());
+    return CelValue::CreateInt64(info_.bytesSent() +
+                                 (headers_.value_ ? headers_.value_->byteSize().value() : 0) +
+                                 (trailers_.value_ ? trailers_.value_->byteSize().value() : 0));
   }
   return {};
 }
diff --git a/test/extensions/filters/common/expr/context_test.cc b/test/extensions/filters/common/expr/context_test.cc
index db5850efd65f0..75971170c0de0 100644
--- a/test/extensions/filters/common/expr/context_test.cc
+++ b/test/extensions/filters/common/expr/context_test.cc
@@ -32,12 +32,14 @@ TEST(Context, EmptyHeadersAttributes) {
 
 TEST(Context, RequestAttributes) {
   NiceMock<StreamInfo::MockStreamInfo> info;
+  NiceMock<StreamInfo::MockStreamInfo> empty_info;
   Http::TestHeaderMapImpl header_map{
       {":method", "POST"},           {":scheme", "http"},      {":path", "/meow?yes=1"},
       {":authority", "kittens.com"}, {"referer", "dogs.com"},  {"user-agent", "envoy-mobile"},
       {"content-length", "10"},      {"x-request-id", "blah"},
   };
   RequestWrapper request(&header_map, info);
+  RequestWrapper empty_request(nullptr, empty_info);
 
   EXPECT_CALL(info, bytesReceived()).WillRepeatedly(Return(10));
   // "2018-04-03T23:06:09.123Z".
@@ -66,6 +68,12 @@ TEST(Context, RequestAttributes) {
     ASSERT_TRUE(value.value().IsString());
     EXPECT_EQ("http", value.value().StringOrDie().value());
   }
+
+  {
+    auto value = empty_request[CelValue::CreateString(Scheme)];
+    EXPECT_FALSE(value.has_value());
+  }
+
   {
     auto value = request[CelValue::CreateString(Host)];
     EXPECT_TRUE(value.has_value());
@@ -130,6 +138,14 @@ TEST(Context, RequestAttributes) {
     EXPECT_EQ(138, value.value().Int64OrDie());
   }
 
+  {
+    auto value = empty_request[CelValue::CreateString(TotalSize)];
+    EXPECT_TRUE(value.has_value());
+    ASSERT_TRUE(value.value().IsInt64());
+    // this includes the headers size
+    EXPECT_EQ(0, value.value().Int64OrDie());
+  }
+
   {
     auto value = request[CelValue::CreateString(Time)];
     EXPECT_TRUE(value.has_value());
@@ -187,11 +203,13 @@ TEST(Context, RequestFallbackAttributes) {
 
 TEST(Context, ResponseAttributes) {
   NiceMock<StreamInfo::MockStreamInfo> info;
+  NiceMock<StreamInfo::MockStreamInfo> empty_info;
   const std::string header_name = "test-header";
   const std::string trailer_name = "test-trailer";
   Http::TestHeaderMapImpl header_map{{header_name, "a"}};
   Http::TestHeaderMapImpl trailer_map{{trailer_name, "b"}};
   ResponseWrapper response(&header_map, &trailer_map, info);
+  ResponseWrapper empty_response(nullptr, nullptr, empty_info);
 
   EXPECT_CALL(info, responseCode()).WillRepeatedly(Return(404));
   EXPECT_CALL(info, bytesSent()).WillRepeatedly(Return(123));
@@ -222,6 +240,13 @@ TEST(Context, ResponseAttributes) {
   }
 
 
+  {
+    auto value = empty_response[CelValue::CreateString(TotalSize)];
+    EXPECT_TRUE(value.has_value());
+    ASSERT_TRUE(value.value().IsInt64());
+    EXPECT_EQ(0, value.value().Int64OrDie());
+  }
+
   {
     auto value = response[CelValue::CreateString(Code)];
     EXPECT_TRUE(value.has_value());
@@ -259,6 +284,7 @@ TEST(Context, ResponseAttributes) {
     ASSERT_TRUE(header.value().IsString());
     EXPECT_EQ("b", header.value().StringOrDie().value());
   }
+
   {
     auto value = response[CelValue::CreateString(Flags)];
     EXPECT_TRUE(value.has_value());