From 73e0729f63311ea0521d1ff4c33c5a15ea107188 Mon Sep 17 00:00:00 2001
From: yhirose <yuji.hirose.bug@gmail.com>
Date: Mon, 19 Apr 2021 17:21:17 -0400
Subject: [PATCH] Change `sink.write()` to return boolean

---
 httplib.h    | 54 +++++++++++++++++++++++++++-------------------------
 test/test.cc |  9 ++++++---
 2 files changed, 34 insertions(+), 29 deletions(-)

diff --git a/httplib.h b/httplib.h
index 60acbb9..9298429 100644
--- a/httplib.h
+++ b/httplib.h
@@ -308,7 +308,7 @@ public:
   DataSink(DataSink &&) = delete;
   DataSink &operator=(DataSink &&) = delete;
 
-  std::function<void(const char *data, size_t data_len)> write;
+  std::function<bool(const char *data, size_t data_len)> write;
   std::function<void()> done;
   std::function<bool()> is_writable;
   std::ostream os;
@@ -2261,8 +2261,7 @@ inline socket_t create_client_socket(
         {
           timeval tv;
           tv.tv_sec = static_cast<long>(write_timeout_sec);
-          tv.tv_usec =
-          static_cast<decltype(tv.tv_usec)>(write_timeout_usec);
+          tv.tv_usec = static_cast<decltype(tv.tv_usec)>(write_timeout_usec);
           setsockopt(sock, SOL_SOCKET, SO_SNDTIMEO, (char *)&tv, sizeof(tv));
         }
 
@@ -3022,7 +3021,7 @@ inline bool write_content(Stream &strm, const ContentProvider &content_provider,
   auto ok = true;
   DataSink data_sink;
 
-  data_sink.write = [&](const char *d, size_t l) {
+  data_sink.write = [&](const char *d, size_t l) -> bool {
     if (ok) {
       if (write_data(strm, d, l)) {
         offset += l;
@@ -3030,6 +3029,7 @@ inline bool write_content(Stream &strm, const ContentProvider &content_provider,
         ok = false;
       }
     }
+    return ok;
   };
 
   data_sink.is_writable = [&](void) { return ok && strm.is_writable(); };
@@ -3068,11 +3068,12 @@ write_content_without_length(Stream &strm,
   auto ok = true;
   DataSink data_sink;
 
-  data_sink.write = [&](const char *d, size_t l) {
+  data_sink.write = [&](const char *d, size_t l) -> bool {
     if (ok) {
       offset += l;
       if (!write_data(strm, d, l)) { ok = false; }
     }
+    return ok;
   };
 
   data_sink.done = [&](void) { data_available = false; };
@@ -3095,30 +3096,30 @@ write_content_chunked(Stream &strm, const ContentProvider &content_provider,
   auto ok = true;
   DataSink data_sink;
 
-  data_sink.write = [&](const char *d, size_t l) {
-    if (!ok) { return; }
+  data_sink.write = [&](const char *d, size_t l) -> bool {
+    if (ok) {
+      data_available = l > 0;
+      offset += l;
 
-    data_available = l > 0;
-    offset += l;
-
-    std::string payload;
-    if (!compressor.compress(d, l, false,
-                             [&](const char *data, size_t data_len) {
-                               payload.append(data, data_len);
-                               return true;
-                             })) {
-      ok = false;
-      return;
-    }
-
-    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())) {
+      std::string payload;
+      if (compressor.compress(d, l, false,
+                              [&](const char *data, size_t data_len) {
+                                payload.append(data, data_len);
+                                return true;
+                              })) {
+        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())) {
+            ok = false;
+          }
+        }
+      } else {
         ok = false;
-        return;
       }
     }
+    return ok;
   };
 
   data_sink.done = [&](void) {
@@ -5747,7 +5748,7 @@ inline std::unique_ptr<Response> ClientImpl::send_with_content_provider(
       size_t offset = 0;
       DataSink data_sink;
 
-      data_sink.write = [&](const char *data, size_t data_len) {
+      data_sink.write = [&](const char *data, size_t data_len) -> bool {
         if (ok) {
           auto last = offset + data_len == content_length;
 
@@ -5763,6 +5764,7 @@ inline std::unique_ptr<Response> ClientImpl::send_with_content_provider(
             ok = false;
           }
         }
+        return ok;
       };
 
       data_sink.is_writable = [&](void) { return ok && true; };
diff --git a/test/test.cc b/test/test.cc
index 91a905c..0694536 100644
--- a/test/test.cc
+++ b/test/test.cc
@@ -1366,7 +1366,8 @@ protected:
                      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);
+                     auto ret = sink.write(&d[static_cast<size_t>(offset)], out_len);
+                     EXPECT_TRUE(ret);
                      return true;
                    },
                    [data] { delete data; });
@@ -2521,7 +2522,8 @@ TEST_F(ServerTest, SlowPost) {
   auto res = cli_.Post(
       "/slowpost", 64 * 1024 * 1024,
       [&](size_t /*offset*/, size_t /*length*/, DataSink &sink) {
-        sink.write(buffer, sizeof(buffer));
+        auto ret = sink.write(buffer, sizeof(buffer));
+        EXPECT_TRUE(ret);
         return true;
       },
       "text/plain");
@@ -3349,7 +3351,8 @@ TEST(ServerStopTest, StopServerWithChunkedTransmission) {
                                                              DataSink &sink) {
       char buffer[27];
       auto size = static_cast<size_t>(sprintf(buffer, "data:%ld\n\n", offset));
-      sink.write(buffer, size);
+      auto ret = sink.write(buffer, size);
+      EXPECT_TRUE(ret);
       std::this_thread::sleep_for(std::chrono::seconds(1));
       return true;
     });