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 checkReadBytesAvailable(len);
140 uint32_t need = len;
141
142 // TODO(dreiss): Skip urbuf on big reads.
143
144 while (true) {
145 // Copy out whatever we have available, then give them the min of
146 // what we have and what they want, then advance indices.
147 int give = (std::min)((uint32_t)readAvail(), need);
148 memcpy(buf, urbuf_ + urpos_, give);
149 need -= give;
150 buf += give;
151 urpos_ += give;
152
153 // If they were satisfied, we are done.
154 if (need == 0) {
155 return len;
156 }
157
158 // If we will need to read from the underlying transport to get more data,
159 // but we already have some data available, return it now. Reading from
160 // the underlying transport may block, and read() is only allowed to block
161 // when no data is available.
162 if (need < len && rstream_->avail_in == 0) {
163 return len - need;
164 }
165
166 // If we get to this point, we need to get some more data.
167
168 // If zlib has reported the end of a stream, we can't really do any more.
169 if (input_ended_) {
170 return len - need;
171 }
172
173 // The uncompressed read buffer is empty, so reset the stream fields.
174 rstream_->next_out = urbuf_;
175 rstream_->avail_out = urbuf_size_;
176 urpos_ = 0;
177
178 // Call inflate() to uncompress some more data
179 if (!readFromZlib()) {
180 // no data available from underlying transport
181 return len - need;
182 }
183
184 // Okay. The read buffer should have whatever we can give it now.
185 // Loop back to the start and try to give some more.
186 }
187 }
188
readFromZlib()189 bool TZlibTransport::readFromZlib() {
190 assert(!input_ended_);
191
192 // If we don't have any more compressed data available,
193 // read some from the underlying transport.
194 if (rstream_->avail_in == 0) {
195 uint32_t got = transport_->read(crbuf_, crbuf_size_);
196 if (got == 0) {
197 return false;
198 }
199 rstream_->next_in = crbuf_;
200 rstream_->avail_in = got;
201 }
202
203 // We have some compressed data now. Uncompress it.
204 int zlib_rv = inflate(rstream_, Z_SYNC_FLUSH);
205
206 if (zlib_rv == Z_STREAM_END) {
207 input_ended_ = true;
208 } else {
209 checkZlibRv(zlib_rv, rstream_->msg);
210 }
211
212 return true;
213 }
214
215 // WRITING STRATEGY
216 //
217 // We buffer up small writes before sending them to zlib, so our logic is:
218 // - Is the write big?
219 // - Send the buffer to zlib.
220 // - Send this data to zlib.
221 // - Is the write small?
222 // - Is there insufficient space in the buffer for it?
223 // - Send the buffer to zlib.
224 // - Copy the data to the buffer.
225 //
226 // We have two buffers for writing also: the uncompressed buffer (mentioned
227 // above) and the compressed buffer. When sending data to zlib we loop over
228 // the following until the source (uncompressed buffer or big write) is empty:
229 // - Is there no more space in the compressed buffer?
230 // - Write the compressed buffer to the underlying transport.
231 // - Deflate from the source into the compressed buffer.
232
write(const uint8_t * buf,uint32_t len)233 void TZlibTransport::write(const uint8_t* buf, uint32_t len) {
234 if (output_finished_) {
235 throw TTransportException(TTransportException::BAD_ARGS, "write() called after finish()");
236 }
237
238 // zlib's "deflate" function has enough logic in it that I think
239 // we're better off (performance-wise) buffering up small writes.
240 if (len > MIN_DIRECT_DEFLATE_SIZE) {
241 flushToZlib(uwbuf_, uwpos_, Z_NO_FLUSH);
242 uwpos_ = 0;
243 flushToZlib(buf, len, Z_NO_FLUSH);
244 } else if (len > 0) {
245 if (uwbuf_size_ - uwpos_ < len) {
246 flushToZlib(uwbuf_, uwpos_, Z_NO_FLUSH);
247 uwpos_ = 0;
248 }
249 memcpy(uwbuf_ + uwpos_, buf, len);
250 uwpos_ += len;
251 }
252 }
253
flush()254 void TZlibTransport::flush() {
255 if (output_finished_) {
256 throw TTransportException(TTransportException::BAD_ARGS, "flush() called after finish()");
257 }
258
259 flushToZlib(uwbuf_, uwpos_, Z_BLOCK);
260 uwpos_ = 0;
261
262 if(wstream_->avail_out < 6){
263 transport_->write(cwbuf_, cwbuf_size_ - wstream_->avail_out);
264 wstream_->next_out = cwbuf_;
265 wstream_->avail_out = cwbuf_size_;
266 }
267
268 flushToTransport(Z_FULL_FLUSH);
269 resetConsumedMessageSize();
270 }
271
finish()272 void TZlibTransport::finish() {
273 if (output_finished_) {
274 throw TTransportException(TTransportException::BAD_ARGS, "finish() called more than once");
275 }
276
277 flushToTransport(Z_FINISH);
278 }
279
flushToTransport(int flush)280 void TZlibTransport::flushToTransport(int flush) {
281 // write pending data in uwbuf_ to zlib
282 flushToZlib(uwbuf_, uwpos_, flush);
283 uwpos_ = 0;
284
285 // write all available data from zlib to the transport
286 transport_->write(cwbuf_, cwbuf_size_ - wstream_->avail_out);
287 wstream_->next_out = cwbuf_;
288 wstream_->avail_out = cwbuf_size_;
289
290 // flush the transport
291 transport_->flush();
292 }
293
flushToZlib(const uint8_t * buf,int len,int flush)294 void TZlibTransport::flushToZlib(const uint8_t* buf, int len, int flush) {
295 wstream_->next_in = const_cast<uint8_t*>(buf);
296 wstream_->avail_in = len;
297
298 while (true) {
299 if ((flush == Z_NO_FLUSH || flush == Z_BLOCK) && wstream_->avail_in == 0) {
300 break;
301 }
302
303 // If our ouput buffer is full, flush to the underlying transport.
304 if (wstream_->avail_out == 0) {
305 transport_->write(cwbuf_, cwbuf_size_);
306 wstream_->next_out = cwbuf_;
307 wstream_->avail_out = cwbuf_size_;
308 }
309
310 int zlib_rv = deflate(wstream_, flush);
311
312 if (flush == Z_FINISH && zlib_rv == Z_STREAM_END) {
313 assert(wstream_->avail_in == 0);
314 output_finished_ = true;
315 break;
316 }
317
318 checkZlibRv(zlib_rv, wstream_->msg);
319
320 if ((flush == Z_SYNC_FLUSH || flush == Z_FULL_FLUSH) && wstream_->avail_in == 0
321 && wstream_->avail_out != 0) {
322 break;
323 }
324 }
325 }
326
borrow(uint8_t * buf,uint32_t * len)327 const uint8_t* TZlibTransport::borrow(uint8_t* buf, uint32_t* len) {
328 (void)buf;
329 // Don't try to be clever with shifting buffers.
330 // If we have enough data, give a pointer to it,
331 // otherwise let the protcol use its slow path.
332 if (readAvail() >= (int)*len) {
333 *len = (uint32_t)readAvail();
334 return urbuf_ + urpos_;
335 }
336 return nullptr;
337 }
338
consume(uint32_t len)339 void TZlibTransport::consume(uint32_t len) {
340 countConsumedMessageBytes(len);
341 if (readAvail() >= (int)len) {
342 urpos_ += len;
343 } else {
344 throw TTransportException(TTransportException::BAD_ARGS, "consume did not follow a borrow.");
345 }
346 }
347
verifyChecksum()348 void TZlibTransport::verifyChecksum() {
349 // If zlib has already reported the end of the stream,
350 // it has verified the checksum.
351 if (input_ended_) {
352 return;
353 }
354
355 // This should only be called when reading is complete.
356 // If the caller still has unread data, throw an exception.
357 if (readAvail() > 0) {
358 throw TTransportException(TTransportException::CORRUPTED_DATA,
359 "verifyChecksum() called before end of zlib stream");
360 }
361
362 // Reset the rstream fields, in case avail_out is 0.
363 // (Since readAvail() is 0, we know there is no unread data in urbuf_)
364 rstream_->next_out = urbuf_;
365 rstream_->avail_out = urbuf_size_;
366 urpos_ = 0;
367
368 // Call inflate()
369 // This will throw an exception if the checksum is bad.
370 bool performed_inflate = readFromZlib();
371 if (!performed_inflate) {
372 // We needed to read from the underlying transport, and the read() call
373 // returned 0.
374 //
375 // Not all TTransport implementations behave the same way here, so we'll
376 // end up with different behavior depending on the underlying transport.
377 //
378 // For some transports (e.g., TFDTransport), read() blocks if no more data
379 // is available. They only return 0 if EOF has been reached, or if the
380 // remote endpoint has closed the connection. For those transports,
381 // verifyChecksum() will block until the checksum becomes available.
382 //
383 // Other transport types (e.g., TMemoryBuffer) always return 0 immediately
384 // if no more data is available. For those transport types, verifyChecksum
385 // will raise the following exception if the checksum is not available from
386 // the underlying transport yet.
387 throw TTransportException(TTransportException::CORRUPTED_DATA,
388 "checksum not available yet in "
389 "verifyChecksum()");
390 }
391
392 // If input_ended_ is true now, the checksum has been verified
393 if (input_ended_) {
394 return;
395 }
396
397 // The caller invoked us before the actual end of the data stream
398 assert(rstream_->avail_out < urbuf_size_);
399 throw TTransportException(TTransportException::CORRUPTED_DATA,
400 "verifyChecksum() called before end of "
401 "zlib stream");
402 }
403
TZlibTransportFactory(std::shared_ptr<TTransportFactory> transportFactory)404 TZlibTransportFactory::TZlibTransportFactory(std::shared_ptr<TTransportFactory> transportFactory)
405 :transportFactory_(transportFactory) {
406 }
407
getTransport(std::shared_ptr<TTransport> trans)408 std::shared_ptr<TTransport> TZlibTransportFactory::getTransport(std::shared_ptr<TTransport> trans) {
409 if (transportFactory_) {
410 return std::shared_ptr<TTransport>(new TZlibTransport(transportFactory_->getTransport(trans)));
411 } else {
412 return std::shared_ptr<TTransport>(new TZlibTransport(trans));
413 }
414 }
415 }
416 }
417 } // apache::thrift::transport
418