envoyproxy · dio · Sep 22, 2020 · Sep 18, 2020 · Sep 18, 2020 · Sep 18, 2020
diff --git a/docs/root/configuration/http/http_filters/lua_filter.rst b/docs/root/configuration/http/http_filters/lua_filter.rst
@@ -89,10 +89,10 @@ on the virtual host, route, or weighted cluster.
 
 LuaPerRoute provides two ways of overriding the `GLOBAL` Lua script:
 
-* By providing a name reference to the defined :ref:`named Lua source codes map 
+* By providing a name reference to the defined :ref:`named Lua source codes map
   <envoy_v3_api_field_extensions.filters.http.lua.v3.Lua.source_codes>`.
-* By providing inline :ref:`source code 
-  <envoy_v3_api_field_extensions.filters.http.lua.v3.LuaPerRoute.source_code>` (This allows the 
+* By providing inline :ref:`source code
+  <envoy_v3_api_field_extensions.filters.http.lua.v3.LuaPerRoute.source_code>` (This allows the
   code to be sent through RDS).
 
 As a concrete example, given the following Lua filter configuration:
@@ -143,7 +143,7 @@ The ``GLOBAL`` Lua script will be overridden by the referenced script:
   <envoy_v3_api_field_extensions.filters.http.lua.v3.Lua.inline_code>`. Therefore, do not use
   ``GLOBAL`` as name for other Lua scripts.
 
-Or we can define a new Lua script in the LuaPerRoute configuration directly to override the `GLOBAL`  
+Or we can define a new Lua script in the LuaPerRoute configuration directly to override the `GLOBAL`
 Lua script as follows:
 
 .. code-block:: yaml
@@ -237,6 +237,40 @@ more details on the supported API.
     response_handle:logInfo("Status: "..response_handle:headers():get(":status"))
   end
 
+A common use-case is to rewrite upstream response body, for example: an upstream sends non-2xx
+response with JSON data, but the application requires HTML page to be sent to browsers.
+
+There are two ways of doing this, the first one is via the `body()` API.
+
+.. code-block:: lua
+
+    function envoy_on_response(response_handle)
+      local content_length = response_handle:body():setBytes("<html><b>Not Found<b></html>")
+      response_handle:headers():replace("content-length", content_length)
+      response_handle:headers():replace("content-type", "text/html")
+    end
+
+
+Or, through `bodyChunks()` API, which let Envoy to skip buffering the upstream response data.
+
+.. code-block:: lua
+
+    function envoy_on_response(response_handle)
+
+      -- Sets the content-length.
+      response_handle:headers():replace("content-length", 28)
+      response_handle:headers():replace("content-type", "text/html")
+
+      local last
+      for chunk in response_handle:bodyChunks() do
+        -- Clears each received chunk.
+        chunk:setBytes("")
+        last = chunk
+      end
+
+      last:setBytes("<html><b>Not Found<b></html>")
+    end
+
 .. _config_http_filters_lua_stream_handle_api:
 
 Complete example
@@ -556,6 +590,17 @@ cause a buffer segment to be copied. *index* is an integer and supplies the buff
 copy. *length* is an integer and supplies the buffer length to copy. *index* + *length* must be
 less than the buffer length.
 
+.. _config_http_filters_lua_buffer_wrapper_api_set_bytes:
+
+setBytes()
+^^^^^^^^^^
+
+.. code-block:: lua
+
+  buffer:setBytes(string)
+
+Set the content of wrapped buffer with the input string.
+
 .. _config_http_filters_lua_metadata_wrapper:
 
 Metadata object API

diff --git a/docs/root/version_history/current.rst b/docs/root/version_history/current.rst
@@ -97,6 +97,7 @@ New Features
 * local_reply config: added :ref:`content_type<envoy_v3_api_field_config.core.v3.SubstitutionFormatString.content_type>` field to set content-type.
 * lua: added Lua APIs to access :ref:`SSL connection info <config_http_filters_lua_ssl_socket_info>` object.
 * lua: added Lua API for :ref:`base64 escaping a string <config_http_filters_lua_stream_handle_api_base64_escape>`.
+* lua: added Lua API for :ref:`setting the current buffer content <config_http_filters_lua_buffer_wrapper_api_set_bytes>`.
 * lua: added new :ref:`source_code <envoy_v3_api_field_extensions.filters.http.lua.v3.LuaPerRoute.source_code>` field to support the dispatching of inline Lua code in per route configuration of Lua filter.
 * overload management: add :ref:`scaling <envoy_v3_api_field_config.overload.v3.Trigger.scaled>` trigger for OverloadManager actions.
 * postgres network filter: :ref:`metadata <config_network_filters_postgres_proxy_dynamic_metadata>` is produced based on SQL query.

diff --git a/source/extensions/filters/common/lua/wrappers.cc b/source/extensions/filters/common/lua/wrappers.cc
@@ -65,6 +65,14 @@ int BufferWrapper::luaGetBytes(lua_State* state) {
   return 1;
 }
 
