diff --git a/httplib.h b/httplib.h index d4ec6641..26fb23e6 100644 --- a/httplib.h +++ b/httplib.h @@ -1429,9 +1429,16 @@ public: DataSink &operator=(DataSink &&) = delete; std::function write; - std::function is_writable; - std::function done; - std::function done_with_trailer; + + // Only `write` is mandatory. The rest are defaulted so that a provider + // calling one on a writer that does not set it gets sensible behaviour + // rather than std::bad_function_call thrown from a worker thread. Capturing + // `this` is safe: DataSink is neither copyable nor movable. + std::function is_writable = []() { return true; }; + std::function done = []() {}; + std::function done_with_trailer = + [this](const Headers & /*trailer*/) { done(); }; + std::ostream os; private: @@ -8332,6 +8339,7 @@ inline bool write_content_with_progress(Stream &strm, size_t end_offset = offset + length; size_t start_offset = offset; auto ok = true; + auto finished = false; DataSink data_sink; data_sink.write = [&](const char *d, size_t l) -> bool { @@ -8355,7 +8363,12 @@ inline bool write_content_with_progress(Stream &strm, data_sink.is_writable = [&]() -> bool { return strm.is_peer_alive(); }; - while (offset < end_offset && !is_shutting_down()) { + // The body is framed by `length`, so a provider that reports itself done + // early has truncated it. Record that and let the short-body check below + // fail the write, rather than calling the provider again forever. + data_sink.done = [&]() { finished = true; }; + + while (offset < end_offset && !finished && !is_shutting_down()) { if (!strm.wait_writable() || !strm.is_peer_alive()) { error = Error::Write; return false; @@ -8368,7 +8381,7 @@ inline bool write_content_with_progress(Stream &strm, } } - if (offset < end_offset) { // exited due to is_shutting_down(), not completion + if (offset < end_offset) { // done() called early, or is_shutting_down() error = Error::Write; return false; } @@ -15380,6 +15393,7 @@ ClientImpl::send_with_content_provider_and_receiver( if (content_provider) { auto ok = true; + auto finished = false; size_t offset = 0; DataSink data_sink; @@ -15403,13 +15417,27 @@ ClientImpl::send_with_content_provider_and_receiver( return ok; }; - while (ok && offset < content_length) { + // As in detail::write_content_with_progress(): the body is framed by + // content_length, so a provider that finishes early has truncated it. + // Stop and report that instead of calling the provider forever. + data_sink.done = [&]() { finished = true; }; + + while (ok && !finished && offset < content_length) { if (!content_provider(offset, content_length - offset, data_sink)) { error = Error::Canceled; output_error_log(error, &req); return nullptr; } } + + // A short body here means either the provider stopped early or the + // compressor gave up. The branch below reports a failing compressor as + // Error::Compression, so keep the two distinguishable. + if (offset < content_length) { + error = ok ? Error::Write : Error::Compression; + output_error_log(error, &req); + return nullptr; + } } else { if (!compressor->compress(body, content_length, true, [&](const char *data, size_t data_len) { @@ -15712,6 +15740,9 @@ inline ContentProviderWithoutLength ClientImpl::get_multipart_content_provider( DataSink cur_sink; auto has_data = true; cur_sink.write = sink.write; + // Forward is_writable so a provider item asking whether it may keep + // going gets the outer sink's answer rather than the default `true`. + cur_sink.is_writable = sink.is_writable; cur_sink.done = [&]() { has_data = false; }; if (!provider_items[cur_item].provider(offset - cur_start, cur_sink)) { diff --git a/test/test.cc b/test/test.cc index c41c4ef1..984cb434 100644 --- a/test/test.cc +++ b/test/test.cc @@ -10608,6 +10608,122 @@ TEST(ClientProblemDetectionTest, ContentProvider) { } } +TEST(DataSinkTest, OptionalCallbacksAreCallableByDefault) { + // A writer only has to assign `write`. The other three used to be left as + // empty std::functions, so a provider calling one threw + // std::bad_function_call out of the thread running it and took the whole + // process down. + DataSink sink; + + std::string written; + sink.write = [&](const char *data, size_t data_len) { + written.append(data, data_len); + return true; + }; + + EXPECT_TRUE(sink.is_writable()); + sink.done(); + sink.done_with_trailer(Headers{{"X-Trailer", "value"}}); + EXPECT_TRUE(written.empty()); +} + +TEST(DataSinkTest, DoneWithTrailerFallsBackToDone) { + // A sink that cannot carry trailers still has to finish. + DataSink sink; + + auto done_count = 0; + sink.done = [&]() { done_count++; }; + + sink.done_with_trailer(Headers{{"X-Trailer", "value"}}); + EXPECT_EQ(1, done_count); +} + +TEST(DataSinkTest, LengthFramedProviderMayCallDoneAfterWritingEverything) { + Server svr; + + const std::string body(4096, 'x'); + + svr.Get("/", [&](const Request & /*req*/, Response &res) { + res.set_content_provider(body.size(), "text/plain", + [&](size_t offset, size_t length, DataSink &sink) { + sink.write(body.data() + offset, length); + sink.done(); + return true; + }); + }); + + auto port = svr.bind_to_any_port(HOST); + auto listen_thread = std::thread([&svr]() { svr.listen_after_bind(); }); + auto se = detail::scope_exit([&] { + svr.stop(); + listen_thread.join(); + ASSERT_FALSE(svr.is_running()); + }); + + svr.wait_until_ready(); + + Client cli(HOST, port); + auto res = cli.Get("/"); + + ASSERT_TRUE(res) << "Error: " << to_string(res.error()); + EXPECT_EQ(StatusCode::OK_200, res->status); + EXPECT_EQ(body, res->body); +} + +TEST(DataSinkTest, LengthFramedProviderThatFinishesEarlyFailsTheResponse) { + Server svr; + + svr.Get("/", [](const Request & /*req*/, Response &res) { + res.set_content_provider( + 1024, "text/plain", + [](size_t /*offset*/, size_t /*length*/, DataSink &sink) { + sink.write("hello", 5); + sink.done(); // short of the 1024 bytes the response promised + return true; + }); + }); + + auto port = svr.bind_to_any_port(HOST); + auto listen_thread = std::thread([&svr]() { svr.listen_after_bind(); }); + auto se = detail::scope_exit([&] { + svr.stop(); + listen_thread.join(); + ASSERT_FALSE(svr.is_running()); + }); + + svr.wait_until_ready(); + + Client cli(HOST, port); + cli.set_read_timeout(5, 0); + + // The response is short of its Content-Length, so the client cannot read a + // complete body. What matters is that this returns at all: a no-op done() + // would leave the writer looping over a provider that never makes progress. + auto res = cli.Get("/"); + EXPECT_FALSE(res); +} + +#ifdef CPPHTTPLIB_ZLIB_SUPPORT +TEST(DataSinkTest, CompressedRequestProviderThatFinishesEarlyFails) { + // The compressed path builds the whole body before connecting, so this + // fails client-side and never reaches a server. + Client cli(HOST, PORT); + cli.set_compress(true); + + auto res = cli.Post( + "/", 1024, + [](size_t /*offset*/, size_t /*length*/, DataSink &sink) { + sink.write("hello", 5); + sink.done(); // short of the 1024 bytes the request announced + return true; + }, + "text/plain"); + + ASSERT_FALSE(res); + EXPECT_EQ(Error::Write, res.error()); +} +#endif + TEST(ErrorHandlerWithContentProviderTest, ErrorHandler) { Server svr;