Zstd.cpp

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <folly/compression/Zstd.h>

#if FOLLY_HAVE_LIBZSTD

#include <stdexcept>
#include <string>

#include <zstd.h>

#include <folly/Conv.h>
#include <folly/Range.h>
#include <folly/ScopeGuard.h>
#include <folly/compression/CompressionContextPoolSingletons.h>
#include <folly/compression/Utils.h>

static_assert(
    ZSTD_VERSION_NUMBER >= 10400,
    "zstd-1.4.0 is the minimum supported zstd version.");

using folly::compression::detail::dataStartsWithLE;
using folly::compression::detail::prefixToStringLE;

using namespace folly::compression::contexts;

namespace folly {
namespace compression {
namespace zstd {
namespace {

size_t zstdThrowIfError(size_t rc) {
  if (!ZSTD_isError(rc)) {
    return rc;
  }
  throw std::runtime_error(
      to<std::string>("ZSTD returned an error: ", ZSTD_getErrorName(rc)));
}

ZSTD_EndDirective zstdTranslateFlush(StreamCodec::FlushOp flush) {
  switch (flush) {
    case StreamCodec::FlushOp::NONE:
      return ZSTD_e_continue;
    case StreamCodec::FlushOp::FLUSH:
      return ZSTD_e_flush;
    case StreamCodec::FlushOp::END:
      return ZSTD_e_end;
    default:
      throw std::invalid_argument("ZSTDStreamCodec: Invalid flush");
  }
}

class ZSTDStreamCodec final : public StreamCodec {
 public:
  explicit ZSTDStreamCodec(Options options);

  std::vector<std::string> validPrefixes() const override;
  bool canUncompress(
      const IOBuf* data, Optional<uint64_t> uncompressedLength) const override;

 private:
  bool doNeedsUncompressedLength() const override;
  uint64_t doMaxCompressedLength(uint64_t uncompressedLength) const override;
  Optional<uint64_t> doGetUncompressedLength(
      IOBuf const* data, Optional<uint64_t> uncompressedLength) const override;

  void doResetStream() override;
  bool doCompressStream(
      ByteRange& input,
      MutableByteRange& output,
      StreamCodec::FlushOp flushOp) override;
  bool doUncompressStream(
      ByteRange& input,
      MutableByteRange& output,
      StreamCodec::FlushOp flushOp) override;

  void resetCCtx();
  void resetDCtx();