+int BufferWrapper::luaSetBytes(lua_State* state) {
+  data_.drain(data_.length());
+  absl::string_view bytes = luaL_checkstring(state, 2);
+  data_.add(bytes);
+  lua_pushnumber(state, data_.length());
+  return 1;
+}
+
 void MetadataMapHelper::setValue(lua_State* state, const ProtobufWkt::Value& value) {
   ProtobufWkt::Value::KindCase kind = value.kind_case();
 

diff --git a/source/extensions/filters/common/lua/wrappers.h b/source/extensions/filters/common/lua/wrappers.h
@@ -13,14 +13,16 @@ namespace Common {
 namespace Lua {
 
 /**
- * A wrapper for a constant buffer which cannot be modified by Lua.
+ * A wrapper for a buffer.
  */
 class BufferWrapper : public BaseLuaObject<BufferWrapper> {
 public:
-  BufferWrapper(const Buffer::Instance& data) : data_(data) {}
+  BufferWrapper(Buffer::Instance& data) : data_(data) {}
 
   static ExportedFunctions exportedFunctions() {
-    return {{"length", static_luaLength}, {"getBytes", static_luaGetBytes}};
+    return {{"length", static_luaLength},
+            {"getBytes", static_luaGetBytes},
+            {"setBytes", static_luaSetBytes}};
   }
 
 private:
@@ -37,7 +39,14 @@ class BufferWrapper : public BaseLuaObject<BufferWrapper> {
    */
   DECLARE_LUA_FUNCTION(BufferWrapper, luaGetBytes);
 
-  const Buffer::Instance& data_;
+  /**
+   * Set the wrapped data with the input string.
+   * @param 1 (string) input string.
+   * @return int the length of the input string.
+   */
+  DECLARE_LUA_FUNCTION(BufferWrapper, luaSetBytes);
+
+  Buffer::Instance& data_;
 };
 
 class MetadataMapWrapper;

diff --git a/source/extensions/filters/http/lua/lua_filter.cc b/source/extensions/filters/http/lua/lua_filter.cc
@@ -428,8 +428,9 @@ int StreamHandleWrapper::luaBody(lua_State* state) {
       if (body_wrapper_.get() != nullptr) {
         body_wrapper_.pushStack();
       } else {
-        body_wrapper_.reset(
-            Filters::Common::Lua::BufferWrapper::create(state, *callbacks_.bufferedBody()), true);
+        body_wrapper_.reset(Filters::Common::Lua::BufferWrapper::create(
+                                state, const_cast<Buffer::Instance&>(*callbacks_.bufferedBody())),
+                            true);
       }
       return 1;
     }

diff --git a/test/extensions/filters/common/lua/wrappers_test.cc b/test/extensions/filters/common/lua/wrappers_test.cc
@@ -76,6 +76,8 @@ TEST_F(LuaBufferWrapperTest, Methods) {
       testPrint(object:length())
       testPrint(object:getBytes(0, 2))
       testPrint(object:getBytes(6, 5))
+      testPrint(object:setBytes("neverland"))
+      testPrint(object:getBytes(0, 5))
     end
   )EOF"};
 
@@ -85,6 +87,8 @@ TEST_F(LuaBufferWrapperTest, Methods) {
   EXPECT_CALL(printer_, testPrint("11"));
   EXPECT_CALL(printer_, testPrint("he"));
   EXPECT_CALL(printer_, testPrint("world"));
+  EXPECT_CALL(printer_, testPrint("9"));
+  EXPECT_CALL(printer_, testPrint("never"));
   start("callMe");
 }
 

diff --git a/test/extensions/filters/http/lua/lua_filter_test.cc b/test/extensions/filters/http/lua/lua_filter_test.cc
@@ -2247,6 +2247,58 @@ TEST_F(LuaHttpFilterTest, LuaFilterBase64Escape) {
   EXPECT_EQ(Http::FilterDataStatus::Continue, filter_->encodeData(response_body, true));
 }
 
+TEST_F(LuaHttpFilterTest, LuaFilterSetResponseBuffer) {
+  const std::string SCRIPT{R"EOF(
+    function envoy_on_response(response_handle)
+      local content_length = response_handle:body():setBytes("1234")
+      response_handle:logTrace(content_length)
+
+      -- It is possible to replace an entry in headers after overridding encoding buffer.
+      response_handle:headers():replace("content-length", content_length)
+    end
+  )EOF"};
+
+  InSequence s;
+  setup(SCRIPT);
+
+  Http::TestRequestHeaderMapImpl request_headers{{":path", "/"}};
+  EXPECT_EQ(Http::FilterHeadersStatus::Continue, filter_->decodeHeaders(request_headers, true));
+
+  Http::TestResponseHeaderMapImpl response_headers{{":status", "200"}};
+  EXPECT_EQ(Http::FilterHeadersStatus::StopIteration,
+            filter_->encodeHeaders(response_headers, false));
+  Buffer::OwnedImpl response_body("1234567890");
+  EXPECT_CALL(*filter_, scriptLog(spdlog::level::trace, StrEq("4")));
+  EXPECT_EQ(Http::FilterDataStatus::Continue, filter_->encodeData(response_body, true));
+  EXPECT_EQ(4, encoder_callbacks_.buffer_->length());
+}
+
+TEST_F(LuaHttpFilterTest, LuaFilterSetResponseBufferChunked) {
+  const std::string SCRIPT{R"EOF(
+    function envoy_on_response(response_handle)
+      local last
+      for chunk in response_handle:bodyChunks() do
+        chunk:setBytes("")
+        last = chunk
+      end
+      response_handle:logTrace(last:setBytes("1234"))
+    end
+  )EOF"};
+
+  InSequence s;
+  setup(SCRIPT);
+
+  Http::TestRequestHeaderMapImpl request_headers{{":path", "/"}};
+  EXPECT_EQ(Http::FilterHeadersStatus::Continue, filter_->decodeHeaders(request_headers, true));
+
+  Http::TestResponseHeaderMapImpl response_headers{{":status", "200"}};
+  EXPECT_EQ(Http::FilterHeadersStatus::Continue, filter_->encodeHeaders(response_headers, false));
+
+  Buffer::OwnedImpl response_body("1234567890");
+  EXPECT_CALL(*filter_, scriptLog(spdlog::level::trace, StrEq("4")));
+  EXPECT_EQ(Http::FilterDataStatus::Continue, filter_->encodeData(response_body, true));
+}
+
 } // namespace
 } // namespace Lua
 } // namespace HttpFilters

diff --git a/test/extensions/filters/http/lua/lua_integration_test.cc b/test/extensions/filters/http/lua/lua_integration_test.cc
@@ -146,6 +146,39 @@ class LuaIntegrationTest : public testing::TestWithParam<Network::Address::IpVer
     registerTestServerPorts({"http"});
   }
 
