folly/io/Compression.cpp

   1 /*
   2  * Copyright 2017 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include <folly/io/Compression.h>
  18
  19 #if FOLLY_HAVE_LIBLZ4
  20 #include <lz4.h>
  21 #include <lz4hc.h>
  22 #if LZ4_VERSION_NUMBER >= 10301
  23 #include <lz4frame.h>
  24 #endif
  25 #endif
  26
  27 #include <glog/logging.h>
  28
  29 #if FOLLY_HAVE_LIBSNAPPY
  30 #include <snappy.h>
  31 #include <snappy-sinksource.h>
  32 #endif
  33
  34 #if FOLLY_HAVE_LIBZ
  35 #include <zlib.h>
  36 #endif
  37
  38 #if FOLLY_HAVE_LIBLZMA
  39 #include <lzma.h>
  40 #endif
  41
  42 #if FOLLY_HAVE_LIBZSTD
  43 #define ZSTD_STATIC_LINKING_ONLY
  44 #include <zstd.h>
  45 #endif
  46
  47 #if FOLLY_HAVE_LIBBZ2
  48 #include <bzlib.h>
  49 #endif
  50
  51 #include <folly/Bits.h>
  52 #include <folly/Conv.h>
  53 #include <folly/Memory.h>
  54 #include <folly/Portability.h>
  55 #include <folly/ScopeGuard.h>
  56 #include <folly/Varint.h>
  57 #include <folly/io/Cursor.h>
  58 #include <algorithm>
  59 #include <unordered_set>
  60
  61 namespace folly { namespace io {
  62
  63 Codec::Codec(CodecType type) : type_(type) { }
  64
  65 // Ensure consistent behavior in the nullptr case
  66 std::unique_ptr<IOBuf> Codec::compress(const IOBuf* data) {
  67   if (data == nullptr) {
  68     throw std::invalid_argument("Codec: data must not be nullptr");
  69   }
  70   uint64_t len = data->computeChainDataLength();
  71   if (len == 0) {
  72     return IOBuf::create(0);
  73   }
  74   if (len > maxUncompressedLength()) {
  75     throw std::runtime_error("Codec: uncompressed length too large");
  76   }
  77
  78   return doCompress(data);
  79 }
  80
  81 std::string Codec::compress(const StringPiece data) {
  82   const uint64_t len = data.size();
  83   if (len == 0) {
  84     return "";
  85   }
  86   if (len > maxUncompressedLength()) {
  87     throw std::runtime_error("Codec: uncompressed length too large");
  88   }
  89
  90   return doCompressString(data);
  91 }
  92
  93 std::unique_ptr<IOBuf> Codec::uncompress(
  94     const IOBuf* data,
  95     Optional<uint64_t> uncompressedLength) {
  96   if (data == nullptr) {
  97     throw std::invalid_argument("Codec: data must not be nullptr");
  98   }
  99   if (!uncompressedLength) {
 100     if (needsUncompressedLength()) {
 101       throw std::invalid_argument("Codec: uncompressed length required");
 102     }
 103   } else if (*uncompressedLength > maxUncompressedLength()) {
 104     throw std::runtime_error("Codec: uncompressed length too large");
 105   }
 106
 107   if (data->empty()) {
 108     if (uncompressedLength.value_or(0) != 0) {
 109       throw std::runtime_error("Codec: invalid uncompressed length");
 110     }
 111     return IOBuf::create(0);
 112   }
 113
 114   return doUncompress(data, uncompressedLength);
 115 }
 116
 117 std::string Codec::uncompress(
 118     const StringPiece data,
 119     Optional<uint64_t> uncompressedLength) {
 120   if (!uncompressedLength) {
 121     if (needsUncompressedLength()) {
 122       throw std::invalid_argument("Codec: uncompressed length required");
 123     }
 124   } else if (*uncompressedLength > maxUncompressedLength()) {
 125     throw std::runtime_error("Codec: uncompressed length too large");
 126   }
 127
 128   if (data.empty()) {
 129     if (uncompressedLength.value_or(0) != 0) {
 130       throw std::runtime_error("Codec: invalid uncompressed length");
 131     }
 132     return "";
 133   }
 134
 135   return doUncompressString(data, uncompressedLength);
 136 }
 137
 138 bool Codec::needsUncompressedLength() const {
 139   return doNeedsUncompressedLength();
 140 }
 141
 142 uint64_t Codec::maxUncompressedLength() const {
 143   return doMaxUncompressedLength();
 144 }
 145
 146 bool Codec::doNeedsUncompressedLength() const {
 147   return false;
 148 }
 149
 150 uint64_t Codec::doMaxUncompressedLength() const {
 151   return UNLIMITED_UNCOMPRESSED_LENGTH;
 152 }
 153
 154 std::vector<std::string> Codec::validPrefixes() const {
 155   return {};
 156 }
 157
 158 bool Codec::canUncompress(const IOBuf*, Optional<uint64_t>) const {
 159   return false;
 160 }
 161
 162 std::string Codec::doCompressString(const StringPiece data) {
 163   const IOBuf inputBuffer{IOBuf::WRAP_BUFFER, data};
 164   auto outputBuffer = doCompress(&inputBuffer);
 165   std::string output;
 166   output.reserve(outputBuffer->computeChainDataLength());
 167   for (auto range : *outputBuffer) {
 168     output.append(reinterpret_cast<const char*>(range.data()), range.size());
 169   }
 170   return output;
 171 }
 172
 173 std::string Codec::doUncompressString(
 174     const StringPiece data,
 175     Optional<uint64_t> uncompressedLength) {
 176   const IOBuf inputBuffer{IOBuf::WRAP_BUFFER, data};
 177   auto outputBuffer = doUncompress(&inputBuffer, uncompressedLength);
 178   std::string output;
 179   output.reserve(outputBuffer->computeChainDataLength());
 180   for (auto range : *outputBuffer) {
 181     output.append(reinterpret_cast<const char*>(range.data()), range.size());
 182   }
 183   return output;
 184 }
 185
 186 uint64_t Codec::maxCompressedLength(uint64_t uncompressedLength) const {
 187   if (uncompressedLength == 0) {
 188     return 0;
 189   }
 190   return doMaxCompressedLength(uncompressedLength);
 191 }
 192
 193 Optional<uint64_t> Codec::getUncompressedLength(
 194     const folly::IOBuf* data,
 195     Optional<uint64_t> uncompressedLength) const {
 196   auto const compressedLength = data->computeChainDataLength();
 197   if (uncompressedLength == uint64_t(0) || compressedLength == 0) {
 198     if (uncompressedLength.value_or(0) != 0 || compressedLength != 0) {
 199       throw std::runtime_error("Invalid uncompressed length");
 200     }
 201     return 0;
 202   }
 203   return doGetUncompressedLength(data, uncompressedLength);
 204 }
 205
 206 Optional<uint64_t> Codec::doGetUncompressedLength(
 207     const folly::IOBuf*,
 208     Optional<uint64_t> uncompressedLength) const {
 209   return uncompressedLength;
 210 }
 211
 212 bool StreamCodec::needsDataLength() const {
 213   return doNeedsDataLength();
 214 }
 215
 216 bool StreamCodec::doNeedsDataLength() const {
 217   return false;
 218 }
 219
 220 void StreamCodec::assertStateIs(State expected) const {
 221   if (state_ != expected) {
 222     throw std::logic_error(folly::to<std::string>(
 223         "Codec: state is ", state_, "; expected state ", expected));
 224   }
 225 }
 226
 227 void StreamCodec::resetStream(Optional<uint64_t> uncompressedLength) {
 228   state_ = State::RESET;
 229   uncompressedLength_ = uncompressedLength;
 230   doResetStream();
 231 }
 232
 233 bool StreamCodec::compressStream(
 234     ByteRange& input,
 235     MutableByteRange& output,
 236     StreamCodec::FlushOp flushOp) {
 237   if (state_ == State::RESET && input.empty()) {
 238     if (flushOp == StreamCodec::FlushOp::NONE) {
 239       return false;
 240     }
 241     if (flushOp == StreamCodec::FlushOp::END &&
 242         uncompressedLength().value_or(0) != 0) {
 243       throw std::runtime_error("Codec: invalid uncompressed length");
 244     }
 245     return true;
 246   }
 247   if (state_ == State::RESET && !input.empty() &&
 248       uncompressedLength() == uint64_t(0)) {
 249     throw std::runtime_error("Codec: invalid uncompressed length");
 250   }
 251   // Handle input state transitions
 252   switch (flushOp) {
 253     case StreamCodec::FlushOp::NONE:
 254       if (state_ == State::RESET) {
 255         state_ = State::COMPRESS;
 256       }
 257       assertStateIs(State::COMPRESS);
 258       break;
 259     case StreamCodec::FlushOp::FLUSH:
 260       if (state_ == State::RESET || state_ == State::COMPRESS) {
 261         state_ = State::COMPRESS_FLUSH;
 262       }
 263       assertStateIs(State::COMPRESS_FLUSH);
 264       break;
 265     case StreamCodec::FlushOp::END:
 266       if (state_ == State::RESET || state_ == State::COMPRESS) {
 267         state_ = State::COMPRESS_END;
 268       }
 269       assertStateIs(State::COMPRESS_END);
 270       break;
 271   }
 272   bool const done = doCompressStream(input, output, flushOp);
 273   // Handle output state transitions
 274   if (done) {
 275     if (state_ == State::COMPRESS_FLUSH) {
 276       state_ = State::COMPRESS;
 277     } else if (state_ == State::COMPRESS_END) {
 278       state_ = State::END;
 279     }
 280     // Check internal invariants
 281     DCHECK(input.empty());
 282     DCHECK(flushOp != StreamCodec::FlushOp::NONE);
 283   }
 284   return done;
 285 }
 286
 287 bool StreamCodec::uncompressStream(
 288     ByteRange& input,
 289     MutableByteRange& output,
 290     StreamCodec::FlushOp flushOp) {
 291   if (state_ == State::RESET && input.empty()) {
 292     if (uncompressedLength().value_or(0) == 0) {
 293       return true;
 294     }
 295     return false;
 296   }
 297   // Handle input state transitions
 298   if (state_ == State::RESET) {
 299     state_ = State::UNCOMPRESS;
 300   }
 301   assertStateIs(State::UNCOMPRESS);
 302   bool const done = doUncompressStream(input, output, flushOp);
 303   // Handle output state transitions
 304   if (done) {
 305     state_ = State::END;
 306   }
 307   return done;
 308 }
 309
 310 static std::unique_ptr<IOBuf> addOutputBuffer(
 311     MutableByteRange& output,
 312     uint64_t size) {
 313   DCHECK(output.empty());
 314   auto buffer = IOBuf::create(size);
 315   buffer->append(buffer->capacity());
 316   output = {buffer->writableData(), buffer->length()};
 317   return buffer;
 318 }
 319
 320 std::unique_ptr<IOBuf> StreamCodec::doCompress(IOBuf const* data) {
 321   uint64_t const uncompressedLength = data->computeChainDataLength();
 322   resetStream(uncompressedLength);
 323   uint64_t const maxCompressedLen = maxCompressedLength(uncompressedLength);
 324
 325   auto constexpr kMaxSingleStepLength = uint64_t(64) << 20; // 64 MB
 326   auto constexpr kDefaultBufferLength = uint64_t(4) << 20; // 4 MB
 327
 328   MutableByteRange output;
 329   auto buffer = addOutputBuffer(
 330       output,
 331       maxCompressedLen <= kMaxSingleStepLength ? maxCompressedLen
 332                                                : kDefaultBufferLength);
 333
 334   // Compress the entire IOBuf chain into the IOBuf chain pointed to by buffer
 335   IOBuf const* current = data;
 336   ByteRange input{current->data(), current->length()};
 337   StreamCodec::FlushOp flushOp = StreamCodec::FlushOp::NONE;
 338   for (;;) {
 339     while (input.empty() && current->next() != data) {
 340       current = current->next();
 341       input = {current->data(), current->length()};
 342     }
 343     if (current->next() == data) {
 344       // This is the last input buffer so end the stream
 345       flushOp = StreamCodec::FlushOp::END;
 346     }
 347     if (output.empty()) {
 348       buffer->prependChain(addOutputBuffer(output, kDefaultBufferLength));
 349     }
 350     size_t const inputSize = input.size();
 351     size_t const outputSize = output.size();
 352     bool const done = compressStream(input, output, flushOp);
 353     if (done) {
 354       DCHECK(input.empty());
 355       DCHECK(flushOp == StreamCodec::FlushOp::END);
 356       DCHECK_EQ(current->next(), data);
 357       break;
 358     }
 359     if (inputSize == input.size() && outputSize == output.size()) {
 360       throw std::runtime_error("Codec: No forward progress made");
 361     }
 362   }
 363   buffer->prev()->trimEnd(output.size());
 364   return buffer;
 365 }
 366
 367 static uint64_t computeBufferLength(
 368     uint64_t const compressedLength,
 369     uint64_t const blockSize) {
 370   uint64_t constexpr kMaxBufferLength = uint64_t(4) << 20; // 4 MiB
 371   uint64_t const goodBufferSize = 4 * std::max(blockSize, compressedLength);
 372   return std::min(goodBufferSize, kMaxBufferLength);
 373 }
 374
 375 std::unique_ptr<IOBuf> StreamCodec::doUncompress(
 376     IOBuf const* data,
 377     Optional<uint64_t> uncompressedLength) {
 378   auto constexpr kMaxSingleStepLength = uint64_t(64) << 20; // 64 MB
 379   auto constexpr kBlockSize = uint64_t(128) << 10;
 380   auto const defaultBufferLength =
 381       computeBufferLength(data->computeChainDataLength(), kBlockSize);
 382
 383   uncompressedLength = getUncompressedLength(data, uncompressedLength);
 384   resetStream(uncompressedLength);
 385
 386   MutableByteRange output;
 387   auto buffer = addOutputBuffer(
 388       output,
 389       (uncompressedLength && *uncompressedLength <= kMaxSingleStepLength
 390            ? *uncompressedLength
 391            : defaultBufferLength));
 392
 393   // Uncompress the entire IOBuf chain into the IOBuf chain pointed to by buffer
 394   IOBuf const* current = data;
 395   ByteRange input{current->data(), current->length()};
 396   StreamCodec::FlushOp flushOp = StreamCodec::FlushOp::NONE;
 397   for (;;) {
 398     while (input.empty() && current->next() != data) {
 399       current = current->next();
 400       input = {current->data(), current->length()};
 401     }
 402     if (current->next() == data) {
 403       // Tell the uncompressor there is no more input (it may optimize)
 404       flushOp = StreamCodec::FlushOp::END;
 405     }
 406     if (output.empty()) {
 407       buffer->prependChain(addOutputBuffer(output, defaultBufferLength));
 408     }
 409     size_t const inputSize = input.size();
 410     size_t const outputSize = output.size();
 411     bool const done = uncompressStream(input, output, flushOp);
 412     if (done) {
 413       break;
 414     }
 415     if (inputSize == input.size() && outputSize == output.size()) {
 416       throw std::runtime_error("Codec: Truncated data");
 417     }
 418   }
 419   if (!input.empty()) {
 420     throw std::runtime_error("Codec: Junk after end of data");
 421   }
 422
 423   buffer->prev()->trimEnd(output.size());
 424   if (uncompressedLength &&
 425       *uncompressedLength != buffer->computeChainDataLength()) {
 426     throw std::runtime_error("Codec: invalid uncompressed length");
 427   }
 428
 429   return buffer;
 430 }
 431
 432 namespace {
 433
 434 /**
 435  * No compression
 436  */
 437 class NoCompressionCodec final : public Codec {
 438  public:
 439   static std::unique_ptr<Codec> create(int level, CodecType type);
 440   explicit NoCompressionCodec(int level, CodecType type);
 441
 442  private:
 443   uint64_t doMaxCompressedLength(uint64_t uncompressedLength) const override;
 444   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
 445   std::unique_ptr<IOBuf> doUncompress(
 446       const IOBuf* data,
 447       Optional<uint64_t> uncompressedLength) override;
 448 };
 449
 450 std::unique_ptr<Codec> NoCompressionCodec::create(int level, CodecType type) {
 451   return std::make_unique<NoCompressionCodec>(level, type);
 452 }
 453
 454 NoCompressionCodec::NoCompressionCodec(int level, CodecType type)
 455   : Codec(type) {
 456   DCHECK(type == CodecType::NO_COMPRESSION);
 457   switch (level) {
 458   case COMPRESSION_LEVEL_DEFAULT:
 459   case COMPRESSION_LEVEL_FASTEST:
 460   case COMPRESSION_LEVEL_BEST:
 461     level = 0;
 462   }
 463   if (level != 0) {
 464     throw std::invalid_argument(to<std::string>(
 465         "NoCompressionCodec: invalid level ", level));
 466   }
 467 }
 468
 469 uint64_t NoCompressionCodec::doMaxCompressedLength(
 470     uint64_t uncompressedLength) const {
 471   return uncompressedLength;
 472 }
 473
 474 std::unique_ptr<IOBuf> NoCompressionCodec::doCompress(
 475     const IOBuf* data) {
 476   return data->clone();
 477 }
 478
 479 std::unique_ptr<IOBuf> NoCompressionCodec::doUncompress(
 480     const IOBuf* data,
 481     Optional<uint64_t> uncompressedLength) {
 482   if (uncompressedLength &&
 483       data->computeChainDataLength() != *uncompressedLength) {
 484     throw std::runtime_error(
 485         to<std::string>("NoCompressionCodec: invalid uncompressed length"));
 486   }
 487   return data->clone();
 488 }
 489
 490 #if (FOLLY_HAVE_LIBLZ4 || FOLLY_HAVE_LIBLZMA)
 491
 492 namespace {
 493
 494 void encodeVarintToIOBuf(uint64_t val, folly::IOBuf* out) {
 495   DCHECK_GE(out->tailroom(), kMaxVarintLength64);
 496   out->append(encodeVarint(val, out->writableTail()));
 497 }
 498
 499 inline uint64_t decodeVarintFromCursor(folly::io::Cursor& cursor) {
 500   uint64_t val = 0;
 501   int8_t b = 0;
 502   for (int shift = 0; shift <= 63; shift += 7) {
 503     b = cursor.read<int8_t>();
 504     val |= static_cast<uint64_t>(b & 0x7f) << shift;
 505     if (b >= 0) {
 506       break;
 507     }
 508   }
 509   if (b < 0) {
 510     throw std::invalid_argument("Invalid varint value. Too big.");
 511   }
 512   return val;
 513 }
 514
 515 }  // namespace
 516
 517 #endif  // FOLLY_HAVE_LIBLZ4 || FOLLY_HAVE_LIBLZMA
 518
 519 namespace {
 520 /**
 521  * Reads sizeof(T) bytes, and returns false if not enough bytes are available.
 522  * Returns true if the first n bytes are equal to prefix when interpreted as
 523  * a little endian T.
 524  */
 525 template <typename T>
 526 typename std::enable_if<std::is_unsigned<T>::value, bool>::type
 527 dataStartsWithLE(const IOBuf* data, T prefix, uint64_t n = sizeof(T)) {
 528   DCHECK_GT(n, 0);
 529   DCHECK_LE(n, sizeof(T));
 530   T value;
 531   Cursor cursor{data};
 532   if (!cursor.tryReadLE(value)) {
 533     return false;
 534   }
 535   const T mask = n == sizeof(T) ? T(-1) : (T(1) << (8 * n)) - 1;
 536   return prefix == (value & mask);
 537 }
 538
 539 template <typename T>
 540 typename std::enable_if<std::is_arithmetic<T>::value, std::string>::type
 541 prefixToStringLE(T prefix, uint64_t n = sizeof(T)) {
 542   DCHECK_GT(n, 0);
 543   DCHECK_LE(n, sizeof(T));
 544   prefix = Endian::little(prefix);
 545   std::string result;
 546   result.resize(n);
 547   memcpy(&result[0], &prefix, n);
 548   return result;
 549 }
 550 } // namespace
 551
 552 #if FOLLY_HAVE_LIBLZ4
 553
 554 /**
 555  * LZ4 compression
 556  */
 557 class LZ4Codec final : public Codec {
 558  public:
 559   static std::unique_ptr<Codec> create(int level, CodecType type);
 560   explicit LZ4Codec(int level, CodecType type);
 561
 562  private:
 563   bool doNeedsUncompressedLength() const override;
 564   uint64_t doMaxUncompressedLength() const override;
 565   uint64_t doMaxCompressedLength(uint64_t uncompressedLength) const override;
 566
 567   bool encodeSize() const { return type() == CodecType::LZ4_VARINT_SIZE; }
 568
 569   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
 570   std::unique_ptr<IOBuf> doUncompress(
 571       const IOBuf* data,
 572       Optional<uint64_t> uncompressedLength) override;
 573
 574   bool highCompression_;
 575 };
 576
 577 std::unique_ptr<Codec> LZ4Codec::create(int level, CodecType type) {
 578   return std::make_unique<LZ4Codec>(level, type);
 579 }
 580
 581 LZ4Codec::LZ4Codec(int level, CodecType type) : Codec(type) {
 582   DCHECK(type == CodecType::LZ4 || type == CodecType::LZ4_VARINT_SIZE);
 583
 584   switch (level) {
 585   case COMPRESSION_LEVEL_FASTEST:
 586   case COMPRESSION_LEVEL_DEFAULT:
 587     level = 1;
 588     break;
 589   case COMPRESSION_LEVEL_BEST:
 590     level = 2;
 591     break;
 592   }
 593   if (level < 1 || level > 2) {
 594     throw std::invalid_argument(to<std::string>(
 595         "LZ4Codec: invalid level: ", level));
 596   }
 597   highCompression_ = (level > 1);
 598 }
 599
 600 bool LZ4Codec::doNeedsUncompressedLength() const {
 601   return !encodeSize();
 602 }
 603
 604 // The value comes from lz4.h in lz4-r117, but older versions of lz4 don't
 605 // define LZ4_MAX_INPUT_SIZE (even though the max size is the same), so do it
 606 // here.
 607 #ifndef LZ4_MAX_INPUT_SIZE
 608 # define LZ4_MAX_INPUT_SIZE 0x7E000000
 609 #endif
 610
 611 uint64_t LZ4Codec::doMaxUncompressedLength() const {
 612   return LZ4_MAX_INPUT_SIZE;
 613 }
 614
 615 uint64_t LZ4Codec::doMaxCompressedLength(uint64_t uncompressedLength) const {
 616   return LZ4_compressBound(uncompressedLength) +
 617       (encodeSize() ? kMaxVarintLength64 : 0);
 618 }
 619
 620 std::unique_ptr<IOBuf> LZ4Codec::doCompress(const IOBuf* data) {
 621   IOBuf clone;
 622   if (data->isChained()) {
 623     // LZ4 doesn't support streaming, so we have to coalesce
 624     clone = data->cloneCoalescedAsValue();
 625     data = &clone;
 626   }
 627
 628   auto out = IOBuf::create(maxCompressedLength(data->length()));
 629   if (encodeSize()) {
 630     encodeVarintToIOBuf(data->length(), out.get());
 631   }
 632
 633   int n;
 634   auto input = reinterpret_cast<const char*>(data->data());
 635   auto output = reinterpret_cast<char*>(out->writableTail());
 636   const auto inputLength = data->length();
 637 #if LZ4_VERSION_NUMBER >= 10700
 638   if (highCompression_) {
 639     n = LZ4_compress_HC(input, output, inputLength, out->tailroom(), 0);
 640   } else {
 641     n = LZ4_compress_default(input, output, inputLength, out->tailroom());
 642   }
 643 #else
 644   if (highCompression_) {
 645     n = LZ4_compressHC(input, output, inputLength);
 646   } else {
 647     n = LZ4_compress(input, output, inputLength);
 648   }
 649 #endif
 650
 651   CHECK_GE(n, 0);
 652   CHECK_LE(n, out->capacity());
 653
 654   out->append(n);
 655   return out;
 656 }
 657
 658 std::unique_ptr<IOBuf> LZ4Codec::doUncompress(
 659     const IOBuf* data,
 660     Optional<uint64_t> uncompressedLength) {
 661   IOBuf clone;
 662   if (data->isChained()) {
 663     // LZ4 doesn't support streaming, so we have to coalesce
 664     clone = data->cloneCoalescedAsValue();
 665     data = &clone;
 666   }
 667
 668   folly::io::Cursor cursor(data);
 669   uint64_t actualUncompressedLength;
 670   if (encodeSize()) {
 671     actualUncompressedLength = decodeVarintFromCursor(cursor);
 672     if (uncompressedLength && *uncompressedLength != actualUncompressedLength) {
 673       throw std::runtime_error("LZ4Codec: invalid uncompressed length");
 674     }
 675   } else {
 676     // Invariants
 677     DCHECK(uncompressedLength.hasValue());
 678     DCHECK(*uncompressedLength <= maxUncompressedLength());
 679     actualUncompressedLength = *uncompressedLength;
 680   }
 681
 682   auto sp = StringPiece{cursor.peekBytes()};
 683   auto out = IOBuf::create(actualUncompressedLength);
 684   int n = LZ4_decompress_safe(
 685       sp.data(),
 686       reinterpret_cast<char*>(out->writableTail()),
 687       sp.size(),
 688       actualUncompressedLength);
 689
 690   if (n < 0 || uint64_t(n) != actualUncompressedLength) {
 691     throw std::runtime_error(to<std::string>(
 692         "LZ4 decompression returned invalid value ", n));
 693   }
 694   out->append(actualUncompressedLength);
 695   return out;
 696 }
 697
 698 #if LZ4_VERSION_NUMBER >= 10301
 699
 700 class LZ4FrameCodec final : public Codec {
 701  public:
 702   static std::unique_ptr<Codec> create(int level, CodecType type);
 703   explicit LZ4FrameCodec(int level, CodecType type);
 704   ~LZ4FrameCodec() override;
 705
 706   std::vector<std::string> validPrefixes() const override;
 707   bool canUncompress(const IOBuf* data, Optional<uint64_t> uncompressedLength)
 708       const override;
 709
 710  private:
 711   uint64_t doMaxCompressedLength(uint64_t uncompressedLength) const override;
 712
 713   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
 714   std::unique_ptr<IOBuf> doUncompress(
 715       const IOBuf* data,
 716       Optional<uint64_t> uncompressedLength) override;
 717
 718   // Reset the dctx_ if it is dirty or null.
 719   void resetDCtx();
 720
 721   int level_;
 722   LZ4F_decompressionContext_t dctx_{nullptr};
 723   bool dirty_{false};
 724 };
 725
 726 /* static */ std::unique_ptr<Codec> LZ4FrameCodec::create(
 727     int level,
 728     CodecType type) {
 729   return std::make_unique<LZ4FrameCodec>(level, type);
 730 }
 731
 732 static constexpr uint32_t kLZ4FrameMagicLE = 0x184D2204;
 733
 734 std::vector<std::string> LZ4FrameCodec::validPrefixes() const {
 735   return {prefixToStringLE(kLZ4FrameMagicLE)};
 736 }
 737
 738 bool LZ4FrameCodec::canUncompress(const IOBuf* data, Optional<uint64_t>) const {
 739   return dataStartsWithLE(data, kLZ4FrameMagicLE);
 740 }
 741
 742 uint64_t LZ4FrameCodec::doMaxCompressedLength(
 743     uint64_t uncompressedLength) const {
 744   LZ4F_preferences_t prefs{};
 745   prefs.compressionLevel = level_;
 746   prefs.frameInfo.contentSize = uncompressedLength;
 747   return LZ4F_compressFrameBound(uncompressedLength, &prefs);
 748 }
 749
 750 static size_t lz4FrameThrowOnError(size_t code) {
 751   if (LZ4F_isError(code)) {
 752     throw std::runtime_error(
 753         to<std::string>("LZ4Frame error: ", LZ4F_getErrorName(code)));
 754   }
 755   return code;
 756 }
 757
 758 void LZ4FrameCodec::resetDCtx() {
 759   if (dctx_ && !dirty_) {
 760     return;
 761   }
 762   if (dctx_) {
 763     LZ4F_freeDecompressionContext(dctx_);
 764   }
 765   lz4FrameThrowOnError(LZ4F_createDecompressionContext(&dctx_, 100));
 766   dirty_ = false;
 767 }
 768
 769 LZ4FrameCodec::LZ4FrameCodec(int level, CodecType type) : Codec(type) {
 770   DCHECK(type == CodecType::LZ4_FRAME);
 771   switch (level) {
 772     case COMPRESSION_LEVEL_FASTEST:
 773     case COMPRESSION_LEVEL_DEFAULT:
 774       level_ = 0;
 775       break;
 776     case COMPRESSION_LEVEL_BEST:
 777       level_ = 16;
 778       break;
 779     default:
 780       level_ = level;
 781       break;
 782   }
 783 }
 784
 785 LZ4FrameCodec::~LZ4FrameCodec() {
 786   if (dctx_) {
 787     LZ4F_freeDecompressionContext(dctx_);
 788   }
 789 }
 790
 791 std::unique_ptr<IOBuf> LZ4FrameCodec::doCompress(const IOBuf* data) {
 792   // LZ4 Frame compression doesn't support streaming so we have to coalesce
 793   IOBuf clone;
 794   if (data->isChained()) {
 795     clone = data->cloneCoalescedAsValue();
 796     data = &clone;
 797   }
 798   // Set preferences
 799   const auto uncompressedLength = data->length();
 800   LZ4F_preferences_t prefs{};
 801   prefs.compressionLevel = level_;
 802   prefs.frameInfo.contentSize = uncompressedLength;
 803   // Compress
 804   auto buf = IOBuf::create(maxCompressedLength(uncompressedLength));
 805   const size_t written = lz4FrameThrowOnError(LZ4F_compressFrame(
 806       buf->writableTail(),
 807       buf->tailroom(),
 808       data->data(),
 809       data->length(),
 810       &prefs));
 811   buf->append(written);
 812   return buf;
 813 }
 814
 815 std::unique_ptr<IOBuf> LZ4FrameCodec::doUncompress(
 816     const IOBuf* data,
 817     Optional<uint64_t> uncompressedLength) {
 818   // Reset the dctx if any errors have occurred
 819   resetDCtx();
 820   // Coalesce the data
 821   ByteRange in = *data->begin();
 822   IOBuf clone;
 823   if (data->isChained()) {
 824     clone = data->cloneCoalescedAsValue();
 825     in = clone.coalesce();
 826   }
 827   data = nullptr;
 828   // Select decompression options
 829   LZ4F_decompressOptions_t options;
 830   options.stableDst = 1;
 831   // Select blockSize and growthSize for the IOBufQueue
 832   IOBufQueue queue(IOBufQueue::cacheChainLength());
 833   auto blockSize = uint64_t{64} << 10;
 834   auto growthSize = uint64_t{4} << 20;
 835   if (uncompressedLength) {
 836     // Allocate uncompressedLength in one chunk (up to 64 MB)
 837     const auto allocateSize = std::min(*uncompressedLength, uint64_t{64} << 20);
 838     queue.preallocate(allocateSize, allocateSize);
 839     blockSize = std::min(*uncompressedLength, blockSize);
 840     growthSize = std::min(*uncompressedLength, growthSize);
 841   } else {
 842     // Reduce growthSize for small data
 843     const auto guessUncompressedLen =
 844         4 * std::max<uint64_t>(blockSize, in.size());
 845     growthSize = std::min(guessUncompressedLen, growthSize);
 846   }
 847   // Once LZ4_decompress() is called, the dctx_ cannot be reused until it
 848   // returns 0
 849   dirty_ = true;
 850   // Decompress until the frame is over
 851   size_t code = 0;
 852   do {
 853     // Allocate enough space to decompress at least a block
 854     void* out;
 855     size_t outSize;
 856     std::tie(out, outSize) = queue.preallocate(blockSize, growthSize);
 857     // Decompress
 858     size_t inSize = in.size();
 859     code = lz4FrameThrowOnError(
 860         LZ4F_decompress(dctx_, out, &outSize, in.data(), &inSize, &options));
 861     if (in.empty() && outSize == 0 && code != 0) {
 862       // We passed no input, no output was produced, and the frame isn't over
 863       // No more forward progress is possible
 864       throw std::runtime_error("LZ4Frame error: Incomplete frame");
 865     }
 866     in.uncheckedAdvance(inSize);
 867     queue.postallocate(outSize);
 868   } while (code != 0);
 869   // At this point the decompression context can be reused
 870   dirty_ = false;
 871   if (uncompressedLength && queue.chainLength() != *uncompressedLength) {
 872     throw std::runtime_error("LZ4Frame error: Invalid uncompressedLength");
 873   }
 874   return queue.move();
 875 }
 876
 877 #endif // LZ4_VERSION_NUMBER >= 10301
 878 #endif // FOLLY_HAVE_LIBLZ4
 879
 880 #if FOLLY_HAVE_LIBSNAPPY
 881
 882 /**
 883  * Snappy compression
 884  */
 885
 886 /**
 887  * Implementation of snappy::Source that reads from a IOBuf chain.
 888  */
 889 class IOBufSnappySource final : public snappy::Source {
 890  public:
 891   explicit IOBufSnappySource(const IOBuf* data);
 892   size_t Available() const override;
 893   const char* Peek(size_t* len) override;
 894   void Skip(size_t n) override;
 895  private:
 896   size_t available_;
 897   io::Cursor cursor_;
 898 };
 899
 900 IOBufSnappySource::IOBufSnappySource(const IOBuf* data)
 901   : available_(data->computeChainDataLength()),
 902     cursor_(data) {
 903 }
 904
 905 size_t IOBufSnappySource::Available() const {
 906   return available_;
 907 }
 908
 909 const char* IOBufSnappySource::Peek(size_t* len) {
 910   auto sp = StringPiece{cursor_.peekBytes()};
 911   *len = sp.size();
 912   return sp.data();
 913 }
 914
 915 void IOBufSnappySource::Skip(size_t n) {
 916   CHECK_LE(n, available_);
 917   cursor_.skip(n);
 918   available_ -= n;
 919 }
 920
 921 class SnappyCodec final : public Codec {
 922  public:
 923   static std::unique_ptr<Codec> create(int level, CodecType type);
 924   explicit SnappyCodec(int level, CodecType type);
 925
 926  private:
 927   uint64_t doMaxUncompressedLength() const override;
 928   uint64_t doMaxCompressedLength(uint64_t uncompressedLength) const override;
 929   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
 930   std::unique_ptr<IOBuf> doUncompress(
 931       const IOBuf* data,
 932       Optional<uint64_t> uncompressedLength) override;
 933 };
 934
 935 std::unique_ptr<Codec> SnappyCodec::create(int level, CodecType type) {
 936   return std::make_unique<SnappyCodec>(level, type);
 937 }
 938
 939 SnappyCodec::SnappyCodec(int level, CodecType type) : Codec(type) {
 940   DCHECK(type == CodecType::SNAPPY);
 941   switch (level) {
 942   case COMPRESSION_LEVEL_FASTEST:
 943   case COMPRESSION_LEVEL_DEFAULT:
 944   case COMPRESSION_LEVEL_BEST:
 945     level = 1;
 946   }
 947   if (level != 1) {
 948     throw std::invalid_argument(to<std::string>(
 949         "SnappyCodec: invalid level: ", level));
 950   }
 951 }
 952
 953 uint64_t SnappyCodec::doMaxUncompressedLength() const {
 954   // snappy.h uses uint32_t for lengths, so there's that.
 955   return std::numeric_limits<uint32_t>::max();
 956 }
 957
 958 uint64_t SnappyCodec::doMaxCompressedLength(uint64_t uncompressedLength) const {
 959   return snappy::MaxCompressedLength(uncompressedLength);
 960 }
 961
 962 std::unique_ptr<IOBuf> SnappyCodec::doCompress(const IOBuf* data) {
 963   IOBufSnappySource source(data);
 964   auto out = IOBuf::create(maxCompressedLength(source.Available()));
 965
 966   snappy::UncheckedByteArraySink sink(reinterpret_cast<char*>(
 967       out->writableTail()));
 968
 969   size_t n = snappy::Compress(&source, &sink);
 970
 971   CHECK_LE(n, out->capacity());
 972   out->append(n);
 973   return out;
 974 }
 975
 976 std::unique_ptr<IOBuf> SnappyCodec::doUncompress(
 977     const IOBuf* data,
 978     Optional<uint64_t> uncompressedLength) {
 979   uint32_t actualUncompressedLength = 0;
 980
 981   {
 982     IOBufSnappySource source(data);
 983     if (!snappy::GetUncompressedLength(&source, &actualUncompressedLength)) {
 984       throw std::runtime_error("snappy::GetUncompressedLength failed");
 985     }
 986     if (uncompressedLength && *uncompressedLength != actualUncompressedLength) {
 987       throw std::runtime_error("snappy: invalid uncompressed length");
 988     }
 989   }
 990
 991   auto out = IOBuf::create(actualUncompressedLength);
 992
 993   {
 994     IOBufSnappySource source(data);
 995     if (!snappy::RawUncompress(&source,
 996                                reinterpret_cast<char*>(out->writableTail()))) {
 997       throw std::runtime_error("snappy::RawUncompress failed");
 998     }
 999   }
1000
1001   out->append(actualUncompressedLength);
1002   return out;
1003 }
1004
1005 #endif  // FOLLY_HAVE_LIBSNAPPY
1006
1007 #if FOLLY_HAVE_LIBZ
1008 /**
1009  * Zlib codec
1010  */
1011 class ZlibStreamCodec final : public StreamCodec {
1012  public:
1013   static std::unique_ptr<Codec> createCodec(int level, CodecType type);
1014   static std::unique_ptr<StreamCodec> createStream(int level, CodecType type);
1015   explicit ZlibStreamCodec(int level, CodecType type);
1016   ~ZlibStreamCodec() override;
1017
1018   std::vector<std::string> validPrefixes() const override;
1019   bool canUncompress(const IOBuf* data, Optional<uint64_t> uncompressedLength)
1020       const override;
1021
1022  private:
1023   uint64_t doMaxCompressedLength(uint64_t uncompressedLength) const override;
1024
1025   void doResetStream() override;
1026   bool doCompressStream(
1027       ByteRange& input,
1028       MutableByteRange& output,
1029       StreamCodec::FlushOp flush) override;
1030   bool doUncompressStream(
1031       ByteRange& input,
1032       MutableByteRange& output,
1033       StreamCodec::FlushOp flush) override;
1034
1035   void resetDeflateStream();
1036   void resetInflateStream();
1037
1038   Optional<z_stream> deflateStream_{};
1039   Optional<z_stream> inflateStream_{};
1040   int level_;
1041   bool needReset_{true};
1042 };
1043
1044 static constexpr uint16_t kGZIPMagicLE = 0x8B1F;
1045
1046 std::vector<std::string> ZlibStreamCodec::validPrefixes() const {
1047   if (type() == CodecType::ZLIB) {
1048     // Zlib streams start with a 2 byte header.
1049     //
1050     //   0   1
1051     // +---+---+
1052     // |CMF|FLG|
1053     // +---+---+
1054     //
1055     // We won't restrict the values of any sub-fields except as described below.
1056     //
1057     // The lowest 4 bits of CMF is the compression method (CM).
1058     // CM == 0x8 is the deflate compression method, which is currently the only
1059     // supported compression method, so any valid prefix must have CM == 0x8.
1060     //
1061     // The lowest 5 bits of FLG is FCHECK.
1062     // FCHECK must be such that the two header bytes are a multiple of 31 when
1063     // interpreted as a big endian 16-bit number.
1064     std::vector<std::string> result;
1065     // 16 values for the first byte, 8 values for the second byte.
1066     // There are also 4 combinations where both 0x00 and 0x1F work as FCHECK.
1067     result.reserve(132);
1068     // Select all values for the CMF byte that use the deflate algorithm 0x8.
1069     for (uint32_t first = 0x0800; first <= 0xF800; first += 0x1000) {
1070       // Select all values for the FLG, but leave FCHECK as 0 since it's fixed.
1071       for (uint32_t second = 0x00; second <= 0xE0; second += 0x20) {
1072         uint16_t prefix = first | second;
1073         // Compute FCHECK.
1074         prefix += 31 - (prefix % 31);
1075         result.push_back(prefixToStringLE(Endian::big(prefix)));
1076         // zlib won't produce this, but it is a valid prefix.
1077         if ((prefix & 0x1F) == 31) {
1078           prefix -= 31;
1079           result.push_back(prefixToStringLE(Endian::big(prefix)));
1080         }
1081       }
1082     }
1083     return result;
1084   } else {
1085     // The gzip frame starts with 2 magic bytes.
1086     return {prefixToStringLE(kGZIPMagicLE)};
1087   }
1088 }
1089
1090 bool ZlibStreamCodec::canUncompress(const IOBuf* data, Optional<uint64_t>)
1091     const {
1092   if (type() == CodecType::ZLIB) {
1093     uint16_t value;
1094     Cursor cursor{data};
1095     if (!cursor.tryReadBE(value)) {
1096       return false;
1097     }
1098     // zlib compressed if using deflate and is a multiple of 31.
1099     return (value & 0x0F00) == 0x0800 && value % 31 == 0;
1100   } else {
1101     return dataStartsWithLE(data, kGZIPMagicLE);
1102   }
1103 }
1104
1105 uint64_t ZlibStreamCodec::doMaxCompressedLength(
1106     uint64_t uncompressedLength) const {
1107   return deflateBound(nullptr, uncompressedLength);
1108 }
1109
1110 std::unique_ptr<Codec> ZlibStreamCodec::createCodec(int level, CodecType type) {
1111   return std::make_unique<ZlibStreamCodec>(level, type);
1112 }
1113
1114 std::unique_ptr<StreamCodec> ZlibStreamCodec::createStream(
1115     int level,
1116     CodecType type) {
1117   return std::make_unique<ZlibStreamCodec>(level, type);
1118 }
1119
1120 ZlibStreamCodec::ZlibStreamCodec(int level, CodecType type)
1121     : StreamCodec(type) {
1122   DCHECK(type == CodecType::ZLIB || type == CodecType::GZIP);
1123   switch (level) {
1124     case COMPRESSION_LEVEL_FASTEST:
1125       level = 1;
1126       break;
1127     case COMPRESSION_LEVEL_DEFAULT:
1128       level = Z_DEFAULT_COMPRESSION;
1129       break;
1130     case COMPRESSION_LEVEL_BEST:
1131       level = 9;
1132       break;
1133   }
1134   if (level != Z_DEFAULT_COMPRESSION && (level < 0 || level > 9)) {
1135     throw std::invalid_argument(
1136         to<std::string>("ZlibStreamCodec: invalid level: ", level));
1137   }
1138   level_ = level;
1139 }
1140
1141 ZlibStreamCodec::~ZlibStreamCodec() {
1142   if (deflateStream_) {
1143     deflateEnd(deflateStream_.get_pointer());
1144     deflateStream_.clear();
1145   }
1146   if (inflateStream_) {
1147     inflateEnd(inflateStream_.get_pointer());
1148     inflateStream_.clear();
1149   }
1150 }
1151
1152 void ZlibStreamCodec::doResetStream() {
1153   needReset_ = true;
1154 }
1155
1156 void ZlibStreamCodec::resetDeflateStream() {
1157   if (deflateStream_) {
1158     int const rc = deflateReset(deflateStream_.get_pointer());
1159     if (rc != Z_OK) {
1160       deflateStream_.clear();
1161       throw std::runtime_error(
1162           to<std::string>("ZlibStreamCodec: deflateReset error: ", rc));
1163     }
1164     return;
1165   }
1166   deflateStream_ = z_stream{};
1167   // Using deflateInit2() to support gzip.  "The windowBits parameter is the
1168   // base two logarithm of the maximum window size (...) The default value is
1169   // 15 (...) Add 16 to windowBits to write a simple gzip header and trailer
1170   // around the compressed data instead of a zlib wrapper. The gzip header
1171   // will have no file name, no extra data, no comment, no modification time
1172   // (set to zero), no header crc, and the operating system will be set to 255
1173   // (unknown)."
1174   int const windowBits = 15 + (type() == CodecType::GZIP ? 16 : 0);
1175   // All other parameters (method, memLevel, strategy) get default values from
1176   // the zlib manual.
1177   int const rc = deflateInit2(
1178       deflateStream_.get_pointer(),
1179       level_,
1180       Z_DEFLATED,
1181       windowBits,
1182       /* memLevel */ 8,
1183       Z_DEFAULT_STRATEGY);
1184   if (rc != Z_OK) {
1185     deflateStream_.clear();
1186     throw std::runtime_error(
1187         to<std::string>("ZlibStreamCodec: deflateInit error: ", rc));
1188   }
1189 }
1190
1191 void ZlibStreamCodec::resetInflateStream() {
1192   if (inflateStream_) {
1193     int const rc = inflateReset(inflateStream_.get_pointer());
1194     if (rc != Z_OK) {
1195       inflateStream_.clear();
1196       throw std::runtime_error(
1197           to<std::string>("ZlibStreamCodec: inflateReset error: ", rc));
1198     }
1199     return;
1200   }
1201   inflateStream_ = z_stream{};
1202   // "The windowBits parameter is the base two logarithm of the maximum window
1203   // size (...) The default value is 15 (...) add 16 to decode only the gzip
1204   // format (the zlib format will return a Z_DATA_ERROR)."
1205   int const windowBits = 15 + (type() == CodecType::GZIP ? 16 : 0);
1206   int const rc = inflateInit2(inflateStream_.get_pointer(), windowBits);
1207   if (rc != Z_OK) {
1208     inflateStream_.clear();
1209     throw std::runtime_error(
1210         to<std::string>("ZlibStreamCodec: inflateInit error: ", rc));
1211   }
1212 }
1213
1214 static int zlibTranslateFlush(StreamCodec::FlushOp flush) {
1215   switch (flush) {
1216     case StreamCodec::FlushOp::NONE:
1217       return Z_NO_FLUSH;
1218     case StreamCodec::FlushOp::FLUSH:
1219       return Z_SYNC_FLUSH;
1220     case StreamCodec::FlushOp::END:
1221       return Z_FINISH;
1222     default:
1223       throw std::invalid_argument("ZlibStreamCodec: Invalid flush");
1224   }
1225 }
1226
1227 static int zlibThrowOnError(int rc) {
1228   switch (rc) {
1229     case Z_OK:
1230     case Z_BUF_ERROR:
1231     case Z_STREAM_END:
1232       return rc;
1233     default:
1234       throw std::runtime_error(to<std::string>("ZlibStreamCodec: error: ", rc));
1235   }
1236 }
1237
1238 bool ZlibStreamCodec::doCompressStream(
1239     ByteRange& input,
1240     MutableByteRange& output,
1241     StreamCodec::FlushOp flush) {
1242   if (needReset_) {
1243     resetDeflateStream();
1244     needReset_ = false;
1245   }
1246   DCHECK(deflateStream_.hasValue());
1247   // zlib will return Z_STREAM_ERROR if output.data() is null.
1248   if (output.data() == nullptr) {
1249     return false;
1250   }
1251   deflateStream_->next_in = const_cast<uint8_t*>(input.data());
1252   deflateStream_->avail_in = input.size();
1253   deflateStream_->next_out = output.data();
1254   deflateStream_->avail_out = output.size();
1255   SCOPE_EXIT {
1256     input.uncheckedAdvance(input.size() - deflateStream_->avail_in);
1257     output.uncheckedAdvance(output.size() - deflateStream_->avail_out);
1258   };
1259   int const rc = zlibThrowOnError(
1260       deflate(deflateStream_.get_pointer(), zlibTranslateFlush(flush)));
1261   switch (flush) {
1262     case StreamCodec::FlushOp::NONE:
1263       return false;
1264     case StreamCodec::FlushOp::FLUSH:
1265       return deflateStream_->avail_in == 0 && deflateStream_->avail_out != 0;
1266     case StreamCodec::FlushOp::END:
1267       return rc == Z_STREAM_END;
1268     default:
1269       throw std::invalid_argument("ZlibStreamCodec: Invalid flush");
1270   }
1271 }
1272
1273 bool ZlibStreamCodec::doUncompressStream(
1274     ByteRange& input,
1275     MutableByteRange& output,
1276     StreamCodec::FlushOp flush) {
1277   if (needReset_) {
1278     resetInflateStream();
1279     needReset_ = false;
1280   }
1281   DCHECK(inflateStream_.hasValue());
1282   // zlib will return Z_STREAM_ERROR if output.data() is null.
1283   if (output.data() == nullptr) {
1284     return false;
1285   }
1286   inflateStream_->next_in = const_cast<uint8_t*>(input.data());
1287   inflateStream_->avail_in = input.size();
1288   inflateStream_->next_out = output.data();
1289   inflateStream_->avail_out = output.size();
1290   SCOPE_EXIT {
1291     input.advance(input.size() - inflateStream_->avail_in);
1292     output.advance(output.size() - inflateStream_->avail_out);
1293   };
1294   int const rc = zlibThrowOnError(
1295       inflate(inflateStream_.get_pointer(), zlibTranslateFlush(flush)));
1296   return rc == Z_STREAM_END;
1297 }
1298
1299 #endif // FOLLY_HAVE_LIBZ
1300
1301 #if FOLLY_HAVE_LIBLZMA
1302
1303 /**
1304  * LZMA2 compression
1305  */
1306 class LZMA2Codec final : public Codec {
1307  public:
1308   static std::unique_ptr<Codec> create(int level, CodecType type);
1309   explicit LZMA2Codec(int level, CodecType type);
1310
1311   std::vector<std::string> validPrefixes() const override;
1312   bool canUncompress(const IOBuf* data, Optional<uint64_t> uncompressedLength)
1313       const override;
1314
1315  private:
1316   bool doNeedsUncompressedLength() const override;
1317   uint64_t doMaxUncompressedLength() const override;
1318   uint64_t doMaxCompressedLength(uint64_t uncompressedLength) const override;
1319
1320   bool encodeSize() const { return type() == CodecType::LZMA2_VARINT_SIZE; }
1321
1322   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
1323   std::unique_ptr<IOBuf> doUncompress(
1324       const IOBuf* data,
1325       Optional<uint64_t> uncompressedLength) override;
1326
1327   std::unique_ptr<IOBuf> addOutputBuffer(lzma_stream* stream, size_t length);
1328   bool doInflate(lzma_stream* stream, IOBuf* head, size_t bufferLength);
1329
1330   int level_;
1331 };
1332
1333 static constexpr uint64_t kLZMA2MagicLE = 0x005A587A37FD;
1334 static constexpr unsigned kLZMA2MagicBytes = 6;
1335
1336 std::vector<std::string> LZMA2Codec::validPrefixes() const {
1337   if (type() == CodecType::LZMA2_VARINT_SIZE) {
1338     return {};
1339   }
1340   return {prefixToStringLE(kLZMA2MagicLE, kLZMA2MagicBytes)};
1341 }
1342
1343 bool LZMA2Codec::canUncompress(const IOBuf* data, Optional<uint64_t>) const {
1344   if (type() == CodecType::LZMA2_VARINT_SIZE) {
1345     return false;
1346   }
1347   // Returns false for all inputs less than 8 bytes.
1348   // This is okay, because no valid LZMA2 streams are less than 8 bytes.
1349   return dataStartsWithLE(data, kLZMA2MagicLE, kLZMA2MagicBytes);
1350 }
1351
1352 std::unique_ptr<Codec> LZMA2Codec::create(int level, CodecType type) {
1353   return std::make_unique<LZMA2Codec>(level, type);
1354 }
1355
1356 LZMA2Codec::LZMA2Codec(int level, CodecType type) : Codec(type) {
1357   DCHECK(type == CodecType::LZMA2 || type == CodecType::LZMA2_VARINT_SIZE);
1358   switch (level) {
1359   case COMPRESSION_LEVEL_FASTEST:
1360     level = 0;
1361     break;
1362   case COMPRESSION_LEVEL_DEFAULT:
1363     level = LZMA_PRESET_DEFAULT;
1364     break;
1365   case COMPRESSION_LEVEL_BEST:
1366     level = 9;
1367     break;
1368   }
1369   if (level < 0 || level > 9) {
1370     throw std::invalid_argument(to<std::string>(
1371         "LZMA2Codec: invalid level: ", level));
1372   }
1373   level_ = level;
1374 }
1375
1376 bool LZMA2Codec::doNeedsUncompressedLength() const {
1377   return false;
1378 }
1379
1380 uint64_t LZMA2Codec::doMaxUncompressedLength() const {
1381   // From lzma/base.h: "Stream is roughly 8 EiB (2^63 bytes)"
1382   return uint64_t(1) << 63;
1383 }
1384
1385 uint64_t LZMA2Codec::doMaxCompressedLength(uint64_t uncompressedLength) const {
1386   return lzma_stream_buffer_bound(uncompressedLength) +
1387       (encodeSize() ? kMaxVarintLength64 : 0);
1388 }
1389
1390 std::unique_ptr<IOBuf> LZMA2Codec::addOutputBuffer(
1391     lzma_stream* stream,
1392     size_t length) {
1393
1394   CHECK_EQ(stream->avail_out, 0);
1395
1396   auto buf = IOBuf::create(length);
1397   buf->append(buf->capacity());
1398
1399   stream->next_out = buf->writableData();
1400   stream->avail_out = buf->length();
1401
1402   return buf;
1403 }
1404
1405 std::unique_ptr<IOBuf> LZMA2Codec::doCompress(const IOBuf* data) {
1406   lzma_ret rc;
1407   lzma_stream stream = LZMA_STREAM_INIT;
1408
1409   rc = lzma_easy_encoder(&stream, level_, LZMA_CHECK_NONE);
1410   if (rc != LZMA_OK) {
1411     throw std::runtime_error(folly::to<std::string>(
1412       "LZMA2Codec: lzma_easy_encoder error: ", rc));
1413   }
1414
1415   SCOPE_EXIT { lzma_end(&stream); };
1416
1417   uint64_t uncompressedLength = data->computeChainDataLength();
1418   uint64_t maxCompressedLength = lzma_stream_buffer_bound(uncompressedLength);
1419
1420   // Max 64MiB in one go
1421   constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20;    // 64MiB
1422   constexpr uint32_t defaultBufferLength = uint32_t(4) << 20;     // 4MiB
1423
1424   auto out = addOutputBuffer(
1425     &stream,
1426     (maxCompressedLength <= maxSingleStepLength ?
1427      maxCompressedLength :
1428      defaultBufferLength));
1429
1430   if (encodeSize()) {
1431     auto size = IOBuf::createCombined(kMaxVarintLength64);
1432     encodeVarintToIOBuf(uncompressedLength, size.get());
1433     size->appendChain(std::move(out));
1434     out = std::move(size);
1435   }
1436
1437   for (auto& range : *data) {
1438     if (range.empty()) {
1439       continue;
1440     }
1441
1442     stream.next_in = const_cast<uint8_t*>(range.data());
1443     stream.avail_in = range.size();
1444
1445     while (stream.avail_in != 0) {
1446       if (stream.avail_out == 0) {
1447         out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
1448       }
1449
1450       rc = lzma_code(&stream, LZMA_RUN);
1451
1452       if (rc != LZMA_OK) {
1453         throw std::runtime_error(folly::to<std::string>(
1454           "LZMA2Codec: lzma_code error: ", rc));
1455       }
1456     }
1457   }
1458
1459   do {
1460     if (stream.avail_out == 0) {
1461       out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
1462     }
1463
1464     rc = lzma_code(&stream, LZMA_FINISH);
1465   } while (rc == LZMA_OK);
1466
1467   if (rc != LZMA_STREAM_END) {
1468     throw std::runtime_error(folly::to<std::string>(
1469       "LZMA2Codec: lzma_code ended with error: ", rc));
1470   }
1471
1472   out->prev()->trimEnd(stream.avail_out);
1473
1474   return out;
1475 }
1476
1477 bool LZMA2Codec::doInflate(lzma_stream* stream,
1478                           IOBuf* head,
1479                           size_t bufferLength) {
1480   if (stream->avail_out == 0) {
1481     head->prependChain(addOutputBuffer(stream, bufferLength));
1482   }
1483
1484   lzma_ret rc = lzma_code(stream, LZMA_RUN);
1485
1486   switch (rc) {
1487   case LZMA_OK:
1488     break;
1489   case LZMA_STREAM_END:
1490     return true;
1491   default:
1492     throw std::runtime_error(to<std::string>(
1493         "LZMA2Codec: lzma_code error: ", rc));
1494   }
1495
1496   return false;
1497 }
1498
1499 std::unique_ptr<IOBuf> LZMA2Codec::doUncompress(
1500     const IOBuf* data,
1501     Optional<uint64_t> uncompressedLength) {
1502   lzma_ret rc;
1503   lzma_stream stream = LZMA_STREAM_INIT;
1504
1505   rc = lzma_auto_decoder(&stream, std::numeric_limits<uint64_t>::max(), 0);
1506   if (rc != LZMA_OK) {
1507     throw std::runtime_error(folly::to<std::string>(
1508       "LZMA2Codec: lzma_auto_decoder error: ", rc));
1509   }
1510
1511   SCOPE_EXIT { lzma_end(&stream); };
1512
1513   // Max 64MiB in one go
1514   constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20; // 64MiB
1515   constexpr uint32_t defaultBufferLength = uint32_t(256) << 10; // 256 KiB
1516
1517   folly::io::Cursor cursor(data);
1518   if (encodeSize()) {
1519     const uint64_t actualUncompressedLength = decodeVarintFromCursor(cursor);
1520     if (uncompressedLength && *uncompressedLength != actualUncompressedLength) {
1521       throw std::runtime_error("LZMA2Codec: invalid uncompressed length");
1522     }
1523     uncompressedLength = actualUncompressedLength;
1524   }
1525
1526   auto out = addOutputBuffer(
1527       &stream,
1528       ((uncompressedLength && *uncompressedLength <= maxSingleStepLength)
1529            ? *uncompressedLength
1530            : defaultBufferLength));
1531
1532   bool streamEnd = false;
1533   auto buf = cursor.peekBytes();
1534   while (!buf.empty()) {
1535     stream.next_in = const_cast<uint8_t*>(buf.data());
1536     stream.avail_in = buf.size();
1537
1538     while (stream.avail_in != 0) {
1539       if (streamEnd) {
1540         throw std::runtime_error(to<std::string>(
1541             "LZMA2Codec: junk after end of data"));
1542       }
1543
1544       streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
1545     }
1546
1547     cursor.skip(buf.size());
1548     buf = cursor.peekBytes();
1549   }
1550
1551   while (!streamEnd) {
1552     streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
1553   }
1554
1555   out->prev()->trimEnd(stream.avail_out);
1556
1557   if (uncompressedLength && *uncompressedLength != stream.total_out) {
1558     throw std::runtime_error(
1559         to<std::string>("LZMA2Codec: invalid uncompressed length"));
1560   }
1561
1562   return out;
1563 }
1564
1565 #endif  // FOLLY_HAVE_LIBLZMA
1566
1567 #ifdef FOLLY_HAVE_LIBZSTD
1568
1569 namespace {
1570 void zstdFreeCStream(ZSTD_CStream* zcs) {
1571   ZSTD_freeCStream(zcs);
1572 }
1573
1574 void zstdFreeDStream(ZSTD_DStream* zds) {
1575   ZSTD_freeDStream(zds);
1576 }
1577 }
1578
1579 /**
1580  * ZSTD compression
1581  */
1582 class ZSTDStreamCodec final : public StreamCodec {
1583  public:
1584   static std::unique_ptr<Codec> createCodec(int level, CodecType);
1585   static std::unique_ptr<StreamCodec> createStream(int level, CodecType);
1586   explicit ZSTDStreamCodec(int level, CodecType type);
1587
1588   std::vector<std::string> validPrefixes() const override;
1589   bool canUncompress(const IOBuf* data, Optional<uint64_t> uncompressedLength)
1590       const override;
1591
1592  private:
1593   bool doNeedsUncompressedLength() const override;
1594   uint64_t doMaxCompressedLength(uint64_t uncompressedLength) const override;
1595   Optional<uint64_t> doGetUncompressedLength(
1596       IOBuf const* data,
1597       Optional<uint64_t> uncompressedLength) const override;
1598
1599   void doResetStream() override;
1600   bool doCompressStream(
1601       ByteRange& input,
1602       MutableByteRange& output,
1603       StreamCodec::FlushOp flushOp) override;
1604   bool doUncompressStream(
1605       ByteRange& input,
1606       MutableByteRange& output,
1607       StreamCodec::FlushOp flushOp) override;
1608
1609   void resetCStream();
1610   void resetDStream();
1611
1612   bool tryBlockCompress(ByteRange& input, MutableByteRange& output) const;
1613   bool tryBlockUncompress(ByteRange& input, MutableByteRange& output) const;
1614
1615   int level_;
1616   bool needReset_{true};
1617   std::unique_ptr<
1618       ZSTD_CStream,
1619       folly::static_function_deleter<ZSTD_CStream, &zstdFreeCStream>>
1620       cstream_{nullptr};
1621   std::unique_ptr<
1622       ZSTD_DStream,
1623       folly::static_function_deleter<ZSTD_DStream, &zstdFreeDStream>>
1624       dstream_{nullptr};
1625 };
1626
1627 static constexpr uint32_t kZSTDMagicLE = 0xFD2FB528;
1628
1629 std::vector<std::string> ZSTDStreamCodec::validPrefixes() const {
1630   return {prefixToStringLE(kZSTDMagicLE)};
1631 }
1632
1633 bool ZSTDStreamCodec::canUncompress(const IOBuf* data, Optional<uint64_t>)
1634     const {
1635   return dataStartsWithLE(data, kZSTDMagicLE);
1636 }
1637
1638 std::unique_ptr<Codec> ZSTDStreamCodec::createCodec(int level, CodecType type) {
1639   return make_unique<ZSTDStreamCodec>(level, type);
1640 }
1641
1642 std::unique_ptr<StreamCodec> ZSTDStreamCodec::createStream(
1643     int level,
1644     CodecType type) {
1645   return make_unique<ZSTDStreamCodec>(level, type);
1646 }
1647
1648 ZSTDStreamCodec::ZSTDStreamCodec(int level, CodecType type)
1649     : StreamCodec(type) {
1650   DCHECK(type == CodecType::ZSTD);
1651   switch (level) {
1652     case COMPRESSION_LEVEL_FASTEST:
1653       level = 1;
1654       break;
1655     case COMPRESSION_LEVEL_DEFAULT:
1656       level = 1;
1657       break;
1658     case COMPRESSION_LEVEL_BEST:
1659       level = 19;
1660       break;
1661   }
1662   if (level < 1 || level > ZSTD_maxCLevel()) {
1663     throw std::invalid_argument(
1664         to<std::string>("ZSTD: invalid level: ", level));
1665   }
1666   level_ = level;
1667 }
1668
1669 bool ZSTDStreamCodec::doNeedsUncompressedLength() const {
1670   return false;
1671 }
1672
1673 uint64_t ZSTDStreamCodec::doMaxCompressedLength(
1674     uint64_t uncompressedLength) const {
1675   return ZSTD_compressBound(uncompressedLength);
1676 }
1677
1678 void zstdThrowIfError(size_t rc) {
1679   if (!ZSTD_isError(rc)) {
1680     return;
1681   }
1682   throw std::runtime_error(
1683       to<std::string>("ZSTD returned an error: ", ZSTD_getErrorName(rc)));
1684 }
1685
1686 Optional<uint64_t> ZSTDStreamCodec::doGetUncompressedLength(
1687     IOBuf const* data,
1688     Optional<uint64_t> uncompressedLength) const {
1689   // Read decompressed size from frame if available in first IOBuf.
1690   auto const decompressedSize =
1691       ZSTD_getDecompressedSize(data->data(), data->length());
1692   if (decompressedSize != 0) {
1693     if (uncompressedLength && *uncompressedLength != decompressedSize) {
1694       throw std::runtime_error("ZSTD: invalid uncompressed length");
1695     }
1696     uncompressedLength = decompressedSize;
1697   }
1698   return uncompressedLength;
1699 }
1700
1701 void ZSTDStreamCodec::doResetStream() {
1702   needReset_ = true;
1703 }
1704
1705 bool ZSTDStreamCodec::tryBlockCompress(
1706     ByteRange& input,
1707     MutableByteRange& output) const {
1708   DCHECK(needReset_);
1709   // We need to know that we have enough output space to use block compression
1710   if (output.size() < ZSTD_compressBound(input.size())) {
1711     return false;
1712   }
1713   size_t const length = ZSTD_compress(
1714       output.data(), output.size(), input.data(), input.size(), level_);
1715   zstdThrowIfError(length);
1716   input.uncheckedAdvance(input.size());
1717   output.uncheckedAdvance(length);
1718   return true;
1719 }
1720
1721 void ZSTDStreamCodec::resetCStream() {
1722   if (!cstream_) {
1723     cstream_.reset(ZSTD_createCStream());
1724     if (!cstream_) {
1725       throw std::bad_alloc{};
1726     }
1727   }
1728   // Advanced API usage works for all supported versions of zstd.
1729   // Required to set contentSizeFlag.
1730   auto params = ZSTD_getParams(level_, uncompressedLength().value_or(0), 0);
1731   params.fParams.contentSizeFlag = uncompressedLength().hasValue();
1732   zstdThrowIfError(ZSTD_initCStream_advanced(
1733       cstream_.get(), nullptr, 0, params, uncompressedLength().value_or(0)));
1734 }
1735
1736 bool ZSTDStreamCodec::doCompressStream(
1737     ByteRange& input,
1738     MutableByteRange& output,
1739     StreamCodec::FlushOp flushOp) {
1740   if (needReset_) {
1741     // If we are given all the input in one chunk try to use block compression
1742     if (flushOp == StreamCodec::FlushOp::END &&
1743         tryBlockCompress(input, output)) {
1744       return true;
1745     }
1746     resetCStream();
1747     needReset_ = false;
1748   }
1749   ZSTD_inBuffer in = {input.data(), input.size(), 0};
1750   ZSTD_outBuffer out = {output.data(), output.size(), 0};
1751   SCOPE_EXIT {
1752     input.uncheckedAdvance(in.pos);
1753     output.uncheckedAdvance(out.pos);
1754   };
1755   if (flushOp == StreamCodec::FlushOp::NONE || !input.empty()) {
1756     zstdThrowIfError(ZSTD_compressStream(cstream_.get(), &out, &in));
1757   }
1758   if (in.pos == in.size && flushOp != StreamCodec::FlushOp::NONE) {
1759     size_t rc;
1760     switch (flushOp) {
1761       case StreamCodec::FlushOp::FLUSH:
1762         rc = ZSTD_flushStream(cstream_.get(), &out);
1763         break;
1764       case StreamCodec::FlushOp::END:
1765         rc = ZSTD_endStream(cstream_.get(), &out);
1766         break;
1767       default:
1768         throw std::invalid_argument("ZSTD: invalid FlushOp");
1769     }
1770     zstdThrowIfError(rc);
1771     if (rc == 0) {
1772       return true;
1773     }
1774   }
1775   return false;
1776 }
1777
1778 bool ZSTDStreamCodec::tryBlockUncompress(
1779     ByteRange& input,
1780     MutableByteRange& output) const {
1781   DCHECK(needReset_);
1782 #if ZSTD_VERSION_NUMBER < 10104
1783   // We require ZSTD_findFrameCompressedSize() to perform this optimization.
1784   return false;
1785 #else
1786   // We need to know the uncompressed length and have enough output space.
1787   if (!uncompressedLength() || output.size() < *uncompressedLength()) {
1788     return false;
1789   }
1790   size_t const compressedLength =
1791       ZSTD_findFrameCompressedSize(input.data(), input.size());
1792   zstdThrowIfError(compressedLength);
1793   size_t const length = ZSTD_decompress(
1794       output.data(), *uncompressedLength(), input.data(), compressedLength);
1795   zstdThrowIfError(length);
1796   if (length != *uncompressedLength()) {
1797     throw std::runtime_error("ZSTDStreamCodec: Incorrect uncompressed length");
1798   }
1799   input.uncheckedAdvance(compressedLength);
1800   output.uncheckedAdvance(length);
1801   return true;
1802 #endif
1803 }
1804
1805 void ZSTDStreamCodec::resetDStream() {
1806   if (!dstream_) {
1807     dstream_.reset(ZSTD_createDStream());
1808     if (!dstream_) {
1809       throw std::bad_alloc{};
1810     }
1811   }
1812   zstdThrowIfError(ZSTD_initDStream(dstream_.get()));
1813 }
1814
1815 bool ZSTDStreamCodec::doUncompressStream(
1816     ByteRange& input,
1817     MutableByteRange& output,
1818     StreamCodec::FlushOp flushOp) {
1819   if (needReset_) {
1820     // If we are given all the input in one chunk try to use block uncompression
1821     if (flushOp == StreamCodec::FlushOp::END &&
1822         tryBlockUncompress(input, output)) {
1823       return true;
1824     }
1825     resetDStream();
1826     needReset_ = false;
1827   }
1828   ZSTD_inBuffer in = {input.data(), input.size(), 0};
1829   ZSTD_outBuffer out = {output.data(), output.size(), 0};
1830   SCOPE_EXIT {
1831     input.uncheckedAdvance(in.pos);
1832     output.uncheckedAdvance(out.pos);
1833   };
1834   size_t const rc = ZSTD_decompressStream(dstream_.get(), &out, &in);
1835   zstdThrowIfError(rc);
1836   return rc == 0;
1837 }
1838
1839 #endif // FOLLY_HAVE_LIBZSTD
1840
1841 #if FOLLY_HAVE_LIBBZ2
1842
1843 class Bzip2Codec final : public Codec {
1844  public:
1845   static std::unique_ptr<Codec> create(int level, CodecType type);
1846   explicit Bzip2Codec(int level, CodecType type);
1847
1848   std::vector<std::string> validPrefixes() const override;
1849   bool canUncompress(IOBuf const* data, Optional<uint64_t> uncompressedLength)
1850       const override;
1851
1852  private:
1853   uint64_t doMaxCompressedLength(uint64_t uncompressedLength) const override;
1854   std::unique_ptr<IOBuf> doCompress(IOBuf const* data) override;
1855   std::unique_ptr<IOBuf> doUncompress(
1856       IOBuf const* data,
1857       Optional<uint64_t> uncompressedLength) override;
1858
1859   int level_;
1860 };
1861
1862 /* static */ std::unique_ptr<Codec> Bzip2Codec::create(
1863     int level,
1864     CodecType type) {
1865   return std::make_unique<Bzip2Codec>(level, type);
1866 }
1867
1868 Bzip2Codec::Bzip2Codec(int level, CodecType type) : Codec(type) {
1869   DCHECK(type == CodecType::BZIP2);
1870   switch (level) {
1871     case COMPRESSION_LEVEL_FASTEST:
1872       level = 1;
1873       break;
1874     case COMPRESSION_LEVEL_DEFAULT:
1875       level = 9;
1876       break;
1877     case COMPRESSION_LEVEL_BEST:
1878       level = 9;
1879       break;
1880   }
1881   if (level < 1 || level > 9) {
1882     throw std::invalid_argument(
1883         to<std::string>("Bzip2: invalid level: ", level));
1884   }
1885   level_ = level;
1886 }
1887
1888 static uint32_t constexpr kBzip2MagicLE = 0x685a42;
1889 static uint64_t constexpr kBzip2MagicBytes = 3;
1890
1891 std::vector<std::string> Bzip2Codec::validPrefixes() const {
1892   return {prefixToStringLE(kBzip2MagicLE, kBzip2MagicBytes)};
1893 }
1894
1895 bool Bzip2Codec::canUncompress(IOBuf const* data, Optional<uint64_t>) const {
1896   return dataStartsWithLE(data, kBzip2MagicLE, kBzip2MagicBytes);
1897 }
1898
1899 uint64_t Bzip2Codec::doMaxCompressedLength(uint64_t uncompressedLength) const {
1900   // http://www.bzip.org/1.0.5/bzip2-manual-1.0.5.html#bzbufftobuffcompress
1901   //   To guarantee that the compressed data will fit in its buffer, allocate an
1902   //   output buffer of size 1% larger than the uncompressed data, plus six
1903   //   hundred extra bytes.
1904   return uncompressedLength + uncompressedLength / 100 + 600;
1905 }
1906
1907 static bz_stream createBzStream() {
1908   bz_stream stream;
1909   stream.bzalloc = nullptr;
1910   stream.bzfree = nullptr;
1911   stream.opaque = nullptr;
1912   stream.next_in = stream.next_out = nullptr;
1913   stream.avail_in = stream.avail_out = 0;
1914   return stream;
1915 }
1916
1917 // Throws on error condition, otherwise returns the code.
1918 static int bzCheck(int const rc) {
1919   switch (rc) {
1920     case BZ_OK:
1921     case BZ_RUN_OK:
1922     case BZ_FLUSH_OK:
1923     case BZ_FINISH_OK:
1924     case BZ_STREAM_END:
1925       return rc;
1926     default:
1927       throw std::runtime_error(to<std::string>("Bzip2 error: ", rc));
1928   }
1929 }
1930
1931 static std::unique_ptr<IOBuf> addOutputBuffer(
1932     bz_stream* stream,
1933     uint64_t const bufferLength) {
1934   DCHECK_LE(bufferLength, std::numeric_limits<unsigned>::max());
1935   DCHECK_EQ(stream->avail_out, 0);
1936
1937   auto buf = IOBuf::create(bufferLength);
1938   buf->append(buf->capacity());
1939
1940   stream->next_out = reinterpret_cast<char*>(buf->writableData());
1941   stream->avail_out = buf->length();
1942
1943   return buf;
1944 }
1945
1946 std::unique_ptr<IOBuf> Bzip2Codec::doCompress(IOBuf const* data) {
1947   bz_stream stream = createBzStream();
1948   bzCheck(BZ2_bzCompressInit(&stream, level_, 0, 0));
1949   SCOPE_EXIT {
1950     bzCheck(BZ2_bzCompressEnd(&stream));
1951   };
1952
1953   uint64_t const uncompressedLength = data->computeChainDataLength();
1954   uint64_t const maxCompressedLen = maxCompressedLength(uncompressedLength);
1955   uint64_t constexpr kMaxSingleStepLength = uint64_t(64) << 20; // 64 MiB
1956   uint64_t constexpr kDefaultBufferLength = uint64_t(4) << 20;
1957
1958   auto out = addOutputBuffer(
1959       &stream,
1960       maxCompressedLen <= kMaxSingleStepLength ? maxCompressedLen
1961                                                : kDefaultBufferLength);
1962
1963   for (auto range : *data) {
1964     while (!range.empty()) {
1965       auto const inSize = std::min<size_t>(range.size(), kMaxSingleStepLength);
1966       stream.next_in =
1967           const_cast<char*>(reinterpret_cast<char const*>(range.data()));
1968       stream.avail_in = inSize;
1969
1970       if (stream.avail_out == 0) {
1971         out->prependChain(addOutputBuffer(&stream, kDefaultBufferLength));
1972       }
1973
1974       bzCheck(BZ2_bzCompress(&stream, BZ_RUN));
1975       range.uncheckedAdvance(inSize - stream.avail_in);
1976     }
1977   }
1978   do {
1979     if (stream.avail_out == 0) {
1980       out->prependChain(addOutputBuffer(&stream, kDefaultBufferLength));
1981     }
1982   } while (bzCheck(BZ2_bzCompress(&stream, BZ_FINISH)) != BZ_STREAM_END);
1983
1984   out->prev()->trimEnd(stream.avail_out);
1985
1986   return out;
1987 }
1988
1989 std::unique_ptr<IOBuf> Bzip2Codec::doUncompress(
1990     const IOBuf* data,
1991     Optional<uint64_t> uncompressedLength) {
1992   bz_stream stream = createBzStream();
1993   bzCheck(BZ2_bzDecompressInit(&stream, 0, 0));
1994   SCOPE_EXIT {
1995     bzCheck(BZ2_bzDecompressEnd(&stream));
1996   };
1997
1998   uint64_t constexpr kMaxSingleStepLength = uint64_t(64) << 20; // 64 MiB
1999   uint64_t const kBlockSize = uint64_t(100) << 10; // 100 KiB
2000   uint64_t const kDefaultBufferLength =
2001       computeBufferLength(data->computeChainDataLength(), kBlockSize);
2002
2003   auto out = addOutputBuffer(
2004       &stream,
2005       ((uncompressedLength && *uncompressedLength <= kMaxSingleStepLength)
2006            ? *uncompressedLength
2007            : kDefaultBufferLength));
2008
2009   int rc = BZ_OK;
2010   for (auto range : *data) {
2011     while (!range.empty()) {
2012       auto const inSize = std::min<size_t>(range.size(), kMaxSingleStepLength);
2013       stream.next_in =
2014           const_cast<char*>(reinterpret_cast<char const*>(range.data()));
2015       stream.avail_in = inSize;
2016
2017       if (stream.avail_out == 0) {
2018         out->prependChain(addOutputBuffer(&stream, kDefaultBufferLength));
2019       }
2020
2021       rc = bzCheck(BZ2_bzDecompress(&stream));
2022       range.uncheckedAdvance(inSize - stream.avail_in);
2023     }
2024   }
2025   while (rc != BZ_STREAM_END) {
2026     if (stream.avail_out == 0) {
2027       out->prependChain(addOutputBuffer(&stream, kDefaultBufferLength));
2028     }
2029     size_t const outputSize = stream.avail_out;
2030     rc = bzCheck(BZ2_bzDecompress(&stream));
2031     if (outputSize == stream.avail_out) {
2032       throw std::runtime_error("Bzip2Codec: Truncated input");
2033     }
2034   }
2035
2036   out->prev()->trimEnd(stream.avail_out);
2037
2038   uint64_t const totalOut =
2039       (uint64_t(stream.total_out_hi32) << 32) + stream.total_out_lo32;
2040   if (uncompressedLength && uncompressedLength != totalOut) {
2041     throw std::runtime_error("Bzip2 error: Invalid uncompressed length");
2042   }
2043
2044   return out;
2045 }
2046
2047 #endif // FOLLY_HAVE_LIBBZ2
2048
2049 /**
2050  * Automatic decompression
2051  */
2052 class AutomaticCodec final : public Codec {
2053  public:
2054   static std::unique_ptr<Codec> create(
2055       std::vector<std::unique_ptr<Codec>> customCodecs);
2056   explicit AutomaticCodec(std::vector<std::unique_ptr<Codec>> customCodecs);
2057
2058   std::vector<std::string> validPrefixes() const override;
2059   bool canUncompress(const IOBuf* data, Optional<uint64_t> uncompressedLength)
2060       const override;
2061
2062  private:
2063   bool doNeedsUncompressedLength() const override;
2064   uint64_t doMaxUncompressedLength() const override;
2065
2066   uint64_t doMaxCompressedLength(uint64_t) const override {
2067     throw std::runtime_error(
2068         "AutomaticCodec error: maxCompressedLength() not supported.");
2069   }
2070   std::unique_ptr<IOBuf> doCompress(const IOBuf*) override {
2071     throw std::runtime_error("AutomaticCodec error: compress() not supported.");
2072   }
2073   std::unique_ptr<IOBuf> doUncompress(
2074       const IOBuf* data,
2075       Optional<uint64_t> uncompressedLength) override;
2076
2077   void addCodecIfSupported(CodecType type);
2078
2079   // Throws iff the codecs aren't compatible (very slow)
2080   void checkCompatibleCodecs() const;
2081
2082   std::vector<std::unique_ptr<Codec>> codecs_;
2083   bool needsUncompressedLength_;
2084   uint64_t maxUncompressedLength_;
2085 };
2086
2087 std::vector<std::string> AutomaticCodec::validPrefixes() const {
2088   std::unordered_set<std::string> prefixes;
2089   for (const auto& codec : codecs_) {
2090     const auto codecPrefixes = codec->validPrefixes();
2091     prefixes.insert(codecPrefixes.begin(), codecPrefixes.end());
2092   }
2093   return std::vector<std::string>{prefixes.begin(), prefixes.end()};
2094 }
2095
2096 bool AutomaticCodec::canUncompress(
2097     const IOBuf* data,
2098     Optional<uint64_t> uncompressedLength) const {
2099   return std::any_of(
2100       codecs_.begin(),
2101       codecs_.end(),
2102       [data, uncompressedLength](std::unique_ptr<Codec> const& codec) {
2103         return codec->canUncompress(data, uncompressedLength);
2104       });
2105 }
2106
2107 void AutomaticCodec::addCodecIfSupported(CodecType type) {
2108   const bool present = std::any_of(
2109       codecs_.begin(),
2110       codecs_.end(),
2111       [&type](std::unique_ptr<Codec> const& codec) {
2112         return codec->type() == type;
2113       });
2114   if (hasCodec(type) && !present) {
2115     codecs_.push_back(getCodec(type));
2116   }
2117 }
2118
2119 /* static */ std::unique_ptr<Codec> AutomaticCodec::create(
2120     std::vector<std::unique_ptr<Codec>> customCodecs) {
2121   return std::make_unique<AutomaticCodec>(std::move(customCodecs));
2122 }
2123
2124 AutomaticCodec::AutomaticCodec(std::vector<std::unique_ptr<Codec>> customCodecs)
2125     : Codec(CodecType::USER_DEFINED), codecs_(std::move(customCodecs)) {
2126   // Fastest -> slowest
2127   addCodecIfSupported(CodecType::LZ4_FRAME);
2128   addCodecIfSupported(CodecType::ZSTD);
2129   addCodecIfSupported(CodecType::ZLIB);
2130   addCodecIfSupported(CodecType::GZIP);
2131   addCodecIfSupported(CodecType::LZMA2);
2132   addCodecIfSupported(CodecType::BZIP2);
2133   if (kIsDebug) {
2134     checkCompatibleCodecs();
2135   }
2136   // Check that none of the codes are are null
2137   DCHECK(std::none_of(
2138       codecs_.begin(), codecs_.end(), [](std::unique_ptr<Codec> const& codec) {
2139         return codec == nullptr;
2140       }));
2141
2142   needsUncompressedLength_ = std::any_of(
2143       codecs_.begin(), codecs_.end(), [](std::unique_ptr<Codec> const& codec) {
2144         return codec->needsUncompressedLength();
2145       });
2146
2147   const auto it = std::max_element(
2148       codecs_.begin(),
2149       codecs_.end(),
2150       [](std::unique_ptr<Codec> const& lhs, std::unique_ptr<Codec> const& rhs) {
2151         return lhs->maxUncompressedLength() < rhs->maxUncompressedLength();
2152       });
2153   DCHECK(it != codecs_.end());
2154   maxUncompressedLength_ = (*it)->maxUncompressedLength();
2155 }
2156
2157 void AutomaticCodec::checkCompatibleCodecs() const {
2158   // Keep track of all the possible headers.
2159   std::unordered_set<std::string> headers;
2160   // The empty header is not allowed.
2161   headers.insert("");
2162   // Step 1:
2163   // Construct a set of headers and check that none of the headers occur twice.
2164   // Eliminate edge cases.
2165   for (auto&& codec : codecs_) {
2166     const auto codecHeaders = codec->validPrefixes();
2167     // Codecs without any valid headers are not allowed.
2168     if (codecHeaders.empty()) {
2169       throw std::invalid_argument{
2170           "AutomaticCodec: validPrefixes() must not be empty."};
2171     }
2172     // Insert all the headers for the current codec.
2173     const size_t beforeSize = headers.size();
2174     headers.insert(codecHeaders.begin(), codecHeaders.end());
2175     // Codecs are not compatible if any header occurred twice.
2176     if (beforeSize + codecHeaders.size() != headers.size()) {
2177       throw std::invalid_argument{
2178           "AutomaticCodec: Two valid prefixes collide."};
2179     }
2180   }
2181   // Step 2:
2182   // Check if any strict non-empty prefix of any header is a header.
2183   for (const auto& header : headers) {
2184     for (size_t i = 1; i < header.size(); ++i) {
2185       if (headers.count(header.substr(0, i))) {
2186         throw std::invalid_argument{
2187             "AutomaticCodec: One valid prefix is a prefix of another valid "
2188             "prefix."};
2189       }
2190     }
2191   }
2192 }
2193
2194 bool AutomaticCodec::doNeedsUncompressedLength() const {
2195   return needsUncompressedLength_;
2196 }
2197
2198 uint64_t AutomaticCodec::doMaxUncompressedLength() const {
2199   return maxUncompressedLength_;
2200 }
2201
2202 std::unique_ptr<IOBuf> AutomaticCodec::doUncompress(
2203     const IOBuf* data,
2204     Optional<uint64_t> uncompressedLength) {
2205   for (auto&& codec : codecs_) {
2206     if (codec->canUncompress(data, uncompressedLength)) {
2207       return codec->uncompress(data, uncompressedLength);
2208     }
2209   }
2210   throw std::runtime_error("AutomaticCodec error: Unknown compressed data");
2211 }
2212
2213 using CodecFactory = std::unique_ptr<Codec> (*)(int, CodecType);
2214 using StreamCodecFactory = std::unique_ptr<StreamCodec> (*)(int, CodecType);
2215 struct Factory {
2216   CodecFactory codec;
2217   StreamCodecFactory stream;
2218 };
2219
2220 constexpr Factory
2221     codecFactories[static_cast<size_t>(CodecType::NUM_CODEC_TYPES)] = {
2222         {}, // USER_DEFINED
2223         {NoCompressionCodec::create, nullptr},
2224
2225 #if FOLLY_HAVE_LIBLZ4
2226         {LZ4Codec::create, nullptr},
2227 #else
2228         {},
2229 #endif
2230
2231 #if FOLLY_HAVE_LIBSNAPPY
2232         {SnappyCodec::create, nullptr},
2233 #else
2234         {},
2235 #endif
2236
2237 #if FOLLY_HAVE_LIBZ
2238         {ZlibStreamCodec::createCodec, ZlibStreamCodec::createStream},
2239 #else
2240         {},
2241 #endif
2242
2243 #if FOLLY_HAVE_LIBLZ4
2244         {LZ4Codec::create, nullptr},
2245 #else
2246         {},
2247 #endif
2248
2249 #if FOLLY_HAVE_LIBLZMA
2250         {LZMA2Codec::create, nullptr},
2251         {LZMA2Codec::create, nullptr},
2252 #else
2253         {},
2254         {},
2255 #endif
2256
2257 #if FOLLY_HAVE_LIBZSTD
2258         {ZSTDStreamCodec::createCodec, ZSTDStreamCodec::createStream},
2259 #else
2260         {},
2261 #endif
2262
2263 #if FOLLY_HAVE_LIBZ
2264         {ZlibStreamCodec::createCodec, ZlibStreamCodec::createStream},
2265 #else
2266         {},
2267 #endif
2268
2269 #if (FOLLY_HAVE_LIBLZ4 && LZ4_VERSION_NUMBER >= 10301)
2270         {LZ4FrameCodec::create, nullptr},
2271 #else
2272         {},
2273 #endif
2274
2275 #if FOLLY_HAVE_LIBBZ2
2276         {Bzip2Codec::create, nullptr},
2277 #else
2278         {},
2279 #endif
2280 };
2281
2282 Factory const& getFactory(CodecType type) {
2283   size_t const idx = static_cast<size_t>(type);
2284   if (idx >= static_cast<size_t>(CodecType::NUM_CODEC_TYPES)) {
2285     throw std::invalid_argument(
2286         to<std::string>("Compression type ", idx, " invalid"));
2287   }
2288   return codecFactories[idx];
2289 }
2290 } // namespace
2291
2292 bool hasCodec(CodecType type) {
2293   return getFactory(type).codec != nullptr;
2294 }
2295
2296 std::unique_ptr<Codec> getCodec(CodecType type, int level) {
2297   auto const factory = getFactory(type).codec;
2298   if (!factory) {
2299     throw std::invalid_argument(
2300         to<std::string>("Compression type ", type, " not supported"));
2301   }
2302   auto codec = (*factory)(level, type);
2303   DCHECK(codec->type() == type);
2304   return codec;
2305 }
2306
2307 bool hasStreamCodec(CodecType type) {
2308   return getFactory(type).stream != nullptr;
2309 }
2310
2311 std::unique_ptr<StreamCodec> getStreamCodec(CodecType type, int level) {
2312   auto const factory = getFactory(type).stream;
2313   if (!factory) {
2314     throw std::invalid_argument(
2315         to<std::string>("Compression type ", type, " not supported"));
2316   }
2317   auto codec = (*factory)(level, type);
2318   DCHECK(codec->type() == type);
2319   return codec;
2320 }
2321
2322 std::unique_ptr<Codec> getAutoUncompressionCodec(
2323     std::vector<std::unique_ptr<Codec>> customCodecs) {
2324   return AutomaticCodec::create(std::move(customCodecs));
2325 }
2326 }}  // namespaces