Fix #472

README.md

@@ -181,6 +181,7 @@ svr.Get("/stream", [&](const Request &req, Response &res) {
     [data](size_t offset, size_t length, DataSink &sink) {
       const auto &d = *data;
       sink.write(&d[offset], std::min(length, DATA_CHUNK_SIZE));
+      return true; // return 'false' if you want to cancel the process.
     },
     [data] { delete data; });
 });
@@ -192,10 +193,11 @@ svr.Get("/stream", [&](const Request &req, Response &res) {
 svr.Get("/chunked", [&](const Request& req, Response& res) {
   res.set_chunked_content_provider(
     [](size_t offset, DataSink &sink) {
-       sink.write("123", 3);
-       sink.write("345", 3);
-       sink.write("789", 3);
-       sink.done();
+      sink.write("123", 3);
+      sink.write("345", 3);
+      sink.write("789", 3);
+      sink.done();
+      return true; // return 'false' if you want to cancel the process.
     }
   );
 });
@@ -363,6 +365,35 @@ res = cli.Options("/resource/foo");
 ```c++
 cli.set_timeout_sec(5); // timeouts in 5 seconds
 ```
+### Receive content with Content receiver
+
+```cpp
+std::string body;
+auto res = cli.Get(
+  "/stream", Headers(),
+  [&](const Response &response) {
+    EXPECT_EQ(200, response.status);
+    return true; // return 'false' if you want to cancel the request.
+  },
+  [&](const char *data, size_t data_length) {
+    body.append(data, data_length);
+    return true; // return 'false' if you want to cancel the request.
+  });
+```
+
+### Send content with Content provider
+
+```cpp
+std::string body = ...;
+auto res = cli_.Post(
+  "/stream", body.size(),
+  [](size_t offset, size_t length, DataSink &sink) {
+    sink.write(body.data() + offset, length);
+    return true; // return 'false' if you want to cancel the request.
+  },
+  "text/plain");
+```
+
 ### With Progress Callback
 
 ```cpp

httplib.h

@@ -227,7 +227,10 @@ public:
 };
 
 using ContentProvider =
-    std::function<void(size_t offset, size_t length, DataSink &sink)>;
+    std::function<bool(size_t offset, size_t length, DataSink &sink)>;
+
+using ChunkedContentProvider =
+    std::function<bool(size_t offset, DataSink &sink)>;
 
 using ContentReceiver =
     std::function<bool(const char *data, size_t data_length)>;
@@ -323,13 +326,11 @@ struct Response {
   void set_content(std::string s, const char *content_type);
 
   void set_content_provider(
-      size_t length,
-      std::function<void(size_t offset, size_t length, DataSink &sink)>
-          provider,
+      size_t length, ContentProvider provider,
       std::function<void()> resource_releaser = [] {});
 
   void set_chunked_content_provider(
-      std::function<void(size_t offset, DataSink &sink)> provider,
+      ChunkedContentProvider provider,
       std::function<void()> resource_releaser = [] {});
 
   Response() = default;
@@ -2074,22 +2075,25 @@ inline ssize_t write_content(Stream &strm, ContentProvider content_provider,
                              size_t offset, size_t length) {
   size_t begin_offset = offset;
   size_t end_offset = offset + length;
-  while (offset < end_offset) {
-    ssize_t written_length = 0;
 
-    DataSink data_sink;
-    data_sink.write = [&](const char *d, size_t l) {
-      offset += l;
-      written_length = strm.write(d, l);
-    };
-    data_sink.done = [&](void) { written_length = -1; };
-    data_sink.is_writable = [&](void) {
-      return strm.is_writable() && written_length >= 0;
-    };
+  ssize_t written_length = 0;
+
+  DataSink data_sink;
+  data_sink.write = [&](const char *d, size_t l) {
+    offset += l;
+    written_length = strm.write(d, l);
+  };
+  data_sink.is_writable = [&](void) {
+    return strm.is_writable() && written_length >= 0;
+  };
 
-    content_provider(offset, end_offset - offset, data_sink);
+  while (offset < end_offset) {
+    if (!content_provider(offset, end_offset - offset, data_sink)) {
+      return -1;
+    }
     if (written_length < 0) { return written_length; }
   }
+
   return static_cast<ssize_t>(offset - begin_offset);
 }
 
@@ -2100,31 +2104,32 @@ inline ssize_t write_content_chunked(Stream &strm,
   size_t offset = 0;
   auto data_available = true;
   ssize_t total_written_length = 0;
-  while (data_available && !is_shutting_down()) {
-    ssize_t written_length = 0;
 
-    DataSink data_sink;
-    data_sink.write = [&](const char *d, size_t l) {
-      data_available = l > 0;
-      offset += l;
+  ssize_t written_length = 0;
 
-      // Emit chunked response header and footer for each chunk
-      auto chunk = from_i_to_hex(l) + "\r\n" + std::string(d, l) + "\r\n";
-      written_length = strm.write(chunk);
-    };
-    data_sink.done = [&](void) {
-      data_available = false;
-      written_length = strm.write("0\r\n\r\n");
-    };
-    data_sink.is_writable = [&](void) {
-      return strm.is_writable() && written_length >= 0;
-    };
+  DataSink data_sink;
+  data_sink.write = [&](const char *d, size_t l) {
+    data_available = l > 0;
+    offset += l;
 
-    content_provider(offset, 0, data_sink);
+    // Emit chunked response header and footer for each chunk
+    auto chunk = from_i_to_hex(l) + "\r\n" + std::string(d, l) + "\r\n";
+    written_length = strm.write(chunk);
+  };
+  data_sink.done = [&](void) {
+    data_available = false;
+    written_length = strm.write("0\r\n\r\n");
+  };
+  data_sink.is_writable = [&](void) {
+    return strm.is_writable() && written_length >= 0;
+  };
 
+  while (data_available && !is_shutting_down()) {
+    if (!content_provider(offset, 0, data_sink)) { return -1; }
     if (written_length < 0) { return written_length; }
     total_written_length += written_length;
   }
+
   return total_written_length;
 }
 
@@ -2904,24 +2909,23 @@ inline void Response::set_content(std::string s, const char *content_type) {
   set_header("Content-Type", content_type);
 }
 
-inline void Response::set_content_provider(
-    size_t in_length,
-    std::function<void(size_t offset, size_t length, DataSink &sink)> provider,
-    std::function<void()> resource_releaser) {
+inline void
+Response::set_content_provider(size_t in_length, ContentProvider provider,
+                               std::function<void()> resource_releaser) {
   assert(in_length > 0);
   content_length = in_length;
   content_provider = [provider](size_t offset, size_t length, DataSink &sink) {
-    provider(offset, length, sink);
+    return provider(offset, length, sink);
   };
   content_provider_resource_releaser = resource_releaser;
 }
 
 inline void Response::set_chunked_content_provider(
-    std::function<void(size_t offset, DataSink &sink)> provider,
+    ChunkedContentProvider provider,
     std::function<void()> resource_releaser) {
   content_length = 0;
   content_provider = [provider](size_t offset, size_t, DataSink &sink) {
-    provider(offset, sink);
+    return provider(offset, sink);
   };
   content_provider_resource_releaser = resource_releaser;
 }
@@ -4106,15 +4110,22 @@ inline bool Client::write_request(Stream &strm, const Request &req,
       size_t offset = 0;
       size_t end_offset = req.content_length;
 
+      ssize_t written_length = 0;
+
       DataSink data_sink;
       data_sink.write = [&](const char *d, size_t l) {
-        auto written_length = strm.write(d, l);
+        written_length = strm.write(d, l);
         offset += static_cast<size_t>(written_length);
       };
-      data_sink.is_writable = [&](void) { return strm.is_writable(); };
+      data_sink.is_writable = [&](void) {
+        return strm.is_writable() && written_length >= 0;
+      };
 
       while (offset < end_offset) {
-        req.content_provider(offset, end_offset - offset, data_sink);
+        if (!req.content_provider(offset, end_offset - offset, data_sink)) {
+          return false;
+        }
+        if (written_length < 0) { return false; }
       }
     }
   } else {
@@ -4148,7 +4159,9 @@ inline std::shared_ptr<Response> Client::send_with_content_provider(
       data_sink.is_writable = [&](void) { return true; };
 
       while (offset < content_length) {
-        content_provider(offset, content_length - offset, data_sink);
+        if (!content_provider(offset, content_length - offset, data_sink)) {
+          return nullptr;
+        }
       }
     } else {
       req.body = body;

test/test.cc

@@ -899,20 +899,21 @@ protected:
         .Get("/streamed-chunked",
              [&](const Request & /*req*/, Response &res) {
                res.set_chunked_content_provider(
-                   [](uint64_t /*offset*/, DataSink &sink) {
-                     ASSERT_TRUE(sink.is_writable());
+                   [](size_t /*offset*/, DataSink &sink) {
+                     EXPECT_TRUE(sink.is_writable());
                      sink.write("123", 3);
                      sink.write("456", 3);
                      sink.write("789", 3);
                      sink.done();
+                     return true;
                    });
              })
         .Get("/streamed-chunked2",
              [&](const Request & /*req*/, Response &res) {
                auto i = new int(0);
                res.set_chunked_content_provider(
-                   [i](uint64_t /*offset*/, DataSink &sink) {
-                     ASSERT_TRUE(sink.is_writable());
+                   [i](size_t /*offset*/, DataSink &sink) {
+                     EXPECT_TRUE(sink.is_writable());
                      switch (*i) {
                      case 0: sink.write("123", 3); break;
                      case 1: sink.write("456", 3); break;
@@ -920,14 +921,16 @@ protected:
                      case 3: sink.done(); break;
                      }
                      (*i)++;
+                     return true;
                    },
                    [i] { delete i; });
              })
         .Get("/streamed",
              [&](const Request & /*req*/, Response &res) {
                res.set_content_provider(
-                   6, [](uint64_t offset, uint64_t /*length*/, DataSink &sink) {
+                   6, [](size_t offset, size_t /*length*/, DataSink &sink) {
                      sink.write(offset < 3 ? "a" : "b", 1);
+                     return true;
                    });
              })
         .Get("/streamed-with-range",
@@ -935,25 +938,27 @@ protected:
                auto data = new std::string("abcdefg");
                res.set_content_provider(
                    data->size(),
-                   [data](uint64_t offset, uint64_t length, DataSink &sink) {
-                     ASSERT_TRUE(sink.is_writable());
+                   [data](size_t offset, size_t length, DataSink &sink) {
+                     EXPECT_TRUE(sink.is_writable());
                      size_t DATA_CHUNK_SIZE = 4;
                      const auto &d = *data;
                      auto out_len =
                          std::min(static_cast<size_t>(length), DATA_CHUNK_SIZE);
                      sink.write(&d[static_cast<size_t>(offset)], out_len);
+                     return true;
                    },
                    [data] { delete data; });
              })
         .Get("/streamed-cancel",
              [&](const Request & /*req*/, Response &res) {
-               res.set_content_provider(size_t(-1), [](uint64_t /*offset*/,
-                                                       uint64_t /*length*/,
-                                                       DataSink &sink) {
-                 ASSERT_TRUE(sink.is_writable());
-                 std::string data = "data_chunk";
-                 sink.write(data.data(), data.size());
-               });
+               res.set_content_provider(
+                   size_t(-1),
+                   [](size_t /*offset*/, size_t /*length*/, DataSink &sink) {
+                     EXPECT_TRUE(sink.is_writable());
+                     std::string data = "data_chunk";
+                     sink.write(data.data(), data.size());
+                     return true;
+                   });
              })
         .Get("/with-range",
              [&](const Request & /*req*/, Response &res) {
@@ -1918,8 +1923,9 @@ TEST_F(ServerTest, PutWithContentProvider) {
   auto res = cli_.Put(
       "/put", 3,
       [](size_t /*offset*/, size_t /*length*/, DataSink &sink) {
-        ASSERT_TRUE(sink.is_writable());
+        EXPECT_TRUE(sink.is_writable());
         sink.write("PUT", 3);
+        return true;
       },
       "text/plain");
 
@@ -1928,14 +1934,26 @@ TEST_F(ServerTest, PutWithContentProvider) {
   EXPECT_EQ("PUT", res->body);
 }
 
+TEST_F(ServerTest, PostWithContentProviderAbort) {
+  auto res = cli_.Post(
+      "/post", 42,
+      [](size_t /*offset*/, size_t /*length*/, DataSink & /*sink*/) {
+        return false;
+      },
+      "text/plain");
+
+  ASSERT_TRUE(res == nullptr);
+}
+
 #ifdef CPPHTTPLIB_ZLIB_SUPPORT
 TEST_F(ServerTest, PutWithContentProviderWithGzip) {
   cli_.set_compress(true);
   auto res = cli_.Put(
       "/put", 3,
       [](size_t /*offset*/, size_t /*length*/, DataSink &sink) {
-        ASSERT_TRUE(sink.is_writable());
+        EXPECT_TRUE(sink.is_writable());
         sink.write("PUT", 3);
+        return true;
       },
       "text/plain");
 
@@ -1944,6 +1962,18 @@ TEST_F(ServerTest, PutWithContentProviderWithGzip) {
   EXPECT_EQ("PUT", res->body);
 }
 
+TEST_F(ServerTest, PostWithContentProviderWithGzipAbort) {
+  cli_.set_compress(true);
+  auto res = cli_.Post(
+      "/post", 42,
+      [](size_t /*offset*/, size_t /*length*/, DataSink & /*sink*/) {
+        return false;
+      },
+      "text/plain");
+
+  ASSERT_TRUE(res == nullptr);
+}
+
 TEST_F(ServerTest, PutLargeFileWithGzip) {
   cli_.set_compress(true);
   auto res = cli_.Put("/put-large", LARGE_DATA, "text/plain");
@@ -2091,8 +2121,8 @@ TEST_F(ServerTest, KeepAlive) {
   Get(requests, "/not-exist");
   Post(requests, "/empty", "", "text/plain");
   Post(
-      requests, "/empty", 0, [&](size_t, size_t, httplib::DataSink &) {},
-      "text/plain");
+      requests, "/empty", 0,
+      [&](size_t, size_t, httplib::DataSink &) { return true; }, "text/plain");
 
   std::vector<Response> responses;
   auto ret = cli_.send(requests, responses);
@@ -2440,6 +2470,7 @@ TEST(ServerStopTest, StopServerWithChunkedTransmission) {
       auto size = static_cast<size_t>(sprintf(buffer, "data:%ld\n\n", offset));
       sink.write(buffer, size);
       std::this_thread::sleep_for(std::chrono::seconds(1));
+      return true;
     });
   });