diff options
Diffstat (limited to 'lib/cpp/src/transport/TZlibTransport.cpp')
-rw-r--r-- | lib/cpp/src/transport/TZlibTransport.cpp | 299 |
1 files changed, 299 insertions, 0 deletions
diff --git a/lib/cpp/src/transport/TZlibTransport.cpp b/lib/cpp/src/transport/TZlibTransport.cpp new file mode 100644 index 000000000..2f14e906b --- /dev/null +++ b/lib/cpp/src/transport/TZlibTransport.cpp @@ -0,0 +1,299 @@ +/* + * 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 <transport/TZlibTransport.h> +#include <zlib.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_, Z_DEFAULT_COMPRESSION); + 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_); + checkZlibRvNothrow(rv, wstream_->msg); + + delete[] urbuf_; + delete[] crbuf_; + delete[] uwbuf_; + delete[] cwbuf_; + delete rstream_; + delete wstream_; +} + +bool TZlibTransport::isOpen() { + return (readAvail() > 0) || transport_->isOpen(); +} + +// 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() { + return urbuf_size_ - rstream_->avail_out - urpos_; +} + +uint32_t TZlibTransport::read(uint8_t* buf, uint32_t len) { + int 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(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 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; + + // 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 len - need; + } + 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) { + if (standalone_) { + input_ended_ = true; + } + } else { + checkZlibRv(zlib_rv, rstream_->msg); + } + + // Okay. The read buffer should have whatever we can give it now. + // Loop back to the start and try to give some more. + } +} + + +// 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) { + // zlib's "deflate" function has enough logic in it that I think + // we're better off (performance-wise) buffering up small writes. + if ((int)len > MIN_DIRECT_DEFLATE_SIZE) { + flushToZlib(uwbuf_, uwpos_); + uwpos_ = 0; + flushToZlib(buf, len); + } else if (len > 0) { + if (uwbuf_size_ - uwpos_ < (int)len) { + flushToZlib(uwbuf_, uwpos_); + uwpos_ = 0; + } + memcpy(uwbuf_ + uwpos_, buf, len); + uwpos_ += len; + } +} + +void TZlibTransport::flush() { + flushToZlib(uwbuf_, uwpos_, true); + assert((int)wstream_->avail_out != cwbuf_size_); + transport_->write(cwbuf_, cwbuf_size_ - wstream_->avail_out); + transport_->flush(); +} + +void TZlibTransport::flushToZlib(const uint8_t* buf, int len, bool finish) { + int flush = (finish ? Z_FINISH : Z_NO_FLUSH); + + wstream_->next_in = const_cast<uint8_t*>(buf); + wstream_->avail_in = len; + + while (wstream_->avail_in > 0 || finish) { + // 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 (finish && zlib_rv == Z_STREAM_END) { + assert(wstream_->avail_in == 0); + break; + } + + checkZlibRv(zlib_rv, wstream_->msg); + } +} + +const uint8_t* TZlibTransport::borrow(uint8_t* buf, uint32_t* len) { + // 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 NULL; +} + +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 (!standalone_) { + throw TTransportException( + TTransportException::BAD_ARGS, + "TZLibTransport can only verify checksums for standalone objects."); + } + + if (!input_ended_) { + // This should only be called when reading is complete, + // but it's possible that the whole checksum has not been fed to zlib yet. + // We try to read an extra byte here to force zlib to finish the stream. + // It might not always be easy to "unread" this byte, + // but we throw an exception if we get it, which is not really + // a recoverable error, so it doesn't matter. + uint8_t buf[1]; + uint32_t got = this->read(buf, sizeof(buf)); + if (got || !input_ended_) { + throw TTransportException( + TTransportException::CORRUPTED_DATA, + "Zlib stream not complete."); + } + } + + // If the checksum had been bad, we would have gotten an error while + // inflating. +} + + +}}} // apache::thrift::transport |