From 5a43bb8149e001625f1a0d81e5300e4e997c5ee8 Mon Sep 17 00:00:00 2001
From: yhirose <yuji.hirose.bug@gmail.com>
Date: Wed, 2 Jun 2021 13:28:49 -0400
Subject: [PATCH] Implemented #946 in a different way

---
 README.md    |  2 +-
 httplib.h    | 12 ++++++++----
 test/test.cc | 48 ++++++++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 55 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 45a0e12..8410be5 100644
--- a/README.md
+++ b/README.md
@@ -266,7 +266,7 @@ svr.Get("/stream", [&](const Request &req, Response &res) {
       sink.write(&d[offset], std::min(length, DATA_CHUNK_SIZE));
       return true; // return 'false' if you want to cancel the process.
     },
-    [data] { delete data; });
+    [data](bool success) { delete data; });
 });
 ```
 
diff --git a/httplib.h b/httplib.h
index d96a7b8..d3cddc9 100644
--- a/httplib.h
+++ b/httplib.h
@@ -337,7 +337,7 @@ using ContentProvider =
 using ContentProviderWithoutLength =
     std::function<bool(size_t offset, DataSink &sink)>;
 
-using ContentProviderResourceReleaser = std::function<void()>;
+using ContentProviderResourceReleaser = std::function<void(bool success)>;
 
 using ContentReceiverWithProgress =
     std::function<bool(const char *data, size_t data_length, uint64_t offset,
@@ -465,7 +465,7 @@ struct Response {
   Response &operator=(Response &&) = default;
   ~Response() {
     if (content_provider_resource_releaser_) {
-      content_provider_resource_releaser_();
+      content_provider_resource_releaser_(content_provider_success_);
     }
   }
 
@@ -474,6 +474,7 @@ struct Response {
   ContentProvider content_provider_;
   ContentProviderResourceReleaser content_provider_resource_releaser_;
   bool is_chunked_content_provider_ = false;
+  bool content_provider_success_ = false;
 };
 
 class Stream {
@@ -4614,8 +4615,11 @@ inline bool Server::write_response_core(Stream &strm, bool close_connection,
     if (!res.body.empty()) {
       if (!strm.write(res.body)) { ret = false; }
     } else if (res.content_provider_) {
-      if (!write_content_with_provider(strm, req, res, boundary,
-                                       content_type)) {
+      if (write_content_with_provider(strm, req, res, boundary,
+                                      content_type)) {
+        res.content_provider_success_ = true;
+      } else {
+        res.content_provider_success_ = false;
         ret = false;
       }
     }
diff --git a/test/test.cc b/test/test.cc
index 2492de3..b93b1a2 100644
--- a/test/test.cc
+++ b/test/test.cc
@@ -1378,7 +1378,10 @@ protected:
                      (*i)++;
                      return true;
                    },
-                   [i] { delete i; });
+                   [i](bool success) {
+                     EXPECT_TRUE(success);
+                     delete i;
+                   });
              })
         .Get("/streamed",
              [&](const Request & /*req*/, Response &res) {
@@ -1405,7 +1408,10 @@ protected:
                      EXPECT_TRUE(ret);
                      return true;
                    },
-                   [data] { delete data; });
+                   [data](bool success) {
+                     EXPECT_TRUE(success);
+                     delete data;
+                   });
              })
         .Get("/streamed-cancel",
              [&](const Request & /*req*/, Response &res) {
@@ -3567,6 +3573,44 @@ TEST(KeepAliveTest, ReadTimeout) {
   ASSERT_FALSE(svr.is_running());
 }
 
+TEST(ClientProblemDetectionTest, ContentProvider) {
+  Server svr;
+
+  size_t content_length = 1024 * 1024;
+
+  svr.Get("/hi", [&](const Request & /*req*/, Response &res) {
+    res.set_content_provider(
+        content_length, "text/plain",
+        [&](size_t offset, size_t length, DataSink &sink) {
+          auto out_len = std::min(length, static_cast<size_t>(1024));
+          std::string out(out_len, '@');
+          sink.write(out.data(), out_len);
+          return offset < 4096;
+        },
+        [](bool success) { ASSERT_FALSE(success); });
+  });
+
+  auto listen_thread = std::thread([&svr]() { svr.listen("localhost", PORT); });
+  while (!svr.is_running()) {
+    std::this_thread::sleep_for(std::chrono::milliseconds(1));
+  }
+
+  // Give GET time to get a few messages.
+  std::this_thread::sleep_for(std::chrono::seconds(1));
+
+  Client cli("localhost", PORT);
+
+  auto res = cli.Get("/hi", [&](const char * /*data*/, size_t /*data_length*/) {
+    return false;
+  });
+
+  ASSERT_FALSE(res);
+
+  svr.stop();
+  listen_thread.join();
+  ASSERT_FALSE(svr.is_running());
+}
+
 TEST(ErrorHandlerWithContentProviderTest, ErrorHandler) {
   Server svr;