#include "duckdb/common/gzip_file_system.hpp" #include "duckdb/common/exception.hpp" #include "duckdb/common/file_system.hpp" #include "miniz.hpp" #include "duckdb/common/limits.hpp" namespace duckdb { /* 0 2 bytes magic header 0x1f, 0x8b (\037 \213) 2 1 byte compression method 0: store (copied) 1: compress 2: pack 3: lzh 4..7: reserved 8: deflate 3 1 byte flags bit 0 set: file probably ascii text bit 1 set: continuation of multi-part gzip file, part number present bit 2 set: extra field present bit 3 set: original file name present bit 4 set: file comment present bit 5 set: file is encrypted, encryption header present bit 6,7: reserved 4 4 bytes file modification time in Unix format 8 1 byte extra flags (depend on compression method) 9 1 byte OS type [ 2 bytes optional part number (second part=1) ]? [ 2 bytes optional extra field length (e) (e)bytes optional extra field ]? [ bytes optional original file name, zero terminated ]? [ bytes optional file comment, zero terminated ]? [ 12 bytes optional encryption header ]? bytes compressed data 4 bytes crc32 4 bytes uncompressed input size modulo 2^32 */ static idx_t GZipConsumeString(FileHandle &input) { idx_t size = 1; // terminator char buffer[1]; while (input.Read(buffer, 1) == 1) { if (buffer[0] == '\0') { break; } size++; } return size; } static constexpr const uint8_t GZIP_COMPRESSION_DEFLATE = 0x08; static constexpr const uint8_t GZIP_FLAG_ASCII = 0x1; static constexpr const uint8_t GZIP_FLAG_MULTIPART = 0x2; static constexpr const uint8_t GZIP_FLAG_EXTRA = 0x4; static constexpr const uint8_t GZIP_FLAG_NAME = 0x8; static constexpr const uint8_t GZIP_FLAG_COMMENT = 0x10; static constexpr const uint8_t GZIP_FLAG_ENCRYPT = 0x20; static constexpr const uint8_t GZIP_HEADER_MINSIZE = 10; static constexpr const unsigned char GZIP_FLAG_UNSUPPORTED = GZIP_FLAG_ASCII | GZIP_FLAG_MULTIPART | GZIP_FLAG_EXTRA | GZIP_FLAG_COMMENT | GZIP_FLAG_ENCRYPT; struct MiniZStreamWrapper { ~MiniZStreamWrapper() { Close(); } duckdb_miniz::mz_stream *mz_stream_ptr = nullptr; public: void Initialize() { Close(); mz_stream_ptr = new duckdb_miniz::mz_stream(); memset(mz_stream_ptr, 0, sizeof(duckdb_miniz::mz_stream)); } void Close() { if (!mz_stream_ptr) { return; } duckdb_miniz::mz_inflateEnd(mz_stream_ptr); delete mz_stream_ptr; mz_stream_ptr = nullptr; } }; class GZipFile : public FileHandle { static constexpr const idx_t BUFFER_SIZE = 1024; public: GZipFile(unique_ptr child_handle_p, const string &path) : FileHandle(gzip_fs, path), child_handle(move(child_handle_p)) { Initialize(); } ~GZipFile() override { Close(); } void Initialize(); int64_t ReadData(void *buffer, int64_t nr_bytes); GZipFileSystem gzip_fs; unique_ptr child_handle; protected: void Close() override { miniz_stream.reset(); in_buff.reset(); out_buff.reset(); } private: idx_t data_start = 0; unique_ptr miniz_stream; // various buffers & pointers unique_ptr in_buff; unique_ptr out_buff; data_ptr_t out_buff_start = nullptr; data_ptr_t out_buff_end = nullptr; data_ptr_t in_buff_start = nullptr; data_ptr_t in_buff_end = nullptr; }; void GZipFile::Initialize() { Close(); D_ASSERT(BUFFER_SIZE >= 3); // found to work fine with 3 uint8_t gzip_hdr[10]; data_start = GZIP_HEADER_MINSIZE; in_buff = unique_ptr(new data_t[BUFFER_SIZE]); in_buff_start = in_buff.get(); in_buff_end = in_buff.get(); out_buff = unique_ptr(new data_t[BUFFER_SIZE]); out_buff_start = out_buff.get(); out_buff_end = out_buff.get(); miniz_stream = make_unique(); miniz_stream->Initialize(); auto &mz_stream_ptr = miniz_stream->mz_stream_ptr; // TODO use custom alloc/free methods in miniz to throw exceptions on OOM auto read_count = child_handle->Read(gzip_hdr, GZIP_HEADER_MINSIZE); // check for incorrectly formatted files // LCOV_EXCL_START if (read_count != GZIP_HEADER_MINSIZE) { throw Exception("Input is not a GZIP stream"); } if (gzip_hdr[0] != 0x1F || gzip_hdr[1] != 0x8B) { // magic header throw Exception("Input is not a GZIP stream"); } if (gzip_hdr[2] != GZIP_COMPRESSION_DEFLATE) { // compression method throw Exception("Unsupported GZIP compression method"); } if (gzip_hdr[3] & GZIP_FLAG_UNSUPPORTED) { throw Exception("Unsupported GZIP archive"); } // LCOV_EXCL_STOP if (gzip_hdr[3] & GZIP_FLAG_NAME) { child_handle->Seek(data_start); data_start += GZipConsumeString(*child_handle); } child_handle->Seek(data_start); // stream is now set to beginning of payload data auto ret = duckdb_miniz::mz_inflateInit2((duckdb_miniz::mz_streamp)mz_stream_ptr, -MZ_DEFAULT_WINDOW_BITS); if (ret != duckdb_miniz::MZ_OK) { throw InternalException("Failed to initialize miniz"); } } unique_ptr GZipFileSystem::OpenCompressedFile(unique_ptr handle) { auto path = handle->path; return make_unique(move(handle), path); } int64_t GZipFile::ReadData(void *buffer, int64_t remaining) { auto &mz_stream_ptr = miniz_stream->mz_stream_ptr; idx_t total_read = 0; while (true) { // first check if there are input bytes available in the output buffers if (out_buff_start != out_buff_end) { // there is! copy it into the output buffer idx_t available = MinValue(remaining, out_buff_end - out_buff_start); memcpy(data_ptr_t(buffer) + total_read, out_buff_start, available); // increment the total read variables as required out_buff_start += available; total_read += available; remaining -= available; if (remaining == 0) { // done! read enough return total_read; } } if (!mz_stream_ptr) { return total_read; } // ran out of buffer: read more data from the child stream out_buff_start = out_buff.get(); out_buff_end = out_buff.get(); D_ASSERT(in_buff_start <= in_buff_end); D_ASSERT(in_buff_end <= in_buff_start + GZipFile::BUFFER_SIZE); // read more input if none available if (in_buff_start == in_buff_end) { // empty input buffer: refill from the start in_buff_start = in_buff.get(); auto sz = child_handle->Read(in_buff.get(), BUFFER_SIZE); if (sz <= 0) { break; } in_buff_end = in_buff_start + sz; } // actually decompress mz_stream_ptr->next_in = (data_ptr_t)in_buff_start; D_ASSERT(in_buff_end - in_buff_start < NumericLimits::Maximum()); mz_stream_ptr->avail_in = (uint32_t)(in_buff_end - in_buff_start); mz_stream_ptr->next_out = (data_ptr_t)out_buff_end; D_ASSERT((out_buff.get() + BUFFER_SIZE) - out_buff_end < NumericLimits::Maximum()); mz_stream_ptr->avail_out = (uint32_t)((out_buff.get() + BUFFER_SIZE) - out_buff_end); auto ret = duckdb_miniz::mz_inflate(mz_stream_ptr, duckdb_miniz::MZ_NO_FLUSH); if (ret != duckdb_miniz::MZ_OK && ret != duckdb_miniz::MZ_STREAM_END) { throw Exception(duckdb_miniz::mz_error(ret)); } // update pointers following inflate() in_buff_start = (data_ptr_t)mz_stream_ptr->next_in; in_buff_end = in_buff_start + mz_stream_ptr->avail_in; out_buff_end = (data_ptr_t)mz_stream_ptr->next_out; D_ASSERT(out_buff_end + mz_stream_ptr->avail_out == out_buff.get() + BUFFER_SIZE); // if stream ended, deallocate inflator if (ret == duckdb_miniz::MZ_STREAM_END) { miniz_stream->Close(); } } return total_read; } int64_t GZipFileSystem::Read(FileHandle &handle, void *buffer, int64_t nr_bytes) { auto &gzip_file = (GZipFile &)handle; return gzip_file.ReadData(buffer, nr_bytes); } void GZipFileSystem::Reset(FileHandle &handle) { auto &gzip_file = (GZipFile &)handle; gzip_file.child_handle->Reset(); gzip_file.Initialize(); } int64_t GZipFileSystem::GetFileSize(FileHandle &handle) { auto &gzip_file = (GZipFile &)handle; return gzip_file.child_handle->GetFileSize(); } bool GZipFileSystem::OnDiskFile(FileHandle &handle) { auto &gzip_file = (GZipFile &)handle; return gzip_file.child_handle->OnDiskFile(); } } // namespace duckdb