  Options options_;
  ZSTD_CCtx_Pool::Ref cctx_{getNULL_ZSTD_CCtx()};
  ZSTD_DCtx_Pool::Ref dctx_{getNULL_ZSTD_DCtx()};
};

constexpr uint32_t kZSTDMagicLE = 0xFD2FB528;

std::vector<std::string> ZSTDStreamCodec::validPrefixes() const {
  return {prefixToStringLE(kZSTDMagicLE)};
}

bool ZSTDStreamCodec::canUncompress(
    const IOBuf* data, Optional<uint64_t>) const {
  return dataStartsWithLE(data, kZSTDMagicLE);
}

CodecType codecType(Options const& options) {
  int const level = options.level();
  DCHECK_NE(level, 0);
  return level > 0 ? CodecType::ZSTD : CodecType::ZSTD_FAST;
}

ZSTDStreamCodec::ZSTDStreamCodec(Options options)
    : StreamCodec(codecType(options), options.level()),
      options_(std::move(options)) {}

bool ZSTDStreamCodec::doNeedsUncompressedLength() const {
  return false;
}

uint64_t ZSTDStreamCodec::doMaxCompressedLength(
    uint64_t uncompressedLength) const {
  return ZSTD_compressBound(uncompressedLength);
}

Optional<uint64_t> ZSTDStreamCodec::doGetUncompressedLength(
    IOBuf const* data, Optional<uint64_t> uncompressedLength) const {
  // Read decompressed size from frame if available in first IOBuf.
  auto const decompressedSize =
      ZSTD_getFrameContentSize(data->data(), data->length());
  if (decompressedSize == ZSTD_CONTENTSIZE_UNKNOWN ||
      decompressedSize == ZSTD_CONTENTSIZE_ERROR) {
    return uncompressedLength;
  }
  if (uncompressedLength && *uncompressedLength != decompressedSize) {
    throw std::runtime_error("ZSTD: invalid uncompressed length");
  }
  return decompressedSize;
}

void ZSTDStreamCodec::doResetStream() {
  cctx_.reset(nullptr);
  dctx_.reset(nullptr);
}

void ZSTDStreamCodec::resetCCtx() {
  DCHECK(cctx_ == nullptr);
  cctx_ = getZSTD_CCtx(); // Gives us a clean context
  DCHECK(cctx_ != nullptr);
  zstdThrowIfError(
      ZSTD_CCtx_setParametersUsingCCtxParams(cctx_.get(), options_.params()));
  zstdThrowIfError(ZSTD_CCtx_setPledgedSrcSize(
      cctx_.get(), uncompressedLength().value_or(ZSTD_CONTENTSIZE_UNKNOWN)));
}

bool ZSTDStreamCodec::doCompressStream(
    ByteRange& input, MutableByteRange& output, StreamCodec::FlushOp flushOp) {
  if (cctx_ == nullptr) {
    resetCCtx();
  }
  ZSTD_inBuffer in = {input.data(), input.size(), 0};
  ZSTD_outBuffer out = {output.data(), output.size(), 0};
  SCOPE_EXIT {
    input.uncheckedAdvance(in.pos);
    output.uncheckedAdvance(out.pos);
  };
  size_t const rc = zstdThrowIfError(ZSTD_compressStream2(
      cctx_.get(), &out, &in, zstdTranslateFlush(flushOp)));
  switch (flushOp) {
    case StreamCodec::FlushOp::NONE:
      return false;
    case StreamCodec::FlushOp::FLUSH:
      return rc == 0;
    case StreamCodec::FlushOp::END:
      if (rc == 0) {
        // Surrender our cctx_
        doResetStream();
      }
      return rc == 0;
    default:
      throw std::invalid_argument("ZSTD: invalid FlushOp");
  }
}

void ZSTDStreamCodec::resetDCtx() {
  DCHECK(dctx_ == nullptr);
  dctx_ = getZSTD_DCtx(); // Gives us a clean context
  DCHECK(dctx_ != nullptr);
  if (options_.maxWindowSize() != 0) {
    zstdThrowIfError(
        ZSTD_DCtx_setMaxWindowSize(dctx_.get(), options_.maxWindowSize()));
  }
}

bool ZSTDStreamCodec::doUncompressStream(
    ByteRange& input, MutableByteRange& output, StreamCodec::FlushOp) {
  if (dctx_ == nullptr) {
    resetDCtx();
  }
  ZSTD_inBuffer in = {input.data(), input.size(), 0};
  ZSTD_outBuffer out = {output.data(), output.size(), 0};
  SCOPE_EXIT {
    input.uncheckedAdvance(in.pos);
    output.uncheckedAdvance(out.pos);
  };
  size_t const rc =
      zstdThrowIfError(ZSTD_decompressStream(dctx_.get(), &out, &in));
  if (rc == 0) {
    // Surrender our dctx_
    doResetStream();
  }
  return rc == 0;
}

} // namespace

Options::Options(int level) : params_(ZSTD_createCCtxParams()), level_(level) {
  if (params_ == nullptr) {
    throw std::bad_alloc{};
  }
  zstdThrowIfError(ZSTD_CCtxParams_init(params_.get(), level));
}

void Options::set(ZSTD_cParameter param, unsigned value) {
  zstdThrowIfError(ZSTD_CCtxParams_setParameter(params_.get(), param, value));
  if (param == ZSTD_c_compressionLevel) {
    level_ = static_cast<int>(value);
  }
}

/* static */ void Options::freeCCtxParams(ZSTD_CCtx_params* params) {
  ZSTD_freeCCtxParams(params);
}

std::unique_ptr<Codec> getCodec(Options options) {
  return std::make_unique<ZSTDStreamCodec>(std::move(options));
}

std::unique_ptr<StreamCodec> getStreamCodec(Options options) {
  return std::make_unique<ZSTDStreamCodec>(std::move(options));
}

} // namespace zstd
} // namespace compression
} // namespace folly

#endif