Fix yhirose#1486

Author: yhirose

README.md

@@ -347,6 +347,27 @@ svr.Get("/chunked", [&](const Request& req, Response& res) {
 });
 ```
 
+With trailer:
+
+```cpp
+svr.Get("/chunked", [&](const Request& req, Response& res) {
+  res.set_header("Trailer", "Dummy1, Dummy2");
+  res.set_chunked_content_provider(
+    "text/plain",
+    [](size_t offset, DataSink &sink) {
+      sink.write("123", 3);
+      sink.write("345", 3);
+      sink.write("789", 3);
+      sink.done_with_trailer({
+        {"Dummy1", "DummyVal1"},
+        {"Dummy2", "DummyVal2"}
+      });
+      return true;
+    }
+  );
+});
+```
+
 ### 'Expect: 100-continue' handler
 
 By default, the server sends a `100 Continue` response for an `Expect: 100-continue` header.

httplib.h

@@ -371,6 +371,7 @@ class DataSink {
 
   std::function<bool(const char *data, size_t data_len)> write;
   std::function<void()> done;
+  std::function<void(const Headers &trailer)> done_with_trailer;
   std::ostream os;
 
 private:
@@ -3525,7 +3526,8 @@ inline bool read_content_without_length(Stream &strm,
   return true;
 }
 
-inline bool read_content_chunked(Stream &strm,
+template <typename T>
+inline bool read_content_chunked(Stream &strm, T &x,
                                  ContentReceiverWithProgress out) {
   const auto bufsiz = 16;
   char buf[bufsiz];
@@ -3551,15 +3553,29 @@ inline bool read_content_chunked(Stream &strm,
 
     if (!line_reader.getline()) { return false; }
 
-    if (strcmp(line_reader.ptr(), "\r\n")) { break; }
+    if (strcmp(line_reader.ptr(), "\r\n")) { return false; }
 
     if (!line_reader.getline()) { return false; }
   }
 
-  if (chunk_len == 0) {
-    // Reader terminator after chunks
-    if (!line_reader.getline() || strcmp(line_reader.ptr(), "\r\n"))
-      return false;
+  assert(chunk_len == 0);
+
+  // Trailer
+  if (!line_reader.getline()) { return false; }
+
+  while (strcmp(line_reader.ptr(), "\r\n")) {
+    if (line_reader.size() > CPPHTTPLIB_HEADER_MAX_LENGTH) { return false; }
+
+    // Exclude line terminator
+    constexpr auto line_terminator_len = 2;
+    auto end = line_reader.ptr() + line_reader.size() - line_terminator_len;
+
+    parse_header(line_reader.ptr(), end,
+                 [&](std::string &&key, std::string &&val) {
+                   x.headers.emplace(std::move(key), std::move(val));
+                 });
+
+    if (!line_reader.getline()) { return false; }
   }
 
   return true;
@@ -3629,7 +3645,7 @@ bool read_content(Stream &strm, T &x, size_t payload_max_length, int &status,
         auto exceed_payload_max_length = false;
 
         if (is_chunked_transfer_encoding(x.headers)) {
-          ret = read_content_chunked(strm, out);
+          ret = read_content_chunked(strm, x, out);
         } else if (!has_header(x.headers, "Content-Length")) {
           ret = read_content_without_length(strm, out);
         } else {
@@ -3785,7 +3801,7 @@ write_content_chunked(Stream &strm, const ContentProvider &content_provider,
     return ok;
   };
 
-  data_sink.done = [&](void) {
+  auto done_with_trailer = [&](const Headers *trailer) {
     if (!ok) { return; }
 
     data_available = false;
@@ -3803,16 +3819,36 @@ write_content_chunked(Stream &strm, const ContentProvider &content_provider,
     if (!payload.empty()) {
       // Emit chunked response header and footer for each chunk
       auto chunk = from_i_to_hex(payload.size()) + "\r\n" + payload + "\r\n";
-      if (!write_data(strm, chunk.data(), chunk.size())) {
+      if (!strm.is_writable() ||
+          !write_data(strm, chunk.data(), chunk.size())) {
         ok = false;
         return;
       }
     }
 
-    static const std::string done_marker("0\r\n\r\n");
+    static const std::string done_marker("0\r\n");
     if (!write_data(strm, done_marker.data(), done_marker.size())) {
       ok = false;
     }
+
+    // Trailer
+    if (trailer) {
+      for (const auto &kv : *trailer) {
+        std::string field_line = kv.first + ": " + kv.second + "\r\n";
+        if (!write_data(strm, field_line.data(), field_line.size())) {
+          ok = false;
+        }
+      }
+    }
+
+    static const std::string crlf("\r\n");
+    if (!write_data(strm, crlf.data(), crlf.size())) { ok = false; }
+  };
+
+  data_sink.done = [&](void) { done_with_trailer(nullptr); };
+
+  data_sink.done_with_trailer = [&](const Headers &trailer) {
+    done_with_trailer(&trailer);
   };
 
   while (data_available && !is_shutting_down()) {

test/test.cc

@@ -186,7 +186,8 @@ TEST(ParseMultipartBoundaryTest, ValueWithQuote) {
 }
 
 TEST(ParseMultipartBoundaryTest, ValueWithCharset) {
-  string content_type = "multipart/mixed; boundary=THIS_STRING_SEPARATES;charset=UTF-8";
+  string content_type =
+      "multipart/mixed; boundary=THIS_STRING_SEPARATES;charset=UTF-8";
   string boundary;
   auto ret = detail::parse_multipart_boundary(content_type, boundary);
   EXPECT_TRUE(ret);
@@ -1710,6 +1711,30 @@ class ServerTest : public ::testing::Test {
                      delete i;
                    });
              })
+        .Get("/streamed-chunked-with-trailer",
+             [&](const Request & /*req*/, Response &res) {
+               auto i = new int(0);
+               res.set_header("Trailer", "Dummy1, Dummy2");
+               res.set_chunked_content_provider(
+                   "text/plain",
+                   [i](size_t /*offset*/, DataSink &sink) {
+                     switch (*i) {
+                     case 0: sink.os << "123"; break;
+                     case 1: sink.os << "456"; break;
+                     case 2: sink.os << "789"; break;
+                     case 3: {
+                       sink.done_with_trailer(
+                           {{"Dummy1", "DummyVal1"}, {"Dummy2", "DummyVal2"}});
+                     } break;
+                     }
+                     (*i)++;
+                     return true;
+                   },
+                   [i](bool success) {
+                     EXPECT_TRUE(success);
+                     delete i;
+                   });
+             })
         .Get("/streamed",
              [&](const Request & /*req*/, Response &res) {
                res.set_content_provider(
@@ -1801,39 +1826,39 @@ class ServerTest : public ::testing::Test {
                 }
               })
         .Post("/multipart/multi_file_values",
-          [&](const Request &req, Response & /*res*/) {
-              EXPECT_EQ(5u, req.files.size());
-              ASSERT_TRUE(!req.has_file("???"));
-              ASSERT_TRUE(req.body.empty());
+              [&](const Request &req, Response & /*res*/) {
+                EXPECT_EQ(5u, req.files.size());
+                ASSERT_TRUE(!req.has_file("???"));
+                ASSERT_TRUE(req.body.empty());
 
-              {
+                {
                   const auto &text_value = req.get_file_values("text");
                   EXPECT_EQ(text_value.size(), 1);
                   auto &text = text_value[0];
                   EXPECT_TRUE(text.filename.empty());
                   EXPECT_EQ("default text", text.content);
-              }
-              {
-                const auto &text1_values = req.get_file_values("multi_text1");
-                EXPECT_EQ(text1_values.size(), 2);
-                EXPECT_EQ("aaaaa", text1_values[0].content);
-                EXPECT_EQ("bbbbb", text1_values[1].content);
-              }
-
-              {
-                const auto &file1_values = req.get_file_values("multi_file1");
-                EXPECT_EQ(file1_values.size(), 2);
-                auto file1 = file1_values[0];
-                EXPECT_EQ(file1.filename, "hello.txt");
-                EXPECT_EQ(file1.content_type, "text/plain");
-                EXPECT_EQ("h\ne\n\nl\nl\no\n", file1.content);
-
-                auto file2 = file1_values[1];
-                EXPECT_EQ(file2.filename, "world.json");
-                EXPECT_EQ(file2.content_type, "application/json");
-                EXPECT_EQ("{\n  \"world\", true\n}\n", file2.content);
-              }
-          })
+                }
+                {
+                  const auto &text1_values = req.get_file_values("multi_text1");
+                  EXPECT_EQ(text1_values.size(), 2);
+                  EXPECT_EQ("aaaaa", text1_values[0].content);
+                  EXPECT_EQ("bbbbb", text1_values[1].content);
+                }
+
+                {
+                  const auto &file1_values = req.get_file_values("multi_file1");
+                  EXPECT_EQ(file1_values.size(), 2);
+                  auto file1 = file1_values[0];
+                  EXPECT_EQ(file1.filename, "hello.txt");
+                  EXPECT_EQ(file1.content_type, "text/plain");
+                  EXPECT_EQ("h\ne\n\nl\nl\no\n", file1.content);
+
+                  auto file2 = file1_values[1];
+                  EXPECT_EQ(file2.filename, "world.json");
+                  EXPECT_EQ(file2.content_type, "application/json");
+                  EXPECT_EQ("{\n  \"world\", true\n}\n", file2.content);
+                }
+              })
         .Post("/empty",
               [&](const Request &req, Response &res) {
                 EXPECT_EQ(req.body, "");
@@ -2680,13 +2705,14 @@ TEST_F(ServerTest, MultipartFormData) {
 
 TEST_F(ServerTest, MultipartFormDataMultiFileValues) {
   MultipartFormDataItems items = {
-    {"text", "default text", "", ""},
+      {"text", "default text", "", ""},
 
-    {"multi_text1", "aaaaa", "", ""},
-    {"multi_text1", "bbbbb", "", ""},
+      {"multi_text1", "aaaaa", "", ""},
+      {"multi_text1", "bbbbb", "", ""},
 
-    {"multi_file1", "h\ne\n\nl\nl\no\n", "hello.txt", "text/plain"},
-    {"multi_file1", "{\n  \"world\", true\n}\n", "world.json", "application/json"},
+      {"multi_file1", "h\ne\n\nl\nl\no\n", "hello.txt", "text/plain"},
+      {"multi_file1", "{\n  \"world\", true\n}\n", "world.json",
+       "application/json"},
   };
 
   auto res = cli_.Post("/multipart/multi_file_values", items);
@@ -2920,6 +2946,15 @@ TEST_F(ServerTest, GetStreamedChunked2) {
   EXPECT_EQ(std::string("123456789"), res->body);
 }
 
+TEST_F(ServerTest, GetStreamedChunkedWithTrailer) {
+  auto res = cli_.Get("/streamed-chunked-with-trailer");
+  ASSERT_TRUE(res);
+  EXPECT_EQ(200, res->status);
+  EXPECT_EQ(std::string("123456789"), res->body);
+  EXPECT_EQ(std::string("DummyVal1"), res->get_header_value("Dummy1"));
+  EXPECT_EQ(std::string("DummyVal2"), res->get_header_value("Dummy2"));
+}
+
 TEST_F(ServerTest, LargeChunkedPost) {
   Request req;
   req.method = "POST";
@@ -3906,9 +3941,8 @@ TEST(ServerStopTest, StopServerWithChunkedTransmission) {
 
 TEST(ServerStopTest, ClientAccessAfterServerDown) {
   httplib::Server svr;
-  svr.Post("/hi", [&](const httplib::Request & /*req*/, httplib::Response &res) {
-    res.status = 200;
-  });
+  svr.Post("/hi", [&](const httplib::Request & /*req*/,
+                      httplib::Response &res) { res.status = 200; });
 
   auto thread = std::thread([&]() { svr.listen(HOST, PORT); });