+  void testRewriteResponse(const std::string& code) {
+    initializeFilter(code);
+    codec_client_ = makeHttpConnection(makeClientConnection(lookupPort("http")));
+    Http::TestRequestHeaderMapImpl request_headers{{":method", "POST"},
+                                                   {":path", "/test/long/url"},
+                                                   {":scheme", "http"},
+                                                   {":authority", "host"},
+                                                   {"x-forwarded-for", "10.0.0.1"}};
+
+    auto encoder_decoder = codec_client_->startRequest(request_headers);
+    Http::StreamEncoder& encoder = encoder_decoder.first;
+    auto response = std::move(encoder_decoder.second);
+    Buffer::OwnedImpl request_data("done");
+    encoder.encodeData(request_data, true);
+
+    waitForNextUpstreamRequest();
+
+    Http::TestResponseHeaderMapImpl response_headers{{":status", "200"}, {"foo", "bar"}};
+    upstream_request_->encodeHeaders(response_headers, false);
+    Buffer::OwnedImpl response_data1("good");
+    upstream_request_->encodeData(response_data1, false);
+    Buffer::OwnedImpl response_data2("bye");
+    upstream_request_->encodeData(response_data2, true);
+
+    response->waitForEndStream();
+
+    EXPECT_EQ(
+        "2",
+        response->headers().get(Http::LowerCaseString("content-length"))->value().getStringView());
+    EXPECT_EQ("ok", response->body());
+    cleanup();
+  }
+
   void cleanup() {
     codec_client_->close();
     if (fake_lua_connection_ != nullptr) {
@@ -900,5 +933,46 @@ TEST_P(LuaIntegrationTest, RdsTestOfLuaPerRoute) {
 #endif
 }
 
+// Rewrite response buffer.
+TEST_P(LuaIntegrationTest, RewriteResponseBuffer) {
+  const std::string FILTER_AND_CODE =
+      R"EOF(
+name: lua
+typed_config:
+  "@type": type.googleapis.com/envoy.config.filter.http.lua.v2.Lua
+  inline_code: |
+    function envoy_on_response(response_handle)
+      local content_length = response_handle:body():setBytes("ok")
+      response_handle:logTrace(content_length)
+
+      response_handle:headers():replace("content-length", content_length)
+    end
+)EOF";
+
+  testRewriteResponse(FILTER_AND_CODE);
+}
+
+// Rewrite chunked response body.
+TEST_P(LuaIntegrationTest, RewriteChunkedBody) {
+  const std::string FILTER_AND_CODE =
+      R"EOF(
+name: lua
+typed_config:
+  "@type": type.googleapis.com/envoy.config.filter.http.lua.v2.Lua
+  inline_code: |
+    function envoy_on_response(response_handle)
+      response_handle:headers():replace("content-length", 2)
+      local last
+      for chunk in response_handle:bodyChunks() do
+        chunk:setBytes("")
+        last = chunk
+      end
+      last:setBytes("ok")
+    end
+)EOF";
+
+  testRewriteResponse(FILTER_AND_CODE);
+}
+
 } // namespace
 } // namespace Envoy