diff --git a/dbms/include/DB/IO/BufferWithOwnMemory.h b/dbms/include/DB/IO/BufferWithOwnMemory.h index 8afa07e66c8..a676ed6b22c 100644 --- a/dbms/include/DB/IO/BufferWithOwnMemory.h +++ b/dbms/include/DB/IO/BufferWithOwnMemory.h @@ -19,7 +19,7 @@ class BufferWithOwnMemory : public Base protected: std::vector memory; public: - BufferWithOwnMemory(size_t size = DBMS_DEFAULT_BUFFER_SIZE) : Base(NULL, size), memory(size) + BufferWithOwnMemory(size_t size = DBMS_DEFAULT_BUFFER_SIZE) : Base(NULL, 0), memory(size) { Base::set(&memory[0], size); } diff --git a/dbms/include/DB/IO/CompressedReadBuffer.h b/dbms/include/DB/IO/CompressedReadBuffer.h index ed158fc6277..b999c5ee3aa 100644 --- a/dbms/include/DB/IO/CompressedReadBuffer.h +++ b/dbms/include/DB/IO/CompressedReadBuffer.h @@ -13,9 +13,6 @@ #include -#define DBMS_COMPRESSED_READ_BUFFER_MAX_COMPRESSED_SIZE 0x40000000ULL /// 1GB - - namespace DB { @@ -38,7 +35,7 @@ private: in.readStrict(&compressed_buffer[0], QUICKLZ_HEADER_SIZE); size_t size_compressed = qlz_size_compressed(&compressed_buffer[0]); - if (size_compressed > DBMS_COMPRESSED_READ_BUFFER_MAX_COMPRESSED_SIZE) + if (size_compressed > DBMS_MAX_COMPRESSED_SIZE) throw Exception("Too large size_compressed. Most likely corrupted data.", ErrorCodes::TOO_LARGE_SIZE_COMPRESSED); size_t size_decompressed = qlz_size_decompressed(&compressed_buffer[0]); diff --git a/dbms/include/DB/IO/CompressedStream.h b/dbms/include/DB/IO/CompressedStream.h index 1c1e0e55eff..b096995e18d 100644 --- a/dbms/include/DB/IO/CompressedStream.h +++ b/dbms/include/DB/IO/CompressedStream.h @@ -3,6 +3,7 @@ /** Общие для CompressingStream.h и DecompressingStream.h дефайны */ +#define DBMS_MAX_COMPRESSED_SIZE 0x40000000ULL /// 1GB #define DBMS_STREAM_BUFFER_SIZE 4096 #define DBMS_COMPRESSING_STREAM_BUFFER_SIZE 1048576 #define QUICKLZ_ADDITIONAL_SPACE 400 diff --git a/dbms/src/IO/CompressedInputStream.cpp b/dbms/src/IO/CompressedInputStream.cpp index 5a8edfe8e2c..d666b2e39f8 100644 --- a/dbms/src/IO/CompressedInputStream.cpp +++ b/dbms/src/IO/CompressedInputStream.cpp @@ -1,7 +1,10 @@ #include +#include #include +#include +#include #include @@ -30,13 +33,25 @@ void DecompressingStreamBuf::getChunk(std::vector & res) void DecompressingStreamBuf::readCompressedChunk() { + /// прочитаем чексумму + uint128 checksum; + p_istr->read(reinterpret_cast(&checksum), sizeof(checksum)); + + if (p_istr->eof()) + return; + if (!p_istr->good()) + throw Exception("Cannot read all data.", ErrorCodes::CANNOT_READ_ALL_DATA); + /// прочитаем заголовок p_istr->read(&compressed_buffer[0], QUICKLZ_HEADER_SIZE); if (!p_istr->good()) - return; + throw Exception("Cannot read all data.", ErrorCodes::CANNOT_READ_ALL_DATA); size_t size_compressed = qlz_size_compressed(&compressed_buffer[0]); + if (size_compressed > DBMS_MAX_COMPRESSED_SIZE) + throw Exception("Too large size_compressed. Most likely corrupted data.", ErrorCodes::TOO_LARGE_SIZE_COMPRESSED); + size_t size_decompressed = qlz_size_decompressed(&compressed_buffer[0]); compressed_buffer.resize(size_compressed); @@ -47,7 +62,10 @@ void DecompressingStreamBuf::readCompressedChunk() p_istr->read(&compressed_buffer[QUICKLZ_HEADER_SIZE], size_compressed - QUICKLZ_HEADER_SIZE); if (!p_istr->good()) - return; + throw Exception("Cannot read all data.", ErrorCodes::CANNOT_READ_ALL_DATA); + + if (checksum != CityHash128(&compressed_buffer[0], size_compressed)) + throw Exception("Checksum doesnt match: corrupted data.", ErrorCodes::CHECKSUM_DOESNT_MATCH); /// разжимаем блок qlz_decompress(&compressed_buffer[0], &uncompressed_buffer[0], &scratch[0]); diff --git a/dbms/src/IO/CompressedOutputStream.cpp b/dbms/src/IO/CompressedOutputStream.cpp index 9877dbcac65..f5217d0c311 100644 --- a/dbms/src/IO/CompressedOutputStream.cpp +++ b/dbms/src/IO/CompressedOutputStream.cpp @@ -1,5 +1,7 @@ #include +#include + #include @@ -40,6 +42,8 @@ int CompressingStreamBuf::writeToDevice(const char * buffer, std::streamsize len length, &scratch[0]); + uint128 checksum = CityHash128(&compressed_buffer[0], compressed_size); + p_ostr->write(reinterpret_cast(&checksum), sizeof(checksum)); p_ostr->write(&compressed_buffer[0], compressed_size); return static_cast(length); }