remove headers in reverse proxy (#653)

include/cinatra/coro_http_server.hpp

@@ -855,6 +855,43 @@ class coro_http_server {
     co_return true;
   }
 
+  template <class T, class Pred>
+  size_t erase_if(std::span<T> &sp, Pred p) {
+    auto it = std::remove_if(sp.begin(), sp.end(), p);
+    size_t count = sp.end() - it;
+    sp = std::span<T>(sp.begin(), it);
+    return count;
+  }
+
+  int remove_result_headers(resp_data &result, std::string_view value) {
+    bool r = false;
+    return erase_if(result.resp_headers, [&](http_header &header) {
+      if (r) {
+        return false;
+      }
+
+      r = (header.value.find(value) != std::string_view::npos);
+
+      return r;
+    });
+  }
+
+  void handle_response_header(resp_data &result, std::string &length) {
+    int r = remove_result_headers(result, "chunked");
+    if (r == 0) {
+      r = remove_result_headers(result, "multipart/form-data");
+      if (r) {
+        length = std::to_string(result.resp_body.size());
+        for (auto &[key, val] : result.resp_headers) {
+          if (key == "Content-Length") {
+            val = length;
+            break;
+          }
+        }
+      }
+    }
+  }
+
   async_simple::coro::Lazy<void> reply(coro_http_client &client,
                                        std::string_view host,
                                        coro_http_request &req,
@@ -880,6 +917,8 @@ class coro_http_server {
         req.full_url(), method_type(req.get_method()), std::move(ctx),
         std::move(req_headers));
 
+    std::string length;
+    handle_response_header(result, length);
     response.add_header_span(result.resp_headers);
 
     response.set_status_and_content_view(

tests/test_coro_http_server.cpp

@@ -210,7 +210,27 @@ TEST_CASE("test multiple download") {
           co_return;
         }
 
-        std::vector<std::string> vec{"hello", " world", " ok"};
+        std::vector<std::string> vec{"hello", " world", " chunked"};
+
+        for (auto &str : vec) {
+          if (ok = co_await resp.get_conn()->write_multipart(str, "text/plain");
+              !ok) {
+            co_return;
+          }
+        }
+
+        ok = co_await resp.get_conn()->end_multipart();
+      });
+  server.set_http_handler<GET>(
+      "/multipart",
+      [](coro_http_request &req,
+         coro_http_response &resp) -> async_simple::coro::Lazy<void> {
+        bool ok;
+        if (ok = co_await resp.get_conn()->begin_multipart(); !ok) {
+          co_return;
+        }
+
+        std::vector<std::string> vec{"hello", " world", " multipart"};
 
         for (auto &str : vec) {
           if (ok = co_await resp.get_conn()->write_multipart(str, "text/plain");
@@ -227,7 +247,10 @@ TEST_CASE("test multiple download") {
   coro_http_client client{};
   auto result = client.get("http://127.0.0.1:9001/");
   CHECK(result.status == 200);
-  CHECK(result.resp_body == "hello world ok");
+  CHECK(result.resp_body == "hello world chunked");
+  result = client.get("http://127.0.0.1:9001/multipart");
+  CHECK(result.status == 200);
+  CHECK(result.resp_body == "hello world multipart");
 }
 
 TEST_CASE("test range download") {
@@ -1515,6 +1538,74 @@ TEST_CASE("test reverse proxy") {
   CHECK(!resp_random.resp_body.empty());
 }
 
+TEST_CASE("test reverse proxy download") {
+  cinatra::coro_http_server server(1, 9001);
+  server.set_http_handler<cinatra::GET, cinatra::POST>(
+      "/test_chunked",
+      [](coro_http_request &req,
+         coro_http_response &resp) -> async_simple::coro::Lazy<void> {
+        resp.set_format_type(format_type::chunked);
+        bool ok;
+        if (ok = co_await resp.get_conn()->begin_chunked(); !ok) {
+          co_return;
+        }
+
+        std::vector<std::string> vec{"hello", " world", " ok"};
+
+        for (auto &str : vec) {
+          if (ok = co_await resp.get_conn()->write_chunked(str); !ok) {
+            co_return;
+          }
+        }
+
+        ok = co_await resp.get_conn()->end_chunked();
+      });
+  server.set_http_handler<cinatra::GET, cinatra::POST>(
+      "/test", [](coro_http_request &req, coro_http_response &resp) {
+        resp.set_status_and_content(status_type::ok, "hello world");
+      });
+  server.set_http_handler<GET>(
+      "/test_multipart",
+      [](coro_http_request &req,
+         coro_http_response &resp) -> async_simple::coro::Lazy<void> {
+        bool ok;
+        if (ok = co_await resp.get_conn()->begin_multipart(); !ok) {
+          co_return;
+        }
+
+        std::vector<std::string> vec{"hello", " world", " multipart"};
+
+        for (auto &str : vec) {
+          if (ok = co_await resp.get_conn()->write_multipart(str, "text/plain");
+              !ok) {
+            co_return;
+          }
+        }
+
+        ok = co_await resp.get_conn()->end_multipart();
+      });
+  server.async_start();
+
+  coro_http_server proxy_rr(2, 8001);
+  proxy_rr.set_http_proxy_handler<GET, POST>(
+      "/([^]+)", {"127.0.0.1:9001"}, coro_io::load_blance_algorithm::RR);
+  proxy_rr.async_start();
+
+  coro_http_client client{};
+  auto result = client.get("http://127.0.0.1:8001/test");
+  CHECK(result.resp_body == "hello world");
+
+  result = client.get("http://127.0.0.1:8001/test_chunked");
+  CHECK(result.status == 200);
+  CHECK(result.resp_body == "hello world ok");
+
+  coro_http_client client1{};
+  result = client1.get("http://127.0.0.1:8001/test_multipart");
+  std::cout << result.net_err.message() << std::endl;
+  CHECK(result.status == 200);
+  CHECK(result.resp_body == "hello world multipart");
+}
+
 TEST_CASE("test reverse proxy websocket") {
   {
     coro_http_server proxy_server(1, 9005);