-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you 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 <cassert>
-#include <cstring>
-#include <algorithm>
-#include <thrift/transport/TZlibTransport.h>
-
-using std::string;
-
-namespace apache {
-namespace thrift {
-namespace transport {
-
-// Don't call this outside of the constructor.
-void TZlibTransport::initZlib() {
- int rv;
- bool r_init = false;
- try {
- rstream_ = new z_stream;
- wstream_ = new z_stream;
-
- rstream_->zalloc = Z_NULL;
- wstream_->zalloc = Z_NULL;
- rstream_->zfree = Z_NULL;
- wstream_->zfree = Z_NULL;
- rstream_->opaque = Z_NULL;
- wstream_->opaque = Z_NULL;
-
- rstream_->next_in = crbuf_;
- wstream_->next_in = uwbuf_;
- rstream_->next_out = urbuf_;
- wstream_->next_out = cwbuf_;
- rstream_->avail_in = 0;
- wstream_->avail_in = 0;
- rstream_->avail_out = urbuf_size_;
- wstream_->avail_out = cwbuf_size_;
-
- rv = inflateInit(rstream_);
- checkZlibRv(rv, rstream_->msg);
-
- // Have to set this flag so we know whether to de-initialize.
- r_init = true;
-
- rv = deflateInit(wstream_, comp_level_);
- checkZlibRv(rv, wstream_->msg);
- }
-
- catch (...) {
- if (r_init) {
- rv = inflateEnd(rstream_);
- checkZlibRvNothrow(rv, rstream_->msg);
- }
- // There is no way we can get here if wstream_ was initialized.
-
- throw;
- }
-}
-
-inline void TZlibTransport::checkZlibRv(int status, const char* message) {
- if (status != Z_OK) {
- throw TZlibTransportException(status, message);
- }
-}
-
-inline void TZlibTransport::checkZlibRvNothrow(int status, const char* message) {
- if (status != Z_OK) {
- string output = "TZlibTransport: zlib failure in destructor: "
- + TZlibTransportException::errorMessage(status, message);
- GlobalOutput(output.c_str());
- }
-}
-
-TZlibTransport::~TZlibTransport() {
- int rv;
- rv = inflateEnd(rstream_);
- checkZlibRvNothrow(rv, rstream_->msg);
-
- rv = deflateEnd(wstream_);
- // Z_DATA_ERROR may be returned if the caller has written data, but not
- // called flush() to actually finish writing the data out to the underlying
- // transport. The defined TTransport behavior in this case is that this data
- // may be discarded, so we ignore the error and silently discard the data.
- // For other erros, log a message.
- if (rv != Z_DATA_ERROR) {
- checkZlibRvNothrow(rv, wstream_->msg);
- }
-
- delete[] urbuf_;
- delete[] crbuf_;
- delete[] uwbuf_;
- delete[] cwbuf_;
- delete rstream_;
- delete wstream_;
-}
-
-bool TZlibTransport::isOpen() const {
- return (readAvail() > 0) || (rstream_->avail_in > 0) || transport_->isOpen();
-}
-
-bool TZlibTransport::peek() {
- return (readAvail() > 0) || (rstream_->avail_in > 0) || transport_->peek();
-}
-
-// READING STRATEGY
-//
-// We have two buffers for reading: one containing the compressed data (crbuf_)
-// and one containing the uncompressed data (urbuf_). When read is called,
-// we repeat the following steps until we have satisfied the request:
-// - Copy data from urbuf_ into the caller's buffer.
-// - If we had enough, return.
-// - If urbuf_ is empty, read some data into it from the underlying transport.
-// - Inflate data from crbuf_ into urbuf_.
-//
-// In standalone objects, we set input_ended_ to true when inflate returns
-// Z_STREAM_END. This allows to make sure that a checksum was verified.
-
-inline int TZlibTransport::readAvail() const {
- return urbuf_size_ - rstream_->avail_out - urpos_;
-}
-
-uint32_t TZlibTransport::read(uint8_t* buf, uint32_t len) {
- uint32_t need = len;
-
- // TODO(dreiss): Skip urbuf on big reads.
-
- while (true) {
- // Copy out whatever we have available, then give them the min of
- // what we have and what they want, then advance indices.
- int give = (std::min)((uint32_t)readAvail(), need);
- memcpy(buf, urbuf_ + urpos_, give);
- need -= give;
- buf += give;
- urpos_ += give;
-
- // If they were satisfied, we are done.
- if (need == 0) {
- return len;
- }
-
- // If we will need to read from the underlying transport to get more data,
- // but we already have some data available, return it now. Reading from
- // the underlying transport may block, and read() is only allowed to block
- // when no data is available.
- if (need < len && rstream_->avail_in == 0) {
- return len - need;
- }
-
- // If we get to this point, we need to get some more data.
-
- // If zlib has reported the end of a stream, we can't really do any more.
- if (input_ended_) {
- return len - need;
- }
-
- // The uncompressed read buffer is empty, so reset the stream fields.
- rstream_->next_out = urbuf_;
- rstream_->avail_out = urbuf_size_;
- urpos_ = 0;
-
- // Call inflate() to uncompress some more data
- if (!readFromZlib()) {
- // no data available from underlying transport
- return len - need;
- }
-
- // Okay. The read buffer should have whatever we can give it now.
- // Loop back to the start and try to give some more.
- }
-}
-
-bool TZlibTransport::readFromZlib() {
- assert(!input_ended_);
-
- // If we don't have any more compressed data available,
- // read some from the underlying transport.
- if (rstream_->avail_in == 0) {
- uint32_t got = transport_->read(crbuf_, crbuf_size_);
- if (got == 0) {
- return false;
- }
- rstream_->next_in = crbuf_;
- rstream_->avail_in = got;
- }
-
- // We have some compressed data now. Uncompress it.
- int zlib_rv = inflate(rstream_, Z_SYNC_FLUSH);
-
- if (zlib_rv == Z_STREAM_END) {
- input_ended_ = true;
- } else {
- checkZlibRv(zlib_rv, rstream_->msg);
- }
-
- return true;
-}
-
-// WRITING STRATEGY
-//
-// We buffer up small writes before sending them to zlib, so our logic is:
-// - Is the write big?
-// - Send the buffer to zlib.
-// - Send this data to zlib.
-// - Is the write small?
-// - Is there insufficient space in the buffer for it?
-// - Send the buffer to zlib.
-// - Copy the data to the buffer.
-//
-// We have two buffers for writing also: the uncompressed buffer (mentioned
-// above) and the compressed buffer. When sending data to zlib we loop over
-// the following until the source (uncompressed buffer or big write) is empty:
-// - Is there no more space in the compressed buffer?
-// - Write the compressed buffer to the underlying transport.
-// - Deflate from the source into the compressed buffer.
-
-void TZlibTransport::write(const uint8_t* buf, uint32_t len) {
- if (output_finished_) {
- throw TTransportException(TTransportException::BAD_ARGS, "write() called after finish()");
- }
-
- // zlib's "deflate" function has enough logic in it that I think
- // we're better off (performance-wise) buffering up small writes.
- if (len > MIN_DIRECT_DEFLATE_SIZE) {
- flushToZlib(uwbuf_, uwpos_, Z_NO_FLUSH);
- uwpos_ = 0;
- flushToZlib(buf, len, Z_NO_FLUSH);
- } else if (len > 0) {
- if (uwbuf_size_ - uwpos_ < len) {
- flushToZlib(uwbuf_, uwpos_, Z_NO_FLUSH);
- uwpos_ = 0;
- }
- memcpy(uwbuf_ + uwpos_, buf, len);
- uwpos_ += len;
- }
-}
-
-void TZlibTransport::flush() {
- if (output_finished_) {
- throw TTransportException(TTransportException::BAD_ARGS, "flush() called after finish()");
- }
-
- flushToZlib(uwbuf_, uwpos_, Z_BLOCK);
- uwpos_ = 0;
-
- if(wstream_->avail_out < 6){
- transport_->write(cwbuf_, cwbuf_size_ - wstream_->avail_out);
- wstream_->next_out = cwbuf_;
- wstream_->avail_out = cwbuf_size_;
- }
-
- flushToTransport(Z_FULL_FLUSH);
-}
-
-void TZlibTransport::finish() {
- if (output_finished_) {
- throw TTransportException(TTransportException::BAD_ARGS, "finish() called more than once");
- }
-
- flushToTransport(Z_FINISH);
-}
-
-void TZlibTransport::flushToTransport(int flush) {
- // write pending data in uwbuf_ to zlib
- flushToZlib(uwbuf_, uwpos_, flush);
- uwpos_ = 0;
-
- // write all available data from zlib to the transport
- transport_->write(cwbuf_, cwbuf_size_ - wstream_->avail_out);
- wstream_->next_out = cwbuf_;
- wstream_->avail_out = cwbuf_size_;
-
- // flush the transport
- transport_->flush();
-}
-
-void TZlibTransport::flushToZlib(const uint8_t* buf, int len, int flush) {
- wstream_->next_in = const_cast<uint8_t*>(buf);
- wstream_->avail_in = len;
-
- while (true) {
- if ((flush == Z_NO_FLUSH || flush == Z_BLOCK) && wstream_->avail_in == 0) {
- break;
- }
-
- // If our ouput buffer is full, flush to the underlying transport.
- if (wstream_->avail_out == 0) {
- transport_->write(cwbuf_, cwbuf_size_);
- wstream_->next_out = cwbuf_;
- wstream_->avail_out = cwbuf_size_;
- }
-
- int zlib_rv = deflate(wstream_, flush);
-
- if (flush == Z_FINISH && zlib_rv == Z_STREAM_END) {
- assert(wstream_->avail_in == 0);
- output_finished_ = true;
- break;
- }
-
- checkZlibRv(zlib_rv, wstream_->msg);
-
- if ((flush == Z_SYNC_FLUSH || flush == Z_FULL_FLUSH) && wstream_->avail_in == 0
- && wstream_->avail_out != 0) {
- break;
- }
- }
-}
-
-const uint8_t* TZlibTransport::borrow(uint8_t* buf, uint32_t* len) {
- (void)buf;
- // Don't try to be clever with shifting buffers.
- // If we have enough data, give a pointer to it,
- // otherwise let the protcol use its slow path.
- if (readAvail() >= (int)*len) {
- *len = (uint32_t)readAvail();
- return urbuf_ + urpos_;
- }
- return nullptr;
-}
-
-void TZlibTransport::consume(uint32_t len) {
- if (readAvail() >= (int)len) {
- urpos_ += len;
- } else {
- throw TTransportException(TTransportException::BAD_ARGS, "consume did not follow a borrow.");
- }
-}
-
-void TZlibTransport::verifyChecksum() {
- // If zlib has already reported the end of the stream,
- // it has verified the checksum.
- if (input_ended_) {
- return;
- }
-
- // This should only be called when reading is complete.
- // If the caller still has unread data, throw an exception.
- if (readAvail() > 0) {
- throw TTransportException(TTransportException::CORRUPTED_DATA,
- "verifyChecksum() called before end of zlib stream");
- }
-
- // Reset the rstream fields, in case avail_out is 0.
- // (Since readAvail() is 0, we know there is no unread data in urbuf_)
- rstream_->next_out = urbuf_;
- rstream_->avail_out = urbuf_size_;
- urpos_ = 0;
-
- // Call inflate()
- // This will throw an exception if the checksum is bad.
- bool performed_inflate = readFromZlib();
- if (!performed_inflate) {
- // We needed to read from the underlying transport, and the read() call
- // returned 0.
- //
- // Not all TTransport implementations behave the same way here, so we'll
- // end up with different behavior depending on the underlying transport.
- //
- // For some transports (e.g., TFDTransport), read() blocks if no more data
- // is available. They only return 0 if EOF has been reached, or if the
- // remote endpoint has closed the connection. For those transports,
- // verifyChecksum() will block until the checksum becomes available.
- //
- // Other transport types (e.g., TMemoryBuffer) always return 0 immediately
- // if no more data is available. For those transport types, verifyChecksum
- // will raise the following exception if the checksum is not available from
- // the underlying transport yet.
- throw TTransportException(TTransportException::CORRUPTED_DATA,
- "checksum not available yet in "
- "verifyChecksum()");
- }
-
- // If input_ended_ is true now, the checksum has been verified
- if (input_ended_) {
- return;
- }
-
- // The caller invoked us before the actual end of the data stream
- assert(rstream_->avail_out < urbuf_size_);
- throw TTransportException(TTransportException::CORRUPTED_DATA,
- "verifyChecksum() called before end of "
- "zlib stream");
-}
-}
-}
-} // apache::thrift::transport