diff --git a/httplib.h b/httplib.h index 2011394..0268232 100644 --- a/httplib.h +++ b/httplib.h @@ -215,7 +215,8 @@ using MultipartFormDataMap = std::multimap; class DataSink { public: - DataSink() = default; + DataSink() : os(&sb_), sb_(*this) {} + DataSink(const DataSink &) = delete; DataSink &operator=(const DataSink &) = delete; DataSink(DataSink &&) = delete; @@ -224,6 +225,24 @@ public: std::function write; std::function done; std::function is_writable; + std::ostream os; + +private: + class data_sink_streambuf : public std::streambuf { + public: + data_sink_streambuf(DataSink &sink) : sink_(sink) {} + + protected: + std::streamsize xsputn(const char *s, std::streamsize n) { + sink_.write(s, static_cast(n)); + return n; + } + + private: + DataSink &sink_; + }; + + data_sink_streambuf sb_; }; using ContentProvider = @@ -2084,22 +2103,20 @@ inline ssize_t write_content(Stream &strm, ContentProvider content_provider, size_t begin_offset = offset; size_t end_offset = offset + length; - ssize_t written_length = 0; + auto ok = true; 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; + if (strm.write(d, l) < 0) { ok = false; } }; + data_sink.is_writable = [&](void) { return strm.is_writable() && ok; }; while (offset < end_offset) { if (!content_provider(offset, end_offset - offset, data_sink)) { return -1; } - if (written_length < 0) { return written_length; } + if (!ok) { return -1; } } return static_cast(offset - begin_offset); diff --git a/test/test.cc b/test/test.cc index c6ee24c..e0464a3 100644 --- a/test/test.cc +++ b/test/test.cc @@ -875,9 +875,9 @@ protected: res.set_chunked_content_provider( [](size_t /*offset*/, DataSink &sink) { EXPECT_TRUE(sink.is_writable()); - sink.write("123", 3); - sink.write("456", 3); - sink.write("789", 3); + sink.os << "123"; + sink.os << "456"; + sink.os << "789"; sink.done(); return true; }); @@ -889,9 +889,9 @@ protected: [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; - case 2: sink.write("789", 3); break; + case 0: sink.os << "123"; break; + case 1: sink.os << "456"; break; + case 2: sink.os << "789"; break; case 3: sink.done(); break; } (*i)++; @@ -903,7 +903,7 @@ protected: [&](const Request & /*req*/, Response &res) { res.set_content_provider( 6, [](size_t offset, size_t /*length*/, DataSink &sink) { - sink.write(offset < 3 ? "a" : "b", 1); + sink.os << (offset < 3 ? "a" : "b"); return true; }); }) @@ -929,8 +929,7 @@ protected: 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()); + sink.os << "data_chunk"; return true; }); }) @@ -1898,7 +1897,7 @@ TEST_F(ServerTest, PutWithContentProvider) { "/put", 3, [](size_t /*offset*/, size_t /*length*/, DataSink &sink) { EXPECT_TRUE(sink.is_writable()); - sink.write("PUT", 3); + sink.os << "PUT"; return true; }, "text/plain"); @@ -1926,7 +1925,7 @@ TEST_F(ServerTest, PutWithContentProviderWithGzip) { "/put", 3, [](size_t /*offset*/, size_t /*length*/, DataSink &sink) { EXPECT_TRUE(sink.is_writable()); - sink.write("PUT", 3); + sink.os << "PUT"; return true; }, "text/plain");