From 7e92ffec48228afecb46639ac0fbfbcd4560595b Mon Sep 17 00:00:00 2001
From: Steffen Schuemann <s.schuemann@pobox.com>
Date: Sat, 14 Sep 2019 14:55:12 +0200
Subject: [PATCH] Added new Client::Get variant that combines a ContentReceiver
 with a new ResponseHandler

While trying to implement streaming of internet radio, where a ContentReceiver is needed to handle the audio data, I had the problem, that important information about the stream data is part of the HTTP header (e.g. size of audio chunks between meta data), so I added a ResponseHandler and a new Get variant, to gain access to the header before handling the first chunk of data.

The ResponseHandler can abort the request by returning false, in the same way as the ContentReceiver.

A test case was also added.
---
 httplib.h    | 42 +++++++++++++++++++++++++++++++++++++++---
 test/test.cc | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 3 deletions(-)

diff --git a/httplib.h b/httplib.h
index f9d59a1..d4b60b5 100644
--- a/httplib.h
+++ b/httplib.h
@@ -153,6 +153,9 @@ typedef std::function<bool(const char *data, size_t data_length, size_t offset,
 
 typedef std::function<bool(uint64_t current, uint64_t total)> Progress;
 
+struct Response;
+typedef std::function<bool(const Response& response)> ResponseHandler;
+
 struct MultipartFile {
   std::string filename;
   std::string content_type;
@@ -188,6 +191,7 @@ struct Request {
 
   // for client
   size_t redirect_count = CPPHTTPLIB_REDIRECT_MAX_COUNT;
+  ResponseHandler response_handler;
   ContentReceiver content_receiver;
   Progress progress;
 
@@ -518,6 +522,15 @@ public:
                                 ContentReceiver content_receiver,
                                 Progress progress);
 
+  std::shared_ptr<Response> Get(const char *path, const Headers &headers,
+                                ResponseHandler response_handler,
+                                ContentReceiver content_receiver);
+
+  std::shared_ptr<Response> Get(const char *path, const Headers &headers,
+                                ResponseHandler response_handler,
+                                ContentReceiver content_receiver,
+                                Progress progress);
+
   std::shared_ptr<Response> Head(const char *path);
 
   std::shared_ptr<Response> Head(const char *path, const Headers &headers);
@@ -2869,6 +2882,12 @@ inline bool Client::process_request(Stream &strm, const Request &req,
     connection_close = true;
   }
 
+  if (req.response_handler) {
+    if(!req.response_handler(res)) {
+      return false;
+    }
+  }
+
   // Body
   if (req.method != "HEAD") {
     detail::ContentReceiverCore out = [&](const char *buf, size_t n) {
@@ -2940,30 +2959,47 @@ Client::Get(const char *path, const Headers &headers, Progress progress) {
 inline std::shared_ptr<Response> Client::Get(const char *path,
                                              ContentReceiver content_receiver) {
   Progress dummy;
-  return Get(path, Headers(), content_receiver, dummy);
+  return Get(path, Headers(), nullptr, content_receiver, dummy);
 }
 
 inline std::shared_ptr<Response> Client::Get(const char *path,
                                              ContentReceiver content_receiver,
                                              Progress progress) {
-  return Get(path, Headers(), content_receiver, progress);
+  return Get(path, Headers(), nullptr, content_receiver, progress);
 }
 
 inline std::shared_ptr<Response> Client::Get(const char *path,
                                              const Headers &headers,
                                              ContentReceiver content_receiver) {
   Progress dummy;
-  return Get(path, headers, content_receiver, dummy);
+  return Get(path, headers, nullptr, content_receiver, dummy);
 }
 
 inline std::shared_ptr<Response> Client::Get(const char *path,
                                              const Headers &headers,
                                              ContentReceiver content_receiver,
                                              Progress progress) {
+  return Get(path, headers, nullptr, content_receiver, progress);
+}
+
+inline std::shared_ptr<Response> Client::Get(const char *path,
+                                             const Headers &headers,
+                                             ResponseHandler response_handler,
+                                             ContentReceiver content_receiver) {
+  Progress dummy;
+  return Get(path, headers, response_handler, content_receiver, dummy);
+}
+
+inline std::shared_ptr<Response> Client::Get(const char *path,
+                                             const Headers &headers,
+                                             ResponseHandler response_handler,
+                                             ContentReceiver content_receiver,
+                                             Progress progress) {
   Request req;
   req.method = "GET";
   req.path = path;
   req.headers = headers;
+  req.response_handler = response_handler;
   req.content_receiver = content_receiver;
   req.progress = progress;
 
diff --git a/test/test.cc b/test/test.cc
index 57049c3..24651a0 100644
--- a/test/test.cc
+++ b/test/test.cc
@@ -242,6 +242,38 @@ TEST(ChunkedEncodingTest, WithContentReceiver) {
   EXPECT_EQ(out, body);
 }
 
+TEST(ChunkedEncodingTest, WithResponseHandlerAndContentReceiver) {
+  auto host = "www.httpwatch.com";
+  auto sec = 2;
+
+#ifdef CPPHTTPLIB_OPENSSL_SUPPORT
+  auto port = 443;
+  httplib::SSLClient cli(host, port, sec);
+#else
+  auto port = 80;
+  httplib::Client cli(host, port, sec);
+#endif
+
+  std::string body;
+  auto res =
+      cli.Get("/httpgallery/chunked/chunkedimage.aspx?0.4153841143030137", Headers(),
+              [&](const Response& response) {
+                EXPECT_EQ(200, response.status);
+                return true;
+              },
+              [&](const char *data, size_t data_length, uint64_t, uint64_t) {
+                body.append(data, data_length);
+                return true;
+              });
+  ASSERT_TRUE(res != nullptr);
+
+  std::string out;
+  httplib::detail::read_file("./image.jpg", out);
+
+  EXPECT_EQ(200, res->status);
+  EXPECT_EQ(out, body);
+}
+
 TEST(RangeTest, FromHTTPBin) {
   auto host = "httpbin.org";
   auto sec = 5;