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 static uint64_t computeBufferLength(uint64_t const compressedLength) {
 850   constexpr uint64_t kMaxBufferLength = uint64_t(4) << 20; // 4 MiB
 851   constexpr uint64_t kBlockSize = uint64_t(32) << 10; // 32 KiB
 852   const uint64_t goodBufferSize = 4 * std::max(kBlockSize, compressedLength);
 853   return std::min(goodBufferSize, kMaxBufferLength);
 854 }
 855
 856 std::unique_ptr<IOBuf> ZlibCodec::doUncompress(const IOBuf* data,
 857                                                uint64_t uncompressedLength) {
 858   z_stream stream;
 859   stream.zalloc = nullptr;
 860   stream.zfree = nullptr;
 861   stream.opaque = nullptr;
 862
 863   // "The windowBits parameter is the base two logarithm of the maximum window
 864   // size (...) The default value is 15 (...) add 16 to decode only the gzip
 865   // format (the zlib format will return a Z_DATA_ERROR)."
 866   int windowBits = 15 + (type() == CodecType::GZIP ? 16 : 0);
 867   int rc = inflateInit2(&stream, windowBits);
 868   if (rc != Z_OK) {
 869     throw std::runtime_error(to<std::string>(
 870         "ZlibCodec: inflateInit error: ", rc, ": ", stream.msg));
 871   }
 872
 873   stream.next_in = stream.next_out = nullptr;
 874   stream.avail_in = stream.avail_out = 0;
 875   stream.total_in = stream.total_out = 0;
 876
 877   bool success = false;
 878
 879   SCOPE_EXIT {
 880     rc = inflateEnd(&stream);
 881     // If we're here because of an exception, it's okay if some data
 882     // got dropped.
 883     CHECK(rc == Z_OK || (!success && rc == Z_DATA_ERROR))
 884       << rc << ": " << stream.msg;
 885   };
 886
 887   // Max 64MiB in one go
 888   constexpr uint64_t maxSingleStepLength = uint64_t(64) << 20; // 64MiB
 889   const uint64_t defaultBufferLength =
 890       computeBufferLength(data->computeChainDataLength());
 891
 892   auto out = addOutputBuffer(
 893       &stream,
 894       ((uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
 895         uncompressedLength <= maxSingleStepLength) ?
 896        uncompressedLength :
 897        defaultBufferLength));
 898
 899   bool streamEnd = false;
 900   for (auto& range : *data) {
 901     if (range.empty()) {
 902       continue;
 903     }
 904
 905     stream.next_in = const_cast<uint8_t*>(range.data());
 906     stream.avail_in = range.size();
 907
 908     while (stream.avail_in != 0) {
 909       if (streamEnd) {
 910         throw std::runtime_error(to<std::string>(
 911             "ZlibCodec: junk after end of data"));
 912       }
 913
 914       streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
 915     }
 916   }
 917
 918   while (!streamEnd) {
 919     streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
 920   }
 921
 922   out->prev()->trimEnd(stream.avail_out);
 923
 924   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
 925       uncompressedLength != stream.total_out) {
 926     throw std::runtime_error(to<std::string>(
 927         "ZlibCodec: invalid uncompressed length"));
 928   }
 929
 930   success = true;  // we survived
 931
 932   return out;
 933 }
 934
 935 #endif  // FOLLY_HAVE_LIBZ
 936
 937 #if FOLLY_HAVE_LIBLZMA
 938
 939 /**
 940  * LZMA2 compression
 941  */
 942 class LZMA2Codec final : public Codec {
 943  public:
 944   static std::unique_ptr<Codec> create(int level, CodecType type);
 945   explicit LZMA2Codec(int level, CodecType type);
 946
 947  private:
 948   bool doNeedsUncompressedLength() const override;
 949   uint64_t doMaxUncompressedLength() const override;
 950
 951   bool encodeSize() const { return type() == CodecType::LZMA2_VARINT_SIZE; }
 952
 953   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
 954   std::unique_ptr<IOBuf> doUncompress(
 955       const IOBuf* data,
 956       uint64_t uncompressedLength) override;
 957
 958   std::unique_ptr<IOBuf> addOutputBuffer(lzma_stream* stream, size_t length);
 959   bool doInflate(lzma_stream* stream, IOBuf* head, size_t bufferLength);
 960
 961   int level_;
 962 };
 963
 964 std::unique_ptr<Codec> LZMA2Codec::create(int level, CodecType type) {
 965   return make_unique<LZMA2Codec>(level, type);
 966 }
 967
 968 LZMA2Codec::LZMA2Codec(int level, CodecType type) : Codec(type) {
 969   DCHECK(type == CodecType::LZMA2 || type == CodecType::LZMA2_VARINT_SIZE);
 970   switch (level) {
 971   case COMPRESSION_LEVEL_FASTEST:
 972     level = 0;
 973     break;
 974   case COMPRESSION_LEVEL_DEFAULT:
 975     level = LZMA_PRESET_DEFAULT;
 976     break;
 977   case COMPRESSION_LEVEL_BEST:
 978     level = 9;
 979     break;
 980   }
 981   if (level < 0 || level > 9) {
 982     throw std::invalid_argument(to<std::string>(
 983         "LZMA2Codec: invalid level: ", level));
 984   }
 985   level_ = level;
 986 }
 987
 988 bool LZMA2Codec::doNeedsUncompressedLength() const {
 989   return false;
 990 }
 991
 992 uint64_t LZMA2Codec::doMaxUncompressedLength() const {
 993   // From lzma/base.h: "Stream is roughly 8 EiB (2^63 bytes)"
 994   return uint64_t(1) << 63;
 995 }
 996
 997 std::unique_ptr<IOBuf> LZMA2Codec::addOutputBuffer(
 998     lzma_stream* stream,
 999     size_t length) {
1000
1001   CHECK_EQ(stream->avail_out, 0);
1002
1003   auto buf = IOBuf::create(length);
1004   buf->append(length);
1005
1006   stream->next_out = buf->writableData();
1007   stream->avail_out = buf->length();
1008
1009   return buf;
1010 }
1011
1012 std::unique_ptr<IOBuf> LZMA2Codec::doCompress(const IOBuf* data) {
1013   lzma_ret rc;
1014   lzma_stream stream = LZMA_STREAM_INIT;
1015
1016   rc = lzma_easy_encoder(&stream, level_, LZMA_CHECK_NONE);
1017   if (rc != LZMA_OK) {
1018     throw std::runtime_error(folly::to<std::string>(
1019       "LZMA2Codec: lzma_easy_encoder error: ", rc));
1020   }
1021
1022   SCOPE_EXIT { lzma_end(&stream); };
1023
1024   uint64_t uncompressedLength = data->computeChainDataLength();
1025   uint64_t maxCompressedLength = lzma_stream_buffer_bound(uncompressedLength);
1026
1027   // Max 64MiB in one go
1028   constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20;    // 64MiB
1029   constexpr uint32_t defaultBufferLength = uint32_t(4) << 20;     // 4MiB
1030
1031   auto out = addOutputBuffer(
1032     &stream,
1033     (maxCompressedLength <= maxSingleStepLength ?
1034      maxCompressedLength :
1035      defaultBufferLength));
1036
1037   if (encodeSize()) {
1038     auto size = IOBuf::createCombined(kMaxVarintLength64);
1039     encodeVarintToIOBuf(uncompressedLength, size.get());
1040     size->appendChain(std::move(out));
1041     out = std::move(size);
1042   }
1043
1044   for (auto& range : *data) {
1045     if (range.empty()) {
1046       continue;
1047     }
1048
1049     stream.next_in = const_cast<uint8_t*>(range.data());
1050     stream.avail_in = range.size();
1051
1052     while (stream.avail_in != 0) {
1053       if (stream.avail_out == 0) {
1054         out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
1055       }
1056
1057       rc = lzma_code(&stream, LZMA_RUN);
1058
1059       if (rc != LZMA_OK) {
1060         throw std::runtime_error(folly::to<std::string>(
1061           "LZMA2Codec: lzma_code error: ", rc));
1062       }
1063     }
1064   }
1065
1066   do {
1067     if (stream.avail_out == 0) {
1068       out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
1069     }
1070
1071     rc = lzma_code(&stream, LZMA_FINISH);
1072   } while (rc == LZMA_OK);
1073
1074   if (rc != LZMA_STREAM_END) {
1075     throw std::runtime_error(folly::to<std::string>(
1076       "LZMA2Codec: lzma_code ended with error: ", rc));
1077   }
1078
1079   out->prev()->trimEnd(stream.avail_out);
1080
1081   return out;
1082 }
1083
1084 bool LZMA2Codec::doInflate(lzma_stream* stream,
1085                           IOBuf* head,
1086                           size_t bufferLength) {
1087   if (stream->avail_out == 0) {
1088     head->prependChain(addOutputBuffer(stream, bufferLength));
1089   }
1090
1091   lzma_ret rc = lzma_code(stream, LZMA_RUN);
1092
1093   switch (rc) {
1094   case LZMA_OK:
1095     break;
1096   case LZMA_STREAM_END:
1097     return true;
1098   default:
1099     throw std::runtime_error(to<std::string>(
1100         "LZMA2Codec: lzma_code error: ", rc));
1101   }
1102
1103   return false;
1104 }
1105
1106 std::unique_ptr<IOBuf> LZMA2Codec::doUncompress(const IOBuf* data,
1107                                                uint64_t uncompressedLength) {
1108   lzma_ret rc;
1109   lzma_stream stream = LZMA_STREAM_INIT;
1110
1111   rc = lzma_auto_decoder(&stream, std::numeric_limits<uint64_t>::max(), 0);
1112   if (rc != LZMA_OK) {
1113     throw std::runtime_error(folly::to<std::string>(
1114       "LZMA2Codec: lzma_auto_decoder error: ", rc));
1115   }
1116
1117   SCOPE_EXIT { lzma_end(&stream); };
1118
1119   // Max 64MiB in one go
1120   constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20; // 64MiB
1121   constexpr uint32_t defaultBufferLength = uint32_t(256) << 10; // 256 KiB
1122
1123   folly::io::Cursor cursor(data);
1124   if (encodeSize()) {
1125     const uint64_t actualUncompressedLength = decodeVarintFromCursor(cursor);
1126     if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
1127         uncompressedLength != actualUncompressedLength) {
1128       throw std::runtime_error("LZMA2Codec: invalid uncompressed length");
1129     }
1130     uncompressedLength = actualUncompressedLength;
1131   }
1132
1133   auto out = addOutputBuffer(
1134       &stream,
1135       ((uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
1136         uncompressedLength <= maxSingleStepLength)
1137            ? uncompressedLength
1138            : defaultBufferLength));
1139
1140   bool streamEnd = false;
1141   auto buf = cursor.peekBytes();
1142   while (!buf.empty()) {
1143     stream.next_in = const_cast<uint8_t*>(buf.data());
1144     stream.avail_in = buf.size();
1145
1146     while (stream.avail_in != 0) {
1147       if (streamEnd) {
1148         throw std::runtime_error(to<std::string>(
1149             "LZMA2Codec: junk after end of data"));
1150       }
1151
1152       streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
1153     }
1154
1155     cursor.skip(buf.size());
1156     buf = cursor.peekBytes();
1157   }
1158
1159   while (!streamEnd) {
1160     streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
1161   }
1162
1163   out->prev()->trimEnd(stream.avail_out);
1164
1165   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
1166       uncompressedLength != stream.total_out) {
1167     throw std::runtime_error(
1168         to<std::string>("LZMA2Codec: invalid uncompressed length"));
1169   }
1170
1171   return out;
1172 }
1173
1174 #endif  // FOLLY_HAVE_LIBLZMA
1175
1176 #ifdef FOLLY_HAVE_LIBZSTD
1177
1178 /**
1179  * ZSTD compression
1180  */
1181 class ZSTDCodec final : public Codec {
1182  public:
1183   static std::unique_ptr<Codec> create(int level, CodecType);
1184   explicit ZSTDCodec(int level, CodecType type);
1185
1186  private:
1187   bool doNeedsUncompressedLength() const override;
1188   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
1189   std::unique_ptr<IOBuf> doUncompress(
1190       const IOBuf* data,
1191       uint64_t uncompressedLength) override;
1192
1193   int level_;
1194 };
1195
1196 std::unique_ptr<Codec> ZSTDCodec::create(int level, CodecType type) {
1197   return make_unique<ZSTDCodec>(level, type);
1198 }
1199
1200 ZSTDCodec::ZSTDCodec(int level, CodecType type) : Codec(type) {
1201   DCHECK(type == CodecType::ZSTD);
1202   switch (level) {
1203     case COMPRESSION_LEVEL_FASTEST:
1204       level = 1;
1205       break;
1206     case COMPRESSION_LEVEL_DEFAULT:
1207       level = 1;
1208       break;
1209     case COMPRESSION_LEVEL_BEST:
1210       level = 19;
1211       break;
1212   }
1213   if (level < 1 || level > ZSTD_maxCLevel()) {
1214     throw std::invalid_argument(
1215         to<std::string>("ZSTD: invalid level: ", level));
1216   }
1217   level_ = level;
1218 }
1219
1220 bool ZSTDCodec::doNeedsUncompressedLength() const {
1221   return false;
1222 }
1223
1224 void zstdThrowIfError(size_t rc) {
1225   if (!ZSTD_isError(rc)) {
1226     return;
1227   }
1228   throw std::runtime_error(
1229       to<std::string>("ZSTD returned an error: ", ZSTD_getErrorName(rc)));
1230 }
1231
1232 std::unique_ptr<IOBuf> ZSTDCodec::doCompress(const IOBuf* data) {
1233   // Support earlier versions of the codec (working with a single IOBuf,
1234   // and using ZSTD_decompress which requires ZSTD frame to contain size,
1235   // which isn't populated by streaming API).
1236   if (!data->isChained()) {
1237     auto out = IOBuf::createCombined(ZSTD_compressBound(data->length()));
1238     const auto rc = ZSTD_compress(
1239         out->writableData(),
1240         out->capacity(),
1241         data->data(),
1242         data->length(),
1243         level_);
1244     zstdThrowIfError(rc);
1245     out->append(rc);
1246     return out;
1247   }
1248
1249   auto zcs = ZSTD_createCStream();
1250   SCOPE_EXIT {
1251     ZSTD_freeCStream(zcs);
1252   };
1253
1254   auto rc = ZSTD_initCStream(zcs, level_);
1255   zstdThrowIfError(rc);
1256
1257   Cursor cursor(data);
1258   auto result = IOBuf::createCombined(ZSTD_compressBound(cursor.totalLength()));
1259
1260   ZSTD_outBuffer out;
1261   out.dst = result->writableTail();
1262   out.size = result->capacity();
1263   out.pos = 0;
1264
1265   for (auto buffer = cursor.peekBytes(); !buffer.empty();) {
1266     ZSTD_inBuffer in;
1267     in.src = buffer.data();
1268     in.size = buffer.size();
1269     for (in.pos = 0; in.pos != in.size;) {
1270       rc = ZSTD_compressStream(zcs, &out, &in);
1271       zstdThrowIfError(rc);
1272     }
1273     cursor.skip(in.size);
1274     buffer = cursor.peekBytes();
1275   }
1276
1277   rc = ZSTD_endStream(zcs, &out);
1278   zstdThrowIfError(rc);
1279   CHECK_EQ(rc, 0);
1280
1281   result->append(out.pos);
1282   return result;
1283 }
1284
1285 static std::unique_ptr<IOBuf> zstdUncompressBuffer(
1286     const IOBuf* data,
1287     uint64_t uncompressedLength) {
1288   // Check preconditions
1289   DCHECK(!data->isChained());
1290   DCHECK(uncompressedLength != Codec::UNKNOWN_UNCOMPRESSED_LENGTH);
1291
1292   auto uncompressed = IOBuf::create(uncompressedLength);
1293   const auto decompressedSize = ZSTD_decompress(
1294       uncompressed->writableTail(),
1295       uncompressed->tailroom(),
1296       data->data(),
1297       data->length());
1298   zstdThrowIfError(decompressedSize);
1299   if (decompressedSize != uncompressedLength) {
1300     throw std::runtime_error("ZSTD: invalid uncompressed length");
1301   }
1302   uncompressed->append(decompressedSize);
1303   return uncompressed;
1304 }
1305
1306 static std::unique_ptr<IOBuf> zstdUncompressStream(
1307     const IOBuf* data,
1308     uint64_t uncompressedLength) {
1309   auto zds = ZSTD_createDStream();
1310   SCOPE_EXIT {
1311     ZSTD_freeDStream(zds);
1312   };
1313
1314   auto rc = ZSTD_initDStream(zds);
1315   zstdThrowIfError(rc);
1316
1317   ZSTD_outBuffer out{};
1318   ZSTD_inBuffer in{};
1319
1320   auto outputSize = ZSTD_DStreamOutSize();
1321   if (uncompressedLength != Codec::UNKNOWN_UNCOMPRESSED_LENGTH) {
1322     outputSize = uncompressedLength;
1323   }
1324
1325   IOBufQueue queue(IOBufQueue::cacheChainLength());
1326
1327   Cursor cursor(data);
1328   for (rc = 0;;) {
1329     if (in.pos == in.size) {
1330       auto buffer = cursor.peekBytes();
1331       in.src = buffer.data();
1332       in.size = buffer.size();
1333       in.pos = 0;
1334       cursor.skip(in.size);
1335       if (rc > 1 && in.size == 0) {
1336         throw std::runtime_error(to<std::string>("ZSTD: incomplete input"));
1337       }
1338     }
1339     if (out.pos == out.size) {
1340       if (out.pos != 0) {
1341         queue.postallocate(out.pos);
1342       }
1343       auto buffer = queue.preallocate(outputSize, outputSize);
1344       out.dst = buffer.first;
1345       out.size = buffer.second;
1346       out.pos = 0;
1347       outputSize = ZSTD_DStreamOutSize();
1348     }
1349     rc = ZSTD_decompressStream(zds, &out, &in);
1350     zstdThrowIfError(rc);
1351     if (rc == 0) {
1352       break;
1353     }
1354   }
1355   if (out.pos != 0) {
1356     queue.postallocate(out.pos);
1357   }
1358   if (in.pos != in.size || !cursor.isAtEnd()) {
1359     throw std::runtime_error("ZSTD: junk after end of data");
1360   }
1361   if (uncompressedLength != Codec::UNKNOWN_UNCOMPRESSED_LENGTH &&
1362       queue.chainLength() != uncompressedLength) {
1363     throw std::runtime_error("ZSTD: invalid uncompressed length");
1364   }
1365
1366   return queue.move();
1367 }
1368
1369 std::unique_ptr<IOBuf> ZSTDCodec::doUncompress(
1370     const IOBuf* data,
1371     uint64_t uncompressedLength) {
1372   {
1373     // Read decompressed size from frame if available in first IOBuf.
1374     const auto decompressedSize =
1375         ZSTD_getDecompressedSize(data->data(), data->length());
1376     if (decompressedSize != 0) {
1377       if (uncompressedLength != Codec::UNKNOWN_UNCOMPRESSED_LENGTH &&
1378           uncompressedLength != decompressedSize) {
1379         throw std::runtime_error("ZSTD: invalid uncompressed length");
1380       }
1381       uncompressedLength = decompressedSize;
1382     }
1383   }
1384   // Faster to decompress using ZSTD_decompress() if we can.
1385   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH && !data->isChained()) {
1386     return zstdUncompressBuffer(data, uncompressedLength);
1387   }
1388   // Fall back to slower streaming decompression.
1389   return zstdUncompressStream(data, uncompressedLength);
1390 }
1391
1392 #endif  // FOLLY_HAVE_LIBZSTD
1393
1394 }  // namespace
1395
1396 typedef std::unique_ptr<Codec> (*CodecFactory)(int, CodecType);
1397 static constexpr CodecFactory
1398     codecFactories[static_cast<size_t>(CodecType::NUM_CODEC_TYPES)] = {
1399         nullptr, // USER_DEFINED
1400         NoCompressionCodec::create,
1401
1402 #if FOLLY_HAVE_LIBLZ4
1403         LZ4Codec::create,
1404 #else
1405         nullptr,
1406 #endif
1407
1408 #if FOLLY_HAVE_LIBSNAPPY
1409         SnappyCodec::create,
1410 #else
1411         nullptr,
1412 #endif
1413
1414 #if FOLLY_HAVE_LIBZ
1415         ZlibCodec::create,
1416 #else
1417         nullptr,
1418 #endif
1419
1420 #if FOLLY_HAVE_LIBLZ4
1421         LZ4Codec::create,
1422 #else
1423         nullptr,
1424 #endif
1425
1426 #if FOLLY_HAVE_LIBLZMA
1427         LZMA2Codec::create,
1428         LZMA2Codec::create,
1429 #else
1430         nullptr,
1431         nullptr,
1432 #endif
1433
1434 #if FOLLY_HAVE_LIBZSTD
1435         ZSTDCodec::create,
1436 #else
1437         nullptr,
1438 #endif
1439
1440 #if FOLLY_HAVE_LIBZ
1441         ZlibCodec::create,
1442 #else
1443         nullptr,
1444 #endif
1445
1446 #if (FOLLY_HAVE_LIBLZ4 && LZ4_VERSION_NUMBER >= 10301)
1447         LZ4FrameCodec::create,
1448 #else
1449         nullptr,
1450 #endif
1451 };
1452
1453 bool hasCodec(CodecType type) {
1454   size_t idx = static_cast<size_t>(type);
1455   if (idx >= static_cast<size_t>(CodecType::NUM_CODEC_TYPES)) {
1456     throw std::invalid_argument(
1457         to<std::string>("Compression type ", idx, " invalid"));
1458   }
1459   return codecFactories[idx] != nullptr;
1460 }
1461
1462 std::unique_ptr<Codec> getCodec(CodecType type, int level) {
1463   size_t idx = static_cast<size_t>(type);
1464   if (idx >= static_cast<size_t>(CodecType::NUM_CODEC_TYPES)) {
1465     throw std::invalid_argument(
1466         to<std::string>("Compression type ", idx, " invalid"));
1467   }
1468   auto factory = codecFactories[idx];
1469   if (!factory) {
1470     throw std::invalid_argument(to<std::string>(
1471         "Compression type ", idx, " not supported"));
1472   }
1473   auto codec = (*factory)(level, type);
1474   DCHECK_EQ(static_cast<size_t>(codec->type()), idx);
1475   return codec;
1476 }
1477
1478 }}  // namespaces