Skip to content

Commit e5dd410

Browse files
yhiroseyhirose
committed
Added set_content_provider without content length
1 parent 951e469 commit e5dd410

File tree

2 files changed

+123
-29
lines changed

2 files changed

+123
-29
lines changed

‎httplib.h‎

Lines changed: 82 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -299,7 +299,7 @@ class DataSink {
299299
using ContentProvider =
300300
std::function<bool(size_t offset, size_t length, DataSink &sink)>;
301301

302-
using ChunkedContentProvider =
302+
using ContentProviderWithoutLength =
303303
std::function<bool(size_t offset, DataSink &sink)>;
304304

305305
using ContentReceiver =
@@ -404,8 +404,12 @@ struct Response {
404404
size_t length, const char *content_type, ContentProvider provider,
405405
std::function<void()> resource_releaser = [] {});
406406

407+
void set_content_provider(
408+
const char *content_type, ContentProviderWithoutLength provider,
409+
std::function<void()> resource_releaser = [] {});
410+
407411
void set_chunked_content_provider(
408-
const char *content_type, ChunkedContentProvider provider,
412+
const char *content_type, ContentProviderWithoutLength provider,
409413
std::function<void()> resource_releaser = [] {});
410414

411415
Response() = default;
@@ -423,6 +427,7 @@ struct Response {
423427
size_t content_length_ = 0;
424428
ContentProvider content_provider_;
425429
std::function<void()> content_provider_resource_releaser_;
430+
bool is_chunked_content_provider = false;
426431
};
427432

428433
class Stream {
@@ -2664,19 +2669,19 @@ inline ssize_t write_content(Stream &strm, ContentProvider content_provider,
26642669
size_t offset, size_t length, T is_shutting_down) {
26652670
size_t begin_offset = offset;
26662671
size_t end_offset = offset + length;
2667-
26682672
auto ok = true;
2669-
26702673
DataSink data_sink;
2674+
26712675
data_sink.write = [&](const char *d, size_t l) {
26722676
if (ok) {
26732677
offset += l;
26742678
if (!write_data(strm, d, l)) { ok = false; }
26752679
}
26762680
};
2681+
26772682
data_sink.is_writable = [&](void) { return ok && strm.is_writable(); };
26782683

2679-
while (ok && offset < end_offset && !is_shutting_down()) {
2684+
while (offset < end_offset && !is_shutting_down()) {
26802685
if (!content_provider(offset, end_offset - offset, data_sink)) {
26812686
return -1;
26822687
}
@@ -2686,14 +2691,41 @@ inline ssize_t write_content(Stream &strm, ContentProvider content_provider,
26862691
return static_cast<ssize_t>(offset - begin_offset);
26872692
}
26882693

2694+
template <typename T>
2695+
inline ssize_t write_content_without_length(Stream &strm,
2696+
ContentProvider content_provider,
2697+
T is_shutting_down) {
2698+
size_t offset = 0;
2699+
auto data_available = true;
2700+
auto ok = true;
2701+
DataSink data_sink;
2702+
2703+
data_sink.write = [&](const char *d, size_t l) {
2704+
if (ok) {
2705+
offset += l;
2706+
if (!write_data(strm, d, l)) { ok = false; }
2707+
}
2708+
};
2709+
2710+
data_sink.done = [&](void) { data_available = false; };
2711+
2712+
data_sink.is_writable = [&](void) { return ok && strm.is_writable(); };
2713+
2714+
while (data_available && !is_shutting_down()) {
2715+
if (!content_provider(offset, 0, data_sink)) { return -1; }
2716+
if (!ok) { return -1; }
2717+
}
2718+
2719+
return static_cast<ssize_t>(offset);
2720+
}
2721+
26892722
template <typename T, typename U>
26902723
inline ssize_t write_content_chunked(Stream &strm,
26912724
ContentProvider content_provider,
26922725
T is_shutting_down, U &compressor) {
26932726
size_t offset = 0;
26942727
auto data_available = true;
26952728
ssize_t total_written_length = 0;
2696-
26972729
auto ok = true;
26982730
DataSink data_sink;
26992731

@@ -3544,17 +3576,31 @@ Response::set_content_provider(size_t in_length, const char *content_type,
35443576
return provider(offset, length, sink);
35453577
};
35463578
content_provider_resource_releaser_ = resource_releaser;
3579+
is_chunked_content_provider = false;
3580+
}
3581+
3582+
inline void Response::set_content_provider(
3583+
const char *content_type, ContentProviderWithoutLength provider,
3584+
std::function<void()> resource_releaser) {
3585+
set_header("Content-Type", content_type);
3586+
content_length_ = 0;
3587+
content_provider_ = [provider](size_t offset, size_t, DataSink &sink) {
3588+
return provider(offset, sink);
3589+
};
3590+
content_provider_resource_releaser_ = resource_releaser;
3591+
is_chunked_content_provider = false;
35473592
}
35483593

35493594
inline void Response::set_chunked_content_provider(
3550-
const char *content_type, ChunkedContentProvider provider,
3595+
const char *content_type, ContentProviderWithoutLength provider,
35513596
std::function<void()> resource_releaser) {
35523597
set_header("Content-Type", content_type);
35533598
content_length_ = 0;
35543599
content_provider_ = [provider](size_t offset, size_t, DataSink &sink) {
35553600
return provider(offset, sink);
35563601
};
35573602
content_provider_resource_releaser_ = resource_releaser;
3603+
is_chunked_content_provider = true;
35583604
}
35593605

35603606
// Rstream implementation
@@ -3893,7 +3939,7 @@ inline bool Server::write_response(Stream &strm, bool close_connection,
38933939
}
38943940

38953941
if (!res.has_header("Content-Type") &&
3896-
(!res.body.empty() || res.content_length_ > 0)) {
3942+
(!res.body.empty() || res.content_length_ > 0 || res.content_provider_)) {
38973943
res.set_header("Content-Type", "text/plain");
38983944
}
38993945

@@ -3939,11 +3985,13 @@ inline bool Server::write_response(Stream &strm, bool close_connection,
39393985
res.set_header("Content-Length", std::to_string(length));
39403986
} else {
39413987
if (res.content_provider_) {
3942-
res.set_header("Transfer-Encoding", "chunked");
3943-
if (type == detail::EncodingType::Gzip) {
3944-
res.set_header("Content-Encoding", "gzip");
3945-
} else if (type == detail::EncodingType::Brotli) {
3946-
res.set_header("Content-Encoding", "br");
3988+
if (res.is_chunked_content_provider) {
3989+
res.set_header("Transfer-Encoding", "chunked");
3990+
if (type == detail::EncodingType::Gzip) {
3991+
res.set_header("Content-Encoding", "gzip");
3992+
} else if (type == detail::EncodingType::Brotli) {
3993+
res.set_header("Content-Encoding", "br");
3994+
}
39473995
}
39483996
} else {
39493997
res.set_header("Content-Length", "0");
@@ -4033,7 +4081,7 @@ Server::write_content_with_provider(Stream &strm, const Request &req,
40334081
return this->svr_sock_ == INVALID_SOCKET;
40344082
};
40354083

4036-
if (res.content_length_) {
4084+
if (res.content_length_ > 0) {
40374085
if (req.ranges.empty()) {
40384086
if (detail::write_content(strm, res.content_provider_, 0,
40394087
res.content_length_, is_shutting_down) < 0) {
@@ -4055,25 +4103,32 @@ Server::write_content_with_provider(Stream &strm, const Request &req,
40554103
}
40564104
}
40574105
} else {
4058-
auto type = detail::encoding_type(req, res);
4106+
if (res.is_chunked_content_provider) {
4107+
auto type = detail::encoding_type(req, res);
40594108

4060-
std::shared_ptr<detail::compressor> compressor;
4061-
if (type == detail::EncodingType::Gzip) {
4109+
std::shared_ptr<detail::compressor> compressor;
4110+
if (type == detail::EncodingType::Gzip) {
40624111
#ifdef CPPHTTPLIB_ZLIB_SUPPORT
4063-
compressor = std::make_shared<detail::gzip_compressor>();
4112+
compressor = std::make_shared<detail::gzip_compressor>();
40644113
#endif
4065-
} else if (type == detail::EncodingType::Brotli) {
4114+
} else if (type == detail::EncodingType::Brotli) {
40664115
#ifdef CPPHTTPLIB_BROTLI_SUPPORT
4067-
compressor = std::make_shared<detail::brotli_compressor>();
4116+
compressor = std::make_shared<detail::brotli_compressor>();
40684117
#endif
4069-
} else {
4070-
compressor = std::make_shared<detail::nocompressor>();
4071-
}
4072-
assert(compressor != nullptr);
4118+
} else {
4119+
compressor = std::make_shared<detail::nocompressor>();
4120+
}
4121+
assert(compressor != nullptr);
40734122

4074-
if (detail::write_content_chunked(strm, res.content_provider_,
4075-
is_shutting_down, *compressor) < 0) {
4076-
return false;
4123+
if (detail::write_content_chunked(strm, res.content_provider_,
4124+
is_shutting_down, *compressor) < 0) {
4125+
return false;
4126+
}
4127+
} else {
4128+
if (detail::write_content_without_length(strm, res.content_provider_,
4129+
is_shutting_down) < 0) {
4130+
return false;
4131+
}
40774132
}
40784133
}
40794134
return true;

‎test/test.cc‎

Lines changed: 41 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1895,8 +1895,7 @@ TEST_F(ServerTest, ClientStop) {
18951895
auto res = cli_.Get("/streamed-cancel",
18961896
[&](const char *, uint64_t) { return true; });
18971897
ASSERT_TRUE(!res);
1898-
EXPECT_TRUE(res.error() == Error::Canceled ||
1899-
res.error() == Error::Read);
1898+
EXPECT_TRUE(res.error() == Error::Canceled || res.error() == Error::Read);
19001899
}));
19011900
}
19021901

@@ -2730,6 +2729,46 @@ TEST(ServerStopTest, StopServerWithChunkedTransmission) {
27302729
ASSERT_FALSE(svr.is_running());
27312730
}
27322731

2732+
TEST(StreamingTest, NoContentLengthStreaming) {
2733+
Server svr;
2734+
2735+
svr.Get("/stream", [](const Request & /*req*/, Response &res) {
2736+
res.set_content_provider(
2737+
"text/plain", [](size_t offset, DataSink &sink) {
2738+
if (offset < 6) {
2739+
sink.os << (offset < 3 ? "a" : "b");
2740+
} else {
2741+
sink.done();
2742+
}
2743+
return true;
2744+
});
2745+
});
2746+
2747+
auto listen_thread = std::thread([&svr]() { svr.listen("localhost", PORT); });
2748+
while (!svr.is_running()) {
2749+
std::this_thread::sleep_for(std::chrono::milliseconds(1));
2750+
}
2751+
2752+
Client client(HOST, PORT);
2753+
2754+
auto get_thread = std::thread([&client]() {
2755+
auto res = client.Get("/stream", [](const char *data, size_t len) -> bool {
2756+
EXPECT_EQ("aaabbb", std::string(data, len));
2757+
return true;
2758+
});
2759+
});
2760+
2761+
// Give GET time to get a few messages.
2762+
std::this_thread::sleep_for(std::chrono::milliseconds(500));
2763+
2764+
svr.stop();
2765+
2766+
listen_thread.join();
2767+
get_thread.join();
2768+
2769+
ASSERT_FALSE(svr.is_running());
2770+
}
2771+
27332772
TEST(MountTest, Unmount) {
27342773
Server svr;
27352774

0 commit comments

Comments
 (0)