#include <glog/logging.h>
#if FOLLY_HAVE_LIBSNAPPY
-#include <snappy.h>
#include <snappy-sinksource.h>
+#include <snappy.h>
#endif
#if FOLLY_HAVE_LIBZ
-#include <zlib.h>
+#include <folly/io/compression/Zlib.h>
#endif
#if FOLLY_HAVE_LIBLZMA
#endif
#if FOLLY_HAVE_LIBZSTD
+#define ZSTD_STATIC_LINKING_ONLY
#include <zstd.h>
#endif
#include <folly/ScopeGuard.h>
#include <folly/Varint.h>
#include <folly/io/Cursor.h>
+#include <folly/io/compression/Utils.h>
#include <algorithm>
#include <unordered_set>
-namespace folly { namespace io {
+using folly::io::compression::detail::dataStartsWithLE;
+using folly::io::compression::detail::prefixToStringLE;
+
+namespace folly {
+namespace io {
Codec::Codec(CodecType type) : type_(type) { }
// Ensure consistent behavior in the nullptr case
std::unique_ptr<IOBuf> Codec::compress(const IOBuf* data) {
- uint64_t len = data->computeChainDataLength();
- if (len == 0) {
- return IOBuf::create(0);
+ if (data == nullptr) {
+ throw std::invalid_argument("Codec: data must not be nullptr");
}
+ uint64_t len = data->computeChainDataLength();
if (len > maxUncompressedLength()) {
throw std::runtime_error("Codec: uncompressed length too large");
}
std::string Codec::compress(const StringPiece data) {
const uint64_t len = data.size();
- if (len == 0) {
- return "";
- }
if (len > maxUncompressedLength()) {
throw std::runtime_error("Codec: uncompressed length too large");
}
std::unique_ptr<IOBuf> Codec::uncompress(
const IOBuf* data,
Optional<uint64_t> uncompressedLength) {
+ if (data == nullptr) {
+ throw std::invalid_argument("Codec: data must not be nullptr");
+ }
if (!uncompressedLength) {
if (needsUncompressedLength()) {
throw std::invalid_argument("Codec: uncompressed length required");
}
uint64_t Codec::maxCompressedLength(uint64_t uncompressedLength) const {
- if (uncompressedLength == 0) {
- return 0;
- }
return doMaxCompressedLength(uncompressedLength);
}
const folly::IOBuf* data,
Optional<uint64_t> uncompressedLength) const {
auto const compressedLength = data->computeChainDataLength();
- if (uncompressedLength == uint64_t(0) || compressedLength == 0) {
- if (uncompressedLength.value_or(0) != 0 || compressedLength != 0) {
+ if (compressedLength == 0) {
+ if (uncompressedLength.value_or(0) != 0) {
throw std::runtime_error("Invalid uncompressed length");
}
return 0;
void StreamCodec::resetStream(Optional<uint64_t> uncompressedLength) {
state_ = State::RESET;
uncompressedLength_ = uncompressedLength;
+ progressMade_ = true;
doResetStream();
}
ByteRange& input,
MutableByteRange& output,
StreamCodec::FlushOp flushOp) {
- if (state_ == State::RESET && input.empty()) {
- if (flushOp == StreamCodec::FlushOp::NONE) {
- return false;
- }
- if (flushOp == StreamCodec::FlushOp::END &&
- uncompressedLength().value_or(0) != 0) {
- throw std::runtime_error("Codec: invalid uncompressed length");
- }
- return true;
+ if (state_ == State::RESET && input.empty() &&
+ flushOp == StreamCodec::FlushOp::END &&
+ uncompressedLength().value_or(0) != 0) {
+ throw std::runtime_error("Codec: invalid uncompressed length");
+ }
+
+ if (!uncompressedLength() && needsDataLength()) {
+ throw std::runtime_error("Codec: uncompressed length required");
}
if (state_ == State::RESET && !input.empty() &&
uncompressedLength() == uint64_t(0)) {
assertStateIs(State::COMPRESS_END);
break;
}
+ size_t const inputSize = input.size();
+ size_t const outputSize = output.size();
bool const done = doCompressStream(input, output, flushOp);
+ if (!done && inputSize == input.size() && outputSize == output.size()) {
+ if (!progressMade_) {
+ throw std::runtime_error("Codec: No forward progress made");
+ }
+ // Throw an exception if there is no progress again next time
+ progressMade_ = false;
+ } else {
+ progressMade_ = true;
+ }
// Handle output state transitions
if (done) {
if (state_ == State::COMPRESS_FLUSH) {
state_ = State::UNCOMPRESS;
}
assertStateIs(State::UNCOMPRESS);
+ size_t const inputSize = input.size();
+ size_t const outputSize = output.size();
bool const done = doUncompressStream(input, output, flushOp);
+ if (!done && inputSize == input.size() && outputSize == output.size()) {
+ if (!progressMade_) {
+ throw std::runtime_error("Codec: no forward progress made");
+ }
+ // Throw an exception if there is no progress again next time
+ progressMade_ = false;
+ } else {
+ progressMade_ = true;
+ }
// Handle output state transitions
if (done) {
state_ = State::END;
IOBuf const* current = data;
ByteRange input{current->data(), current->length()};
StreamCodec::FlushOp flushOp = StreamCodec::FlushOp::NONE;
- for (;;) {
+ bool done = false;
+ while (!done) {
while (input.empty() && current->next() != data) {
current = current->next();
input = {current->data(), current->length()};
if (output.empty()) {
buffer->prependChain(addOutputBuffer(output, kDefaultBufferLength));
}
- bool const done = compressStream(input, output, flushOp);
+ done = compressStream(input, output, flushOp);
if (done) {
DCHECK(input.empty());
DCHECK(flushOp == StreamCodec::FlushOp::END);
DCHECK_EQ(current->next(), data);
- break;
}
}
buffer->prev()->trimEnd(output.size());
IOBuf const* current = data;
ByteRange input{current->data(), current->length()};
StreamCodec::FlushOp flushOp = StreamCodec::FlushOp::NONE;
- for (;;) {
+ bool done = false;
+ while (!done) {
while (input.empty() && current->next() != data) {
current = current->next();
input = {current->data(), current->length()};
if (output.empty()) {
buffer->prependChain(addOutputBuffer(output, defaultBufferLength));
}
- bool const done = uncompressStream(input, output, flushOp);
- if (done) {
- break;
- }
+ done = uncompressStream(input, output, flushOp);
}
if (!input.empty()) {
throw std::runtime_error("Codec: Junk after end of data");
}
NoCompressionCodec::NoCompressionCodec(int level, CodecType type)
- : Codec(type) {
+ : Codec(type) {
DCHECK(type == CodecType::NO_COMPRESSION);
switch (level) {
- case COMPRESSION_LEVEL_DEFAULT:
- case COMPRESSION_LEVEL_FASTEST:
- case COMPRESSION_LEVEL_BEST:
- level = 0;
+ case COMPRESSION_LEVEL_DEFAULT:
+ case COMPRESSION_LEVEL_FASTEST:
+ case COMPRESSION_LEVEL_BEST:
+ level = 0;
}
if (level != 0) {
throw std::invalid_argument(to<std::string>(
return val;
}
-} // namespace
+} // namespace
#endif // FOLLY_HAVE_LIBLZ4 || FOLLY_HAVE_LIBLZMA
-namespace {
-/**
- * Reads sizeof(T) bytes, and returns false if not enough bytes are available.
- * Returns true if the first n bytes are equal to prefix when interpreted as
- * a little endian T.
- */
-template <typename T>
-typename std::enable_if<std::is_unsigned<T>::value, bool>::type
-dataStartsWithLE(const IOBuf* data, T prefix, uint64_t n = sizeof(T)) {
- DCHECK_GT(n, 0);
- DCHECK_LE(n, sizeof(T));
- T value;
- Cursor cursor{data};
- if (!cursor.tryReadLE(value)) {
- return false;
- }
- const T mask = n == sizeof(T) ? T(-1) : (T(1) << (8 * n)) - 1;
- return prefix == (value & mask);
-}
-
-template <typename T>
-typename std::enable_if<std::is_arithmetic<T>::value, std::string>::type
-prefixToStringLE(T prefix, uint64_t n = sizeof(T)) {
- DCHECK_GT(n, 0);
- DCHECK_LE(n, sizeof(T));
- prefix = Endian::little(prefix);
- std::string result;
- result.resize(n);
- memcpy(&result[0], &prefix, n);
- return result;
-}
-} // namespace
-
#if FOLLY_HAVE_LIBLZ4
/**
DCHECK(type == CodecType::LZ4 || type == CodecType::LZ4_VARINT_SIZE);
switch (level) {
- case COMPRESSION_LEVEL_FASTEST:
- case COMPRESSION_LEVEL_DEFAULT:
- level = 1;
- break;
- case COMPRESSION_LEVEL_BEST:
- level = 2;
- break;
+ case COMPRESSION_LEVEL_FASTEST:
+ case COMPRESSION_LEVEL_DEFAULT:
+ level = 1;
+ break;
+ case COMPRESSION_LEVEL_BEST:
+ level = 2;
+ break;
}
if (level < 1 || level > 2) {
throw std::invalid_argument(to<std::string>(
SnappyCodec::SnappyCodec(int level, CodecType type) : Codec(type) {
DCHECK(type == CodecType::SNAPPY);
switch (level) {
- case COMPRESSION_LEVEL_FASTEST:
- case COMPRESSION_LEVEL_DEFAULT:
- case COMPRESSION_LEVEL_BEST:
- level = 1;
+ case COMPRESSION_LEVEL_FASTEST:
+ case COMPRESSION_LEVEL_DEFAULT:
+ case COMPRESSION_LEVEL_BEST:
+ level = 1;
}
if (level != 1) {
throw std::invalid_argument(to<std::string>(
#endif // FOLLY_HAVE_LIBSNAPPY
-#if FOLLY_HAVE_LIBZ
-/**
- * Zlib codec
- */
-class ZlibCodec final : public Codec {
- public:
- static std::unique_ptr<Codec> create(int level, CodecType type);
- explicit ZlibCodec(int level, CodecType type);
-
- std::vector<std::string> validPrefixes() const override;
- bool canUncompress(const IOBuf* data, Optional<uint64_t> uncompressedLength)
- const override;
-
- private:
- uint64_t doMaxCompressedLength(uint64_t uncompressedLength) const override;
- std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
- std::unique_ptr<IOBuf> doUncompress(
- const IOBuf* data,
- Optional<uint64_t> uncompressedLength) override;
-
- std::unique_ptr<IOBuf> addOutputBuffer(z_stream* stream, uint32_t length);
- bool doInflate(z_stream* stream, IOBuf* head, uint32_t bufferLength);
-
- int level_;
-};
-
-static constexpr uint16_t kGZIPMagicLE = 0x8B1F;
-
-std::vector<std::string> ZlibCodec::validPrefixes() const {
- if (type() == CodecType::ZLIB) {
- // Zlib streams start with a 2 byte header.
- //
- // 0 1
- // +---+---+
- // |CMF|FLG|
- // +---+---+
- //
- // We won't restrict the values of any sub-fields except as described below.
- //
- // The lowest 4 bits of CMF is the compression method (CM).
- // CM == 0x8 is the deflate compression method, which is currently the only
- // supported compression method, so any valid prefix must have CM == 0x8.
- //
- // The lowest 5 bits of FLG is FCHECK.
- // FCHECK must be such that the two header bytes are a multiple of 31 when
- // interpreted as a big endian 16-bit number.
- std::vector<std::string> result;
- // 16 values for the first byte, 8 values for the second byte.
- // There are also 4 combinations where both 0x00 and 0x1F work as FCHECK.
- result.reserve(132);
- // Select all values for the CMF byte that use the deflate algorithm 0x8.
- for (uint32_t first = 0x0800; first <= 0xF800; first += 0x1000) {
- // Select all values for the FLG, but leave FCHECK as 0 since it's fixed.
- for (uint32_t second = 0x00; second <= 0xE0; second += 0x20) {
- uint16_t prefix = first | second;
- // Compute FCHECK.
- prefix += 31 - (prefix % 31);
- result.push_back(prefixToStringLE(Endian::big(prefix)));
- // zlib won't produce this, but it is a valid prefix.
- if ((prefix & 0x1F) == 31) {
- prefix -= 31;
- result.push_back(prefixToStringLE(Endian::big(prefix)));
- }
- }
- }
- return result;
- } else {
- // The gzip frame starts with 2 magic bytes.
- return {prefixToStringLE(kGZIPMagicLE)};
- }
-}
-
-bool ZlibCodec::canUncompress(const IOBuf* data, Optional<uint64_t>) const {
- if (type() == CodecType::ZLIB) {
- uint16_t value;
- Cursor cursor{data};
- if (!cursor.tryReadBE(value)) {
- return false;
- }
- // zlib compressed if using deflate and is a multiple of 31.
- return (value & 0x0F00) == 0x0800 && value % 31 == 0;
- } else {
- return dataStartsWithLE(data, kGZIPMagicLE);
- }
-}
-
-uint64_t ZlibCodec::doMaxCompressedLength(uint64_t uncompressedLength) const {
- return deflateBound(nullptr, uncompressedLength);
-}
-
-std::unique_ptr<Codec> ZlibCodec::create(int level, CodecType type) {
- return std::make_unique<ZlibCodec>(level, type);
-}
-
-ZlibCodec::ZlibCodec(int level, CodecType type) : Codec(type) {
- DCHECK(type == CodecType::ZLIB || type == CodecType::GZIP);
- switch (level) {
- case COMPRESSION_LEVEL_FASTEST:
- level = 1;
- break;
- case COMPRESSION_LEVEL_DEFAULT:
- level = Z_DEFAULT_COMPRESSION;
- break;
- case COMPRESSION_LEVEL_BEST:
- level = 9;
- break;
- }
- if (level != Z_DEFAULT_COMPRESSION && (level < 0 || level > 9)) {
- throw std::invalid_argument(to<std::string>(
- "ZlibCodec: invalid level: ", level));
- }
- level_ = level;
-}
-
-std::unique_ptr<IOBuf> ZlibCodec::addOutputBuffer(z_stream* stream,
- uint32_t length) {
- CHECK_EQ(stream->avail_out, 0);
-
- auto buf = IOBuf::create(length);
- buf->append(buf->capacity());
-
- stream->next_out = buf->writableData();
- stream->avail_out = buf->length();
-
- return buf;
-}
-
-bool ZlibCodec::doInflate(z_stream* stream,
- IOBuf* head,
- uint32_t bufferLength) {
- if (stream->avail_out == 0) {
- head->prependChain(addOutputBuffer(stream, bufferLength));
- }
-
- int rc = inflate(stream, Z_NO_FLUSH);
-
- switch (rc) {
- case Z_OK:
- break;
- case Z_STREAM_END:
- return true;
- case Z_BUF_ERROR:
- case Z_NEED_DICT:
- case Z_DATA_ERROR:
- case Z_MEM_ERROR:
- throw std::runtime_error(to<std::string>(
- "ZlibCodec: inflate error: ", rc, ": ", stream->msg));
- default:
- CHECK(false) << rc << ": " << stream->msg;
- }
-
- return false;
-}
-
-std::unique_ptr<IOBuf> ZlibCodec::doCompress(const IOBuf* data) {
- z_stream stream;
- stream.zalloc = nullptr;
- stream.zfree = nullptr;
- stream.opaque = nullptr;
-
- // Using deflateInit2() to support gzip. "The windowBits parameter is the
- // base two logarithm of the maximum window size (...) The default value is
- // 15 (...) Add 16 to windowBits to write a simple gzip header and trailer
- // around the compressed data instead of a zlib wrapper. The gzip header
- // will have no file name, no extra data, no comment, no modification time
- // (set to zero), no header crc, and the operating system will be set to 255
- // (unknown)."
- int windowBits = 15 + (type() == CodecType::GZIP ? 16 : 0);
- // All other parameters (method, memLevel, strategy) get default values from
- // the zlib manual.
- int rc = deflateInit2(&stream,
- level_,
- Z_DEFLATED,
- windowBits,
- /* memLevel */ 8,
- Z_DEFAULT_STRATEGY);
- if (rc != Z_OK) {
- throw std::runtime_error(to<std::string>(
- "ZlibCodec: deflateInit error: ", rc, ": ", stream.msg));
- }
-
- stream.next_in = stream.next_out = nullptr;
- stream.avail_in = stream.avail_out = 0;
- stream.total_in = stream.total_out = 0;
-
- bool success = false;
-
- SCOPE_EXIT {
- rc = deflateEnd(&stream);
- // If we're here because of an exception, it's okay if some data
- // got dropped.
- CHECK(rc == Z_OK || (!success && rc == Z_DATA_ERROR))
- << rc << ": " << stream.msg;
- };
-
- uint64_t uncompressedLength = data->computeChainDataLength();
- uint64_t maxCompressedLength = deflateBound(&stream, uncompressedLength);
-
- // Max 64MiB in one go
- constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20; // 64MiB
- constexpr uint32_t defaultBufferLength = uint32_t(4) << 20; // 4MiB
-
- auto out = addOutputBuffer(
- &stream,
- (maxCompressedLength <= maxSingleStepLength ?
- maxCompressedLength :
- defaultBufferLength));
-
- for (auto& range : *data) {
- uint64_t remaining = range.size();
- uint64_t written = 0;
- while (remaining) {
- uint32_t step = (remaining > maxSingleStepLength ?
- maxSingleStepLength : remaining);
- stream.next_in = const_cast<uint8_t*>(range.data() + written);
- stream.avail_in = step;
- remaining -= step;
- written += step;
-
- while (stream.avail_in != 0) {
- if (stream.avail_out == 0) {
- out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
- }
-
- rc = deflate(&stream, Z_NO_FLUSH);
-
- CHECK_EQ(rc, Z_OK) << stream.msg;
- }
- }
- }
-
- do {
- if (stream.avail_out == 0) {
- out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
- }
-
- rc = deflate(&stream, Z_FINISH);
- } while (rc == Z_OK);
-
- CHECK_EQ(rc, Z_STREAM_END) << stream.msg;
-
- out->prev()->trimEnd(stream.avail_out);
-
- success = true; // we survived
-
- return out;
-}
-
-std::unique_ptr<IOBuf> ZlibCodec::doUncompress(
- const IOBuf* data,
- Optional<uint64_t> uncompressedLength) {
- z_stream stream;
- stream.zalloc = nullptr;
- stream.zfree = nullptr;
- stream.opaque = nullptr;
-
- // "The windowBits parameter is the base two logarithm of the maximum window
- // size (...) The default value is 15 (...) add 16 to decode only the gzip
- // format (the zlib format will return a Z_DATA_ERROR)."
- int windowBits = 15 + (type() == CodecType::GZIP ? 16 : 0);
- int rc = inflateInit2(&stream, windowBits);
- if (rc != Z_OK) {
- throw std::runtime_error(to<std::string>(
- "ZlibCodec: inflateInit error: ", rc, ": ", stream.msg));
- }
-
- stream.next_in = stream.next_out = nullptr;
- stream.avail_in = stream.avail_out = 0;
- stream.total_in = stream.total_out = 0;
-
- bool success = false;
-
- SCOPE_EXIT {
- rc = inflateEnd(&stream);
- // If we're here because of an exception, it's okay if some data
- // got dropped.
- CHECK(rc == Z_OK || (!success && rc == Z_DATA_ERROR))
- << rc << ": " << stream.msg;
- };
-
- // Max 64MiB in one go
- constexpr uint64_t maxSingleStepLength = uint64_t(64) << 20; // 64MiB
- constexpr uint64_t kBlockSize = uint64_t(32) << 10; // 32 KiB
- const uint64_t defaultBufferLength =
- computeBufferLength(data->computeChainDataLength(), kBlockSize);
-
- auto out = addOutputBuffer(
- &stream,
- ((uncompressedLength && *uncompressedLength <= maxSingleStepLength)
- ? *uncompressedLength
- : defaultBufferLength));
-
- bool streamEnd = false;
- for (auto& range : *data) {
- if (range.empty()) {
- continue;
- }
-
- stream.next_in = const_cast<uint8_t*>(range.data());
- stream.avail_in = range.size();
-
- while (stream.avail_in != 0) {
- if (streamEnd) {
- throw std::runtime_error(to<std::string>(
- "ZlibCodec: junk after end of data"));
- }
-
- streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
- }
- }
-
- while (!streamEnd) {
- streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
- }
-
- out->prev()->trimEnd(stream.avail_out);
-
- if (uncompressedLength && *uncompressedLength != stream.total_out) {
- throw std::runtime_error(
- to<std::string>("ZlibCodec: invalid uncompressed length"));
- }
-
- success = true; // we survived
-
- return out;
-}
-
-#endif // FOLLY_HAVE_LIBZ
-
#if FOLLY_HAVE_LIBLZMA
/**
* LZMA2 compression
*/
-class LZMA2Codec final : public Codec {
+class LZMA2StreamCodec final : public StreamCodec {
public:
- static std::unique_ptr<Codec> create(int level, CodecType type);
- explicit LZMA2Codec(int level, CodecType type);
+ static std::unique_ptr<Codec> createCodec(int level, CodecType type);
+ static std::unique_ptr<StreamCodec> createStream(int level, CodecType type);
+ explicit LZMA2StreamCodec(int level, CodecType type);
+ ~LZMA2StreamCodec() override;
std::vector<std::string> validPrefixes() const override;
bool canUncompress(const IOBuf* data, Optional<uint64_t> uncompressedLength)
const override;
private:
- bool doNeedsUncompressedLength() const override;
+ bool doNeedsDataLength() const override;
uint64_t doMaxUncompressedLength() const override;
uint64_t doMaxCompressedLength(uint64_t uncompressedLength) const override;
- bool encodeSize() const { return type() == CodecType::LZMA2_VARINT_SIZE; }
+ bool encodeSize() const {
+ return type() == CodecType::LZMA2_VARINT_SIZE;
+ }
- std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
- std::unique_ptr<IOBuf> doUncompress(
- const IOBuf* data,
- Optional<uint64_t> uncompressedLength) override;
+ void doResetStream() override;
+ bool doCompressStream(
+ ByteRange& input,
+ MutableByteRange& output,
+ StreamCodec::FlushOp flushOp) override;
+ bool doUncompressStream(
+ ByteRange& input,
+ MutableByteRange& output,
+ StreamCodec::FlushOp flushOp) override;
+
+ void resetCStream();
+ void resetDStream();
- std::unique_ptr<IOBuf> addOutputBuffer(lzma_stream* stream, size_t length);
- bool doInflate(lzma_stream* stream, IOBuf* head, size_t bufferLength);
+ bool decodeAndCheckVarint(ByteRange& input);
+ bool flushVarintBuffer(MutableByteRange& output);
+ void resetVarintBuffer();
+
+ Optional<lzma_stream> cstream_{};
+ Optional<lzma_stream> dstream_{};
+
+ std::array<uint8_t, kMaxVarintLength64> varintBuffer_;
+ ByteRange varintToEncode_;
+ size_t varintBufferPos_{0};
int level_;
+ bool needReset_{true};
+ bool needDecodeSize_{false};
};
static constexpr uint64_t kLZMA2MagicLE = 0x005A587A37FD;
static constexpr unsigned kLZMA2MagicBytes = 6;
-std::vector<std::string> LZMA2Codec::validPrefixes() const {
+std::vector<std::string> LZMA2StreamCodec::validPrefixes() const {
if (type() == CodecType::LZMA2_VARINT_SIZE) {
return {};
}
return {prefixToStringLE(kLZMA2MagicLE, kLZMA2MagicBytes)};
}
-bool LZMA2Codec::canUncompress(const IOBuf* data, Optional<uint64_t>) const {
+bool LZMA2StreamCodec::doNeedsDataLength() const {
+ return encodeSize();
+}
+
+bool LZMA2StreamCodec::canUncompress(const IOBuf* data, Optional<uint64_t>)
+ const {
if (type() == CodecType::LZMA2_VARINT_SIZE) {
return false;
}
return dataStartsWithLE(data, kLZMA2MagicLE, kLZMA2MagicBytes);
}
-std::unique_ptr<Codec> LZMA2Codec::create(int level, CodecType type) {
- return std::make_unique<LZMA2Codec>(level, type);
+std::unique_ptr<Codec> LZMA2StreamCodec::createCodec(
+ int level,
+ CodecType type) {
+ return make_unique<LZMA2StreamCodec>(level, type);
+}
+
+std::unique_ptr<StreamCodec> LZMA2StreamCodec::createStream(
+ int level,
+ CodecType type) {
+ return make_unique<LZMA2StreamCodec>(level, type);
}
-LZMA2Codec::LZMA2Codec(int level, CodecType type) : Codec(type) {
+LZMA2StreamCodec::LZMA2StreamCodec(int level, CodecType type)
+ : StreamCodec(type) {
DCHECK(type == CodecType::LZMA2 || type == CodecType::LZMA2_VARINT_SIZE);
switch (level) {
- case COMPRESSION_LEVEL_FASTEST:
- level = 0;
- break;
- case COMPRESSION_LEVEL_DEFAULT:
- level = LZMA_PRESET_DEFAULT;
- break;
- case COMPRESSION_LEVEL_BEST:
- level = 9;
- break;
+ case COMPRESSION_LEVEL_FASTEST:
+ level = 0;
+ break;
+ case COMPRESSION_LEVEL_DEFAULT:
+ level = LZMA_PRESET_DEFAULT;
+ break;
+ case COMPRESSION_LEVEL_BEST:
+ level = 9;
+ break;
}
if (level < 0 || level > 9) {
- throw std::invalid_argument(to<std::string>(
- "LZMA2Codec: invalid level: ", level));
+ throw std::invalid_argument(
+ to<std::string>("LZMA2Codec: invalid level: ", level));
}
level_ = level;
}
-bool LZMA2Codec::doNeedsUncompressedLength() const {
- return false;
+LZMA2StreamCodec::~LZMA2StreamCodec() {
+ if (cstream_) {
+ lzma_end(cstream_.get_pointer());
+ cstream_.clear();
+ }
+ if (dstream_) {
+ lzma_end(dstream_.get_pointer());
+ dstream_.clear();
+ }
}
-uint64_t LZMA2Codec::doMaxUncompressedLength() const {
+uint64_t LZMA2StreamCodec::doMaxUncompressedLength() const {
// From lzma/base.h: "Stream is roughly 8 EiB (2^63 bytes)"
return uint64_t(1) << 63;
}
-uint64_t LZMA2Codec::doMaxCompressedLength(uint64_t uncompressedLength) const {
+uint64_t LZMA2StreamCodec::doMaxCompressedLength(
+ uint64_t uncompressedLength) const {
return lzma_stream_buffer_bound(uncompressedLength) +
(encodeSize() ? kMaxVarintLength64 : 0);
}
-std::unique_ptr<IOBuf> LZMA2Codec::addOutputBuffer(
- lzma_stream* stream,
- size_t length) {
-
- CHECK_EQ(stream->avail_out, 0);
-
- auto buf = IOBuf::create(length);
- buf->append(buf->capacity());
-
- stream->next_out = buf->writableData();
- stream->avail_out = buf->length();
-
- return buf;
+void LZMA2StreamCodec::doResetStream() {
+ needReset_ = true;
}
-std::unique_ptr<IOBuf> LZMA2Codec::doCompress(const IOBuf* data) {
- lzma_ret rc;
- lzma_stream stream = LZMA_STREAM_INIT;
-
- rc = lzma_easy_encoder(&stream, level_, LZMA_CHECK_NONE);
+void LZMA2StreamCodec::resetCStream() {
+ if (!cstream_) {
+ cstream_.assign(LZMA_STREAM_INIT);
+ }
+ lzma_ret const rc =
+ lzma_easy_encoder(cstream_.get_pointer(), level_, LZMA_CHECK_NONE);
if (rc != LZMA_OK) {
throw std::runtime_error(folly::to<std::string>(
- "LZMA2Codec: lzma_easy_encoder error: ", rc));
- }
-
- SCOPE_EXIT { lzma_end(&stream); };
-
- uint64_t uncompressedLength = data->computeChainDataLength();
- uint64_t maxCompressedLength = lzma_stream_buffer_bound(uncompressedLength);
-
- // Max 64MiB in one go
- constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20; // 64MiB
- constexpr uint32_t defaultBufferLength = uint32_t(4) << 20; // 4MiB
-
- auto out = addOutputBuffer(
- &stream,
- (maxCompressedLength <= maxSingleStepLength ?
- maxCompressedLength :
- defaultBufferLength));
-
- if (encodeSize()) {
- auto size = IOBuf::createCombined(kMaxVarintLength64);
- encodeVarintToIOBuf(uncompressedLength, size.get());
- size->appendChain(std::move(out));
- out = std::move(size);
+ "LZMA2StreamCodec: lzma_easy_encoder error: ", rc));
}
+}
- for (auto& range : *data) {
- if (range.empty()) {
- continue;
- }
-
- stream.next_in = const_cast<uint8_t*>(range.data());
- stream.avail_in = range.size();
-
- while (stream.avail_in != 0) {
- if (stream.avail_out == 0) {
- out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
- }
-
- rc = lzma_code(&stream, LZMA_RUN);
-
- if (rc != LZMA_OK) {
- throw std::runtime_error(folly::to<std::string>(
- "LZMA2Codec: lzma_code error: ", rc));
- }
- }
+void LZMA2StreamCodec::resetDStream() {
+ if (!dstream_) {
+ dstream_.assign(LZMA_STREAM_INIT);
}
-
- do {
- if (stream.avail_out == 0) {
- out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
- }
-
- rc = lzma_code(&stream, LZMA_FINISH);
- } while (rc == LZMA_OK);
-
- if (rc != LZMA_STREAM_END) {
+ lzma_ret const rc = lzma_auto_decoder(
+ dstream_.get_pointer(), std::numeric_limits<uint64_t>::max(), 0);
+ if (rc != LZMA_OK) {
throw std::runtime_error(folly::to<std::string>(
- "LZMA2Codec: lzma_code ended with error: ", rc));
+ "LZMA2StreamCodec: lzma_auto_decoder error: ", rc));
}
-
- out->prev()->trimEnd(stream.avail_out);
-
- return out;
}
-bool LZMA2Codec::doInflate(lzma_stream* stream,
- IOBuf* head,
- size_t bufferLength) {
- if (stream->avail_out == 0) {
- head->prependChain(addOutputBuffer(stream, bufferLength));
+static lzma_ret lzmaThrowOnError(lzma_ret const rc) {
+ switch (rc) {
+ case LZMA_OK:
+ case LZMA_STREAM_END:
+ case LZMA_BUF_ERROR: // not fatal: returned if no progress was made twice
+ return rc;
+ default:
+ throw std::runtime_error(
+ to<std::string>("LZMA2StreamCodec: error: ", rc));
}
+}
- lzma_ret rc = lzma_code(stream, LZMA_RUN);
+static lzma_action lzmaTranslateFlush(StreamCodec::FlushOp flush) {
+ switch (flush) {
+ case StreamCodec::FlushOp::NONE:
+ return LZMA_RUN;
+ case StreamCodec::FlushOp::FLUSH:
+ return LZMA_SYNC_FLUSH;
+ case StreamCodec::FlushOp::END:
+ return LZMA_FINISH;
+ default:
+ throw std::invalid_argument("LZMA2StreamCodec: Invalid flush");
+ }
+}
- switch (rc) {
- case LZMA_OK:
- break;
- case LZMA_STREAM_END:
+/**
+ * Flushes the varint buffer.
+ * Advances output by the number of bytes written.
+ * Returns true when flushing is complete.
+ */
+bool LZMA2StreamCodec::flushVarintBuffer(MutableByteRange& output) {
+ if (varintToEncode_.empty()) {
return true;
- default:
- throw std::runtime_error(to<std::string>(
- "LZMA2Codec: lzma_code error: ", rc));
}
-
- return false;
+ const size_t numBytesToCopy = std::min(varintToEncode_.size(), output.size());
+ if (numBytesToCopy > 0) {
+ memcpy(output.data(), varintToEncode_.data(), numBytesToCopy);
+ }
+ varintToEncode_.advance(numBytesToCopy);
+ output.advance(numBytesToCopy);
+ return varintToEncode_.empty();
}
-std::unique_ptr<IOBuf> LZMA2Codec::doUncompress(
- const IOBuf* data,
- Optional<uint64_t> uncompressedLength) {
- lzma_ret rc;
- lzma_stream stream = LZMA_STREAM_INIT;
-
- rc = lzma_auto_decoder(&stream, std::numeric_limits<uint64_t>::max(), 0);
- if (rc != LZMA_OK) {
- throw std::runtime_error(folly::to<std::string>(
- "LZMA2Codec: lzma_auto_decoder error: ", rc));
+bool LZMA2StreamCodec::doCompressStream(
+ ByteRange& input,
+ MutableByteRange& output,
+ StreamCodec::FlushOp flushOp) {
+ if (needReset_) {
+ resetCStream();
+ if (encodeSize()) {
+ varintBufferPos_ = 0;
+ size_t const varintSize =
+ encodeVarint(*uncompressedLength(), varintBuffer_.data());
+ varintToEncode_ = {varintBuffer_.data(), varintSize};
+ }
+ needReset_ = false;
}
- SCOPE_EXIT { lzma_end(&stream); };
-
- // Max 64MiB in one go
- constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20; // 64MiB
- constexpr uint32_t defaultBufferLength = uint32_t(256) << 10; // 256 KiB
-
- folly::io::Cursor cursor(data);
- if (encodeSize()) {
- const uint64_t actualUncompressedLength = decodeVarintFromCursor(cursor);
- if (uncompressedLength && *uncompressedLength != actualUncompressedLength) {
- throw std::runtime_error("LZMA2Codec: invalid uncompressed length");
- }
- uncompressedLength = actualUncompressedLength;
+ if (!flushVarintBuffer(output)) {
+ return false;
}
- auto out = addOutputBuffer(
- &stream,
- ((uncompressedLength && *uncompressedLength <= maxSingleStepLength)
- ? *uncompressedLength
- : defaultBufferLength));
+ cstream_->next_in = const_cast<uint8_t*>(input.data());
+ cstream_->avail_in = input.size();
+ cstream_->next_out = output.data();
+ cstream_->avail_out = output.size();
+ SCOPE_EXIT {
+ input.uncheckedAdvance(input.size() - cstream_->avail_in);
+ output.uncheckedAdvance(output.size() - cstream_->avail_out);
+ };
+ lzma_ret const rc = lzmaThrowOnError(
+ lzma_code(cstream_.get_pointer(), lzmaTranslateFlush(flushOp)));
+ switch (flushOp) {
+ case StreamCodec::FlushOp::NONE:
+ return false;
+ case StreamCodec::FlushOp::FLUSH:
+ return cstream_->avail_in == 0 && cstream_->avail_out != 0;
+ case StreamCodec::FlushOp::END:
+ return rc == LZMA_STREAM_END;
+ default:
+ throw std::invalid_argument("LZMA2StreamCodec: invalid FlushOp");
+ }
+}
- bool streamEnd = false;
- auto buf = cursor.peekBytes();
- while (!buf.empty()) {
- stream.next_in = const_cast<uint8_t*>(buf.data());
- stream.avail_in = buf.size();
+/**
+ * Attempts to decode a varint from input.
+ * The function advances input by the number of bytes read.
+ *
+ * If there are too many bytes and the varint is not valid, throw a
+ * runtime_error.
+ *
+ * If the uncompressed length was provided and a decoded varint does not match
+ * the provided length, throw a runtime_error.
+ *
+ * Returns true if the varint was successfully decoded and matches the
+ * uncompressed length if provided, and false if more bytes are needed.
+ */
+bool LZMA2StreamCodec::decodeAndCheckVarint(ByteRange& input) {
+ if (input.empty()) {
+ return false;
+ }
+ size_t const numBytesToCopy =
+ std::min(kMaxVarintLength64 - varintBufferPos_, input.size());
+ memcpy(varintBuffer_.data() + varintBufferPos_, input.data(), numBytesToCopy);
- while (stream.avail_in != 0) {
- if (streamEnd) {
- throw std::runtime_error(to<std::string>(
- "LZMA2Codec: junk after end of data"));
- }
+ size_t const rangeSize = varintBufferPos_ + numBytesToCopy;
+ ByteRange range{varintBuffer_.data(), rangeSize};
+ auto const ret = tryDecodeVarint(range);
- streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
+ if (ret.hasValue()) {
+ size_t const varintSize = rangeSize - range.size();
+ input.advance(varintSize - varintBufferPos_);
+ if (uncompressedLength() && *uncompressedLength() != ret.value()) {
+ throw std::runtime_error("LZMA2StreamCodec: invalid uncompressed length");
}
+ return true;
+ } else if (ret.error() == DecodeVarintError::TooManyBytes) {
+ throw std::runtime_error("LZMA2StreamCodec: invalid uncompressed length");
+ } else {
+ // Too few bytes
+ input.advance(numBytesToCopy);
+ varintBufferPos_ += numBytesToCopy;
+ return false;
+ }
+}
- cursor.skip(buf.size());
- buf = cursor.peekBytes();
+bool LZMA2StreamCodec::doUncompressStream(
+ ByteRange& input,
+ MutableByteRange& output,
+ StreamCodec::FlushOp flushOp) {
+ if (needReset_) {
+ resetDStream();
+ needReset_ = false;
+ needDecodeSize_ = encodeSize();
+ if (encodeSize()) {
+ // Reset buffer
+ varintBufferPos_ = 0;
+ }
}
- while (!streamEnd) {
- streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
+ if (needDecodeSize_) {
+ // Try decoding the varint. If the input does not contain the entire varint,
+ // buffer the input. If the varint can not be decoded, fail.
+ if (!decodeAndCheckVarint(input)) {
+ return false;
+ }
+ needDecodeSize_ = false;
}
- out->prev()->trimEnd(stream.avail_out);
+ dstream_->next_in = const_cast<uint8_t*>(input.data());
+ dstream_->avail_in = input.size();
+ dstream_->next_out = output.data();
+ dstream_->avail_out = output.size();
+ SCOPE_EXIT {
+ input.advance(input.size() - dstream_->avail_in);
+ output.advance(output.size() - dstream_->avail_out);
+ };
- if (uncompressedLength && *uncompressedLength != stream.total_out) {
- throw std::runtime_error(
- to<std::string>("LZMA2Codec: invalid uncompressed length"));
+ lzma_ret rc;
+ switch (flushOp) {
+ case StreamCodec::FlushOp::NONE:
+ case StreamCodec::FlushOp::FLUSH:
+ rc = lzmaThrowOnError(lzma_code(dstream_.get_pointer(), LZMA_RUN));
+ break;
+ case StreamCodec::FlushOp::END:
+ rc = lzmaThrowOnError(lzma_code(dstream_.get_pointer(), LZMA_FINISH));
+ break;
+ default:
+ throw std::invalid_argument("LZMA2StreamCodec: invalid flush");
}
-
- return out;
+ return rc == LZMA_STREAM_END;
}
-
-#endif // FOLLY_HAVE_LIBLZMA
+#endif // FOLLY_HAVE_LIBLZMA
#ifdef FOLLY_HAVE_LIBZSTD
+namespace {
+void zstdFreeCStream(ZSTD_CStream* zcs) {
+ ZSTD_freeCStream(zcs);
+}
+
+void zstdFreeDStream(ZSTD_DStream* zds) {
+ ZSTD_freeDStream(zds);
+}
+}
+
/**
* ZSTD compression
*/
-class ZSTDCodec final : public Codec {
+class ZSTDStreamCodec final : public StreamCodec {
public:
- static std::unique_ptr<Codec> create(int level, CodecType);
- explicit ZSTDCodec(int level, CodecType type);
+ static std::unique_ptr<Codec> createCodec(int level, CodecType);
+ static std::unique_ptr<StreamCodec> createStream(int level, CodecType);
+ explicit ZSTDStreamCodec(int level, CodecType type);
std::vector<std::string> validPrefixes() const override;
bool canUncompress(const IOBuf* data, Optional<uint64_t> uncompressedLength)
private:
bool doNeedsUncompressedLength() const override;
uint64_t doMaxCompressedLength(uint64_t uncompressedLength) const override;
- std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
- std::unique_ptr<IOBuf> doUncompress(
- const IOBuf* data,
- Optional<uint64_t> uncompressedLength) override;
+ Optional<uint64_t> doGetUncompressedLength(
+ IOBuf const* data,
+ Optional<uint64_t> uncompressedLength) const override;
+
+ void doResetStream() override;
+ bool doCompressStream(
+ ByteRange& input,
+ MutableByteRange& output,
+ StreamCodec::FlushOp flushOp) override;
+ bool doUncompressStream(
+ ByteRange& input,
+ MutableByteRange& output,
+ StreamCodec::FlushOp flushOp) override;
+
+ void resetCStream();
+ void resetDStream();
+
+ bool tryBlockCompress(ByteRange& input, MutableByteRange& output) const;
+ bool tryBlockUncompress(ByteRange& input, MutableByteRange& output) const;
int level_;
+ bool needReset_{true};
+ std::unique_ptr<
+ ZSTD_CStream,
+ folly::static_function_deleter<ZSTD_CStream, &zstdFreeCStream>>
+ cstream_{nullptr};
+ std::unique_ptr<
+ ZSTD_DStream,
+ folly::static_function_deleter<ZSTD_DStream, &zstdFreeDStream>>
+ dstream_{nullptr};
};
static constexpr uint32_t kZSTDMagicLE = 0xFD2FB528;
-std::vector<std::string> ZSTDCodec::validPrefixes() const {
+std::vector<std::string> ZSTDStreamCodec::validPrefixes() const {
return {prefixToStringLE(kZSTDMagicLE)};
}
-bool ZSTDCodec::canUncompress(const IOBuf* data, Optional<uint64_t>) const {
+bool ZSTDStreamCodec::canUncompress(const IOBuf* data, Optional<uint64_t>)
+ const {
return dataStartsWithLE(data, kZSTDMagicLE);
}
-std::unique_ptr<Codec> ZSTDCodec::create(int level, CodecType type) {
- return std::make_unique<ZSTDCodec>(level, type);
+std::unique_ptr<Codec> ZSTDStreamCodec::createCodec(int level, CodecType type) {
+ return make_unique<ZSTDStreamCodec>(level, type);
+}
+
+std::unique_ptr<StreamCodec> ZSTDStreamCodec::createStream(
+ int level,
+ CodecType type) {
+ return make_unique<ZSTDStreamCodec>(level, type);
}
-ZSTDCodec::ZSTDCodec(int level, CodecType type) : Codec(type) {
+ZSTDStreamCodec::ZSTDStreamCodec(int level, CodecType type)
+ : StreamCodec(type) {
DCHECK(type == CodecType::ZSTD);
switch (level) {
case COMPRESSION_LEVEL_FASTEST:
level_ = level;
}
-bool ZSTDCodec::doNeedsUncompressedLength() const {
+bool ZSTDStreamCodec::doNeedsUncompressedLength() const {
return false;
}
-uint64_t ZSTDCodec::doMaxCompressedLength(uint64_t uncompressedLength) const {
+uint64_t ZSTDStreamCodec::doMaxCompressedLength(
+ uint64_t uncompressedLength) const {
return ZSTD_compressBound(uncompressedLength);
}
to<std::string>("ZSTD returned an error: ", ZSTD_getErrorName(rc)));
}
-std::unique_ptr<IOBuf> ZSTDCodec::doCompress(const IOBuf* data) {
- // Support earlier versions of the codec (working with a single IOBuf,
- // and using ZSTD_decompress which requires ZSTD frame to contain size,
- // which isn't populated by streaming API).
- if (!data->isChained()) {
- auto out = IOBuf::createCombined(ZSTD_compressBound(data->length()));
- const auto rc = ZSTD_compress(
- out->writableData(),
- out->capacity(),
- data->data(),
- data->length(),
- level_);
- zstdThrowIfError(rc);
- out->append(rc);
- return out;
- }
-
- auto zcs = ZSTD_createCStream();
- SCOPE_EXIT {
- ZSTD_freeCStream(zcs);
- };
-
- auto rc = ZSTD_initCStream(zcs, level_);
- zstdThrowIfError(rc);
-
- Cursor cursor(data);
- auto result =
- IOBuf::createCombined(maxCompressedLength(cursor.totalLength()));
-
- ZSTD_outBuffer out;
- out.dst = result->writableTail();
- out.size = result->capacity();
- out.pos = 0;
-
- for (auto buffer = cursor.peekBytes(); !buffer.empty();) {
- ZSTD_inBuffer in;
- in.src = buffer.data();
- in.size = buffer.size();
- for (in.pos = 0; in.pos != in.size;) {
- rc = ZSTD_compressStream(zcs, &out, &in);
- zstdThrowIfError(rc);
+Optional<uint64_t> ZSTDStreamCodec::doGetUncompressedLength(
+ IOBuf const* data,
+ Optional<uint64_t> uncompressedLength) const {
+ // Read decompressed size from frame if available in first IOBuf.
+ auto const decompressedSize =
+ ZSTD_getDecompressedSize(data->data(), data->length());
+ if (decompressedSize != 0) {
+ if (uncompressedLength && *uncompressedLength != decompressedSize) {
+ throw std::runtime_error("ZSTD: invalid uncompressed length");
}
- cursor.skip(in.size);
- buffer = cursor.peekBytes();
+ uncompressedLength = decompressedSize;
}
+ return uncompressedLength;
+}
- rc = ZSTD_endStream(zcs, &out);
- zstdThrowIfError(rc);
- CHECK_EQ(rc, 0);
+void ZSTDStreamCodec::doResetStream() {
+ needReset_ = true;
+}
- result->append(out.pos);
- return result;
+bool ZSTDStreamCodec::tryBlockCompress(
+ ByteRange& input,
+ MutableByteRange& output) const {
+ DCHECK(needReset_);
+ // We need to know that we have enough output space to use block compression
+ if (output.size() < ZSTD_compressBound(input.size())) {
+ return false;
+ }
+ size_t const length = ZSTD_compress(
+ output.data(), output.size(), input.data(), input.size(), level_);
+ zstdThrowIfError(length);
+ input.uncheckedAdvance(input.size());
+ output.uncheckedAdvance(length);
+ return true;
}
-static std::unique_ptr<IOBuf> zstdUncompressBuffer(
- const IOBuf* data,
- Optional<uint64_t> uncompressedLength) {
- // Check preconditions
- DCHECK(!data->isChained());
- DCHECK(uncompressedLength.hasValue());
-
- auto uncompressed = IOBuf::create(*uncompressedLength);
- const auto decompressedSize = ZSTD_decompress(
- uncompressed->writableTail(),
- uncompressed->tailroom(),
- data->data(),
- data->length());
- zstdThrowIfError(decompressedSize);
- if (decompressedSize != uncompressedLength) {
- throw std::runtime_error("ZSTD: invalid uncompressed length");
+void ZSTDStreamCodec::resetCStream() {
+ if (!cstream_) {
+ cstream_.reset(ZSTD_createCStream());
+ if (!cstream_) {
+ throw std::bad_alloc{};
+ }
}
- uncompressed->append(decompressedSize);
- return uncompressed;
+ // Advanced API usage works for all supported versions of zstd.
+ // Required to set contentSizeFlag.
+ auto params = ZSTD_getParams(level_, uncompressedLength().value_or(0), 0);
+ params.fParams.contentSizeFlag = uncompressedLength().hasValue();
+ zstdThrowIfError(ZSTD_initCStream_advanced(
+ cstream_.get(), nullptr, 0, params, uncompressedLength().value_or(0)));
}
-static std::unique_ptr<IOBuf> zstdUncompressStream(
- const IOBuf* data,
- Optional<uint64_t> uncompressedLength) {
- auto zds = ZSTD_createDStream();
+bool ZSTDStreamCodec::doCompressStream(
+ ByteRange& input,
+ MutableByteRange& output,
+ StreamCodec::FlushOp flushOp) {
+ if (needReset_) {
+ // If we are given all the input in one chunk try to use block compression
+ if (flushOp == StreamCodec::FlushOp::END &&
+ tryBlockCompress(input, output)) {
+ return true;
+ }
+ resetCStream();
+ needReset_ = false;
+ }
+ ZSTD_inBuffer in = {input.data(), input.size(), 0};
+ ZSTD_outBuffer out = {output.data(), output.size(), 0};
SCOPE_EXIT {
- ZSTD_freeDStream(zds);
+ input.uncheckedAdvance(in.pos);
+ output.uncheckedAdvance(out.pos);
};
-
- auto rc = ZSTD_initDStream(zds);
- zstdThrowIfError(rc);
-
- ZSTD_outBuffer out{};
- ZSTD_inBuffer in{};
-
- auto outputSize = uncompressedLength.value_or(ZSTD_DStreamOutSize());
-
- IOBufQueue queue(IOBufQueue::cacheChainLength());
-
- Cursor cursor(data);
- for (rc = 0;;) {
- if (in.pos == in.size) {
- auto buffer = cursor.peekBytes();
- in.src = buffer.data();
- in.size = buffer.size();
- in.pos = 0;
- cursor.skip(in.size);
- if (rc > 1 && in.size == 0) {
- throw std::runtime_error(to<std::string>("ZSTD: incomplete input"));
- }
- }
- if (out.pos == out.size) {
- if (out.pos != 0) {
- queue.postallocate(out.pos);
- }
- auto buffer = queue.preallocate(outputSize, outputSize);
- out.dst = buffer.first;
- out.size = buffer.second;
- out.pos = 0;
- outputSize = ZSTD_DStreamOutSize();
+ if (flushOp == StreamCodec::FlushOp::NONE || !input.empty()) {
+ zstdThrowIfError(ZSTD_compressStream(cstream_.get(), &out, &in));
+ }
+ if (in.pos == in.size && flushOp != StreamCodec::FlushOp::NONE) {
+ size_t rc;
+ switch (flushOp) {
+ case StreamCodec::FlushOp::FLUSH:
+ rc = ZSTD_flushStream(cstream_.get(), &out);
+ break;
+ case StreamCodec::FlushOp::END:
+ rc = ZSTD_endStream(cstream_.get(), &out);
+ break;
+ default:
+ throw std::invalid_argument("ZSTD: invalid FlushOp");
}
- rc = ZSTD_decompressStream(zds, &out, &in);
zstdThrowIfError(rc);
if (rc == 0) {
- break;
+ return true;
}
}
- if (out.pos != 0) {
- queue.postallocate(out.pos);
- }
- if (in.pos != in.size || !cursor.isAtEnd()) {
- throw std::runtime_error("ZSTD: junk after end of data");
- }
- if (uncompressedLength && queue.chainLength() != *uncompressedLength) {
- throw std::runtime_error("ZSTD: invalid uncompressed length");
- }
+ return false;
+}
- return queue.move();
+bool ZSTDStreamCodec::tryBlockUncompress(
+ ByteRange& input,
+ MutableByteRange& output) const {
+ DCHECK(needReset_);
+#if ZSTD_VERSION_NUMBER < 10104
+ // We require ZSTD_findFrameCompressedSize() to perform this optimization.
+ return false;
+#else
+ // We need to know the uncompressed length and have enough output space.
+ if (!uncompressedLength() || output.size() < *uncompressedLength()) {
+ return false;
+ }
+ size_t const compressedLength =
+ ZSTD_findFrameCompressedSize(input.data(), input.size());
+ zstdThrowIfError(compressedLength);
+ size_t const length = ZSTD_decompress(
+ output.data(), *uncompressedLength(), input.data(), compressedLength);
+ zstdThrowIfError(length);
+ if (length != *uncompressedLength()) {
+ throw std::runtime_error("ZSTDStreamCodec: Incorrect uncompressed length");
+ }
+ input.uncheckedAdvance(compressedLength);
+ output.uncheckedAdvance(length);
+ return true;
+#endif
}
-std::unique_ptr<IOBuf> ZSTDCodec::doUncompress(
- const IOBuf* data,
- Optional<uint64_t> uncompressedLength) {
- {
- // Read decompressed size from frame if available in first IOBuf.
- const auto decompressedSize =
- ZSTD_getDecompressedSize(data->data(), data->length());
- if (decompressedSize != 0) {
- if (uncompressedLength && *uncompressedLength != decompressedSize) {
- throw std::runtime_error("ZSTD: invalid uncompressed length");
- }
- uncompressedLength = decompressedSize;
+void ZSTDStreamCodec::resetDStream() {
+ if (!dstream_) {
+ dstream_.reset(ZSTD_createDStream());
+ if (!dstream_) {
+ throw std::bad_alloc{};
}
}
- // Faster to decompress using ZSTD_decompress() if we can.
- if (uncompressedLength && !data->isChained()) {
- return zstdUncompressBuffer(data, uncompressedLength);
+ zstdThrowIfError(ZSTD_initDStream(dstream_.get()));
+}
+
+bool ZSTDStreamCodec::doUncompressStream(
+ ByteRange& input,
+ MutableByteRange& output,
+ StreamCodec::FlushOp flushOp) {
+ if (needReset_) {
+ // If we are given all the input in one chunk try to use block uncompression
+ if (flushOp == StreamCodec::FlushOp::END &&
+ tryBlockUncompress(input, output)) {
+ return true;
+ }
+ resetDStream();
+ needReset_ = false;
}
- // Fall back to slower streaming decompression.
- return zstdUncompressStream(data, uncompressedLength);
+ ZSTD_inBuffer in = {input.data(), input.size(), 0};
+ ZSTD_outBuffer out = {output.data(), output.size(), 0};
+ SCOPE_EXIT {
+ input.uncheckedAdvance(in.pos);
+ output.uncheckedAdvance(out.pos);
+ };
+ size_t const rc = ZSTD_decompressStream(dstream_.get(), &out, &in);
+ zstdThrowIfError(rc);
+ return rc == 0;
}
-#endif // FOLLY_HAVE_LIBZSTD
+#endif // FOLLY_HAVE_LIBZSTD
#if FOLLY_HAVE_LIBBZ2
if (stream.avail_out == 0) {
out->prependChain(addOutputBuffer(&stream, kDefaultBufferLength));
}
-
+ size_t const outputSize = stream.avail_out;
rc = bzCheck(BZ2_bzDecompress(&stream));
+ if (outputSize == stream.avail_out) {
+ throw std::runtime_error("Bzip2Codec: Truncated input");
+ }
}
out->prev()->trimEnd(stream.avail_out);
#endif // FOLLY_HAVE_LIBBZ2
+#if FOLLY_HAVE_LIBZ
+
+zlib::Options getZlibOptions(CodecType type) {
+ DCHECK(type == CodecType::GZIP || type == CodecType::ZLIB);
+ return type == CodecType::GZIP ? zlib::defaultGzipOptions()
+ : zlib::defaultZlibOptions();
+}
+
+std::unique_ptr<Codec> getZlibCodec(int level, CodecType type) {
+ return zlib::getCodec(getZlibOptions(type), level);
+}
+
+std::unique_ptr<StreamCodec> getZlibStreamCodec(int level, CodecType type) {
+ return zlib::getStreamCodec(getZlibOptions(type), level);
+}
+
+#endif // FOLLY_HAVE_LIBZ
+
/**
* Automatic decompression
*/
#endif
#if FOLLY_HAVE_LIBZ
- {ZlibCodec::create, nullptr},
+ {getZlibCodec, getZlibStreamCodec},
#else
{},
#endif
#endif
#if FOLLY_HAVE_LIBLZMA
- {LZMA2Codec::create, nullptr},
- {LZMA2Codec::create, nullptr},
+ {LZMA2StreamCodec::createCodec, LZMA2StreamCodec::createStream},
+ {LZMA2StreamCodec::createCodec, LZMA2StreamCodec::createStream},
#else
{},
{},
#endif
#if FOLLY_HAVE_LIBZSTD
- {ZSTDCodec::create, nullptr},
+ {ZSTDStreamCodec::createCodec, ZSTDStreamCodec::createStream},
#else
{},
#endif
#if FOLLY_HAVE_LIBZ
- {ZlibCodec::create, nullptr},
+ {getZlibCodec, getZlibStreamCodec},
#else
{},
#endif
std::vector<std::unique_ptr<Codec>> customCodecs) {
return AutomaticCodec::create(std::move(customCodecs));
}
-}} // namespaces
+} // namespace io
+} // namespace folly