diff --git a/dbms/programs/server/HTTPHandler.cpp b/dbms/programs/server/HTTPHandler.cpp index 2d531b3840f..8c69573dd36 100644 --- a/dbms/programs/server/HTTPHandler.cpp +++ b/dbms/programs/server/HTTPHandler.cpp @@ -397,25 +397,23 @@ void HTTPHandler::processQuery( String http_request_compression_method_str = request.get("Content-Encoding", ""); if (!http_request_compression_method_str.empty()) { - do { - if (http_request_compression_method_str == "gzip") - { - in_post = std::make_unique(*in_post_raw, ZlibCompressionMethod::Gzip); - break; - } - if (http_request_compression_method_str == "deflate") - { - in_post = std::make_unique(*in_post_raw, ZlibCompressionMethod::Zlib); - break; - } - if (http_request_compression_method_str == "br") - { - in_post = std::make_unique(*in_post_raw); - break; - } + if (http_request_compression_method_str == "gzip") + { + in_post = std::make_unique(*in_post_raw, ZlibCompressionMethod::Gzip); + } + else if (http_request_compression_method_str == "deflate") + { + in_post = std::make_unique(*in_post_raw, ZlibCompressionMethod::Zlib); + } + else if (http_request_compression_method_str == "br") + { + in_post = std::make_unique(*in_post_raw); + } + else + { throw Exception("Unknown Content-Encoding of HTTP request: " + http_request_compression_method_str, - ErrorCodes::UNKNOWN_COMPRESSION_METHOD); - } while(0); + ErrorCodes::UNKNOWN_COMPRESSION_METHOD); + } } else in_post = std::move(in_post_raw); diff --git a/dbms/src/IO/BrotliReadBuffer.cpp b/dbms/src/IO/BrotliReadBuffer.cpp index f67e9ffc78d..3ddd2cb165e 100644 --- a/dbms/src/IO/BrotliReadBuffer.cpp +++ b/dbms/src/IO/BrotliReadBuffer.cpp @@ -5,14 +5,19 @@ namespace DB BrotliReadBuffer::BrotliReadBuffer(ReadBuffer &in_, size_t buf_size, char *existing_memory, size_t alignment) : BufferWithOwnMemory(buf_size, existing_memory, alignment) , in(in_) + , bstate(BrotliDecoderCreateInstance(nullptr, nullptr, nullptr)) + , bresult(BROTLI_DECODER_RESULT_NEEDS_MORE_INPUT) + , in_available(0) + , in_data(nullptr) + , out_capacity(0) + , out_data(nullptr) , eof(false) { - bstate_ = BrotliDecoderCreateInstance(NULL,NULL,NULL); } BrotliReadBuffer::~BrotliReadBuffer() { - BrotliDecoderDestroyInstance(bstate_); + BrotliDecoderDestroyInstance(bstate); } bool BrotliReadBuffer::nextImpl() @@ -20,20 +25,42 @@ bool BrotliReadBuffer::nextImpl() if (eof) return false; - auto ptr_in = reinterpret_cast(in.position()); - size_t size_in = in.buffer().end() - in.position(); + if (!in_available) + { + in.nextIfAtEnd(); + in_available = in.buffer().end() - in.position(); + in_data = reinterpret_cast(in.position()); + } - auto ptr_out = reinterpret_cast(internal_buffer.begin()); - size_t size_out = internal_buffer.size(); + if (bresult == BROTLI_DECODER_RESULT_NEEDS_MORE_INPUT && (!in_available || in.eof())) + { + throw Exception(std::string("brotli decode error"), ErrorCodes::CANNOT_READ_ALL_DATA); + } - BrotliDecoderDecompressStream(bstate_, &size_in, &ptr_in, &size_out, &ptr_out, nullptr); + out_capacity = internal_buffer.size(); + out_data = reinterpret_cast(internal_buffer.begin()); - in.position() = in.buffer().end() - size_in; - working_buffer.resize(internal_buffer.size() - size_out); + bresult = BrotliDecoderDecompressStream(bstate, &in_available, &in_data, &out_capacity, &out_data, nullptr); - if (in.eof()) { - eof = true; - return working_buffer.size() != 0; + in.position() = in.buffer().end() - in_available; + working_buffer.resize(internal_buffer.size() - out_capacity); + + if (bresult == BROTLI_DECODER_RESULT_SUCCESS) + { + if (in.eof()) + { + eof = true; + return working_buffer.size() != 0; + } + else + { + throw Exception(std::string("brotli decode error"), ErrorCodes::CANNOT_READ_ALL_DATA); + } + } + + if (bresult == BROTLI_DECODER_RESULT_ERROR) + { + throw Exception(std::string("brotli decode error"), ErrorCodes::CANNOT_READ_ALL_DATA); } return true; diff --git a/dbms/src/IO/BrotliReadBuffer.h b/dbms/src/IO/BrotliReadBuffer.h index cc06587add4..400d343984c 100644 --- a/dbms/src/IO/BrotliReadBuffer.h +++ b/dbms/src/IO/BrotliReadBuffer.h @@ -5,9 +5,11 @@ #include -namespace DB { +namespace DB +{ -class BrotliReadBuffer : public BufferWithOwnMemory { +class BrotliReadBuffer : public BufferWithOwnMemory +{ public: BrotliReadBuffer( ReadBuffer &in_, @@ -21,7 +23,16 @@ private: bool nextImpl() override; ReadBuffer ∈ - BrotliDecoderState * bstate_; + + BrotliDecoderState * bstate; + BrotliDecoderResult bresult; + + size_t in_available; + const uint8_t * in_data; + + size_t out_capacity; + uint8_t * out_data; + bool eof; }; } \ No newline at end of file