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 #include <zstd.h>
  44 #endif
  45
  46 #include <folly/Conv.h>
  47 #include <folly/Memory.h>
  48 #include <folly/Portability.h>
  49 #include <folly/ScopeGuard.h>
  50 #include <folly/Varint.h>
  51 #include <folly/io/Cursor.h>
  52
  53 namespace folly { namespace io {
  54
  55 Codec::Codec(CodecType type) : type_(type) { }
  56
  57 // Ensure consistent behavior in the nullptr case
  58 std::unique_ptr<IOBuf> Codec::compress(const IOBuf* data) {
  59   uint64_t len = data->computeChainDataLength();
  60   if (len == 0) {
  61     return IOBuf::create(0);
  62   }
  63   if (len > maxUncompressedLength()) {
  64     throw std::runtime_error("Codec: uncompressed length too large");
  65   }
  66
  67   return doCompress(data);
  68 }
  69
  70 std::string Codec::compress(const StringPiece data) {
  71   const uint64_t len = data.size();
  72   if (len == 0) {
  73     return "";
  74   }
  75   if (len > maxUncompressedLength()) {
  76     throw std::runtime_error("Codec: uncompressed length too large");
  77   }
  78
  79   return doCompressString(data);
  80 }
  81
  82 std::unique_ptr<IOBuf> Codec::uncompress(const IOBuf* data,
  83                                          uint64_t uncompressedLength) {
  84   if (uncompressedLength == UNKNOWN_UNCOMPRESSED_LENGTH) {
  85     if (needsUncompressedLength()) {
  86       throw std::invalid_argument("Codec: uncompressed length required");
  87     }
  88   } else if (uncompressedLength > maxUncompressedLength()) {
  89     throw std::runtime_error("Codec: uncompressed length too large");
  90   }
  91
  92   if (data->empty()) {
  93     if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
  94         uncompressedLength != 0) {
  95       throw std::runtime_error("Codec: invalid uncompressed length");
  96     }
  97     return IOBuf::create(0);
  98   }
  99
 100   return doUncompress(data, uncompressedLength);
 101 }
 102
 103 std::string Codec::uncompress(
 104     const StringPiece data,
 105     uint64_t uncompressedLength) {
 106   if (uncompressedLength == UNKNOWN_UNCOMPRESSED_LENGTH) {
 107     if (needsUncompressedLength()) {
 108       throw std::invalid_argument("Codec: uncompressed length required");
 109     }
 110   } else if (uncompressedLength > maxUncompressedLength()) {
 111     throw std::runtime_error("Codec: uncompressed length too large");
 112   }
 113
 114   if (data.empty()) {
 115     if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
 116         uncompressedLength != 0) {
 117       throw std::runtime_error("Codec: invalid uncompressed length");
 118     }
 119     return "";
 120   }
 121
 122   return doUncompressString(data, uncompressedLength);
 123 }
 124
 125 bool Codec::needsUncompressedLength() const {
 126   return doNeedsUncompressedLength();
 127 }
 128
 129 uint64_t Codec::maxUncompressedLength() const {
 130   return doMaxUncompressedLength();
 131 }
 132
 133 bool Codec::doNeedsUncompressedLength() const {
 134   return false;
 135 }
 136
 137 uint64_t Codec::doMaxUncompressedLength() const {
 138   return UNLIMITED_UNCOMPRESSED_LENGTH;
 139 }
 140
 141 std::string Codec::doCompressString(const StringPiece data) {
 142   const IOBuf inputBuffer{IOBuf::WRAP_BUFFER, data};
 143   auto outputBuffer = doCompress(&inputBuffer);
 144   std::string output;
 145   output.reserve(outputBuffer->computeChainDataLength());
 146   for (auto range : *outputBuffer) {
 147     output.append(reinterpret_cast<const char*>(range.data()), range.size());
 148   }
 149   return output;
 150 }
 151
 152 std::string Codec::doUncompressString(
 153     const StringPiece data,
 154     uint64_t uncompressedLength) {
 155   const IOBuf inputBuffer{IOBuf::WRAP_BUFFER, data};
 156   auto outputBuffer = doUncompress(&inputBuffer, uncompressedLength);
 157   std::string output;
 158   output.reserve(outputBuffer->computeChainDataLength());
 159   for (auto range : *outputBuffer) {
 160     output.append(reinterpret_cast<const char*>(range.data()), range.size());
 161   }
 162   return output;
 163 }
 164
 165 namespace {
 166
 167 /**
 168  * No compression
 169  */
 170 class NoCompressionCodec final : public Codec {
 171  public:
 172   static std::unique_ptr<Codec> create(int level, CodecType type);
 173   explicit NoCompressionCodec(int level, CodecType type);
 174
 175  private:
 176   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
 177   std::unique_ptr<IOBuf> doUncompress(
 178       const IOBuf* data,
 179       uint64_t uncompressedLength) override;
 180 };
 181
 182 std::unique_ptr<Codec> NoCompressionCodec::create(int level, CodecType type) {
 183   return make_unique<NoCompressionCodec>(level, type);
 184 }
 185
 186 NoCompressionCodec::NoCompressionCodec(int level, CodecType type)
 187   : Codec(type) {
 188   DCHECK(type == CodecType::NO_COMPRESSION);
 189   switch (level) {
 190   case COMPRESSION_LEVEL_DEFAULT:
 191   case COMPRESSION_LEVEL_FASTEST:
 192   case COMPRESSION_LEVEL_BEST:
 193     level = 0;
 194   }
 195   if (level != 0) {
 196     throw std::invalid_argument(to<std::string>(
 197         "NoCompressionCodec: invalid level ", level));
 198   }
 199 }
 200
 201 std::unique_ptr<IOBuf> NoCompressionCodec::doCompress(
 202     const IOBuf* data) {
 203   return data->clone();
 204 }
 205
 206 std::unique_ptr<IOBuf> NoCompressionCodec::doUncompress(
 207     const IOBuf* data,
 208     uint64_t uncompressedLength) {
 209   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
 210       data->computeChainDataLength() != uncompressedLength) {
 211     throw std::runtime_error(to<std::string>(
 212         "NoCompressionCodec: invalid uncompressed length"));
 213   }
 214   return data->clone();
 215 }
 216
 217 #if (FOLLY_HAVE_LIBLZ4 || FOLLY_HAVE_LIBLZMA)
 218
 219 namespace {
 220
 221 void encodeVarintToIOBuf(uint64_t val, folly::IOBuf* out) {
 222   DCHECK_GE(out->tailroom(), kMaxVarintLength64);
 223   out->append(encodeVarint(val, out->writableTail()));
 224 }
 225
 226 inline uint64_t decodeVarintFromCursor(folly::io::Cursor& cursor) {
 227   uint64_t val = 0;
 228   int8_t b = 0;
 229   for (int shift = 0; shift <= 63; shift += 7) {
 230     b = cursor.read<int8_t>();
 231     val |= static_cast<uint64_t>(b & 0x7f) << shift;
 232     if (b >= 0) {
 233       break;
 234     }
 235   }
 236   if (b < 0) {
 237     throw std::invalid_argument("Invalid varint value. Too big.");
 238   }
 239   return val;
 240 }
 241
 242 }  // namespace
 243
 244 #endif  // FOLLY_HAVE_LIBLZ4 || FOLLY_HAVE_LIBLZMA
 245
 246 #if FOLLY_HAVE_LIBLZ4
 247
 248 /**
 249  * LZ4 compression
 250  */
 251 class LZ4Codec final : public Codec {
 252  public:
 253   static std::unique_ptr<Codec> create(int level, CodecType type);
 254   explicit LZ4Codec(int level, CodecType type);
 255
 256  private:
 257   bool doNeedsUncompressedLength() const override;
 258   uint64_t doMaxUncompressedLength() const override;
 259
 260   bool encodeSize() const { return type() == CodecType::LZ4_VARINT_SIZE; }
 261
 262   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
 263   std::unique_ptr<IOBuf> doUncompress(
 264       const IOBuf* data,
 265       uint64_t uncompressedLength) override;
 266
 267   bool highCompression_;
 268 };
 269
 270 std::unique_ptr<Codec> LZ4Codec::create(int level, CodecType type) {
 271   return make_unique<LZ4Codec>(level, type);
 272 }
 273
 274 LZ4Codec::LZ4Codec(int level, CodecType type) : Codec(type) {
 275   DCHECK(type == CodecType::LZ4 || type == CodecType::LZ4_VARINT_SIZE);
 276
 277   switch (level) {
 278   case COMPRESSION_LEVEL_FASTEST:
 279   case COMPRESSION_LEVEL_DEFAULT:
 280     level = 1;
 281     break;
 282   case COMPRESSION_LEVEL_BEST:
 283     level = 2;
 284     break;
 285   }
 286   if (level < 1 || level > 2) {
 287     throw std::invalid_argument(to<std::string>(
 288         "LZ4Codec: invalid level: ", level));
 289   }
 290   highCompression_ = (level > 1);
 291 }
 292
 293 bool LZ4Codec::doNeedsUncompressedLength() const {
 294   return !encodeSize();
 295 }
 296
 297 // The value comes from lz4.h in lz4-r117, but older versions of lz4 don't
 298 // define LZ4_MAX_INPUT_SIZE (even though the max size is the same), so do it
 299 // here.
 300 #ifndef LZ4_MAX_INPUT_SIZE
 301 # define LZ4_MAX_INPUT_SIZE 0x7E000000
 302 #endif
 303
 304 uint64_t LZ4Codec::doMaxUncompressedLength() const {
 305   return LZ4_MAX_INPUT_SIZE;
 306 }
 307
 308 std::unique_ptr<IOBuf> LZ4Codec::doCompress(const IOBuf* data) {
 309   IOBuf clone;
 310   if (data->isChained()) {
 311     // LZ4 doesn't support streaming, so we have to coalesce
 312     clone = data->cloneCoalescedAsValue();
 313     data = &clone;
 314   }
 315
 316   uint32_t extraSize = encodeSize() ? kMaxVarintLength64 : 0;
 317   auto out = IOBuf::create(extraSize + LZ4_compressBound(data->length()));
 318   if (encodeSize()) {
 319     encodeVarintToIOBuf(data->length(), out.get());
 320   }
 321
 322   int n;
 323   auto input = reinterpret_cast<const char*>(data->data());
 324   auto output = reinterpret_cast<char*>(out->writableTail());
 325   const auto inputLength = data->length();
 326 #if LZ4_VERSION_NUMBER >= 10700
 327   if (highCompression_) {
 328     n = LZ4_compress_HC(input, output, inputLength, out->tailroom(), 0);
 329   } else {
 330     n = LZ4_compress_default(input, output, inputLength, out->tailroom());
 331   }
 332 #else
 333   if (highCompression_) {
 334     n = LZ4_compressHC(input, output, inputLength);
 335   } else {
 336     n = LZ4_compress(input, output, inputLength);
 337   }
 338 #endif
 339
 340   CHECK_GE(n, 0);
 341   CHECK_LE(n, out->capacity());
 342
 343   out->append(n);
 344   return out;
 345 }
 346
 347 std::unique_ptr<IOBuf> LZ4Codec::doUncompress(
 348     const IOBuf* data,
 349     uint64_t uncompressedLength) {
 350   IOBuf clone;
 351   if (data->isChained()) {
 352     // LZ4 doesn't support streaming, so we have to coalesce
 353     clone = data->cloneCoalescedAsValue();
 354     data = &clone;
 355   }
 356
 357   folly::io::Cursor cursor(data);
 358   uint64_t actualUncompressedLength;
 359   if (encodeSize()) {
 360     actualUncompressedLength = decodeVarintFromCursor(cursor);
 361     if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
 362         uncompressedLength != actualUncompressedLength) {
 363       throw std::runtime_error("LZ4Codec: invalid uncompressed length");
 364     }
 365   } else {
 366     actualUncompressedLength = uncompressedLength;
 367     if (actualUncompressedLength == UNKNOWN_UNCOMPRESSED_LENGTH ||
 368         actualUncompressedLength > maxUncompressedLength()) {
 369       throw std::runtime_error("LZ4Codec: invalid uncompressed length");
 370     }
 371   }
 372
 373   auto sp = StringPiece{cursor.peekBytes()};
 374   auto out = IOBuf::create(actualUncompressedLength);
 375   int n = LZ4_decompress_safe(
 376       sp.data(),
 377       reinterpret_cast<char*>(out->writableTail()),
 378       sp.size(),
 379       actualUncompressedLength);
 380
 381   if (n < 0 || uint64_t(n) != actualUncompressedLength) {
 382     throw std::runtime_error(to<std::string>(
 383         "LZ4 decompression returned invalid value ", n));
 384   }
 385   out->append(actualUncompressedLength);
 386   return out;
 387 }
 388
 389 #if LZ4_VERSION_NUMBER >= 10301
 390
 391 class LZ4FrameCodec final : public Codec {
 392  public:
 393   static std::unique_ptr<Codec> create(int level, CodecType type);
 394   explicit LZ4FrameCodec(int level, CodecType type);
 395   ~LZ4FrameCodec();
 396
 397  private:
 398   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
 399   std::unique_ptr<IOBuf> doUncompress(
 400       const IOBuf* data,
 401       uint64_t uncompressedLength) override;
 402
 403   // Reset the dctx_ if it is dirty or null.
 404   void resetDCtx();
 405
 406   int level_;
 407   LZ4F_decompressionContext_t dctx_{nullptr};
 408   bool dirty_{false};
 409 };
 410
 411 /* static */ std::unique_ptr<Codec> LZ4FrameCodec::create(
 412     int level,
 413     CodecType type) {
 414   return make_unique<LZ4FrameCodec>(level, type);
 415 }
 416
 417 static size_t lz4FrameThrowOnError(size_t code) {
 418   if (LZ4F_isError(code)) {
 419     throw std::runtime_error(
 420         to<std::string>("LZ4Frame error: ", LZ4F_getErrorName(code)));
 421   }
 422   return code;
 423 }
 424
 425 void LZ4FrameCodec::resetDCtx() {
 426   if (dctx_ && !dirty_) {
 427     return;
 428   }
 429   if (dctx_) {
 430     LZ4F_freeDecompressionContext(dctx_);
 431   }
 432   lz4FrameThrowOnError(LZ4F_createDecompressionContext(&dctx_, 100));
 433   dirty_ = false;
 434 }
 435
 436 LZ4FrameCodec::LZ4FrameCodec(int level, CodecType type) : Codec(type) {
 437   DCHECK(type == CodecType::LZ4_FRAME);
 438   switch (level) {
 439     case COMPRESSION_LEVEL_FASTEST:
 440     case COMPRESSION_LEVEL_DEFAULT:
 441       level_ = 0;
 442       break;
 443     case COMPRESSION_LEVEL_BEST:
 444       level_ = 16;
 445       break;
 446     default:
 447       level_ = level;
 448       break;
 449   }
 450 }
 451
 452 LZ4FrameCodec::~LZ4FrameCodec() {
 453   if (dctx_) {
 454     LZ4F_freeDecompressionContext(dctx_);
 455   }
 456 }
 457
 458 std::unique_ptr<IOBuf> LZ4FrameCodec::doCompress(const IOBuf* data) {
 459   // LZ4 Frame compression doesn't support streaming so we have to coalesce
 460   IOBuf clone;
 461   if (data->isChained()) {
 462     clone = data->cloneCoalescedAsValue();
 463     data = &clone;
 464   }
 465   // Set preferences
 466   const auto uncompressedLength = data->length();
 467   LZ4F_preferences_t prefs{};
 468   prefs.compressionLevel = level_;
 469   prefs.frameInfo.contentSize = uncompressedLength;
 470   // Compress
 471   auto buf = IOBuf::create(LZ4F_compressFrameBound(uncompressedLength, &prefs));
 472   const size_t written = lz4FrameThrowOnError(LZ4F_compressFrame(
 473       buf->writableTail(),
 474       buf->tailroom(),
 475       data->data(),
 476       data->length(),
 477       &prefs));
 478   buf->append(written);
 479   return buf;
 480 }
 481
 482 std::unique_ptr<IOBuf> LZ4FrameCodec::doUncompress(
 483     const IOBuf* data,
 484     uint64_t uncompressedLength) {
 485   // Reset the dctx if any errors have occurred
 486   resetDCtx();
 487   // Coalesce the data
 488   ByteRange in = *data->begin();
 489   IOBuf clone;
 490   if (data->isChained()) {
 491     clone = data->cloneCoalescedAsValue();
 492     in = clone.coalesce();
 493   }
 494   data = nullptr;
 495   // Select decompression options
 496   LZ4F_decompressOptions_t options;
 497   options.stableDst = 1;
 498   // Select blockSize and growthSize for the IOBufQueue
 499   IOBufQueue queue(IOBufQueue::cacheChainLength());
 500   auto blockSize = uint64_t{64} << 10;
 501   auto growthSize = uint64_t{4} << 20;
 502   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH) {
 503     // Allocate uncompressedLength in one chunk (up to 64 MB)
 504     const auto allocateSize = std::min(uncompressedLength, uint64_t{64} << 20);
 505     queue.preallocate(allocateSize, allocateSize);
 506     blockSize = std::min(uncompressedLength, blockSize);
 507     growthSize = std::min(uncompressedLength, growthSize);
 508   } else {
 509     // Reduce growthSize for small data
 510     const auto guessUncompressedLen = 4 * std::max(blockSize, in.size());
 511     growthSize = std::min(guessUncompressedLen, growthSize);
 512   }
 513   // Once LZ4_decompress() is called, the dctx_ cannot be reused until it
 514   // returns 0
 515   dirty_ = true;
 516   // Decompress until the frame is over
 517   size_t code = 0;
 518   do {
 519     // Allocate enough space to decompress at least a block
 520     void* out;
 521     size_t outSize;
 522     std::tie(out, outSize) = queue.preallocate(blockSize, growthSize);
 523     // Decompress
 524     size_t inSize = in.size();
 525     code = lz4FrameThrowOnError(
 526         LZ4F_decompress(dctx_, out, &outSize, in.data(), &inSize, &options));
 527     if (in.empty() && outSize == 0 && code != 0) {
 528       // We passed no input, no output was produced, and the frame isn't over
 529       // No more forward progress is possible
 530       throw std::runtime_error("LZ4Frame error: Incomplete frame");
 531     }
 532     in.uncheckedAdvance(inSize);
 533     queue.postallocate(outSize);
 534   } while (code != 0);
 535   // At this point the decompression context can be reused
 536   dirty_ = false;
 537   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
 538       queue.chainLength() != uncompressedLength) {
 539     throw std::runtime_error("LZ4Frame error: Invalid uncompressedLength");
 540   }
 541   return queue.move();
 542 }
 543
 544 #endif // LZ4_VERSION_NUMBER >= 10301
 545 #endif // FOLLY_HAVE_LIBLZ4
 546
 547 #if FOLLY_HAVE_LIBSNAPPY
 548
 549 /**
 550  * Snappy compression
 551  */
 552
 553 /**
 554  * Implementation of snappy::Source that reads from a IOBuf chain.
 555  */
 556 class IOBufSnappySource final : public snappy::Source {
 557  public:
 558   explicit IOBufSnappySource(const IOBuf* data);
 559   size_t Available() const override;
 560   const char* Peek(size_t* len) override;
 561   void Skip(size_t n) override;
 562  private:
 563   size_t available_;
 564   io::Cursor cursor_;
 565 };
 566
 567 IOBufSnappySource::IOBufSnappySource(const IOBuf* data)
 568   : available_(data->computeChainDataLength()),
 569     cursor_(data) {
 570 }
 571
 572 size_t IOBufSnappySource::Available() const {
 573   return available_;
 574 }
 575
 576 const char* IOBufSnappySource::Peek(size_t* len) {
 577   auto sp = StringPiece{cursor_.peekBytes()};
 578   *len = sp.size();
 579   return sp.data();
 580 }
 581
 582 void IOBufSnappySource::Skip(size_t n) {
 583   CHECK_LE(n, available_);
 584   cursor_.skip(n);
 585   available_ -= n;
 586 }
 587
 588 class SnappyCodec final : public Codec {
 589  public:
 590   static std::unique_ptr<Codec> create(int level, CodecType type);
 591   explicit SnappyCodec(int level, CodecType type);
 592
 593  private:
 594   uint64_t doMaxUncompressedLength() const override;
 595   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
 596   std::unique_ptr<IOBuf> doUncompress(
 597       const IOBuf* data,
 598       uint64_t uncompressedLength) override;
 599 };
 600
 601 std::unique_ptr<Codec> SnappyCodec::create(int level, CodecType type) {
 602   return make_unique<SnappyCodec>(level, type);
 603 }
 604
 605 SnappyCodec::SnappyCodec(int level, CodecType type) : Codec(type) {
 606   DCHECK(type == CodecType::SNAPPY);
 607   switch (level) {
 608   case COMPRESSION_LEVEL_FASTEST:
 609   case COMPRESSION_LEVEL_DEFAULT:
 610   case COMPRESSION_LEVEL_BEST:
 611     level = 1;
 612   }
 613   if (level != 1) {
 614     throw std::invalid_argument(to<std::string>(
 615         "SnappyCodec: invalid level: ", level));
 616   }
 617 }
 618
 619 uint64_t SnappyCodec::doMaxUncompressedLength() const {
 620   // snappy.h uses uint32_t for lengths, so there's that.
 621   return std::numeric_limits<uint32_t>::max();
 622 }
 623
 624 std::unique_ptr<IOBuf> SnappyCodec::doCompress(const IOBuf* data) {
 625   IOBufSnappySource source(data);
 626   auto out =
 627     IOBuf::create(snappy::MaxCompressedLength(source.Available()));
 628
 629   snappy::UncheckedByteArraySink sink(reinterpret_cast<char*>(
 630       out->writableTail()));
 631
 632   size_t n = snappy::Compress(&source, &sink);
 633
 634   CHECK_LE(n, out->capacity());
 635   out->append(n);
 636   return out;
 637 }
 638
 639 std::unique_ptr<IOBuf> SnappyCodec::doUncompress(const IOBuf* data,
 640                                                  uint64_t uncompressedLength) {
 641   uint32_t actualUncompressedLength = 0;
 642
 643   {
 644     IOBufSnappySource source(data);
 645     if (!snappy::GetUncompressedLength(&source, &actualUncompressedLength)) {
 646       throw std::runtime_error("snappy::GetUncompressedLength failed");
 647     }
 648     if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
 649         uncompressedLength != actualUncompressedLength) {
 650       throw std::runtime_error("snappy: invalid uncompressed length");
 651     }
 652   }
 653
 654   auto out = IOBuf::create(actualUncompressedLength);
 655
 656   {
 657     IOBufSnappySource source(data);
 658     if (!snappy::RawUncompress(&source,
 659                                reinterpret_cast<char*>(out->writableTail()))) {
 660       throw std::runtime_error("snappy::RawUncompress failed");
 661     }
 662   }
 663
 664   out->append(actualUncompressedLength);
 665   return out;
 666 }
 667
 668 #endif  // FOLLY_HAVE_LIBSNAPPY
 669
 670 #if FOLLY_HAVE_LIBZ
 671 /**
 672  * Zlib codec
 673  */
 674 class ZlibCodec final : public Codec {
 675  public:
 676   static std::unique_ptr<Codec> create(int level, CodecType type);
 677   explicit ZlibCodec(int level, CodecType type);
 678
 679  private:
 680   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
 681   std::unique_ptr<IOBuf> doUncompress(
 682       const IOBuf* data,
 683       uint64_t uncompressedLength) override;
 684
 685   std::unique_ptr<IOBuf> addOutputBuffer(z_stream* stream, uint32_t length);
 686   bool doInflate(z_stream* stream, IOBuf* head, uint32_t bufferLength);
 687
 688   int level_;
 689 };
 690
 691 std::unique_ptr<Codec> ZlibCodec::create(int level, CodecType type) {
 692   return make_unique<ZlibCodec>(level, type);
 693 }
 694
 695 ZlibCodec::ZlibCodec(int level, CodecType type) : Codec(type) {
 696   DCHECK(type == CodecType::ZLIB || type == CodecType::GZIP);
 697   switch (level) {
 698   case COMPRESSION_LEVEL_FASTEST:
 699     level = 1;
 700     break;
 701   case COMPRESSION_LEVEL_DEFAULT:
 702     level = Z_DEFAULT_COMPRESSION;
 703     break;
 704   case COMPRESSION_LEVEL_BEST:
 705     level = 9;
 706     break;
 707   }
 708   if (level != Z_DEFAULT_COMPRESSION && (level < 0 || level > 9)) {
 709     throw std::invalid_argument(to<std::string>(
 710         "ZlibCodec: invalid level: ", level));
 711   }
 712   level_ = level;
 713 }
 714
 715 std::unique_ptr<IOBuf> ZlibCodec::addOutputBuffer(z_stream* stream,
 716                                                   uint32_t length) {
 717   CHECK_EQ(stream->avail_out, 0);
 718
 719   auto buf = IOBuf::create(length);
 720   buf->append(length);
 721
 722   stream->next_out = buf->writableData();
 723   stream->avail_out = buf->length();
 724
 725   return buf;
 726 }
 727
 728 bool ZlibCodec::doInflate(z_stream* stream,
 729                           IOBuf* head,
 730                           uint32_t bufferLength) {
 731   if (stream->avail_out == 0) {
 732     head->prependChain(addOutputBuffer(stream, bufferLength));
 733   }
 734
 735   int rc = inflate(stream, Z_NO_FLUSH);
 736
 737   switch (rc) {
 738   case Z_OK:
 739     break;
 740   case Z_STREAM_END:
 741     return true;
 742   case Z_BUF_ERROR:
 743   case Z_NEED_DICT:
 744   case Z_DATA_ERROR:
 745   case Z_MEM_ERROR:
 746     throw std::runtime_error(to<std::string>(
 747         "ZlibCodec: inflate error: ", rc, ": ", stream->msg));
 748   default:
 749     CHECK(false) << rc << ": " << stream->msg;
 750   }
 751
 752   return false;
 753 }
 754
 755 std::unique_ptr<IOBuf> ZlibCodec::doCompress(const IOBuf* data) {
 756   z_stream stream;
 757   stream.zalloc = nullptr;
 758   stream.zfree = nullptr;
 759   stream.opaque = nullptr;
 760
 761   // Using deflateInit2() to support gzip.  "The windowBits parameter is the
 762   // base two logarithm of the maximum window size (...) The default value is
 763   // 15 (...) Add 16 to windowBits to write a simple gzip header and trailer
 764   // around the compressed data instead of a zlib wrapper. The gzip header
 765   // will have no file name, no extra data, no comment, no modification time
 766   // (set to zero), no header crc, and the operating system will be set to 255
 767   // (unknown)."
 768   int windowBits = 15 + (type() == CodecType::GZIP ? 16 : 0);
 769   // All other parameters (method, memLevel, strategy) get default values from
 770   // the zlib manual.
 771   int rc = deflateInit2(&stream,
 772                         level_,
 773                         Z_DEFLATED,
 774                         windowBits,
 775                         /* memLevel */ 8,
 776                         Z_DEFAULT_STRATEGY);
 777   if (rc != Z_OK) {
 778     throw std::runtime_error(to<std::string>(
 779         "ZlibCodec: deflateInit error: ", rc, ": ", stream.msg));
 780   }
 781
 782   stream.next_in = stream.next_out = nullptr;
 783   stream.avail_in = stream.avail_out = 0;
 784   stream.total_in = stream.total_out = 0;
 785
 786   bool success = false;
 787
 788   SCOPE_EXIT {
 789     rc = deflateEnd(&stream);
 790     // If we're here because of an exception, it's okay if some data
 791     // got dropped.
 792     CHECK(rc == Z_OK || (!success && rc == Z_DATA_ERROR))
 793       << rc << ": " << stream.msg;
 794   };
 795
 796   uint64_t uncompressedLength = data->computeChainDataLength();
 797   uint64_t maxCompressedLength = deflateBound(&stream, uncompressedLength);
 798
 799   // Max 64MiB in one go
 800   constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20;    // 64MiB
 801   constexpr uint32_t defaultBufferLength = uint32_t(4) << 20;     // 4MiB
 802
 803   auto out = addOutputBuffer(
 804       &stream,
 805       (maxCompressedLength <= maxSingleStepLength ?
 806        maxCompressedLength :
 807        defaultBufferLength));
 808
 809   for (auto& range : *data) {
 810     uint64_t remaining = range.size();
 811     uint64_t written = 0;
 812     while (remaining) {
 813       uint32_t step = (remaining > maxSingleStepLength ?
 814                        maxSingleStepLength : remaining);
 815       stream.next_in = const_cast<uint8_t*>(range.data() + written);
 816       stream.avail_in = step;
 817       remaining -= step;
 818       written += step;
 819
 820       while (stream.avail_in != 0) {
 821         if (stream.avail_out == 0) {
 822           out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
 823         }
 824
 825         rc = deflate(&stream, Z_NO_FLUSH);
 826
 827         CHECK_EQ(rc, Z_OK) << stream.msg;
 828       }
 829     }
 830   }
 831
 832   do {
 833     if (stream.avail_out == 0) {
 834       out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
 835     }
 836
 837     rc = deflate(&stream, Z_FINISH);
 838   } while (rc == Z_OK);
 839
 840   CHECK_EQ(rc, Z_STREAM_END) << stream.msg;
 841
 842   out->prev()->trimEnd(stream.avail_out);
 843
 844   success = true;  // we survived
 845
 846   return out;
 847 }
 848
 849 std::unique_ptr<IOBuf> ZlibCodec::doUncompress(const IOBuf* data,
 850                                                uint64_t uncompressedLength) {
 851   z_stream stream;
 852   stream.zalloc = nullptr;
 853   stream.zfree = nullptr;
 854   stream.opaque = nullptr;
 855
 856   // "The windowBits parameter is the base two logarithm of the maximum window
 857   // size (...) The default value is 15 (...) add 16 to decode only the gzip
 858   // format (the zlib format will return a Z_DATA_ERROR)."
 859   int windowBits = 15 + (type() == CodecType::GZIP ? 16 : 0);
 860   int rc = inflateInit2(&stream, windowBits);
 861   if (rc != Z_OK) {
 862     throw std::runtime_error(to<std::string>(
 863         "ZlibCodec: inflateInit error: ", rc, ": ", stream.msg));
 864   }
 865
 866   stream.next_in = stream.next_out = nullptr;
 867   stream.avail_in = stream.avail_out = 0;
 868   stream.total_in = stream.total_out = 0;
 869
 870   bool success = false;
 871
 872   SCOPE_EXIT {
 873     rc = inflateEnd(&stream);
 874     // If we're here because of an exception, it's okay if some data
 875     // got dropped.
 876     CHECK(rc == Z_OK || (!success && rc == Z_DATA_ERROR))
 877       << rc << ": " << stream.msg;
 878   };
 879
 880   // Max 64MiB in one go
 881   constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20;    // 64MiB
 882   constexpr uint32_t defaultBufferLength = uint32_t(4) << 20;     // 4MiB
 883
 884   auto out = addOutputBuffer(
 885       &stream,
 886       ((uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
 887         uncompressedLength <= maxSingleStepLength) ?
 888        uncompressedLength :
 889        defaultBufferLength));
 890
 891   bool streamEnd = false;
 892   for (auto& range : *data) {
 893     if (range.empty()) {
 894       continue;
 895     }
 896
 897     stream.next_in = const_cast<uint8_t*>(range.data());
 898     stream.avail_in = range.size();
 899
 900     while (stream.avail_in != 0) {
 901       if (streamEnd) {
 902         throw std::runtime_error(to<std::string>(
 903             "ZlibCodec: junk after end of data"));
 904       }
 905
 906       streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
 907     }
 908   }
 909
 910   while (!streamEnd) {
 911     streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
 912   }
 913
 914   out->prev()->trimEnd(stream.avail_out);
 915
 916   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
 917       uncompressedLength != stream.total_out) {
 918     throw std::runtime_error(to<std::string>(
 919         "ZlibCodec: invalid uncompressed length"));
 920   }
 921
 922   success = true;  // we survived
 923
 924   return out;
 925 }
 926
 927 #endif  // FOLLY_HAVE_LIBZ
 928
 929 #if FOLLY_HAVE_LIBLZMA
 930
 931 /**
 932  * LZMA2 compression
 933  */
 934 class LZMA2Codec final : public Codec {
 935  public:
 936   static std::unique_ptr<Codec> create(int level, CodecType type);
 937   explicit LZMA2Codec(int level, CodecType type);
 938
 939  private:
 940   bool doNeedsUncompressedLength() const override;
 941   uint64_t doMaxUncompressedLength() const override;
 942
 943   bool encodeSize() const { return type() == CodecType::LZMA2_VARINT_SIZE; }
 944
 945   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
 946   std::unique_ptr<IOBuf> doUncompress(
 947       const IOBuf* data,
 948       uint64_t uncompressedLength) override;
 949
 950   std::unique_ptr<IOBuf> addOutputBuffer(lzma_stream* stream, size_t length);
 951   bool doInflate(lzma_stream* stream, IOBuf* head, size_t bufferLength);
 952
 953   int level_;
 954 };
 955
 956 std::unique_ptr<Codec> LZMA2Codec::create(int level, CodecType type) {
 957   return make_unique<LZMA2Codec>(level, type);
 958 }
 959
 960 LZMA2Codec::LZMA2Codec(int level, CodecType type) : Codec(type) {
 961   DCHECK(type == CodecType::LZMA2 || type == CodecType::LZMA2_VARINT_SIZE);
 962   switch (level) {
 963   case COMPRESSION_LEVEL_FASTEST:
 964     level = 0;
 965     break;
 966   case COMPRESSION_LEVEL_DEFAULT:
 967     level = LZMA_PRESET_DEFAULT;
 968     break;
 969   case COMPRESSION_LEVEL_BEST:
 970     level = 9;
 971     break;
 972   }
 973   if (level < 0 || level > 9) {
 974     throw std::invalid_argument(to<std::string>(
 975         "LZMA2Codec: invalid level: ", level));
 976   }
 977   level_ = level;
 978 }
 979
 980 bool LZMA2Codec::doNeedsUncompressedLength() const {
 981   return false;
 982 }
 983
 984 uint64_t LZMA2Codec::doMaxUncompressedLength() const {
 985   // From lzma/base.h: "Stream is roughly 8 EiB (2^63 bytes)"
 986   return uint64_t(1) << 63;
 987 }
 988
 989 std::unique_ptr<IOBuf> LZMA2Codec::addOutputBuffer(
 990     lzma_stream* stream,
 991     size_t length) {
 992
 993   CHECK_EQ(stream->avail_out, 0);
 994
 995   auto buf = IOBuf::create(length);
 996   buf->append(length);
 997
 998   stream->next_out = buf->writableData();
 999   stream->avail_out = buf->length();
1000
1001   return buf;
1002 }
1003
1004 std::unique_ptr<IOBuf> LZMA2Codec::doCompress(const IOBuf* data) {
1005   lzma_ret rc;
1006   lzma_stream stream = LZMA_STREAM_INIT;
1007
1008   rc = lzma_easy_encoder(&stream, level_, LZMA_CHECK_NONE);
1009   if (rc != LZMA_OK) {
1010     throw std::runtime_error(folly::to<std::string>(
1011       "LZMA2Codec: lzma_easy_encoder error: ", rc));
1012   }
1013
1014   SCOPE_EXIT { lzma_end(&stream); };
1015
1016   uint64_t uncompressedLength = data->computeChainDataLength();
1017   uint64_t maxCompressedLength = lzma_stream_buffer_bound(uncompressedLength);
1018
1019   // Max 64MiB in one go
1020   constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20;    // 64MiB
1021   constexpr uint32_t defaultBufferLength = uint32_t(4) << 20;     // 4MiB
1022
1023   auto out = addOutputBuffer(
1024     &stream,
1025     (maxCompressedLength <= maxSingleStepLength ?
1026      maxCompressedLength :
1027      defaultBufferLength));
1028
1029   if (encodeSize()) {
1030     auto size = IOBuf::createCombined(kMaxVarintLength64);
1031     encodeVarintToIOBuf(uncompressedLength, size.get());
1032     size->appendChain(std::move(out));
1033     out = std::move(size);
1034   }
1035
1036   for (auto& range : *data) {
1037     if (range.empty()) {
1038       continue;
1039     }
1040
1041     stream.next_in = const_cast<uint8_t*>(range.data());
1042     stream.avail_in = range.size();
1043
1044     while (stream.avail_in != 0) {
1045       if (stream.avail_out == 0) {
1046         out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
1047       }
1048
1049       rc = lzma_code(&stream, LZMA_RUN);
1050
1051       if (rc != LZMA_OK) {
1052         throw std::runtime_error(folly::to<std::string>(
1053           "LZMA2Codec: lzma_code error: ", rc));
1054       }
1055     }
1056   }
1057
1058   do {
1059     if (stream.avail_out == 0) {
1060       out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
1061     }
1062
1063     rc = lzma_code(&stream, LZMA_FINISH);
1064   } while (rc == LZMA_OK);
1065
1066   if (rc != LZMA_STREAM_END) {
1067     throw std::runtime_error(folly::to<std::string>(
1068       "LZMA2Codec: lzma_code ended with error: ", rc));
1069   }
1070
1071   out->prev()->trimEnd(stream.avail_out);
1072
1073   return out;
1074 }
1075
1076 bool LZMA2Codec::doInflate(lzma_stream* stream,
1077                           IOBuf* head,
1078                           size_t bufferLength) {
1079   if (stream->avail_out == 0) {
1080     head->prependChain(addOutputBuffer(stream, bufferLength));
1081   }
1082
1083   lzma_ret rc = lzma_code(stream, LZMA_RUN);
1084
1085   switch (rc) {
1086   case LZMA_OK:
1087     break;
1088   case LZMA_STREAM_END:
1089     return true;
1090   default:
1091     throw std::runtime_error(to<std::string>(
1092         "LZMA2Codec: lzma_code error: ", rc));
1093   }
1094
1095   return false;
1096 }
1097
1098 std::unique_ptr<IOBuf> LZMA2Codec::doUncompress(const IOBuf* data,
1099                                                uint64_t uncompressedLength) {
1100   lzma_ret rc;
1101   lzma_stream stream = LZMA_STREAM_INIT;
1102
1103   rc = lzma_auto_decoder(&stream, std::numeric_limits<uint64_t>::max(), 0);
1104   if (rc != LZMA_OK) {
1105     throw std::runtime_error(folly::to<std::string>(
1106       "LZMA2Codec: lzma_auto_decoder error: ", rc));
1107   }
1108
1109   SCOPE_EXIT { lzma_end(&stream); };
1110
1111   // Max 64MiB in one go
1112   constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20; // 64MiB
1113   constexpr uint32_t defaultBufferLength = uint32_t(256) << 10; // 256 KiB
1114
1115   folly::io::Cursor cursor(data);
1116   if (encodeSize()) {
1117     const uint64_t actualUncompressedLength = decodeVarintFromCursor(cursor);
1118     if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
1119         uncompressedLength != actualUncompressedLength) {
1120       throw std::runtime_error("LZMA2Codec: invalid uncompressed length");
1121     }
1122     uncompressedLength = actualUncompressedLength;
1123   }
1124
1125   auto out = addOutputBuffer(
1126       &stream,
1127       ((uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
1128         uncompressedLength <= maxSingleStepLength)
1129            ? uncompressedLength
1130            : defaultBufferLength));
1131
1132   bool streamEnd = false;
1133   auto buf = cursor.peekBytes();
1134   while (!buf.empty()) {
1135     stream.next_in = const_cast<uint8_t*>(buf.data());
1136     stream.avail_in = buf.size();
1137
1138     while (stream.avail_in != 0) {
1139       if (streamEnd) {
1140         throw std::runtime_error(to<std::string>(
1141             "LZMA2Codec: junk after end of data"));
1142       }
1143
1144       streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
1145     }
1146
1147     cursor.skip(buf.size());
1148     buf = cursor.peekBytes();
1149   }
1150
1151   while (!streamEnd) {
1152     streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
1153   }
1154
1155   out->prev()->trimEnd(stream.avail_out);
1156
1157   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
1158       uncompressedLength != stream.total_out) {
1159     throw std::runtime_error(
1160         to<std::string>("LZMA2Codec: invalid uncompressed length"));
1161   }
1162
1163   return out;
1164 }
1165
1166 #endif  // FOLLY_HAVE_LIBLZMA
1167
1168 #ifdef FOLLY_HAVE_LIBZSTD
1169
1170 /**
1171  * ZSTD compression
1172  */
1173 class ZSTDCodec final : public Codec {
1174  public:
1175   static std::unique_ptr<Codec> create(int level, CodecType);
1176   explicit ZSTDCodec(int level, CodecType type);
1177
1178  private:
1179   bool doNeedsUncompressedLength() const override;
1180   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
1181   std::unique_ptr<IOBuf> doUncompress(
1182       const IOBuf* data,
1183       uint64_t uncompressedLength) override;
1184
1185   int level_;
1186 };
1187
1188 std::unique_ptr<Codec> ZSTDCodec::create(int level, CodecType type) {
1189   return make_unique<ZSTDCodec>(level, type);
1190 }
1191
1192 ZSTDCodec::ZSTDCodec(int level, CodecType type) : Codec(type) {
1193   DCHECK(type == CodecType::ZSTD);
1194   switch (level) {
1195     case COMPRESSION_LEVEL_FASTEST:
1196       level = 1;
1197       break;
1198     case COMPRESSION_LEVEL_DEFAULT:
1199       level = 1;
1200       break;
1201     case COMPRESSION_LEVEL_BEST:
1202       level = 19;
1203       break;
1204   }
1205   if (level < 1 || level > ZSTD_maxCLevel()) {
1206     throw std::invalid_argument(
1207         to<std::string>("ZSTD: invalid level: ", level));
1208   }
1209   level_ = level;
1210 }
1211
1212 bool ZSTDCodec::doNeedsUncompressedLength() const {
1213   return false;
1214 }
1215
1216 void zstdThrowIfError(size_t rc) {
1217   if (!ZSTD_isError(rc)) {
1218     return;
1219   }
1220   throw std::runtime_error(
1221       to<std::string>("ZSTD returned an error: ", ZSTD_getErrorName(rc)));
1222 }
1223
1224 std::unique_ptr<IOBuf> ZSTDCodec::doCompress(const IOBuf* data) {
1225   // Support earlier versions of the codec (working with a single IOBuf,
1226   // and using ZSTD_decompress which requires ZSTD frame to contain size,
1227   // which isn't populated by streaming API).
1228   if (!data->isChained()) {
1229     auto out = IOBuf::createCombined(ZSTD_compressBound(data->length()));
1230     const auto rc = ZSTD_compress(
1231         out->writableData(),
1232         out->capacity(),
1233         data->data(),
1234         data->length(),
1235         level_);
1236     zstdThrowIfError(rc);
1237     out->append(rc);
1238     return out;
1239   }
1240
1241   auto zcs = ZSTD_createCStream();
1242   SCOPE_EXIT {
1243     ZSTD_freeCStream(zcs);
1244   };
1245
1246   auto rc = ZSTD_initCStream(zcs, level_);
1247   zstdThrowIfError(rc);
1248
1249   Cursor cursor(data);
1250   auto result = IOBuf::createCombined(ZSTD_compressBound(cursor.totalLength()));
1251
1252   ZSTD_outBuffer out;
1253   out.dst = result->writableTail();
1254   out.size = result->capacity();
1255   out.pos = 0;
1256
1257   for (auto buffer = cursor.peekBytes(); !buffer.empty();) {
1258     ZSTD_inBuffer in;
1259     in.src = buffer.data();
1260     in.size = buffer.size();
1261     for (in.pos = 0; in.pos != in.size;) {
1262       rc = ZSTD_compressStream(zcs, &out, &in);
1263       zstdThrowIfError(rc);
1264     }
1265     cursor.skip(in.size);
1266     buffer = cursor.peekBytes();
1267   }
1268
1269   rc = ZSTD_endStream(zcs, &out);
1270   zstdThrowIfError(rc);
1271   CHECK_EQ(rc, 0);
1272
1273   result->append(out.pos);
1274   return result;
1275 }
1276
1277 static std::unique_ptr<IOBuf> zstdUncompressBuffer(
1278     const IOBuf* data,
1279     uint64_t uncompressedLength) {
1280   // Check preconditions
1281   DCHECK(!data->isChained());
1282   DCHECK(uncompressedLength != Codec::UNKNOWN_UNCOMPRESSED_LENGTH);
1283
1284   auto uncompressed = IOBuf::create(uncompressedLength);
1285   const auto decompressedSize = ZSTD_decompress(
1286       uncompressed->writableTail(),
1287       uncompressed->tailroom(),
1288       data->data(),
1289       data->length());
1290   zstdThrowIfError(decompressedSize);
1291   if (decompressedSize != uncompressedLength) {
1292     throw std::runtime_error("ZSTD: invalid uncompressed length");
1293   }
1294   uncompressed->append(decompressedSize);
1295   return uncompressed;
1296 }
1297
1298 static std::unique_ptr<IOBuf> zstdUncompressStream(
1299     const IOBuf* data,
1300     uint64_t uncompressedLength) {
1301   auto zds = ZSTD_createDStream();
1302   SCOPE_EXIT {
1303     ZSTD_freeDStream(zds);
1304   };
1305
1306   auto rc = ZSTD_initDStream(zds);
1307   zstdThrowIfError(rc);
1308
1309   ZSTD_outBuffer out{};
1310   ZSTD_inBuffer in{};
1311
1312   auto outputSize = ZSTD_DStreamOutSize();
1313   if (uncompressedLength != Codec::UNKNOWN_UNCOMPRESSED_LENGTH) {
1314     outputSize = uncompressedLength;
1315   }
1316
1317   IOBufQueue queue(IOBufQueue::cacheChainLength());
1318
1319   Cursor cursor(data);
1320   for (rc = 0;;) {
1321     if (in.pos == in.size) {
1322       auto buffer = cursor.peekBytes();
1323       in.src = buffer.data();
1324       in.size = buffer.size();
1325       in.pos = 0;
1326       cursor.skip(in.size);
1327       if (rc > 1 && in.size == 0) {
1328         throw std::runtime_error(to<std::string>("ZSTD: incomplete input"));
1329       }
1330     }
1331     if (out.pos == out.size) {
1332       if (out.pos != 0) {
1333         queue.postallocate(out.pos);
1334       }
1335       auto buffer = queue.preallocate(outputSize, outputSize);
1336       out.dst = buffer.first;
1337       out.size = buffer.second;
1338       out.pos = 0;
1339       outputSize = ZSTD_DStreamOutSize();
1340     }
1341     rc = ZSTD_decompressStream(zds, &out, &in);
1342     zstdThrowIfError(rc);
1343     if (rc == 0) {
1344       break;
1345     }
1346   }
1347   if (out.pos != 0) {
1348     queue.postallocate(out.pos);
1349   }
1350   if (in.pos != in.size || !cursor.isAtEnd()) {
1351     throw std::runtime_error("ZSTD: junk after end of data");
1352   }
1353   if (uncompressedLength != Codec::UNKNOWN_UNCOMPRESSED_LENGTH &&
1354       queue.chainLength() != uncompressedLength) {
1355     throw std::runtime_error("ZSTD: invalid uncompressed length");
1356   }
1357
1358   return queue.move();
1359 }
1360
1361 std::unique_ptr<IOBuf> ZSTDCodec::doUncompress(
1362     const IOBuf* data,
1363     uint64_t uncompressedLength) {
1364   {
1365     // Read decompressed size from frame if available in first IOBuf.
1366     const auto decompressedSize =
1367         ZSTD_getDecompressedSize(data->data(), data->length());
1368     if (decompressedSize != 0) {
1369       if (uncompressedLength != Codec::UNKNOWN_UNCOMPRESSED_LENGTH &&
1370           uncompressedLength != decompressedSize) {
1371         throw std::runtime_error("ZSTD: invalid uncompressed length");
1372       }
1373       uncompressedLength = decompressedSize;
1374     }
1375   }
1376   // Faster to decompress using ZSTD_decompress() if we can.
1377   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH && !data->isChained()) {
1378     return zstdUncompressBuffer(data, uncompressedLength);
1379   }
1380   // Fall back to slower streaming decompression.
1381   return zstdUncompressStream(data, uncompressedLength);
1382 }
1383
1384 #endif  // FOLLY_HAVE_LIBZSTD
1385
1386 }  // namespace
1387
1388 typedef std::unique_ptr<Codec> (*CodecFactory)(int, CodecType);
1389 static constexpr CodecFactory
1390     codecFactories[static_cast<size_t>(CodecType::NUM_CODEC_TYPES)] = {
1391         nullptr, // USER_DEFINED
1392         NoCompressionCodec::create,
1393
1394 #if FOLLY_HAVE_LIBLZ4
1395         LZ4Codec::create,
1396 #else
1397         nullptr,
1398 #endif
1399
1400 #if FOLLY_HAVE_LIBSNAPPY
1401         SnappyCodec::create,
1402 #else
1403         nullptr,
1404 #endif
1405
1406 #if FOLLY_HAVE_LIBZ
1407         ZlibCodec::create,
1408 #else
1409         nullptr,
1410 #endif
1411
1412 #if FOLLY_HAVE_LIBLZ4
1413         LZ4Codec::create,
1414 #else
1415         nullptr,
1416 #endif
1417
1418 #if FOLLY_HAVE_LIBLZMA
1419         LZMA2Codec::create,
1420         LZMA2Codec::create,
1421 #else
1422         nullptr,
1423         nullptr,
1424 #endif
1425
1426 #if FOLLY_HAVE_LIBZSTD
1427         ZSTDCodec::create,
1428 #else
1429         nullptr,
1430 #endif
1431
1432 #if FOLLY_HAVE_LIBZ
1433         ZlibCodec::create,
1434 #else
1435         nullptr,
1436 #endif
1437
1438 #if (FOLLY_HAVE_LIBLZ4 && LZ4_VERSION_NUMBER >= 10301)
1439         LZ4FrameCodec::create,
1440 #else
1441         nullptr,
1442 #endif
1443 };
1444
1445 bool hasCodec(CodecType type) {
1446   size_t idx = static_cast<size_t>(type);
1447   if (idx >= static_cast<size_t>(CodecType::NUM_CODEC_TYPES)) {
1448     throw std::invalid_argument(
1449         to<std::string>("Compression type ", idx, " invalid"));
1450   }
1451   return codecFactories[idx] != nullptr;
1452 }
1453
1454 std::unique_ptr<Codec> getCodec(CodecType type, int level) {
1455   size_t idx = static_cast<size_t>(type);
1456   if (idx >= static_cast<size_t>(CodecType::NUM_CODEC_TYPES)) {
1457     throw std::invalid_argument(
1458         to<std::string>("Compression type ", idx, " invalid"));
1459   }
1460   auto factory = codecFactories[idx];
1461   if (!factory) {
1462     throw std::invalid_argument(to<std::string>(
1463         "Compression type ", idx, " not supported"));
1464   }
1465   auto codec = (*factory)(level, type);
1466   DCHECK_EQ(static_cast<size_t>(codec->type()), idx);
1467   return codec;
1468 }
1469
1470 }}  // namespaces