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.");
}
}
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.
}
}
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");
}
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.
}
}
inline int readAvail(String s) { return readAvail(s->data(), s->size()); }
bool TZlibTransport::isOpen() {
return (readAvail() > 0) || transport_->isOpen();
}
bool TZlibTransport::peek() {
return (readAvail() > 0) || (rstream_->avail_in > 0) || transport_->peek();
}
bool TZlibTransport::isOpen() {
return (readAvail() > 0) || (rstream_->avail_in > 0) || transport_->isOpen();
}
