folly/io/Compression.cpp

   1 /*
   2  * Copyright 2014 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 #endif
  23
  24 #include <glog/logging.h>
  25
  26 #if FOLLY_HAVE_LIBSNAPPY
  27 #include <snappy.h>
  28 #include <snappy-sinksource.h>
  29 #endif
  30
  31 #if FOLLY_HAVE_LIBZ
  32 #include <zlib.h>
  33 #endif
  34
  35 #if FOLLY_HAVE_LIBLZMA
  36 #include <lzma.h>
  37 #endif
  38
  39 #include <folly/Conv.h>
  40 #include <folly/Memory.h>
  41 #include <folly/Portability.h>
  42 #include <folly/ScopeGuard.h>
  43 #include <folly/Varint.h>
  44 #include <folly/io/Cursor.h>
  45
  46 namespace folly { namespace io {
  47
  48 Codec::Codec(CodecType type) : type_(type) { }
  49
  50 // Ensure consistent behavior in the nullptr case
  51 std::unique_ptr<IOBuf> Codec::compress(const IOBuf* data) {
  52   uint64_t len = data->computeChainDataLength();
  53   if (len == 0) {
  54     return IOBuf::create(0);
  55   } else if (len > maxUncompressedLength()) {
  56     throw std::runtime_error("Codec: uncompressed length too large");
  57   }
  58
  59   return doCompress(data);
  60 }
  61
  62 std::unique_ptr<IOBuf> Codec::uncompress(const IOBuf* data,
  63                                          uint64_t uncompressedLength) {
  64   if (uncompressedLength == UNKNOWN_UNCOMPRESSED_LENGTH) {
  65     if (needsUncompressedLength()) {
  66       throw std::invalid_argument("Codec: uncompressed length required");
  67     }
  68   } else if (uncompressedLength > maxUncompressedLength()) {
  69     throw std::runtime_error("Codec: uncompressed length too large");
  70   }
  71
  72   if (data->empty()) {
  73     if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
  74         uncompressedLength != 0) {
  75       throw std::runtime_error("Codec: invalid uncompressed length");
  76     }
  77     return IOBuf::create(0);
  78   }
  79
  80   return doUncompress(data, uncompressedLength);
  81 }
  82
  83 bool Codec::needsUncompressedLength() const {
  84   return doNeedsUncompressedLength();
  85 }
  86
  87 uint64_t Codec::maxUncompressedLength() const {
  88   return doMaxUncompressedLength();
  89 }
  90
  91 bool Codec::doNeedsUncompressedLength() const {
  92   return false;
  93 }
  94
  95 uint64_t Codec::doMaxUncompressedLength() const {
  96   return UNLIMITED_UNCOMPRESSED_LENGTH;
  97 }
  98
  99 namespace {
 100
 101 /**
 102  * No compression
 103  */
 104 class NoCompressionCodec FOLLY_FINAL : public Codec {
 105  public:
 106   static std::unique_ptr<Codec> create(int level, CodecType type);
 107   explicit NoCompressionCodec(int level, CodecType type);
 108
 109  private:
 110   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) FOLLY_OVERRIDE;
 111   std::unique_ptr<IOBuf> doUncompress(
 112       const IOBuf* data,
 113       uint64_t uncompressedLength) FOLLY_OVERRIDE;
 114 };
 115
 116 std::unique_ptr<Codec> NoCompressionCodec::create(int level, CodecType type) {
 117   return make_unique<NoCompressionCodec>(level, type);
 118 }
 119
 120 NoCompressionCodec::NoCompressionCodec(int level, CodecType type)
 121   : Codec(type) {
 122   DCHECK(type == CodecType::NO_COMPRESSION);
 123   switch (level) {
 124   case COMPRESSION_LEVEL_DEFAULT:
 125   case COMPRESSION_LEVEL_FASTEST:
 126   case COMPRESSION_LEVEL_BEST:
 127     level = 0;
 128   }
 129   if (level != 0) {
 130     throw std::invalid_argument(to<std::string>(
 131         "NoCompressionCodec: invalid level ", level));
 132   }
 133 }
 134
 135 std::unique_ptr<IOBuf> NoCompressionCodec::doCompress(
 136     const IOBuf* data) {
 137   return data->clone();
 138 }
 139
 140 std::unique_ptr<IOBuf> NoCompressionCodec::doUncompress(
 141     const IOBuf* data,
 142     uint64_t uncompressedLength) {
 143   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
 144       data->computeChainDataLength() != uncompressedLength) {
 145     throw std::runtime_error(to<std::string>(
 146         "NoCompressionCodec: invalid uncompressed length"));
 147   }
 148   return data->clone();
 149 }
 150
 151 namespace {
 152
 153 void encodeVarintToIOBuf(uint64_t val, folly::IOBuf* out) {
 154   DCHECK_GE(out->tailroom(), kMaxVarintLength64);
 155   out->append(encodeVarint(val, out->writableTail()));
 156 }
 157
 158 uint64_t decodeVarintFromCursor(folly::io::Cursor& cursor) {
 159   // Must have enough room in *this* buffer.
 160   auto p = cursor.peek();
 161   folly::ByteRange range(p.first, p.second);
 162   uint64_t val = decodeVarint(range);
 163   cursor.skip(range.data() - p.first);
 164   return val;
 165 }
 166
 167 }  // namespace
 168
 169 #if FOLLY_HAVE_LIBLZ4
 170
 171 /**
 172  * LZ4 compression
 173  */
 174 class LZ4Codec FOLLY_FINAL : public Codec {
 175  public:
 176   static std::unique_ptr<Codec> create(int level, CodecType type);
 177   explicit LZ4Codec(int level, CodecType type);
 178
 179  private:
 180   bool doNeedsUncompressedLength() const FOLLY_OVERRIDE;
 181   uint64_t doMaxUncompressedLength() const FOLLY_OVERRIDE;
 182
 183   bool encodeSize() const { return type() == CodecType::LZ4_VARINT_SIZE; }
 184
 185   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) FOLLY_OVERRIDE;
 186   std::unique_ptr<IOBuf> doUncompress(
 187       const IOBuf* data,
 188       uint64_t uncompressedLength) FOLLY_OVERRIDE;
 189
 190   bool highCompression_;
 191 };
 192
 193 std::unique_ptr<Codec> LZ4Codec::create(int level, CodecType type) {
 194   return make_unique<LZ4Codec>(level, type);
 195 }
 196
 197 LZ4Codec::LZ4Codec(int level, CodecType type) : Codec(type) {
 198   DCHECK(type == CodecType::LZ4 || type == CodecType::LZ4_VARINT_SIZE);
 199
 200   switch (level) {
 201   case COMPRESSION_LEVEL_FASTEST:
 202   case COMPRESSION_LEVEL_DEFAULT:
 203     level = 1;
 204     break;
 205   case COMPRESSION_LEVEL_BEST:
 206     level = 2;
 207     break;
 208   }
 209   if (level < 1 || level > 2) {
 210     throw std::invalid_argument(to<std::string>(
 211         "LZ4Codec: invalid level: ", level));
 212   }
 213   highCompression_ = (level > 1);
 214 }
 215
 216 bool LZ4Codec::doNeedsUncompressedLength() const {
 217   return !encodeSize();
 218 }
 219
 220 // The value comes from lz4.h in lz4-r117, but older versions of lz4 don't
 221 // define LZ4_MAX_INPUT_SIZE (even though the max size is the same), so do it
 222 // here.
 223 #ifndef LZ4_MAX_INPUT_SIZE
 224 # define LZ4_MAX_INPUT_SIZE 0x7E000000
 225 #endif
 226
 227 uint64_t LZ4Codec::doMaxUncompressedLength() const {
 228   return LZ4_MAX_INPUT_SIZE;
 229 }
 230
 231 std::unique_ptr<IOBuf> LZ4Codec::doCompress(const IOBuf* data) {
 232   std::unique_ptr<IOBuf> clone;
 233   if (data->isChained()) {
 234     // LZ4 doesn't support streaming, so we have to coalesce
 235     clone = data->clone();
 236     clone->coalesce();
 237     data = clone.get();
 238   }
 239
 240   uint32_t extraSize = encodeSize() ? kMaxVarintLength64 : 0;
 241   auto out = IOBuf::create(extraSize + LZ4_compressBound(data->length()));
 242   if (encodeSize()) {
 243     encodeVarintToIOBuf(data->length(), out.get());
 244   }
 245
 246   int n;
 247   if (highCompression_) {
 248     n = LZ4_compressHC(reinterpret_cast<const char*>(data->data()),
 249                        reinterpret_cast<char*>(out->writableTail()),
 250                        data->length());
 251   } else {
 252     n = LZ4_compress(reinterpret_cast<const char*>(data->data()),
 253                      reinterpret_cast<char*>(out->writableTail()),
 254                      data->length());
 255   }
 256
 257   CHECK_GE(n, 0);
 258   CHECK_LE(n, out->capacity());
 259
 260   out->append(n);
 261   return out;
 262 }
 263
 264 std::unique_ptr<IOBuf> LZ4Codec::doUncompress(
 265     const IOBuf* data,
 266     uint64_t uncompressedLength) {
 267   std::unique_ptr<IOBuf> clone;
 268   if (data->isChained()) {
 269     // LZ4 doesn't support streaming, so we have to coalesce
 270     clone = data->clone();
 271     clone->coalesce();
 272     data = clone.get();
 273   }
 274
 275   folly::io::Cursor cursor(data);
 276   uint64_t actualUncompressedLength;
 277   if (encodeSize()) {
 278     actualUncompressedLength = decodeVarintFromCursor(cursor);
 279     if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
 280         uncompressedLength != actualUncompressedLength) {
 281       throw std::runtime_error("LZ4Codec: invalid uncompressed length");
 282     }
 283   } else {
 284     actualUncompressedLength = uncompressedLength;
 285     if (actualUncompressedLength == UNKNOWN_UNCOMPRESSED_LENGTH ||
 286         actualUncompressedLength > maxUncompressedLength()) {
 287       throw std::runtime_error("LZ4Codec: invalid uncompressed length");
 288     }
 289   }
 290
 291   auto p = cursor.peek();
 292   auto out = IOBuf::create(actualUncompressedLength);
 293   int n = LZ4_decompress_safe(reinterpret_cast<const char*>(p.first),
 294                               reinterpret_cast<char*>(out->writableTail()),
 295                               p.second,
 296                               actualUncompressedLength);
 297
 298   if (n < 0 || uint64_t(n) != actualUncompressedLength) {
 299     throw std::runtime_error(to<std::string>(
 300         "LZ4 decompression returned invalid value ", n));
 301   }
 302   out->append(actualUncompressedLength);
 303   return out;
 304 }
 305
 306 #endif  // FOLLY_HAVE_LIBLZ4
 307
 308 #if FOLLY_HAVE_LIBSNAPPY
 309
 310 /**
 311  * Snappy compression
 312  */
 313
 314 /**
 315  * Implementation of snappy::Source that reads from a IOBuf chain.
 316  */
 317 class IOBufSnappySource FOLLY_FINAL : public snappy::Source {
 318  public:
 319   explicit IOBufSnappySource(const IOBuf* data);
 320   size_t Available() const FOLLY_OVERRIDE;
 321   const char* Peek(size_t* len) FOLLY_OVERRIDE;
 322   void Skip(size_t n) FOLLY_OVERRIDE;
 323  private:
 324   size_t available_;
 325   io::Cursor cursor_;
 326 };
 327
 328 IOBufSnappySource::IOBufSnappySource(const IOBuf* data)
 329   : available_(data->computeChainDataLength()),
 330     cursor_(data) {
 331 }
 332
 333 size_t IOBufSnappySource::Available() const {
 334   return available_;
 335 }
 336
 337 const char* IOBufSnappySource::Peek(size_t* len) {
 338   auto p = cursor_.peek();
 339   *len = p.second;
 340   return reinterpret_cast<const char*>(p.first);
 341 }
 342
 343 void IOBufSnappySource::Skip(size_t n) {
 344   CHECK_LE(n, available_);
 345   cursor_.skip(n);
 346   available_ -= n;
 347 }
 348
 349 class SnappyCodec FOLLY_FINAL : public Codec {
 350  public:
 351   static std::unique_ptr<Codec> create(int level, CodecType type);
 352   explicit SnappyCodec(int level, CodecType type);
 353
 354  private:
 355   uint64_t doMaxUncompressedLength() const FOLLY_OVERRIDE;
 356   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) FOLLY_OVERRIDE;
 357   std::unique_ptr<IOBuf> doUncompress(
 358       const IOBuf* data,
 359       uint64_t uncompressedLength) FOLLY_OVERRIDE;
 360 };
 361
 362 std::unique_ptr<Codec> SnappyCodec::create(int level, CodecType type) {
 363   return make_unique<SnappyCodec>(level, type);
 364 }
 365
 366 SnappyCodec::SnappyCodec(int level, CodecType type) : Codec(type) {
 367   DCHECK(type == CodecType::SNAPPY);
 368   switch (level) {
 369   case COMPRESSION_LEVEL_FASTEST:
 370   case COMPRESSION_LEVEL_DEFAULT:
 371   case COMPRESSION_LEVEL_BEST:
 372     level = 1;
 373   }
 374   if (level != 1) {
 375     throw std::invalid_argument(to<std::string>(
 376         "SnappyCodec: invalid level: ", level));
 377   }
 378 }
 379
 380 uint64_t SnappyCodec::doMaxUncompressedLength() const {
 381   // snappy.h uses uint32_t for lengths, so there's that.
 382   return std::numeric_limits<uint32_t>::max();
 383 }
 384
 385 std::unique_ptr<IOBuf> SnappyCodec::doCompress(const IOBuf* data) {
 386   IOBufSnappySource source(data);
 387   auto out =
 388     IOBuf::create(snappy::MaxCompressedLength(source.Available()));
 389
 390   snappy::UncheckedByteArraySink sink(reinterpret_cast<char*>(
 391       out->writableTail()));
 392
 393   size_t n = snappy::Compress(&source, &sink);
 394
 395   CHECK_LE(n, out->capacity());
 396   out->append(n);
 397   return out;
 398 }
 399
 400 std::unique_ptr<IOBuf> SnappyCodec::doUncompress(const IOBuf* data,
 401                                                  uint64_t uncompressedLength) {
 402   uint32_t actualUncompressedLength = 0;
 403
 404   {
 405     IOBufSnappySource source(data);
 406     if (!snappy::GetUncompressedLength(&source, &actualUncompressedLength)) {
 407       throw std::runtime_error("snappy::GetUncompressedLength failed");
 408     }
 409     if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
 410         uncompressedLength != actualUncompressedLength) {
 411       throw std::runtime_error("snappy: invalid uncompressed length");
 412     }
 413   }
 414
 415   auto out = IOBuf::create(actualUncompressedLength);
 416
 417   {
 418     IOBufSnappySource source(data);
 419     if (!snappy::RawUncompress(&source,
 420                                reinterpret_cast<char*>(out->writableTail()))) {
 421       throw std::runtime_error("snappy::RawUncompress failed");
 422     }
 423   }
 424
 425   out->append(actualUncompressedLength);
 426   return out;
 427 }
 428
 429 #endif  // FOLLY_HAVE_LIBSNAPPY
 430
 431 #if FOLLY_HAVE_LIBZ
 432 /**
 433  * Zlib codec
 434  */
 435 class ZlibCodec FOLLY_FINAL : public Codec {
 436  public:
 437   static std::unique_ptr<Codec> create(int level, CodecType type);
 438   explicit ZlibCodec(int level, CodecType type);
 439
 440  private:
 441   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) FOLLY_OVERRIDE;
 442   std::unique_ptr<IOBuf> doUncompress(
 443       const IOBuf* data,
 444       uint64_t uncompressedLength) FOLLY_OVERRIDE;
 445
 446   std::unique_ptr<IOBuf> addOutputBuffer(z_stream* stream, uint32_t length);
 447   bool doInflate(z_stream* stream, IOBuf* head, uint32_t bufferLength);
 448
 449   int level_;
 450 };
 451
 452 std::unique_ptr<Codec> ZlibCodec::create(int level, CodecType type) {
 453   return make_unique<ZlibCodec>(level, type);
 454 }
 455
 456 ZlibCodec::ZlibCodec(int level, CodecType type) : Codec(type) {
 457   DCHECK(type == CodecType::ZLIB);
 458   switch (level) {
 459   case COMPRESSION_LEVEL_FASTEST:
 460     level = 1;
 461     break;
 462   case COMPRESSION_LEVEL_DEFAULT:
 463     level = Z_DEFAULT_COMPRESSION;
 464     break;
 465   case COMPRESSION_LEVEL_BEST:
 466     level = 9;
 467     break;
 468   }
 469   if (level != Z_DEFAULT_COMPRESSION && (level < 0 || level > 9)) {
 470     throw std::invalid_argument(to<std::string>(
 471         "ZlibCodec: invalid level: ", level));
 472   }
 473   level_ = level;
 474 }
 475
 476 std::unique_ptr<IOBuf> ZlibCodec::addOutputBuffer(z_stream* stream,
 477                                                   uint32_t length) {
 478   CHECK_EQ(stream->avail_out, 0);
 479
 480   auto buf = IOBuf::create(length);
 481   buf->append(length);
 482
 483   stream->next_out = buf->writableData();
 484   stream->avail_out = buf->length();
 485
 486   return buf;
 487 }
 488
 489 bool ZlibCodec::doInflate(z_stream* stream,
 490                           IOBuf* head,
 491                           uint32_t bufferLength) {
 492   if (stream->avail_out == 0) {
 493     head->prependChain(addOutputBuffer(stream, bufferLength));
 494   }
 495
 496   int rc = inflate(stream, Z_NO_FLUSH);
 497
 498   switch (rc) {
 499   case Z_OK:
 500     break;
 501   case Z_STREAM_END:
 502     return true;
 503   case Z_BUF_ERROR:
 504   case Z_NEED_DICT:
 505   case Z_DATA_ERROR:
 506   case Z_MEM_ERROR:
 507     throw std::runtime_error(to<std::string>(
 508         "ZlibCodec: inflate error: ", rc, ": ", stream->msg));
 509   default:
 510     CHECK(false) << rc << ": " << stream->msg;
 511   }
 512
 513   return false;
 514 }
 515
 516 std::unique_ptr<IOBuf> ZlibCodec::doCompress(const IOBuf* data) {
 517   z_stream stream;
 518   stream.zalloc = nullptr;
 519   stream.zfree = nullptr;
 520   stream.opaque = nullptr;
 521
 522   int rc = deflateInit(&stream, level_);
 523   if (rc != Z_OK) {
 524     throw std::runtime_error(to<std::string>(
 525         "ZlibCodec: deflateInit error: ", rc, ": ", stream.msg));
 526   }
 527
 528   stream.next_in = stream.next_out = nullptr;
 529   stream.avail_in = stream.avail_out = 0;
 530   stream.total_in = stream.total_out = 0;
 531
 532   bool success = false;
 533
 534   SCOPE_EXIT {
 535     int rc = deflateEnd(&stream);
 536     // If we're here because of an exception, it's okay if some data
 537     // got dropped.
 538     CHECK(rc == Z_OK || (!success && rc == Z_DATA_ERROR))
 539       << rc << ": " << stream.msg;
 540   };
 541
 542   uint64_t uncompressedLength = data->computeChainDataLength();
 543   uint64_t maxCompressedLength = deflateBound(&stream, uncompressedLength);
 544
 545   // Max 64MiB in one go
 546   constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20;    // 64MiB
 547   constexpr uint32_t defaultBufferLength = uint32_t(4) << 20;     // 4MiB
 548
 549   auto out = addOutputBuffer(
 550       &stream,
 551       (maxCompressedLength <= maxSingleStepLength ?
 552        maxCompressedLength :
 553        defaultBufferLength));
 554
 555   for (auto& range : *data) {
 556     if (range.empty()) {
 557       continue;
 558     }
 559
 560     stream.next_in = const_cast<uint8_t*>(range.data());
 561     stream.avail_in = range.size();
 562
 563     while (stream.avail_in != 0) {
 564       if (stream.avail_out == 0) {
 565         out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
 566       }
 567
 568       rc = deflate(&stream, Z_NO_FLUSH);
 569
 570       CHECK_EQ(rc, Z_OK) << stream.msg;
 571     }
 572   }
 573
 574   do {
 575     if (stream.avail_out == 0) {
 576       out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
 577     }
 578
 579     rc = deflate(&stream, Z_FINISH);
 580   } while (rc == Z_OK);
 581
 582   CHECK_EQ(rc, Z_STREAM_END) << stream.msg;
 583
 584   out->prev()->trimEnd(stream.avail_out);
 585
 586   success = true;  // we survived
 587
 588   return out;
 589 }
 590
 591 std::unique_ptr<IOBuf> ZlibCodec::doUncompress(const IOBuf* data,
 592                                                uint64_t uncompressedLength) {
 593   z_stream stream;
 594   stream.zalloc = nullptr;
 595   stream.zfree = nullptr;
 596   stream.opaque = nullptr;
 597
 598   int rc = inflateInit(&stream);
 599   if (rc != Z_OK) {
 600     throw std::runtime_error(to<std::string>(
 601         "ZlibCodec: inflateInit error: ", rc, ": ", stream.msg));
 602   }
 603
 604   stream.next_in = stream.next_out = nullptr;
 605   stream.avail_in = stream.avail_out = 0;
 606   stream.total_in = stream.total_out = 0;
 607
 608   bool success = false;
 609
 610   SCOPE_EXIT {
 611     int rc = inflateEnd(&stream);
 612     // If we're here because of an exception, it's okay if some data
 613     // got dropped.
 614     CHECK(rc == Z_OK || (!success && rc == Z_DATA_ERROR))
 615       << rc << ": " << stream.msg;
 616   };
 617
 618   // Max 64MiB in one go
 619   constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20;    // 64MiB
 620   constexpr uint32_t defaultBufferLength = uint32_t(4) << 20;     // 4MiB
 621
 622   auto out = addOutputBuffer(
 623       &stream,
 624       ((uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
 625         uncompressedLength <= maxSingleStepLength) ?
 626        uncompressedLength :
 627        defaultBufferLength));
 628
 629   bool streamEnd = false;
 630   for (auto& range : *data) {
 631     if (range.empty()) {
 632       continue;
 633     }
 634
 635     stream.next_in = const_cast<uint8_t*>(range.data());
 636     stream.avail_in = range.size();
 637
 638     while (stream.avail_in != 0) {
 639       if (streamEnd) {
 640         throw std::runtime_error(to<std::string>(
 641             "ZlibCodec: junk after end of data"));
 642       }
 643
 644       streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
 645     }
 646   }
 647
 648   while (!streamEnd) {
 649     streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
 650   }
 651
 652   out->prev()->trimEnd(stream.avail_out);
 653
 654   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
 655       uncompressedLength != stream.total_out) {
 656     throw std::runtime_error(to<std::string>(
 657         "ZlibCodec: invalid uncompressed length"));
 658   }
 659
 660   success = true;  // we survived
 661
 662   return out;
 663 }
 664
 665 #endif  // FOLLY_HAVE_LIBZ
 666
 667 #if FOLLY_HAVE_LIBLZMA
 668
 669 /**
 670  * LZMA2 compression
 671  */
 672 class LZMA2Codec FOLLY_FINAL : public Codec {
 673  public:
 674   static std::unique_ptr<Codec> create(int level, CodecType type);
 675   explicit LZMA2Codec(int level, CodecType type);
 676
 677  private:
 678   bool doNeedsUncompressedLength() const FOLLY_OVERRIDE;
 679   uint64_t doMaxUncompressedLength() const FOLLY_OVERRIDE;
 680
 681   bool encodeSize() const { return type() == CodecType::LZMA2_VARINT_SIZE; }
 682
 683   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) FOLLY_OVERRIDE;
 684   std::unique_ptr<IOBuf> doUncompress(
 685       const IOBuf* data,
 686       uint64_t uncompressedLength) FOLLY_OVERRIDE;
 687
 688   std::unique_ptr<IOBuf> addOutputBuffer(lzma_stream* stream, size_t length);
 689   bool doInflate(lzma_stream* stream, IOBuf* head, size_t bufferLength);
 690
 691   int level_;
 692 };
 693
 694 std::unique_ptr<Codec> LZMA2Codec::create(int level, CodecType type) {
 695   return make_unique<LZMA2Codec>(level, type);
 696 }
 697
 698 LZMA2Codec::LZMA2Codec(int level, CodecType type) : Codec(type) {
 699   DCHECK(type == CodecType::LZMA2 || type == CodecType::LZMA2_VARINT_SIZE);
 700   switch (level) {
 701   case COMPRESSION_LEVEL_FASTEST:
 702     level = 0;
 703     break;
 704   case COMPRESSION_LEVEL_DEFAULT:
 705     level = LZMA_PRESET_DEFAULT;
 706     break;
 707   case COMPRESSION_LEVEL_BEST:
 708     level = 9;
 709     break;
 710   }
 711   if (level < 0 || level > 9) {
 712     throw std::invalid_argument(to<std::string>(
 713         "LZMA2Codec: invalid level: ", level));
 714   }
 715   level_ = level;
 716 }
 717
 718 bool LZMA2Codec::doNeedsUncompressedLength() const {
 719   return !encodeSize();
 720 }
 721
 722 uint64_t LZMA2Codec::doMaxUncompressedLength() const {
 723   // From lzma/base.h: "Stream is roughly 8 EiB (2^63 bytes)"
 724   return uint64_t(1) << 63;
 725 }
 726
 727 std::unique_ptr<IOBuf> LZMA2Codec::addOutputBuffer(
 728     lzma_stream* stream,
 729     size_t length) {
 730
 731   CHECK_EQ(stream->avail_out, 0);
 732
 733   auto buf = IOBuf::create(length);
 734   buf->append(length);
 735
 736   stream->next_out = buf->writableData();
 737   stream->avail_out = buf->length();
 738
 739   return buf;
 740 }
 741
 742 std::unique_ptr<IOBuf> LZMA2Codec::doCompress(const IOBuf* data) {
 743   lzma_ret rc;
 744   lzma_stream stream = LZMA_STREAM_INIT;
 745
 746   rc = lzma_easy_encoder(&stream, level_, LZMA_CHECK_NONE);
 747   if (rc != LZMA_OK) {
 748     throw std::runtime_error(folly::to<std::string>(
 749       "LZMA2Codec: lzma_easy_encoder error: ", rc));
 750   }
 751
 752   SCOPE_EXIT { lzma_end(&stream); };
 753
 754   uint64_t uncompressedLength = data->computeChainDataLength();
 755   uint64_t maxCompressedLength = lzma_stream_buffer_bound(uncompressedLength);
 756
 757   // Max 64MiB in one go
 758   constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20;    // 64MiB
 759   constexpr uint32_t defaultBufferLength = uint32_t(4) << 20;     // 4MiB
 760
 761   auto out = addOutputBuffer(
 762     &stream,
 763     (maxCompressedLength <= maxSingleStepLength ?
 764      maxCompressedLength :
 765      defaultBufferLength));
 766
 767   if (encodeSize()) {
 768     auto size = IOBuf::createCombined(kMaxVarintLength64);
 769     encodeVarintToIOBuf(uncompressedLength, size.get());
 770     size->appendChain(std::move(out));
 771     out = std::move(size);
 772   }
 773
 774   for (auto& range : *data) {
 775     if (range.empty()) {
 776       continue;
 777     }
 778
 779     stream.next_in = const_cast<uint8_t*>(range.data());
 780     stream.avail_in = range.size();
 781
 782     while (stream.avail_in != 0) {
 783       if (stream.avail_out == 0) {
 784         out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
 785       }
 786
 787       rc = lzma_code(&stream, LZMA_RUN);
 788
 789       if (rc != LZMA_OK) {
 790         throw std::runtime_error(folly::to<std::string>(
 791           "LZMA2Codec: lzma_code error: ", rc));
 792       }
 793     }
 794   }
 795
 796   do {
 797     if (stream.avail_out == 0) {
 798       out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
 799     }
 800
 801     rc = lzma_code(&stream, LZMA_FINISH);
 802   } while (rc == LZMA_OK);
 803
 804   if (rc != LZMA_STREAM_END) {
 805     throw std::runtime_error(folly::to<std::string>(
 806       "LZMA2Codec: lzma_code ended with error: ", rc));
 807   }
 808
 809   out->prev()->trimEnd(stream.avail_out);
 810
 811   return out;
 812 }
 813
 814 bool LZMA2Codec::doInflate(lzma_stream* stream,
 815                           IOBuf* head,
 816                           size_t bufferLength) {
 817   if (stream->avail_out == 0) {
 818     head->prependChain(addOutputBuffer(stream, bufferLength));
 819   }
 820
 821   lzma_ret rc = lzma_code(stream, LZMA_RUN);
 822
 823   switch (rc) {
 824   case LZMA_OK:
 825     break;
 826   case LZMA_STREAM_END:
 827     return true;
 828   default:
 829     throw std::runtime_error(to<std::string>(
 830         "LZMA2Codec: lzma_code error: ", rc));
 831   }
 832
 833   return false;
 834 }
 835
 836 std::unique_ptr<IOBuf> LZMA2Codec::doUncompress(const IOBuf* data,
 837                                                uint64_t uncompressedLength) {
 838   lzma_ret rc;
 839   lzma_stream stream = LZMA_STREAM_INIT;
 840
 841   rc = lzma_auto_decoder(&stream, std::numeric_limits<uint64_t>::max(), 0);
 842   if (rc != LZMA_OK) {
 843     throw std::runtime_error(folly::to<std::string>(
 844       "LZMA2Codec: lzma_auto_decoder error: ", rc));
 845   }
 846
 847   SCOPE_EXIT { lzma_end(&stream); };
 848
 849   // Max 64MiB in one go
 850   constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20;    // 64MiB
 851   constexpr uint32_t defaultBufferLength = uint32_t(4) << 20;     // 4MiB
 852
 853   folly::io::Cursor cursor(data);
 854   uint64_t actualUncompressedLength;
 855   if (encodeSize()) {
 856     actualUncompressedLength = decodeVarintFromCursor(cursor);
 857     if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
 858         uncompressedLength != actualUncompressedLength) {
 859       throw std::runtime_error("LZMA2Codec: invalid uncompressed length");
 860     }
 861   } else {
 862     actualUncompressedLength = uncompressedLength;
 863     DCHECK_NE(actualUncompressedLength, UNKNOWN_UNCOMPRESSED_LENGTH);
 864   }
 865
 866   auto out = addOutputBuffer(
 867       &stream,
 868       (actualUncompressedLength <= maxSingleStepLength ?
 869        actualUncompressedLength :
 870        defaultBufferLength));
 871
 872   bool streamEnd = false;
 873   auto buf = cursor.peek();
 874   while (buf.second != 0) {
 875     stream.next_in = const_cast<uint8_t*>(buf.first);
 876     stream.avail_in = buf.second;
 877
 878     while (stream.avail_in != 0) {
 879       if (streamEnd) {
 880         throw std::runtime_error(to<std::string>(
 881             "LZMA2Codec: junk after end of data"));
 882       }
 883
 884       streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
 885     }
 886
 887     cursor.skip(buf.second);
 888     buf = cursor.peek();
 889   }
 890
 891   while (!streamEnd) {
 892     streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
 893   }
 894
 895   out->prev()->trimEnd(stream.avail_out);
 896
 897   if (actualUncompressedLength != stream.total_out) {
 898     throw std::runtime_error(to<std::string>(
 899         "LZMA2Codec: invalid uncompressed length"));
 900   }
 901
 902   return out;
 903 }
 904
 905 #endif  // FOLLY_HAVE_LIBLZMA
 906
 907 typedef std::unique_ptr<Codec> (*CodecFactory)(int, CodecType);
 908
 909 CodecFactory gCodecFactories[
 910     static_cast<size_t>(CodecType::NUM_CODEC_TYPES)] = {
 911   nullptr,  // USER_DEFINED
 912   NoCompressionCodec::create,
 913
 914 #if FOLLY_HAVE_LIBLZ4
 915   LZ4Codec::create,
 916 #else
 917   nullptr,
 918 #endif
 919
 920 #if FOLLY_HAVE_LIBSNAPPY
 921   SnappyCodec::create,
 922 #else
 923   nullptr,
 924 #endif
 925
 926 #if FOLLY_HAVE_LIBZ
 927   ZlibCodec::create,
 928 #else
 929   nullptr,
 930 #endif
 931
 932 #if FOLLY_HAVE_LIBLZ4
 933   LZ4Codec::create,
 934 #else
 935   nullptr,
 936 #endif
 937
 938 #if FOLLY_HAVE_LIBLZMA
 939   LZMA2Codec::create,
 940   LZMA2Codec::create,
 941 #else
 942   nullptr,
 943   nullptr,
 944 #endif
 945 };
 946
 947 }  // namespace
 948
 949 std::unique_ptr<Codec> getCodec(CodecType type, int level) {
 950   size_t idx = static_cast<size_t>(type);
 951   if (idx >= static_cast<size_t>(CodecType::NUM_CODEC_TYPES)) {
 952     throw std::invalid_argument(to<std::string>(
 953         "Compression type ", idx, " not supported"));
 954   }
 955   auto factory = gCodecFactories[idx];
 956   if (!factory) {
 957     throw std::invalid_argument(to<std::string>(
 958         "Compression type ", idx, " not supported"));
 959   }
 960   auto codec = (*factory)(level, type);
 961   DCHECK_EQ(static_cast<size_t>(codec->type()), idx);
 962   return codec;
 963 }
 964
 965 }}  // namespaces