1 /*
2  * Copyright (C) 2018 The Android Open Source Project
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 "utils/zlib/zlib.h"
18 
19 #include "utils/base/logging.h"
20 #include "utils/flatbuffers/flatbuffers.h"
21 
22 namespace libtextclassifier3 {
23 
Instance()24 std::unique_ptr<ZlibDecompressor> ZlibDecompressor::Instance() {
25   std::unique_ptr<ZlibDecompressor> result(new ZlibDecompressor());
26   if (!result->initialized_) {
27     result.reset();
28   }
29   return result;
30 }
31 
ZlibDecompressor()32 ZlibDecompressor::ZlibDecompressor() {
33   memset(&stream_, 0, sizeof(stream_));
34   stream_.zalloc = Z_NULL;
35   stream_.zfree = Z_NULL;
36   initialized_ = false;
37   if (inflateInit(&stream_) != Z_OK) {
38     TC3_LOG(ERROR) << "Could not initialize decompressor.";
39     return;
40   }
41   initialized_ = true;
42 }
43 
~ZlibDecompressor()44 ZlibDecompressor::~ZlibDecompressor() {
45   if (initialized_) {
46     inflateEnd(&stream_);
47   }
48 }
49 
Decompress(const uint8 * buffer,const int buffer_size,const int uncompressed_size,std::string * out)50 bool ZlibDecompressor::Decompress(const uint8* buffer, const int buffer_size,
51                                   const int uncompressed_size,
52                                   std::string* out) {
53   if (out == nullptr) {
54     return false;
55   }
56   out->resize(uncompressed_size);
57   stream_.next_in =
58       const_cast<z_const Bytef*>(reinterpret_cast<const Bytef*>(buffer));
59   stream_.avail_in = buffer_size;
60   stream_.next_out = reinterpret_cast<Bytef*>(const_cast<char*>(out->c_str()));
61   stream_.avail_out = uncompressed_size;
62   return (inflate(&stream_, Z_SYNC_FLUSH) == Z_OK);
63 }
64 
MaybeDecompress(const CompressedBuffer * compressed_buffer,std::string * out)65 bool ZlibDecompressor::MaybeDecompress(
66     const CompressedBuffer* compressed_buffer, std::string* out) {
67   if (!compressed_buffer) {
68     return true;
69   }
70   return Decompress(compressed_buffer->buffer()->Data(),
71                     compressed_buffer->buffer()->size(),
72                     compressed_buffer->uncompressed_size(), out);
73 }
74 
MaybeDecompress(const CompressedBufferT * compressed_buffer,std::string * out)75 bool ZlibDecompressor::MaybeDecompress(
76     const CompressedBufferT* compressed_buffer, std::string* out) {
77   if (!compressed_buffer) {
78     return true;
79   }
80   return Decompress(compressed_buffer->buffer.data(),
81                     compressed_buffer->buffer.size(),
82                     compressed_buffer->uncompressed_size, out);
83 }
84 
MaybeDecompressOptionallyCompressedBuffer(const flatbuffers::String * uncompressed_buffer,const CompressedBuffer * compressed_buffer,std::string * out)85 bool ZlibDecompressor::MaybeDecompressOptionallyCompressedBuffer(
86     const flatbuffers::String* uncompressed_buffer,
87     const CompressedBuffer* compressed_buffer, std::string* out) {
88   if (uncompressed_buffer != nullptr) {
89     *out = uncompressed_buffer->str();
90     return true;
91   }
92   return MaybeDecompress(compressed_buffer, out);
93 }
94 
MaybeDecompressOptionallyCompressedBuffer(const flatbuffers::Vector<uint8> * uncompressed_buffer,const CompressedBuffer * compressed_buffer,std::string * out)95 bool ZlibDecompressor::MaybeDecompressOptionallyCompressedBuffer(
96     const flatbuffers::Vector<uint8>* uncompressed_buffer,
97     const CompressedBuffer* compressed_buffer, std::string* out) {
98   if (uncompressed_buffer != nullptr) {
99     *out =
100         std::string(reinterpret_cast<const char*>(uncompressed_buffer->data()),
101                     uncompressed_buffer->size());
102     return true;
103   }
104   return MaybeDecompress(compressed_buffer, out);
105 }
106 
Instance()107 std::unique_ptr<ZlibCompressor> ZlibCompressor::Instance() {
108   std::unique_ptr<ZlibCompressor> result(new ZlibCompressor());
109   if (!result->initialized_) {
110     result.reset();
111   }
112   return result;
113 }
114 
ZlibCompressor(const int level,const int tmp_buffer_size)115 ZlibCompressor::ZlibCompressor(const int level, const int tmp_buffer_size) {
116   memset(&stream_, 0, sizeof(stream_));
117   stream_.zalloc = Z_NULL;
118   stream_.zfree = Z_NULL;
119   buffer_size_ = tmp_buffer_size;
120   buffer_.reset(new Bytef[buffer_size_]);
121   initialized_ = false;
122   if (deflateInit(&stream_, level) != Z_OK) {
123     TC3_LOG(ERROR) << "Could not initialize compressor.";
124     return;
125   }
126   initialized_ = true;
127 }
128 
~ZlibCompressor()129 ZlibCompressor::~ZlibCompressor() { deflateEnd(&stream_); }
130 
Compress(const std::string & uncompressed_content,CompressedBufferT * out)131 void ZlibCompressor::Compress(const std::string& uncompressed_content,
132                               CompressedBufferT* out) {
133   out->uncompressed_size = uncompressed_content.size();
134   out->buffer.clear();
135   stream_.next_in = const_cast<z_const Bytef*>(
136       reinterpret_cast<const Bytef*>(uncompressed_content.c_str()));
137   stream_.avail_in = uncompressed_content.size();
138   stream_.next_out = buffer_.get();
139   stream_.avail_out = buffer_size_;
140   unsigned char* buffer_deflate_start_position =
141       reinterpret_cast<unsigned char*>(buffer_.get());
142   int status;
143   do {
144     // Deflate chunk-wise.
145     // Z_SYNC_FLUSH causes all pending output to be flushed, but doesn't
146     // reset the compression state.
147     // As we do not know how big the compressed buffer will be, we compress
148     // chunk wise and append the flushed content to the output string buffer.
149     // As we store the uncompressed size, we do not have to do this during
150     // decompression.
151     status = deflate(&stream_, Z_SYNC_FLUSH);
152     unsigned char* buffer_deflate_end_position =
153         reinterpret_cast<unsigned char*>(stream_.next_out);
154     if (buffer_deflate_end_position != buffer_deflate_start_position) {
155       out->buffer.insert(out->buffer.end(), buffer_deflate_start_position,
156                          buffer_deflate_end_position);
157       stream_.next_out = buffer_deflate_start_position;
158       stream_.avail_out = buffer_size_;
159     } else {
160       break;
161     }
162   } while (status == Z_OK);
163 }
164 
165 }  // namespace libtextclassifier3
166