1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/core/lib/io/zlib_inputstream.h"
17
18 #include <zlib.h>
19
20 #include "tensorflow/core/platform/logging.h"
21 #include "tensorflow/core/platform/strcat.h"
22
23 namespace tensorflow {
24 namespace io {
25
26 struct ZStreamDef {
ZStreamDeftensorflow::io::ZStreamDef27 ZStreamDef(size_t input_buffer_capacity, size_t output_buffer_capacity)
28 : input(new Bytef[input_buffer_capacity]),
29 output(new Bytef[output_buffer_capacity]),
30 stream(new z_stream) {}
31
32 // Buffer for storing contents read from compressed stream.
33 // TODO(srbs): Consider using circular buffers. That would greatly simplify
34 // the implementation.
35 std::unique_ptr<Bytef[]> input;
36
37 // Buffer for storing inflated contents of `input_stream_`.
38 std::unique_ptr<Bytef[]> output;
39
40 // Configuration passed to `inflate`.
41 //
42 // z_stream_def_->stream->next_in:
43 // Next byte to de-compress. Points to some byte in
44 // z_stream_def_->streamdef_.input buffer.
45 // z_stream_def_->stream->avail_in:
46 // Number of bytes available to be decompressed at this time.
47 // z_stream_def_->stream->next_out:
48 // Next byte to write de-compressed data to. Points to some byte in
49 // z_stream_def_->streamdef_.output buffer.
50 // z_stream_def_->stream->avail_out:
51 // Number of free bytes available at write location.
52 std::unique_ptr<z_stream> stream;
53 };
54
ZlibInputStream(InputStreamInterface * input_stream,size_t input_buffer_bytes,size_t output_buffer_bytes,const ZlibCompressionOptions & zlib_options,bool owns_input_stream)55 ZlibInputStream::ZlibInputStream(
56 InputStreamInterface* input_stream,
57 size_t input_buffer_bytes, // size of z_stream.next_in buffer
58 size_t output_buffer_bytes, // size of z_stream.next_out buffer
59 const ZlibCompressionOptions& zlib_options, bool owns_input_stream)
60 : owns_input_stream_(owns_input_stream),
61 input_stream_(input_stream),
62 input_buffer_capacity_(input_buffer_bytes),
63 output_buffer_capacity_(output_buffer_bytes),
64 zlib_options_(zlib_options),
65 z_stream_def_(
66 new ZStreamDef(input_buffer_capacity_, output_buffer_capacity_)),
67 bytes_read_(0) {
68 InitZlibBuffer();
69 }
70
ZlibInputStream(InputStreamInterface * input_stream,size_t input_buffer_bytes,size_t output_buffer_bytes,const ZlibCompressionOptions & zlib_options)71 ZlibInputStream::ZlibInputStream(InputStreamInterface* input_stream,
72 size_t input_buffer_bytes,
73 size_t output_buffer_bytes,
74 const ZlibCompressionOptions& zlib_options)
75 : ZlibInputStream(input_stream, input_buffer_bytes, output_buffer_bytes,
76 zlib_options, false) {}
77
~ZlibInputStream()78 ZlibInputStream::~ZlibInputStream() {
79 if (z_stream_def_->stream && !init_error_) {
80 inflateEnd(z_stream_def_->stream.get());
81 }
82 if (owns_input_stream_) {
83 delete input_stream_;
84 }
85 }
86
Reset()87 Status ZlibInputStream::Reset() {
88 if (init_error_) {
89 return errors::DataLoss("unable to reset stream, cannot decompress.");
90 }
91 TF_RETURN_IF_ERROR(input_stream_->Reset());
92 inflateEnd(z_stream_def_->stream.get());
93 InitZlibBuffer();
94 bytes_read_ = 0;
95 return Status::OK();
96 }
97
InitZlibBuffer()98 void ZlibInputStream::InitZlibBuffer() {
99 memset(z_stream_def_->stream.get(), 0, sizeof(z_stream));
100
101 z_stream_def_->stream->zalloc = Z_NULL;
102 z_stream_def_->stream->zfree = Z_NULL;
103 z_stream_def_->stream->opaque = Z_NULL;
104 z_stream_def_->stream->next_in = Z_NULL;
105 z_stream_def_->stream->avail_in = 0;
106
107 int status =
108 inflateInit2(z_stream_def_->stream.get(), zlib_options_.window_bits);
109
110 if (zlib_options_.soft_fail_on_error && status != Z_OK) {
111 init_error_ = true;
112 return;
113 }
114 CHECK_EQ(status, Z_OK) << "inflateInit failed with status " << status;
115
116 z_stream_def_->stream->next_in = z_stream_def_->input.get();
117 z_stream_def_->stream->next_out = z_stream_def_->output.get();
118 next_unread_byte_ = reinterpret_cast<char*>(z_stream_def_->output.get());
119 z_stream_def_->stream->avail_in = 0;
120 z_stream_def_->stream->avail_out = output_buffer_capacity_;
121 }
122
ReadFromStream()123 Status ZlibInputStream::ReadFromStream() {
124 int bytes_to_read = input_buffer_capacity_;
125 char* read_location = reinterpret_cast<char*>(z_stream_def_->input.get());
126
127 // If there are unread bytes in the input stream we move them to the head
128 // of the stream to maximize the space available to read new data into.
129 if (z_stream_def_->stream->avail_in > 0) {
130 uLong read_bytes =
131 z_stream_def_->stream->next_in - z_stream_def_->input.get();
132 // Remove `read_bytes` from the head of the input stream.
133 // Move unread bytes to the head of the input stream.
134 if (read_bytes > 0) {
135 memmove(z_stream_def_->input.get(), z_stream_def_->stream->next_in,
136 z_stream_def_->stream->avail_in);
137 }
138
139 bytes_to_read -= z_stream_def_->stream->avail_in;
140 read_location += z_stream_def_->stream->avail_in;
141 }
142 tstring data;
143 // Try to read enough data to fill up z_stream_def_->input.
144 // TODO(rohanj): Add a char* version of ReadNBytes to InputStreamInterface
145 // and use that instead to make this more efficient.
146 Status s = input_stream_->ReadNBytes(bytes_to_read, &data);
147 memcpy(read_location, data.data(), data.size());
148
149 // Since we moved unread data to the head of the input stream we can point
150 // next_in to the head of the input stream.
151 z_stream_def_->stream->next_in = z_stream_def_->input.get();
152
153 // Note: data.size() could be different from bytes_to_read.
154 z_stream_def_->stream->avail_in += data.size();
155
156 if (!s.ok() && !errors::IsOutOfRange(s)) {
157 return s;
158 }
159
160 // We throw OutOfRange error iff no new data has been read from stream.
161 // Since we never check how much data is remaining in the stream, it is
162 // possible that on the last read there isn't enough data in the stream to
163 // fill up the buffer in which case input_stream_->ReadNBytes would return an
164 // OutOfRange error.
165 if (data.empty()) {
166 return errors::OutOfRange("EOF reached");
167 }
168 if (errors::IsOutOfRange(s)) {
169 return Status::OK();
170 }
171
172 return s;
173 }
174
ReadBytesFromCache(size_t bytes_to_read,tstring * result)175 size_t ZlibInputStream::ReadBytesFromCache(size_t bytes_to_read,
176 tstring* result) {
177 size_t unread_bytes =
178 reinterpret_cast<char*>(z_stream_def_->stream->next_out) -
179 next_unread_byte_;
180 size_t can_read_bytes = std::min(bytes_to_read, unread_bytes);
181 if (can_read_bytes > 0) {
182 result->append(next_unread_byte_, can_read_bytes);
183 next_unread_byte_ += can_read_bytes;
184 }
185 bytes_read_ += can_read_bytes;
186 return can_read_bytes;
187 }
188
NumUnreadBytes() const189 size_t ZlibInputStream::NumUnreadBytes() const {
190 size_t read_bytes =
191 next_unread_byte_ - reinterpret_cast<char*>(z_stream_def_->output.get());
192 return output_buffer_capacity_ - z_stream_def_->stream->avail_out -
193 read_bytes;
194 }
195
ReadNBytes(int64 bytes_to_read,tstring * result)196 Status ZlibInputStream::ReadNBytes(int64 bytes_to_read, tstring* result) {
197 if (init_error_) {
198 return errors::DataLoss("Unable to decompress Zlib file.");
199 }
200
201 result->clear();
202 // Read as many bytes as possible from cache.
203 bytes_to_read -= ReadBytesFromCache(bytes_to_read, result);
204
205 while (bytes_to_read > 0) {
206 // At this point we can be sure that cache has been emptied.
207 DCHECK_EQ(NumUnreadBytes(), 0);
208
209 // Now that the cache is empty we need to inflate more data.
210
211 // Step 1. Setup output stream.
212 z_stream_def_->stream->next_out = z_stream_def_->output.get();
213 next_unread_byte_ = reinterpret_cast<char*>(z_stream_def_->output.get());
214 z_stream_def_->stream->avail_out = output_buffer_capacity_;
215
216 // Step 2. Try to inflate some input data.
217 TF_RETURN_IF_ERROR(Inflate());
218
219 // Step 3. Read any data produced by inflate. If no progress was made by
220 // inflate, read more compressed data from the input stream.
221 if (NumUnreadBytes() == 0) {
222 TF_RETURN_IF_ERROR(ReadFromStream());
223 } else {
224 bytes_to_read -= ReadBytesFromCache(bytes_to_read, result);
225 }
226 }
227
228 return Status::OK();
229 }
230
231 #if defined(TF_CORD_SUPPORT)
ReadNBytes(int64 bytes_to_read,absl::Cord * result)232 Status ZlibInputStream::ReadNBytes(int64 bytes_to_read, absl::Cord* result) {
233 // TODO(frankchn): Optimize this instead of bouncing through the buffer.
234 tstring buf;
235 TF_RETURN_IF_ERROR(ReadNBytes(bytes_to_read, &buf));
236 result->Clear();
237 result->Append(buf.data());
238 return Status::OK();
239 }
240 #endif
241
Tell() const242 int64 ZlibInputStream::Tell() const { return bytes_read_; }
243
Inflate()244 Status ZlibInputStream::Inflate() {
245 int error = inflate(z_stream_def_->stream.get(), zlib_options_.flush_mode);
246 // Source: http://zlib.net/manual.html
247 // Z_BUF_ERROR: `inflate` returns Z_BUF_ERROR if no progress was made. This is
248 // not fatal and `inflate` can be called again with more input and output
249 // space to continue inflating.
250 if (error != Z_OK && error != Z_STREAM_END && error != Z_BUF_ERROR) {
251 string error_string =
252 strings::StrCat("inflate() failed with error ", error);
253 if (z_stream_def_->stream->msg != nullptr) {
254 strings::StrAppend(&error_string, ": ", z_stream_def_->stream->msg);
255 }
256 return errors::DataLoss(error_string);
257 }
258 if (error == Z_STREAM_END && zlib_options_.window_bits == MAX_WBITS + 16) {
259 inflateReset(z_stream_def_->stream.get());
260 }
261 return Status::OK();
262 }
263
264 } // namespace io
265 } // namespace tensorflow
266