1 /*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20 #include <cassert>
21 #include <cstring>
22 #include <algorithm>
23 #include <thrift/transport/TZlibTransport.h>
24
25 using std::string;
26
27 namespace apache {
28 namespace thrift {
29 namespace transport {
30
31 // Don't call this outside of the constructor.
initZlib()32 void TZlibTransport::initZlib() {
33 int rv;
34 bool r_init = false;
35 try {
36 rstream_ = new z_stream;
37 wstream_ = new z_stream;
38
39 rstream_->zalloc = Z_NULL;
40 wstream_->zalloc = Z_NULL;
41 rstream_->zfree = Z_NULL;
42 wstream_->zfree = Z_NULL;
43 rstream_->opaque = Z_NULL;
44 wstream_->opaque = Z_NULL;
45
46 rstream_->next_in = crbuf_;
47 wstream_->next_in = uwbuf_;
48 rstream_->next_out = urbuf_;
49 wstream_->next_out = cwbuf_;
50 rstream_->avail_in = 0;
51 wstream_->avail_in = 0;
52 rstream_->avail_out = urbuf_size_;
53 wstream_->avail_out = cwbuf_size_;
54
55 rv = inflateInit(rstream_);
56 checkZlibRv(rv, rstream_->msg);
57
58 // Have to set this flag so we know whether to de-initialize.
59 r_init = true;
60
61 rv = deflateInit(wstream_, comp_level_);
62 checkZlibRv(rv, wstream_->msg);
63 }
64
65 catch (...) {
66 if (r_init) {
67 rv = inflateEnd(rstream_);
68 checkZlibRvNothrow(rv, rstream_->msg);
69 }
70 // There is no way we can get here if wstream_ was initialized.
71
72 throw;
73 }
74 }
75
checkZlibRv(int status,const char * message)76 inline void TZlibTransport::checkZlibRv(int status, const char* message) {
77 if (status != Z_OK) {
78 throw TZlibTransportException(status, message);
79 }
80 }
81
checkZlibRvNothrow(int status,const char * message)82 inline void TZlibTransport::checkZlibRvNothrow(int status, const char* message) {
83 if (status != Z_OK) {
84 string output = "TZlibTransport: zlib failure in destructor: "
85 + TZlibTransportException::errorMessage(status, message);
86 GlobalOutput(output.c_str());
87 }
88 }
89
~TZlibTransport()90 TZlibTransport::~TZlibTransport() {
91 int rv;
92 rv = inflateEnd(rstream_);
93 checkZlibRvNothrow(rv, rstream_->msg);
94
95 rv = deflateEnd(wstream_);
96 // Z_DATA_ERROR may be returned if the caller has written data, but not
97 // called flush() to actually finish writing the data out to the underlying
98 // transport. The defined TTransport behavior in this case is that this data
99 // may be discarded, so we ignore the error and silently discard the data.
100 // For other erros, log a message.
101 if (rv != Z_DATA_ERROR) {
102 checkZlibRvNothrow(rv, wstream_->msg);
103 }
104
105 delete[] urbuf_;
106 delete[] crbuf_;
107 delete[] uwbuf_;
108 delete[] cwbuf_;
109 delete rstream_;
110 delete wstream_;
111 }
112
isOpen() const113 bool TZlibTransport::isOpen() const {
114 return (readAvail() > 0) || (rstream_->avail_in > 0) || transport_->isOpen();
115 }
116
peek()117 bool TZlibTransport::peek() {
118 return (readAvail() > 0) || (rstream_->avail_in > 0) || transport_->peek();
119 }
120
121 // READING STRATEGY
122 //
123 // We have two buffers for reading: one containing the compressed data (crbuf_)
124 // and one containing the uncompressed data (urbuf_). When read is called,
125 // we repeat the following steps until we have satisfied the request:
126 // - Copy data from urbuf_ into the caller's buffer.
127 // - If we had enough, return.
128 // - If urbuf_ is empty, read some data into it from the underlying transport.
129 // - Inflate data from crbuf_ into urbuf_.
130 //
131 // In standalone objects, we set input_ended_ to true when inflate returns
132 // Z_STREAM_END. This allows to make sure that a checksum was verified.
133
readAvail() const134 inline int TZlibTransport::readAvail() const {
135 return urbuf_size_ - rstream_->avail_out - urpos_;
136 }
137
read(uint8_t * buf,uint32_t len)138 uint32_t TZlibTransport::read(uint8_t* buf, uint32_t len) {
139 uint32_t need = len;
140
141 // TODO(dreiss): Skip urbuf on big reads.
142
143 while (true) {
144 // Copy out whatever we have available, then give them the min of
145 // what we have and what they want, then advance indices.
146 int give = (std::min)((uint32_t)readAvail(), need);
147 memcpy(buf, urbuf_ + urpos_, give);
148 need -= give;
149 buf += give;
150 urpos_ += give;
151
152 // If they were satisfied, we are done.
153 if (need == 0) {
154 return len;
155 }
156
157 // If we will need to read from the underlying transport to get more data,
158 // but we already have some data available, return it now. Reading from
159 // the underlying transport may block, and read() is only allowed to block
160 // when no data is available.
161 if (need < len && rstream_->avail_in == 0) {
162 return len - need;
163 }
164
165 // If we get to this point, we need to get some more data.
166
167 // If zlib has reported the end of a stream, we can't really do any more.
168 if (input_ended_) {
169 return len - need;
170 }
171
172 // The uncompressed read buffer is empty, so reset the stream fields.
173 rstream_->next_out = urbuf_;
174 rstream_->avail_out = urbuf_size_;
175 urpos_ = 0;
176
177 // Call inflate() to uncompress some more data
178 if (!readFromZlib()) {
179 // no data available from underlying transport
180 return len - need;
181 }
182
183 // Okay. The read buffer should have whatever we can give it now.
184 // Loop back to the start and try to give some more.
185 }
186 }
187
readFromZlib()188 bool TZlibTransport::readFromZlib() {
189 assert(!input_ended_);
190
191 // If we don't have any more compressed data available,
192 // read some from the underlying transport.
193 if (rstream_->avail_in == 0) {
194 uint32_t got = transport_->read(crbuf_, crbuf_size_);
195 if (got == 0) {
196 return false;
197 }
198 rstream_->next_in = crbuf_;
199 rstream_->avail_in = got;
200 }
201
202 // We have some compressed data now. Uncompress it.
203 int zlib_rv = inflate(rstream_, Z_SYNC_FLUSH);
204
205 if (zlib_rv == Z_STREAM_END) {
206 input_ended_ = true;
207 } else {
208 checkZlibRv(zlib_rv, rstream_->msg);
209 }
210
211 return true;
212 }
213
214 // WRITING STRATEGY
215 //
216 // We buffer up small writes before sending them to zlib, so our logic is:
217 // - Is the write big?
218 // - Send the buffer to zlib.
219 // - Send this data to zlib.
220 // - Is the write small?
221 // - Is there insufficient space in the buffer for it?
222 // - Send the buffer to zlib.
223 // - Copy the data to the buffer.
224 //
225 // We have two buffers for writing also: the uncompressed buffer (mentioned
226 // above) and the compressed buffer. When sending data to zlib we loop over
227 // the following until the source (uncompressed buffer or big write) is empty:
228 // - Is there no more space in the compressed buffer?
229 // - Write the compressed buffer to the underlying transport.
230 // - Deflate from the source into the compressed buffer.
231
write(const uint8_t * buf,uint32_t len)232 void TZlibTransport::write(const uint8_t* buf, uint32_t len) {
233 if (output_finished_) {
234 throw TTransportException(TTransportException::BAD_ARGS, "write() called after finish()");
235 }
236
237 // zlib's "deflate" function has enough logic in it that I think
238 // we're better off (performance-wise) buffering up small writes.
239 if (len > MIN_DIRECT_DEFLATE_SIZE) {
240 flushToZlib(uwbuf_, uwpos_, Z_NO_FLUSH);
241 uwpos_ = 0;
242 flushToZlib(buf, len, Z_NO_FLUSH);
243 } else if (len > 0) {
244 if (uwbuf_size_ - uwpos_ < len) {
245 flushToZlib(uwbuf_, uwpos_, Z_NO_FLUSH);
246 uwpos_ = 0;
247 }
248 memcpy(uwbuf_ + uwpos_, buf, len);
249 uwpos_ += len;
250 }
251 }
252
flush()253 void TZlibTransport::flush() {
254 if (output_finished_) {
255 throw TTransportException(TTransportException::BAD_ARGS, "flush() called after finish()");
256 }
257
258 flushToZlib(uwbuf_, uwpos_, Z_BLOCK);
259 uwpos_ = 0;
260
261 if(wstream_->avail_out < 6){
262 transport_->write(cwbuf_, cwbuf_size_ - wstream_->avail_out);
263 wstream_->next_out = cwbuf_;
264 wstream_->avail_out = cwbuf_size_;
265 }
266
267 flushToTransport(Z_FULL_FLUSH);
268 }
269
finish()270 void TZlibTransport::finish() {
271 if (output_finished_) {
272 throw TTransportException(TTransportException::BAD_ARGS, "finish() called more than once");
273 }
274
275 flushToTransport(Z_FINISH);
276 }
277
flushToTransport(int flush)278 void TZlibTransport::flushToTransport(int flush) {
279 // write pending data in uwbuf_ to zlib
280 flushToZlib(uwbuf_, uwpos_, flush);
281 uwpos_ = 0;
282
283 // write all available data from zlib to the transport
284 transport_->write(cwbuf_, cwbuf_size_ - wstream_->avail_out);
285 wstream_->next_out = cwbuf_;
286 wstream_->avail_out = cwbuf_size_;
287
288 // flush the transport
289 transport_->flush();
290 }
291
flushToZlib(const uint8_t * buf,int len,int flush)292 void TZlibTransport::flushToZlib(const uint8_t* buf, int len, int flush) {
293 wstream_->next_in = const_cast<uint8_t*>(buf);
294 wstream_->avail_in = len;
295
296 while (true) {
297 if ((flush == Z_NO_FLUSH || flush == Z_BLOCK) && wstream_->avail_in == 0) {
298 break;
299 }
300
301 // If our ouput buffer is full, flush to the underlying transport.
302 if (wstream_->avail_out == 0) {
303 transport_->write(cwbuf_, cwbuf_size_);
304 wstream_->next_out = cwbuf_;
305 wstream_->avail_out = cwbuf_size_;
306 }
307
308 int zlib_rv = deflate(wstream_, flush);
309
310 if (flush == Z_FINISH && zlib_rv == Z_STREAM_END) {
311 assert(wstream_->avail_in == 0);
312 output_finished_ = true;
313 break;
314 }
315
316 checkZlibRv(zlib_rv, wstream_->msg);
317
318 if ((flush == Z_SYNC_FLUSH || flush == Z_FULL_FLUSH) && wstream_->avail_in == 0
319 && wstream_->avail_out != 0) {
320 break;
321 }
322 }
323 }
324
borrow(uint8_t * buf,uint32_t * len)325 const uint8_t* TZlibTransport::borrow(uint8_t* buf, uint32_t* len) {
326 (void)buf;
327 // Don't try to be clever with shifting buffers.
328 // If we have enough data, give a pointer to it,
329 // otherwise let the protcol use its slow path.
330 if (readAvail() >= (int)*len) {
331 *len = (uint32_t)readAvail();
332 return urbuf_ + urpos_;
333 }
334 return nullptr;
335 }
336
consume(uint32_t len)337 void TZlibTransport::consume(uint32_t len) {
338 if (readAvail() >= (int)len) {
339 urpos_ += len;
340 } else {
341 throw TTransportException(TTransportException::BAD_ARGS, "consume did not follow a borrow.");
342 }
343 }
344
verifyChecksum()345 void TZlibTransport::verifyChecksum() {
346 // If zlib has already reported the end of the stream,
347 // it has verified the checksum.
348 if (input_ended_) {
349 return;
350 }
351
352 // This should only be called when reading is complete.
353 // If the caller still has unread data, throw an exception.
354 if (readAvail() > 0) {
355 throw TTransportException(TTransportException::CORRUPTED_DATA,
356 "verifyChecksum() called before end of zlib stream");
357 }
358
359 // Reset the rstream fields, in case avail_out is 0.
360 // (Since readAvail() is 0, we know there is no unread data in urbuf_)
361 rstream_->next_out = urbuf_;
362 rstream_->avail_out = urbuf_size_;
363 urpos_ = 0;
364
365 // Call inflate()
366 // This will throw an exception if the checksum is bad.
367 bool performed_inflate = readFromZlib();
368 if (!performed_inflate) {
369 // We needed to read from the underlying transport, and the read() call
370 // returned 0.
371 //
372 // Not all TTransport implementations behave the same way here, so we'll
373 // end up with different behavior depending on the underlying transport.
374 //
375 // For some transports (e.g., TFDTransport), read() blocks if no more data
376 // is available. They only return 0 if EOF has been reached, or if the
377 // remote endpoint has closed the connection. For those transports,
378 // verifyChecksum() will block until the checksum becomes available.
379 //
380 // Other transport types (e.g., TMemoryBuffer) always return 0 immediately
381 // if no more data is available. For those transport types, verifyChecksum
382 // will raise the following exception if the checksum is not available from
383 // the underlying transport yet.
384 throw TTransportException(TTransportException::CORRUPTED_DATA,
385 "checksum not available yet in "
386 "verifyChecksum()");
387 }
388
389 // If input_ended_ is true now, the checksum has been verified
390 if (input_ended_) {
391 return;
392 }
393
394 // The caller invoked us before the actual end of the data stream
395 assert(rstream_->avail_out < urbuf_size_);
396 throw TTransportException(TTransportException::CORRUPTED_DATA,
397 "verifyChecksum() called before end of "
398 "zlib stream");
399 }
400 }
401 }
402 } // apache::thrift::transport
403