From f537a1c3ea5acace844e42cda37477cf7f435413 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BE=84=E6=BD=AD?= <zty98751@alibaba-inc.com>
Date: Wed, 21 Feb 2024 14:00:43 +0800
Subject: [PATCH] fix: content-type not work in custom response plugin (#833)

---
 .../extensions/custom_response/plugin.cc      | 14 ++---
 .../extensions/custom_response/plugin.h       |  1 -
 .../extensions/custom_response/plugin_test.cc | 58 ++++++++++++++++++-
 3 files changed, 60 insertions(+), 13 deletions(-)

diff --git a/plugins/wasm-cpp/extensions/custom_response/plugin.cc b/plugins/wasm-cpp/extensions/custom_response/plugin.cc
index fc62034e3..1f420a1d2 100644
--- a/plugins/wasm-cpp/extensions/custom_response/plugin.cc
+++ b/plugins/wasm-cpp/extensions/custom_response/plugin.cc
@@ -74,9 +74,7 @@ bool PluginRootContext::parsePluginConfig(const json& configuration,
             }
             if (absl::AsciiStrToLower(pair[0]) ==
                 Wasm::Common::Http::Header::ContentType) {
-              rule.content_type = pair[1];
               has_content_type = true;
-              return true;
             }
             rule.headers.emplace_back(pair[0], pair[1]);
             return true;
@@ -105,13 +103,11 @@ bool PluginRootContext::parsePluginConfig(const json& configuration,
   if (!rule.body.empty() && !has_content_type) {
     auto try_decode_json = Wasm::Common::JsonParse(rule.body);
     if (try_decode_json.has_value()) {
-      rule.content_type = "application/json; charset=utf-8";
-      //      rule.headers.emplace_back(Wasm::Common::Http::Header::ContentType,
-      //                                "application/json; charset=utf-8");
+      rule.headers.emplace_back(Wasm::Common::Http::Header::ContentType,
+                                "application/json; charset=utf-8");
     } else {
-      rule.content_type = "text/plain; charset=utf-8";
-      //      rule.headers.emplace_back(Wasm::Common::Http::Header::ContentType,
-      //                                "text/plain; charset=utf-8");
+      rule.headers.emplace_back(Wasm::Common::Http::Header::ContentType,
+                                "text/plain; charset=utf-8");
     }
   }
   return true;
@@ -139,8 +135,6 @@ FilterHeadersStatus PluginRootContext::onResponse(
   if (!hit) {
     return FilterHeadersStatus::Continue;
   }
-  replaceResponseHeader(Wasm::Common::Http::Header::ContentType,
-                        rule.content_type);
   sendLocalResponse(rule.status_code, "", rule.body, rule.headers);
   return FilterHeadersStatus::StopIteration;
 }
diff --git a/plugins/wasm-cpp/extensions/custom_response/plugin.h b/plugins/wasm-cpp/extensions/custom_response/plugin.h
index b37af5a4b..e6b7a17bd 100644
--- a/plugins/wasm-cpp/extensions/custom_response/plugin.h
+++ b/plugins/wasm-cpp/extensions/custom_response/plugin.h
@@ -40,7 +40,6 @@ namespace custom_response {
 struct CustomResponseConfigRule {
   std::vector<std::string> enable_on_status;
   std::vector<std::pair<std::string, std::string>> headers;
-  std::string content_type;
   int32_t status_code = 200;
   std::string body;
 };
diff --git a/plugins/wasm-cpp/extensions/custom_response/plugin_test.cc b/plugins/wasm-cpp/extensions/custom_response/plugin_test.cc
index 8c5f06acb..025d4af68 100644
--- a/plugins/wasm-cpp/extensions/custom_response/plugin_test.cc
+++ b/plugins/wasm-cpp/extensions/custom_response/plugin_test.cc
@@ -146,12 +146,66 @@ TEST_F(CustomResponseTest, EnableOnStatus) {
   status_code_ = "429";
   EXPECT_EQ(context_->onRequestHeaders(0, false),
             FilterHeadersStatus::Continue);
-  EXPECT_CALL(*mock_context_, sendLocalResponse(233, testing::_, testing::_,
-                                                testing::_, testing::_));
+  EXPECT_CALL(
+      *mock_context_,
+      sendLocalResponse(
+          233, testing::_,
+          testing::ElementsAre(
+              testing::Pair("abc", "123"), testing::Pair("zty", "test"),
+              testing::Pair("content-type", "application/json; charset=utf-8")),
+          testing::_, testing::_));
   EXPECT_EQ(context_->onResponseHeaders(0, false),
             FilterHeadersStatus::StopIteration);
 }
 
+TEST_F(CustomResponseTest, ContentTypePlain) {
+  std::string configuration = R"(
+{
+   "status_code": 200,
+   "body": "abc"
+})";
+
+  BufferBase buffer;
+  buffer.set({configuration.data(), configuration.size()});
+
+  EXPECT_CALL(*mock_context_, getBuffer(WasmBufferType::PluginConfiguration))
+      .WillOnce([&buffer](WasmBufferType) { return &buffer; });
+  EXPECT_TRUE(root_context_->configure(configuration.size()));
+
+  EXPECT_CALL(
+      *mock_context_,
+      sendLocalResponse(200, testing::_,
+                        testing::ElementsAre(testing::Pair(
+                            "content-type", "text/plain; charset=utf-8")),
+                        testing::_, testing::_));
+  EXPECT_EQ(context_->onRequestHeaders(0, false),
+            FilterHeadersStatus::StopIteration);
+}
+
+TEST_F(CustomResponseTest, ContentTypeCustom) {
+  std::string configuration = R"(
+{
+   "status_code": 200,
+   "headers": ["content-type=application/custom"],
+   "body": "abc"
+})";
+
+  BufferBase buffer;
+  buffer.set({configuration.data(), configuration.size()});
+
+  EXPECT_CALL(*mock_context_, getBuffer(WasmBufferType::PluginConfiguration))
+      .WillOnce([&buffer](WasmBufferType) { return &buffer; });
+  EXPECT_TRUE(root_context_->configure(configuration.size()));
+
+  EXPECT_CALL(*mock_context_,
+              sendLocalResponse(200, testing::_,
+                                testing::ElementsAre(testing::Pair(
+                                    "content-type", "application/custom")),
+                                testing::_, testing::_));
+  EXPECT_EQ(context_->onRequestHeaders(0, false),
+            FilterHeadersStatus::StopIteration);
+}
+
 TEST_F(CustomResponseTest, NoGlobalRule) {
   std::string configuration = R"(
 {