Add optional user defined header writer (#1683)

* Add optional user defined header writer

* Fix errors and add test

httplib.h

@@ -8,7 +8,7 @@
 #ifndef CPPHTTPLIB_HTTPLIB_H
 #define CPPHTTPLIB_HTTPLIB_H
 
-#define CPPHTTPLIB_VERSION "0.13.3"
+#define CPPHTTPLIB_VERSION "0.14.0"
 
 /*
  * Configuration
@@ -487,8 +487,7 @@ struct Request {
 
   bool has_header(const std::string &key) const;
   std::string get_header_value(const std::string &key, size_t id = 0) const;
-  template <typename T>
-  T get_header_value(const std::string &key, size_t id = 0) const;
+  uint64_t get_header_value_u64(const std::string &key, size_t id = 0) const;
   size_t get_header_value_count(const std::string &key) const;
   void set_header(const std::string &key, const std::string &val);
 
@@ -520,8 +519,7 @@ struct Response {
 
   bool has_header(const std::string &key) const;
   std::string get_header_value(const std::string &key, size_t id = 0) const;
-  template <typename T>
-  T get_header_value(const std::string &key, size_t id = 0) const;
+  uint64_t get_header_value_u64(const std::string &key, size_t id = 0) const;
   size_t get_header_value_count(const std::string &key) const;
   void set_header(const std::string &key, const std::string &val);
 
@@ -739,6 +737,8 @@ private:
   std::regex regex_;
 };
 
+ssize_t write_headers(Stream &strm, const Headers &headers);
+
 } // namespace detail
 
 class Server {
@@ -802,6 +802,8 @@ public:
   Server &set_socket_options(SocketOptions socket_options);
 
   Server &set_default_headers(Headers headers);
+  Server &
+  set_header_writer(std::function<ssize_t(Stream &, Headers &)> const &writer);
 
   Server &set_keep_alive_max_count(size_t count);
   Server &set_keep_alive_timeout(time_t sec);
@@ -936,6 +938,8 @@ private:
   SocketOptions socket_options_ = default_socket_options;
 
   Headers default_headers_;
+  std::function<ssize_t(Stream &, Headers &)> header_writer_ =
+      detail::write_headers;
 };
 
 enum class Error {
@@ -988,8 +992,8 @@ public:
   bool has_request_header(const std::string &key) const;
   std::string get_request_header_value(const std::string &key,
                                        size_t id = 0) const;
-  template <typename T>
-  T get_request_header_value(const std::string &key, size_t id = 0) const;
+  uint64_t get_request_header_value_u64(const std::string &key,
+                                        size_t id = 0) const;
   size_t get_request_header_value_count(const std::string &key) const;
 
 private:
@@ -1166,6 +1170,9 @@ public:
 
   void set_default_headers(Headers headers);
 
+  void
+  set_header_writer(std::function<ssize_t(Stream &, Headers &)> const &writer);
+
   void set_address_family(int family);
   void set_tcp_nodelay(bool on);
   void set_socket_options(SocketOptions socket_options);
@@ -1275,6 +1282,10 @@ protected:
   // Default headers
   Headers default_headers_;
 
+  // Header writer
+  std::function<ssize_t(Stream &, Headers &)> header_writer_ =
+      detail::write_headers;
+
   // Settings
   std::string client_cert_path_;
   std::string client_key_path_;
@@ -1541,6 +1552,9 @@ public:
 
   void set_default_headers(Headers headers);
 
+  void
+  set_header_writer(std::function<ssize_t(Stream &, Headers &)> const &writer);
+
   void set_address_family(int family);
   void set_tcp_nodelay(bool on);
   void set_socket_options(SocketOptions socket_options);
@@ -1710,15 +1724,9 @@ inline void duration_to_sec_and_usec(const T &duration, U callback) {
   callback(static_cast<time_t>(sec), static_cast<time_t>(usec));
 }
 
-template <typename T>
-inline T get_header_value(const Headers & /*headers*/,
-                          const std::string & /*key*/, size_t /*id*/ = 0,
-                          uint64_t /*def*/ = 0) {}
-
-template <>
-inline uint64_t get_header_value<uint64_t>(const Headers &headers,
-                                           const std::string &key, size_t id,
-                                           uint64_t def) {
+inline uint64_t get_header_value_u64(const Headers &headers,
+                                     const std::string &key, size_t id,
+                                     uint64_t def) {
   auto rng = headers.equal_range(key);
   auto it = rng.first;
   std::advance(it, static_cast<ssize_t>(id));
@@ -1730,14 +1738,14 @@ inline uint64_t get_header_value<uint64_t>(const Headers &headers,
 
 } // namespace detail
 
-template <typename T>
-inline T Request::get_header_value(const std::string &key, size_t id) const {
-  return detail::get_header_value<T>(headers, key, id, 0);
+inline uint64_t Request::get_header_value_u64(const std::string &key,
+                                              size_t id) const {
+  return detail::get_header_value_u64(headers, key, id, 0);
 }
 
-template <typename T>
-inline T Response::get_header_value(const std::string &key, size_t id) const {
-  return detail::get_header_value<T>(headers, key, id, 0);
+inline uint64_t Response::get_header_value_u64(const std::string &key,
+                                               size_t id) const {
+  return detail::get_header_value_u64(headers, key, id, 0);
 }
 
 template <typename... Args>
@@ -1906,10 +1914,9 @@ inline std::ostream &operator<<(std::ostream &os, const Error &obj) {
   return os;
 }
 
-template <typename T>
-inline T Result::get_request_header_value(const std::string &key,
-                                          size_t id) const {
-  return detail::get_header_value<T>(request_headers_, key, id, 0);
+inline uint64_t Result::get_request_header_value_u64(const std::string &key,
+                                                     size_t id) const {
+  return detail::get_header_value_u64(request_headers_, key, id, 0);
 }
 
 template <class Rep, class Period>
@@ -3894,7 +3901,7 @@ bool read_content(Stream &strm, T &x, size_t payload_max_length, int &status,
         } else if (!has_header(x.headers, "Content-Length")) {
           ret = read_content_without_length(strm, out);
         } else {
-          auto len = get_header_value<uint64_t>(x.headers, "Content-Length");
+          auto len = get_header_value_u64(x.headers, "Content-Length", 0, 0);
           if (len > payload_max_length) {
             exceed_payload_max_length = true;
             skip_content_with_length(strm, len);
@@ -4986,7 +4993,7 @@ inline bool parse_www_authenticate(const Response &res,
         s = s.substr(pos + 1);
         auto beg = std::sregex_iterator(s.begin(), s.end(), re);
         for (auto i = beg; i != std::sregex_iterator(); ++i) {
-          auto m = *i;
+          const auto &m = *i;
           auto key = s.substr(static_cast<size_t>(m.position(1)),
                               static_cast<size_t>(m.length(1)));
           auto val = m.length(2) > 0
@@ -5681,6 +5688,12 @@ inline Server &Server::set_default_headers(Headers headers) {
   return *this;
 }
 
+inline Server &Server::set_header_writer(
+    std::function<ssize_t(Stream &, Headers &)> const &writer) {
+  header_writer_ = writer;
+  return *this;
+}
+
 inline Server &Server::set_keep_alive_max_count(size_t count) {
   keep_alive_max_count_ = count;
   return *this;
@@ -5875,7 +5888,7 @@ inline bool Server::write_response_core(Stream &strm, bool close_connection,
       return false;
     }
 
-    if (!detail::write_headers(bstrm, res.headers)) { return false; }
+    if (!header_writer_(bstrm, res.headers)) { return false; }
 
     // Flush buffer
     auto &data = bstrm.get_buffer();
@@ -6723,9 +6736,9 @@ inline bool ClientImpl::read_response_line(Stream &strm, const Request &req,
   if (!line_reader.getline()) { return false; }
 
 #ifdef CPPHTTPLIB_ALLOW_LF_AS_LINE_TERMINATOR
-  const static std::regex re("(HTTP/1\\.[01]) (\\d{3})(?: (.*?))?\r\n");
-#else
   const static std::regex re("(HTTP/1\\.[01]) (\\d{3})(?: (.*?))?\r?\n");
+#else
+  const static std::regex re("(HTTP/1\\.[01]) (\\d{3})(?: (.*?))?\r\n");
 #endif
 
   std::cmatch m;
@@ -7114,7 +7127,7 @@ inline bool ClientImpl::write_request(Stream &strm, Request &req,
     const auto &path = url_encode_ ? detail::encode_url(req.path) : req.path;
     bstrm.write_format("%s %s HTTP/1.1\r\n", req.method.c_str(), path.c_str());
 
-    detail::write_headers(bstrm, req.headers);
+    header_writer_(bstrm, req.headers);
 
     // Flush buffer
     auto &data = bstrm.get_buffer();
@@ -7925,6 +7938,11 @@ inline void ClientImpl::set_default_headers(Headers headers) {
   default_headers_ = std::move(headers);
 }
 
+inline void ClientImpl::set_header_writer(
+    std::function<ssize_t(Stream &, Headers &)> const &writer) {
+  header_writer_ = writer;
+}
+
 inline void ClientImpl::set_address_family(int family) {
   address_family_ = family;
 }
@@ -9119,6 +9137,11 @@ inline void Client::set_default_headers(Headers headers) {
   cli_->set_default_headers(std::move(headers));
 }
 
+inline void Client::set_header_writer(
+    std::function<ssize_t(Stream &, Headers &)> const &writer) {
+  cli_->set_header_writer(writer);
+}
+
 inline void Client::set_address_family(int family) {
   cli_->set_address_family(family);
 }

test/test.cc

@@ -211,8 +211,7 @@ TEST(GetHeaderValueTest, DefaultValue) {
 
 TEST(GetHeaderValueTest, DefaultValueInt) {
   Headers headers = {{"Dummy", "Dummy"}};
-  auto val =
-      detail::get_header_value<uint64_t>(headers, "Content-Length", 0, 100);
+  auto val = detail::get_header_value_u64(headers, "Content-Length", 0, 100);
   EXPECT_EQ(100ull, val);
 }
 
@@ -241,8 +240,7 @@ TEST(GetHeaderValueTest, SetContent) {
 
 TEST(GetHeaderValueTest, RegularValueInt) {
   Headers headers = {{"Content-Length", "100"}, {"Dummy", "Dummy"}};
-  auto val =
-      detail::get_header_value<uint64_t>(headers, "Content-Length", 0, 0);
+  auto val = detail::get_header_value_u64(headers, "Content-Length", 0, 0);
   EXPECT_EQ(100ull, val);
 }
 
@@ -444,7 +442,7 @@ TEST(HostnameToIPConversionTest, HTTPWatch_Online) {
   auto host = "www.httpwatch.com";
 
   auto ip = hosted_at(host);
-  EXPECT_EQ("191.236.16.12", ip);
+  EXPECT_EQ("23.96.13.243", ip);
 
   std::vector<std::string> addrs;
   hosted_at(host, addrs);
@@ -1016,7 +1014,7 @@ TEST(UrlWithSpace, Redirect_Online) {
   auto res = cli.Get("/files/2595/310/Neat 1.4-17.jar");
   ASSERT_TRUE(res);
   EXPECT_EQ(200, res->status);
-  EXPECT_EQ(18527U, res->get_header_value<uint64_t>("Content-Length"));
+  EXPECT_EQ(18527U, res->get_header_value_u64("Content-Length"));
 }
 
 #endif
@@ -1594,6 +1592,46 @@ TEST(URLFragmentTest, WithFragment) {
   }
 }
 
+TEST(HeaderWriter, SetHeaderWriter) {
+  Server svr;
+
+  svr.set_header_writer([](Stream &strm, Headers &hdrs) {
+    hdrs.emplace("CustomServerHeader", "CustomServerValue");
+    return detail::write_headers(strm, hdrs);
+  });
+  svr.Get("/hi", [](const Request &req, Response &res) {
+    auto it = req.headers.find("CustomClientHeader");
+    EXPECT_TRUE(it != req.headers.end());
+    EXPECT_EQ(it->second, "CustomClientValue");
+    res.set_content("Hello World!\n", "text/plain");
+  });
+
+  auto thread = std::thread([&]() { svr.listen(HOST, PORT); });
+  auto se = detail::scope_exit([&] {
+    svr.stop();
+    thread.join();
+    ASSERT_FALSE(svr.is_running());
+  });
+
+  std::this_thread::sleep_for(std::chrono::seconds(1));
+
+  {
+    Client cli(HOST, PORT);
+    cli.set_header_writer([](Stream &strm, Headers &hdrs) {
+      hdrs.emplace("CustomClientHeader", "CustomClientValue");
+      return detail::write_headers(strm, hdrs);
+    });
+
+    auto res = cli.Get("/hi");
+    EXPECT_TRUE(res);
+    EXPECT_EQ(200, res->status);
+
+    auto it = res->headers.find("CustomServerHeader");
+    EXPECT_TRUE(it != res->headers.end());
+    EXPECT_EQ(it->second, "CustomServerValue");
+  }
+}
+
 class ServerTest : public ::testing::Test {
 protected:
   ServerTest()
@@ -3296,7 +3334,7 @@ TEST_F(ServerTest, PutLargeFileWithGzip2) {
   ASSERT_TRUE(res);
   EXPECT_EQ(200, res->status);
   EXPECT_EQ(LARGE_DATA, res->body);
-  EXPECT_EQ(101942u, res.get_request_header_value<uint64_t>("Content-Length"));
+  EXPECT_EQ(101942u, res.get_request_header_value_u64("Content-Length"));
   EXPECT_EQ("gzip", res.get_request_header_value("Content-Encoding"));
 }