utils/zlib/zlib.cc

/*
 * Copyright (C) 2018 The Android Open Source Project
 *
 * 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 "utils/zlib/zlib.h"

#include <memory>

#include "utils/base/logging.h"
#include "utils/flatbuffers.h"

namespace libtextclassifier3 {

std::unique_ptr<ZlibDecompressor> ZlibDecompressor::Instance() {
  std::unique_ptr<ZlibDecompressor> result(new ZlibDecompressor());
  if (!result->initialized_) {
    result.reset();
  }
  return result;
}

ZlibDecompressor::ZlibDecompressor() {
  memset(&stream_, 0, sizeof(stream_));
  stream_.zalloc = Z_NULL;
  stream_.zfree = Z_NULL;
  initialized_ = (inflateInit(&stream_) == Z_OK);
}

ZlibDecompressor::~ZlibDecompressor() {
  if (initialized_) {
    inflateEnd(&stream_);
  }
}

bool ZlibDecompressor::Decompress(const uint8* buffer, const int buffer_size,
                                  const int uncompressed_size,
                                  std::string* out) {
  if (out == nullptr) {
    return false;
  }
  out->resize(uncompressed_size);
  stream_.next_in = reinterpret_cast<const Bytef*>(buffer);
  stream_.avail_in = buffer_size;
  stream_.next_out = reinterpret_cast<Bytef*>(const_cast<char*>(out->c_str()));
  stream_.avail_out = uncompressed_size;
  return (inflate(&stream_, Z_SYNC_FLUSH) == Z_OK);
}

bool ZlibDecompressor::MaybeDecompress(
    const CompressedBuffer* compressed_buffer, std::string* out) {
  if (!compressed_buffer) {
    return true;
  }
  return Decompress(compressed_buffer->buffer()->Data(),
                    compressed_buffer->buffer()->size(),
                    compressed_buffer->uncompressed_size(), out);
}

bool ZlibDecompressor::MaybeDecompress(
    const CompressedBufferT* compressed_buffer, std::string* out) {
  if (!compressed_buffer) {
    return true;
  }
  return Decompress(compressed_buffer->buffer.data(),
                    compressed_buffer->buffer.size(),
                    compressed_buffer->uncompressed_size, out);
}

std::unique_ptr<ZlibCompressor> ZlibCompressor::Instance() {
  std::unique_ptr<ZlibCompressor> result(new ZlibCompressor());
  if (!result->initialized_) {
    result.reset();
  }
  return result;
}

ZlibCompressor::ZlibCompressor(int level, int tmp_buffer_size) {
  memset(&stream_, 0, sizeof(stream_));
  stream_.zalloc = Z_NULL;
  stream_.zfree = Z_NULL;
  buffer_size_ = tmp_buffer_size;
  buffer_.reset(new Bytef[buffer_size_]);
  initialized_ = (deflateInit(&stream_, level) == Z_OK);
}

ZlibCompressor::~ZlibCompressor() { deflateEnd(&stream_); }

void ZlibCompressor::Compress(const std::string& uncompressed_content,
                              CompressedBufferT* out) {
  out->uncompressed_size = uncompressed_content.size();
  out->buffer.clear();
  stream_.next_in =
      reinterpret_cast<const Bytef*>(uncompressed_content.c_str());
  stream_.avail_in = uncompressed_content.size();
  stream_.next_out = buffer_.get();
  stream_.avail_out = buffer_size_;
  unsigned char* buffer_deflate_start_position =
      reinterpret_cast<unsigned char*>(buffer_.get());
  int status;
  do {
    // Deflate chunk-wise.
    // Z_SYNC_FLUSH causes all pending output to be flushed, but doesn't
    // reset the compression state.
    // As we do not know how big the compressed buffer will be, we compress
    // chunk wise and append the flushed content to the output string buffer.
    // As we store the uncompressed size, we do not have to do this during
    // decompression.
    status = deflate(&stream_, Z_SYNC_FLUSH);
    unsigned char* buffer_deflate_end_position =
        reinterpret_cast<unsigned char*>(stream_.next_out);
    if (buffer_deflate_end_position != buffer_deflate_start_position) {
      out->buffer.insert(out->buffer.end(), buffer_deflate_start_position,
                         buffer_deflate_end_position);
      stream_.next_out = buffer_deflate_start_position;
      stream_.avail_out = buffer_size_;
    } else {
      break;
    }
  } while (status == Z_OK);
}

std::unique_ptr<UniLib::RegexPattern> UncompressMakeRegexPattern(
    const UniLib& unilib, const flatbuffers::String* uncompressed_pattern,
    const CompressedBuffer* compressed_pattern, ZlibDecompressor* decompressor,
    std::string* result_pattern_text) {
  UnicodeText unicode_regex_pattern;
  std::string decompressed_pattern;
  if (compressed_pattern != nullptr &&
      compressed_pattern->buffer() != nullptr) {
    if (decompressor == nullptr ||
        !decompressor->MaybeDecompress(compressed_pattern,
                                       &decompressed_pattern)) {
      TC3_LOG(ERROR) << "Cannot decompress pattern.";
      return nullptr;
    }
    unicode_regex_pattern =
        UTF8ToUnicodeText(decompressed_pattern.data(),
                          decompressed_pattern.size(), /*do_copy=*/false);
  } else {
    if (uncompressed_pattern == nullptr) {
      TC3_LOG(ERROR) << "Cannot load uncompressed pattern.";
      return nullptr;
    }
    unicode_regex_pattern =
        UTF8ToUnicodeText(uncompressed_pattern->c_str(),
                          uncompressed_pattern->Length(), /*do_copy=*/false);
  }

  if (result_pattern_text != nullptr) {
    *result_pattern_text = unicode_regex_pattern.ToUTF8String();
  }

  std::unique_ptr<UniLib::RegexPattern> regex_pattern =
      unilib.CreateRegexPattern(unicode_regex_pattern);
  if (!regex_pattern) {
    TC3_LOG(ERROR) << "Could not create pattern: "
                   << unicode_regex_pattern.ToUTF8String();
  }
  return regex_pattern;
}

}  // namespace libtextclassifier3