1 files changed, 96 insertions, 7 deletions

diff --git a/src/httpaction.cc b/src/httpaction.cc
index 2cce21d..a4dd07e 100644
--- a/src/httpaction.cc
+++ b/src/httpaction.cc
@@ -405,16 +405,17 @@ private:
 
     /* Scan and Replace */
     bool __scan_and_replace(enum edit_zone, std::string & raw);
+    bool __is_replaced{false};
 };
 
 std::array<std::string, HttpActionReplace::kZoneMax> HttpActionReplace::map_edit_zone_to_str =
-{
-    "http_req_uri",
-    "http_req_headers",
-    "http_req_body",
-    "http_resp_headers",
-    "http_resp_body"
-};
+    {
+        "http_req_uri",
+        "http_req_headers",
+        "http_req_body",
+        "http_resp_headers",
+        "http_resp_body"
+    };
 
 void HttpActionReplace::Construct(const std::string & str_kv)
 {
@@ -572,11 +573,99 @@ void HttpActionReplace::__on_request_body(HttpSession * session)
 
 void HttpActionReplace::__on_response_header(HttpSession * session)
 {
+    auto & response = session->response();
+
+    /* 对Headers替换 */
+    response.Headers().ForEachHeader([this, &response]
+        (const std::string & str_field, const std::string str_value) -> bool
+    {
+        /* Field, Value组合成字符串，整体调用正则扫描 */
+        std::string __combine_header = str_field + ":" + str_value;
+
+        /* 没有命中，继续调用本函数处理后面的Headers */
+        if (!__scan_and_replace(kZoneResponseHeader, __combine_header))
+            return true;
+
+        /* 替换标志 */
+        __is_replaced = true;
+
+        /* 命中，若替换后的长度为0，删除该头部 */
+        if (__combine_header.length() == 0)
+        {
+            response.Headers().Remove(str_field);
+            return true;
+        }
+
+        /* 按分号拆分成Field和Value */
+        auto first_comma_pos = __combine_header.find_first_of(':');
+        if (first_comma_pos == std::string::npos)
+        {
+            throw std::runtime_error("Invalid regex replacement for http headers, no comma found.");
+        }
+
+        auto __replaced_field = __combine_header.substr(0, first_comma_pos);
+        auto __replaced_value = __combine_header.substr(first_comma_pos + 1);
+
+        response.Headers().Set(__replaced_field, __replaced_value);
+        return true;
+    });
+
     return;
 }
 
 void HttpActionReplace::__on_response_body(HttpSession * session)
 {
+    auto & response = session->response();
+    auto resp_bodys = response.StolenBody();
+
+    size_t total_body_bytes = 0;
+    std::vector<HttpResponse::body_content_ptr_t> body_replaced_segments;
+
+    for (auto & body_raw_ptr : resp_bodys)
+    {
+        std::string str_body_raw = std::string(body_raw_ptr->begin(), body_raw_ptr->end());
+        std::unique_ptr<std::vector<char>> body_replaced_ptr;
+
+        if (__scan_and_replace(kZoneResponseBody, str_body_raw))
+        {
+            body_replaced_ptr = std::make_unique<std::vector<char>>();
+            body_replaced_ptr->insert(body_replaced_ptr->end(), str_body_raw.begin(), str_body_raw.end());
+            __is_replaced = true;
+        }
+        else
+        {
+            body_replaced_ptr = std::move(body_raw_ptr);
+        }
+
+        total_body_bytes += body_replaced_ptr->size();
+        body_replaced_segments.push_back(std::move(body_replaced_ptr));
+    }
+
+    response.Body(std::move(body_replaced_segments));
+
+    /* 发生替换后，重新构造应答内容 */
+    if (__is_replaced)
+    {
+        bool is_content_length_set = false;
+
+        /* 原始头部是否有Content-Length字段 */
+        response.cHeaders().ForEachValueOfHeader("Content-Length", [&is_content_length_set](
+            const std::string _, const std::string __)
+        {
+            is_content_length_set = true;
+            return false;
+        });
+
+        /* 如果有，需要根据替换后的内容校正该数值 */
+        if (is_content_length_set)
+        {
+            response.Headers().Set("Content-Length", std::to_string(total_body_bytes));
+        }
+
+        /* 构建应答内容 */
+        response.Construct();
+    }
+
     return;
 }