From 2c0613f211268ad80833264ca2c26bb4dba701af Mon Sep 17 00:00:00 2001 From: yhirose Date: Wed, 13 May 2020 21:48:14 -0400 Subject: [PATCH] Fix #472 --- README.md | 39 ++++++++++++++++-- httplib.h | 113 ++++++++++++++++++++++++++++----------------------- test/test.cc | 67 ++++++++++++++++++++++-------- 3 files changed, 147 insertions(+), 72 deletions(-) diff --git a/README.md b/README.md index 3c161bc..c8c5ecd 100644 --- a/README.md +++ b/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 diff --git a/httplib.h b/httplib.h index 84c8c16..b647ce5 100644 --- a/httplib.h +++ b/httplib.h @@ -227,7 +227,10 @@ public: }; using ContentProvider = - std::function; + std::function; + +using ChunkedContentProvider = + std::function; using ContentReceiver = std::function; @@ -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 - provider, + size_t length, ContentProvider provider, std::function resource_releaser = [] {}); void set_chunked_content_provider( - std::function provider, + ChunkedContentProvider provider, std::function 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; + + 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; + }; + 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; - }; - - content_provider(offset, end_offset - offset, data_sink); + if (!content_provider(offset, end_offset - offset, data_sink)) { + return -1; + } if (written_length < 0) { return written_length; } } + return static_cast(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; + + ssize_t written_length = 0; + + DataSink data_sink; + data_sink.write = [&](const char *d, size_t l) { + data_available = l > 0; + offset += l; + + // 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()) { - ssize_t written_length = 0; - - DataSink data_sink; - data_sink.write = [&](const char *d, size_t l) { - data_available = l > 0; - offset += l; - - // 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; - }; - - content_provider(offset, 0, data_sink); - + 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 provider, - std::function resource_releaser) { +inline void +Response::set_content_provider(size_t in_length, ContentProvider provider, + std::function 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 provider, + ChunkedContentProvider provider, std::function 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(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 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; diff --git a/test/test.cc b/test/test.cc index f5ad332..4ebb0c3 100644 --- a/test/test.cc +++ b/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(length), DATA_CHUNK_SIZE); sink.write(&d[static_cast(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 responses; auto ret = cli_.send(requests, responses); @@ -2440,6 +2470,7 @@ TEST(ServerStopTest, StopServerWithChunkedTransmission) { auto size = static_cast(sprintf(buffer, "data:%ld\n\n", offset)); sink.write(buffer, size); std::this_thread::sleep_for(std::chrono::seconds(1)); + return true; }); });