diff --git a/httplib.h b/httplib.h index e5798f72..76960c80 100644 --- a/httplib.h +++ b/httplib.h @@ -3797,6 +3797,8 @@ template bool check_and_write_headers(Stream &strm, Headers &headers, T header_writer, Error &error); +bool write_data(Stream &strm, const char *d, size_t l); + std::string trim_copy(const std::string &s); void divide( @@ -5815,32 +5817,40 @@ inline bool write_websocket_frame(Stream &strm, ws::Opcode opcode, if (len < 126) { header[1] = static_cast(len); if (mask) { header[1] |= 0x80; } - if (strm.write(reinterpret_cast(header), 2) < 0) { return false; } + if (!write_data(strm, reinterpret_cast(header), 2)) { + return false; + } } else if (len <= 0xFFFF) { header[1] = 126; if (mask) { header[1] |= 0x80; } - if (strm.write(reinterpret_cast(header), 2) < 0) { return false; } + if (!write_data(strm, reinterpret_cast(header), 2)) { + return false; + } uint8_t ext[2]; ext[0] = static_cast((len >> 8) & 0xFF); ext[1] = static_cast(len & 0xFF); - if (strm.write(reinterpret_cast(ext), 2) < 0) { return false; } + if (!write_data(strm, reinterpret_cast(ext), 2)) { return false; } } else { header[1] = 127; if (mask) { header[1] |= 0x80; } - if (strm.write(reinterpret_cast(header), 2) < 0) { return false; } + if (!write_data(strm, reinterpret_cast(header), 2)) { + return false; + } uint8_t ext[8]; for (int i = 7; i >= 0; i--) { ext[7 - i] = static_cast((static_cast(len) >> (i * 8)) & 0xFF); } - if (strm.write(reinterpret_cast(ext), 8) < 0) { return false; } + if (!write_data(strm, reinterpret_cast(ext), 8)) { return false; } } if (mask) { // Generate random mask key uint8_t mask_key[4]; fill_random_bytes(mask_key, sizeof(mask_key)); - if (strm.write(reinterpret_cast(mask_key), 4) < 0) { return false; } + if (!write_data(strm, reinterpret_cast(mask_key), 4)) { + return false; + } // Write masked payload in chunks const size_t chunk_size = 4096; @@ -5851,11 +5861,11 @@ inline bool write_websocket_frame(Stream &strm, ws::Opcode opcode, buf[i] = data[offset + i] ^ static_cast(mask_key[(offset + i) % 4]); } - if (strm.write(buf.data(), n) < 0) { return false; } + if (!write_data(strm, buf.data(), n)) { return false; } } } else { if (len > 0) { - if (strm.write(data, len) < 0) { return false; } + if (!write_data(strm, data, len)) { return false; } } } @@ -15969,7 +15979,7 @@ inline ClientImpl::StreamHandle ClientImpl::open_stream_(Request &req) { } if (!req.body.empty()) { - if (strm.write(req.body.data(), req.body.size()) < 0) { + if (!detail::write_data(strm, req.body.data(), req.body.size())) { handle.error = Error::Write; handle.response.reset(); return handle; diff --git a/test/test.cc b/test/test.cc index dc1917fd..fa0bbb2c 100644 --- a/test/test.cc +++ b/test/test.cc @@ -23859,6 +23859,55 @@ TEST(OpenStreamMalformedContentLength, OutOfRange) { server_thread.join(); } +#ifndef _WIN32 +// As above for the body of a streamed request: a send() cut short by the write +// timeout must not leave the rest of the body unsent. +TEST(OpenStreamPartialWriteTest, BodyIsWrittenWhole) { + signal(SIGPIPE, SIG_IGN); + + Server svr; + svr.set_socket_options([](socket_t sock) { + int size = 8192; + setsockopt(sock, SOL_SOCKET, SO_RCVBUF, + reinterpret_cast(&size), sizeof(size)); + }); + svr.Post("/size", [](const Request &, Response &res, + const ContentReader &content_reader) { + // Outlast the client's write timeout before reading the body. + std::this_thread::sleep_for(std::chrono::milliseconds(750)); + size_t received = 0; + content_reader([&](const char *, size_t len) { + received += len; + return true; + }); + res.set_content(std::to_string(received), "text/plain"); + }); + + auto port = svr.bind_to_any_port(HOST); + std::thread t([&]() { svr.listen_after_bind(); }); + auto se = detail::scope_exit([&] { + svr.stop(); + t.join(); + }); + svr.wait_until_ready(); + + Client cli(HOST, port); + cli.set_write_timeout(0, 500000); + cli.set_socket_options([](socket_t sock) { + int size = 8192; + setsockopt(sock, SOL_SOCKET, SO_SNDBUF, + reinterpret_cast(&size), sizeof(size)); + }); + + const std::string body(4 * 1024 * 1024, 'A'); + auto handle = cli.open_stream("POST", "/size", {}, {}, body, + "application/octet-stream"); + ASSERT_TRUE(handle.is_valid()) << to_string(handle.error); + EXPECT_EQ(StatusCode::OK_200, handle.response->status); + EXPECT_EQ(std::to_string(body.size()), read_all(handle)); +} +#endif + // Serves `response` to the single request `fn` makes with a fresh client. template static void with_single_response(const std::string &response, Fn fn) { @@ -28195,6 +28244,64 @@ TEST_F(WebSocketIntegrationTest, ChronoTimeoutSetters) { EXPECT_LT(elapsed, 30000); } +#ifndef _WIN32 +// A send() that times out after part of the data went out returns the count +// written so far. Every byte of the frame must still be sent, or the peer reads +// the next frame's header and payload as the rest of this one. +TEST(WebSocketPartialWriteTest, FrameIsWrittenWhole) { + signal(SIGPIPE, SIG_IGN); + + const std::string large(4 * 1024 * 1024, 'A'); + std::atomic large_sent{false}; + std::atomic small_sent{false}; + + Server svr; + svr.set_write_timeout(0, 500000); + svr.set_socket_options([](socket_t sock) { + int size = 8192; + setsockopt(sock, SOL_SOCKET, SO_SNDBUF, + reinterpret_cast(&size), sizeof(size)); + }); + svr.WebSocket("/ws", [&](const Request &, ws::WebSocket &ws) { + large_sent = ws.send(large.data(), large.size()); + small_sent = ws.send("after"); + std::string msg; + while (ws.read(msg)) {} + }); + + auto port = svr.bind_to_any_port(HOST); + std::thread t([&]() { svr.listen_after_bind(); }); + auto se = detail::scope_exit([&] { + svr.stop(); + t.join(); + }); + svr.wait_until_ready(); + + ws::WebSocketClient client("ws://" + std::string(HOST) + ":" + + std::to_string(port) + "/ws"); + client.set_read_timeout(5); + client.set_socket_options([](socket_t sock) { + int size = 8192; + setsockopt(sock, SOL_SOCKET, SO_RCVBUF, + reinterpret_cast(&size), sizeof(size)); + }); + ASSERT_TRUE(client.connect()); + + // Outlast the server's write timeout, then drain. + std::this_thread::sleep_for(std::chrono::milliseconds(750)); + + std::string msg; + EXPECT_EQ(ws::ReadResult::Binary, client.read(msg)); + EXPECT_TRUE(msg == large); + EXPECT_EQ(ws::ReadResult::Text, client.read(msg)); + EXPECT_EQ("after", msg); + EXPECT_TRUE(large_sent); + EXPECT_TRUE(small_sent); + + client.close(); +} +#endif + TEST(WebSocketPreRoutingTest, RejectWithoutAuth) { Server svr;