folly/compression/Compression.cpp

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