1 | #include <IO/ZlibInflatingReadBuffer.h> |
2 | |
3 | |
4 | namespace DB |
5 | { |
6 | |
7 | ZlibInflatingReadBuffer::ZlibInflatingReadBuffer( |
8 | std::unique_ptr<ReadBuffer> in_, |
9 | CompressionMethod compression_method, |
10 | size_t buf_size, |
11 | char * existing_memory, |
12 | size_t alignment) |
13 | : BufferWithOwnMemory<ReadBuffer>(buf_size, existing_memory, alignment) |
14 | , in(std::move(in_)) |
15 | , eof(false) |
16 | { |
17 | zstr.zalloc = nullptr; |
18 | zstr.zfree = nullptr; |
19 | zstr.opaque = nullptr; |
20 | zstr.next_in = nullptr; |
21 | zstr.avail_in = 0; |
22 | zstr.next_out = nullptr; |
23 | zstr.avail_out = 0; |
24 | |
25 | int window_bits = 15; |
26 | if (compression_method == CompressionMethod::Gzip) |
27 | { |
28 | window_bits += 16; |
29 | } |
30 | |
31 | #pragma GCC diagnostic push |
32 | #pragma GCC diagnostic ignored "-Wold-style-cast" |
33 | int rc = inflateInit2(&zstr, window_bits); |
34 | #pragma GCC diagnostic pop |
35 | |
36 | if (rc != Z_OK) |
37 | throw Exception(std::string("inflateInit2 failed: " ) + zError(rc) + "; zlib version: " + ZLIB_VERSION, ErrorCodes::ZLIB_INFLATE_FAILED); |
38 | } |
39 | |
40 | ZlibInflatingReadBuffer::~ZlibInflatingReadBuffer() |
41 | { |
42 | inflateEnd(&zstr); |
43 | } |
44 | |
45 | bool ZlibInflatingReadBuffer::nextImpl() |
46 | { |
47 | if (eof) |
48 | return false; |
49 | |
50 | if (!zstr.avail_in) |
51 | { |
52 | in->nextIfAtEnd(); |
53 | zstr.next_in = reinterpret_cast<unsigned char *>(in->position()); |
54 | zstr.avail_in = in->buffer().end() - in->position(); |
55 | } |
56 | zstr.next_out = reinterpret_cast<unsigned char *>(internal_buffer.begin()); |
57 | zstr.avail_out = internal_buffer.size(); |
58 | |
59 | int rc = inflate(&zstr, Z_NO_FLUSH); |
60 | |
61 | in->position() = in->buffer().end() - zstr.avail_in; |
62 | working_buffer.resize(internal_buffer.size() - zstr.avail_out); |
63 | |
64 | if (rc == Z_STREAM_END) |
65 | { |
66 | if (in->eof()) |
67 | { |
68 | eof = true; |
69 | return working_buffer.size() != 0; |
70 | } |
71 | else |
72 | { |
73 | rc = inflateReset(&zstr); |
74 | if (rc != Z_OK) |
75 | throw Exception(std::string("inflateReset failed: " ) + zError(rc), ErrorCodes::ZLIB_INFLATE_FAILED); |
76 | return true; |
77 | } |
78 | } |
79 | if (rc != Z_OK) |
80 | throw Exception(std::string("inflate failed: " ) + zError(rc), ErrorCodes::ZLIB_INFLATE_FAILED); |
81 | |
82 | return true; |
83 | } |
84 | |
85 | } |
86 | |