1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
|
#include "decompress.hh"
#include "buffer.hh"
#define ZLIB_CONST
#include <algorithm>
#include <cstddef>
#include <expected>
#include <limits>
#include <memory>
#include <optional>
#include <utility>
#include <zlib.h>
namespace decompress {
namespace {
const size_t kBufferSizeZ = static_cast<size_t>(1024) * 1024;
class DecompressReader : public io::Reader {
public:
DecompressReader(std::unique_ptr<io::Reader> reader, bool gzip)
: reader_(std::move(reader)), gzip_(gzip) {}
~DecompressReader() override {
if (initialized_)
inflateEnd(&stream_);
}
std::expected<size_t, io::ReadError> read(void* dst, size_t max) override {
auto err = fill();
if (err.has_value())
return std::unexpected(err.value());
// NOLINTNEXTLINE(misc-include-cleaner)
stream_.next_out = reinterpret_cast<Bytef*>(dst);
stream_.avail_out = max;
if (!initialized_) {
if (in_eof_ && buffer_->empty())
return std::unexpected(io::ReadError::Eof);
stream_.zalloc = Z_NULL;
stream_.zfree = Z_NULL;
stream_.opaque = Z_NULL;
if (inflateInit2(&stream_, gzip_ ? 16 : 0) != Z_OK) {
return std::unexpected(io::ReadError::Error);
}
initialized_ = true;
}
auto* const rptr = stream_.next_in;
auto ret = inflate(&stream_, in_eof_ ? Z_FINISH : Z_NO_FLUSH);
auto got = max - stream_.avail_out;
if (ret == Z_STREAM_END) {
inflateEnd(&stream_);
initialized_ = false;
buffer_->consume(stream_.next_in - rptr);
} else if (ret == Z_OK) {
if (!in_eof_)
buffer_->consume(stream_.next_in - rptr);
} else {
switch (ret) {
case Z_DATA_ERROR:
return std::unexpected(io::ReadError::InvalidData);
case Z_BUF_ERROR:
// Tricky, could also be because of too little input.
// TODO: Improve logic here if we can
return std::unexpected(io::ReadError::MaxTooSmall);
default:
return std::unexpected(io::ReadError::Error);
}
}
return got;
}
std::expected<size_t, io::ReadError> skip(size_t max) override {
auto tmp = std::make_unique_for_overwrite<char[]>(max);
return read(tmp.get(), max);
}
private:
std::optional<io::ReadError> fill() {
size_t avail;
auto* rptr = buffer_->rptr(avail);
if (!in_eof_ && avail < kBufferSizeZ / 2) {
auto* wptr = buffer_->wptr(avail);
auto got = reader_->read(wptr, avail);
if (got.has_value()) {
buffer_->commit(got.value());
} else if (got.error() == io::ReadError::Eof) {
in_eof_ = true;
} else {
return got.error();
}
rptr = buffer_->rptr(avail);
}
// NOLINTNEXTLINE(misc-include-cleaner)
stream_.next_in = reinterpret_cast<z_const Bytef*>(rptr);
stream_.avail_in = std::min(
// NOLINTNEXTLINE(misc-include-cleaner)
static_cast<size_t>(std::numeric_limits<uInt>::max()), avail);
return std::nullopt;
}
std::unique_ptr<io::Reader> reader_;
bool const gzip_;
bool in_eof_{false};
std::unique_ptr<Buffer> buffer_{Buffer::fixed(kBufferSizeZ)};
bool initialized_{false};
z_stream stream_;
};
} // namespace
std::unique_ptr<io::Reader> zlib(std::unique_ptr<io::Reader> reader) {
return std::make_unique<DecompressReader>(std::move(reader),
/* gzip = */ false);
}
std::unique_ptr<io::Reader> gzip(std::unique_ptr<io::Reader> reader) {
return std::make_unique<DecompressReader>(std::move(reader),
/* gzip = */ true);
}
} // namespace decompress
|
