1 | #include <Common/config.h> |
2 | #if USE_BROTLI |
3 | |
4 | #include "BrotliReadBuffer.h" |
5 | #include <brotli/decode.h> |
6 | |
7 | namespace DB |
8 | { |
9 | |
10 | namespace ErrorCodes |
11 | { |
12 | extern const int BROTLI_READ_FAILED; |
13 | } |
14 | |
15 | |
16 | class BrotliReadBuffer::BrotliStateWrapper |
17 | { |
18 | public: |
19 | BrotliStateWrapper() |
20 | : state(BrotliDecoderCreateInstance(nullptr, nullptr, nullptr)) |
21 | , result(BROTLI_DECODER_RESULT_NEEDS_MORE_INPUT) |
22 | { |
23 | } |
24 | |
25 | ~BrotliStateWrapper() |
26 | { |
27 | BrotliDecoderDestroyInstance(state); |
28 | } |
29 | |
30 | public: |
31 | BrotliDecoderState * state; |
32 | BrotliDecoderResult result; |
33 | }; |
34 | |
35 | BrotliReadBuffer::BrotliReadBuffer(std::unique_ptr<ReadBuffer> in_, size_t buf_size, char *existing_memory, size_t alignment) |
36 | : BufferWithOwnMemory<ReadBuffer>(buf_size, existing_memory, alignment) |
37 | , in(std::move(in_)) |
38 | , brotli(std::make_unique<BrotliStateWrapper>()) |
39 | , in_available(0) |
40 | , in_data(nullptr) |
41 | , out_capacity(0) |
42 | , out_data(nullptr) |
43 | , eof(false) |
44 | { |
45 | } |
46 | |
47 | BrotliReadBuffer::~BrotliReadBuffer() |
48 | { |
49 | } |
50 | |
51 | bool BrotliReadBuffer::nextImpl() |
52 | { |
53 | if (eof) |
54 | return false; |
55 | |
56 | if (!in_available) |
57 | { |
58 | in->nextIfAtEnd(); |
59 | in_available = in->buffer().end() - in->position(); |
60 | in_data = reinterpret_cast<uint8_t *>(in->position()); |
61 | } |
62 | |
63 | if (brotli->result == BROTLI_DECODER_RESULT_NEEDS_MORE_INPUT && (!in_available || in->eof())) |
64 | { |
65 | throw Exception("brotli decode error" , ErrorCodes::BROTLI_READ_FAILED); |
66 | } |
67 | |
68 | out_capacity = internal_buffer.size(); |
69 | out_data = reinterpret_cast<uint8_t *>(internal_buffer.begin()); |
70 | |
71 | brotli->result = BrotliDecoderDecompressStream(brotli->state, &in_available, &in_data, &out_capacity, &out_data, nullptr); |
72 | |
73 | in->position() = in->buffer().end() - in_available; |
74 | working_buffer.resize(internal_buffer.size() - out_capacity); |
75 | |
76 | if (brotli->result == BROTLI_DECODER_RESULT_SUCCESS) |
77 | { |
78 | if (in->eof()) |
79 | { |
80 | eof = true; |
81 | return working_buffer.size() != 0; |
82 | } |
83 | else |
84 | { |
85 | throw Exception("brotli decode error" , ErrorCodes::BROTLI_READ_FAILED); |
86 | } |
87 | } |
88 | |
89 | if (brotli->result == BROTLI_DECODER_RESULT_ERROR) |
90 | { |
91 | throw Exception("brotli decode error" , ErrorCodes::BROTLI_READ_FAILED); |
92 | } |
93 | |
94 | return true; |
95 | } |
96 | } |
97 | |
98 | #endif |
99 | |