diff --git a/base/poco/Net/include/Poco/Net/HTTPChunkedStream.h b/base/poco/Net/include/Poco/Net/HTTPChunkedStream.h index a6576aa561d0..604e66da5786 100644 --- a/base/poco/Net/include/Poco/Net/HTTPChunkedStream.h +++ b/base/poco/Net/include/Poco/Net/HTTPChunkedStream.h @@ -45,7 +45,7 @@ namespace Net ~HTTPChunkedStreamBuf(); void close(); - bool isComplete() const { return _chunk == std::char_traits::eof(); } + bool isComplete(bool read_from_device_to_check_eof = false) noexcept; protected: int readFromDevice(char * buffer, std::streamsize length); @@ -70,8 +70,6 @@ namespace Net ~HTTPChunkedIOS(); HTTPChunkedStreamBuf * rdbuf(); - bool isComplete() const { return _buf.isComplete(); } - protected: HTTPChunkedStreamBuf _buf; }; @@ -83,6 +81,8 @@ namespace Net public: HTTPChunkedInputStream(HTTPSession & session); ~HTTPChunkedInputStream(); + + bool isComplete() { return _buf.isComplete(/*read_from_device_to_check_eof*/ true); } }; @@ -92,6 +92,8 @@ namespace Net public: HTTPChunkedOutputStream(HTTPSession & session); ~HTTPChunkedOutputStream(); + + bool isComplete() { return _buf.isComplete(); } }; diff --git a/base/poco/Net/src/HTTPChunkedStream.cpp b/base/poco/Net/src/HTTPChunkedStream.cpp index 16ed1e71c314..ca0b73128ff4 100644 --- a/base/poco/Net/src/HTTPChunkedStream.cpp +++ b/base/poco/Net/src/HTTPChunkedStream.cpp @@ -54,7 +54,7 @@ void HTTPChunkedStreamBuf::close() sync(); _session.write("0\r\n\r\n", 5); - _chunk = std::char_traits::eof(); + _chunk = std::char_traits::eof(); } } @@ -90,7 +90,7 @@ unsigned int HTTPChunkedStreamBuf::parseChunkLen() if (size_t pos = line.find(';'); pos != std::string::npos) line.resize(pos); - unsigned chunkLen; + unsigned chunkLen = 0; if (NumberParser::tryParseHex(line, chunkLen)) return chunkLen; else @@ -118,6 +118,9 @@ int HTTPChunkedStreamBuf::readFromDevice(char* buffer, std::streamsize length) if (_chunk > 0) { + if (length == 0) + return 0; + if (length > _chunk) length = _chunk; int n = _session.read(buffer, length); if (n > 0) @@ -128,7 +131,7 @@ int HTTPChunkedStreamBuf::readFromDevice(char* buffer, std::streamsize length) if (_chunk == 0) skipCRLF(); return n; } - else + else { skipCRLF(); _chunk = eof; @@ -137,6 +140,26 @@ int HTTPChunkedStreamBuf::readFromDevice(char* buffer, std::streamsize length) } +bool HTTPChunkedStreamBuf::isComplete(bool read_from_device_to_check_eof) noexcept +{ + if (read_from_device_to_check_eof) + { + try + { + /// If the stream is closed without final last chunk + /// "Unexpected EOF" exception would be thrown + readFromDevice(nullptr, 0); + } + catch (...) + { + return false; + } + } + + return _chunk == std::char_traits::eof(); +} + + int HTTPChunkedStreamBuf::writeToDevice(const char* buffer, std::streamsize length) { _chunkBuffer.clear(); diff --git a/src/Common/HTTPConnectionPool.cpp b/src/Common/HTTPConnectionPool.cpp index cf272e612b03..4443fe4e15a6 100644 --- a/src/Common/HTTPConnectionPool.cpp +++ b/src/Common/HTTPConnectionPool.cpp @@ -456,6 +456,8 @@ class EndpointConnectionPool : public std::enable_shared_from_thisatConnectionDestroy(); diff --git a/src/Common/tests/gtest_http_chunked_stream.cpp b/src/Common/tests/gtest_http_chunked_stream.cpp new file mode 100644 index 000000000000..c737f4794105 --- /dev/null +++ b/src/Common/tests/gtest_http_chunked_stream.cpp @@ -0,0 +1,16 @@ +#include + +#include +#include + + +TEST(HTTPChunkedStreamBuf, IsCompleteHandlesInvalidSocketException) +{ + Poco::Net::HTTPClientSession session; + Poco::Net::HTTPChunkedStreamBuf buf(session, std::ios::in); + + /// Default-initialized socket throws InvalidSocketException in SocketImpl::receiveBytes, + /// which HTTPChunkedStreamBuf::isComplete should swallow and return false. + bool complete = buf.isComplete(true); + ASSERT_FALSE(complete); +} diff --git a/tests/queries/0_stateless/02537_distributed_loosing_files_after_exception.reference b/tests/queries/0_stateless/02537_distributed_losing_files_after_exception.reference similarity index 100% rename from tests/queries/0_stateless/02537_distributed_loosing_files_after_exception.reference rename to tests/queries/0_stateless/02537_distributed_losing_files_after_exception.reference diff --git a/tests/queries/0_stateless/02537_distributed_loosing_files_after_exception.sql.j2 b/tests/queries/0_stateless/02537_distributed_losing_files_after_exception.sql.j2 similarity index 100% rename from tests/queries/0_stateless/02537_distributed_loosing_files_after_exception.sql.j2 rename to tests/queries/0_stateless/02537_distributed_losing_files_after_exception.sql.j2