IOBuf::IOBuf(CopyBufferOp op, const void* buf, uint64_t size,
uint64_t headroom, uint64_t minTailroom)
: IOBuf(CREATE, headroom + size + minTailroom) {
advance(headroom);
memcpy(writableData(), buf, size);
append(size);
}
void CompressionCorruptionTest::runSimpleTest(const DataHolder& dh) {
constexpr uint64_t uncompressedLength = 42;
auto original = IOBuf::wrapBuffer(dh.data(uncompressedLength));
auto compressed = codec_->compress(original.get());
if (!codec_->needsUncompressedLength()) {
auto uncompressed = codec_->uncompress(compressed.get());
EXPECT_EQ(uncompressedLength, uncompressed->computeChainDataLength());
EXPECT_EQ(dh.hash(uncompressedLength), hashIOBuf(uncompressed.get()));
}
{
auto uncompressed = codec_->uncompress(compressed.get(),
uncompressedLength);
EXPECT_EQ(uncompressedLength, uncompressed->computeChainDataLength());
EXPECT_EQ(dh.hash(uncompressedLength), hashIOBuf(uncompressed.get()));
}
EXPECT_THROW(codec_->uncompress(compressed.get(), uncompressedLength + 1),
std::runtime_error);
// Corrupt the first character
++(compressed->writableData()[0]);
if (!codec_->needsUncompressedLength()) {
EXPECT_THROW(codec_->uncompress(compressed.get()),
std::runtime_error);
}
EXPECT_THROW(codec_->uncompress(compressed.get(), uncompressedLength),
std::runtime_error);
}
fbstring IOBuf::moveToFbString() {
// malloc-allocated buffers are just fine, everything else needs
// to be turned into one.
if (!sharedInfo() || // user owned, not ours to give up
sharedInfo()->freeFn || // not malloc()-ed
headroom() != 0 || // malloc()-ed block doesn't start at beginning
tailroom() == 0 || // no room for NUL terminator
isShared() || // shared
isChained()) { // chained
// We might as well get rid of all head and tailroom if we're going
// to reallocate; we need 1 byte for NUL terminator.
coalesceAndReallocate(0, computeChainDataLength(), this, 1);
}
// Ensure NUL terminated
*writableTail() = 0;
fbstring str(reinterpret_cast<char*>(writableData()),
length(), capacity(),
AcquireMallocatedString());
if (flags() & kFlagFreeSharedInfo) {
delete sharedInfo();
}
// Reset to a state where we can be deleted cleanly
flagsAndSharedInfo_ = 0;
buf_ = nullptr;
clear();
return str;
}
std::unique_ptr<folly::IOBuf> FILEReader::operator()(size_t n) {
auto buf = folly::IOBuf::create(n);
size_t bytesRead = fread(buf->writableData(), 1, n, fp_);
if (bytesRead < n) {
folly::throwSystemError("FILEReader: fread");
}
buf->append(n);
return buf;
}
std::unique_ptr<IOBuf> ZlibCodec::addOutputBuffer(z_stream* stream,
uint32_t length) {
CHECK_EQ(stream->avail_out, 0);
auto buf = IOBuf::create(length);
buf->append(length);
stream->next_out = buf->writableData();
stream->avail_out = buf->length();
return buf;
}
std::unique_ptr<IOBuf> ZSTDCodec::doCompress(const IOBuf* data) {
// Support earlier versions of the codec (working with a single IOBuf,
// and using ZSTD_decompress which requires ZSTD frame to contain size,
// which isn't populated by streaming API).
if (!data->isChained()) {
auto out = IOBuf::createCombined(ZSTD_compressBound(data->length()));
const auto rc = ZSTD_compress(
out->writableData(),
out->capacity(),
data->data(),
data->length(),
level_);
zstdThrowIfError(rc);
out->append(rc);
return out;
}
auto zcs = ZSTD_createCStream();
SCOPE_EXIT {
ZSTD_freeCStream(zcs);
};
auto rc = ZSTD_initCStream(zcs, level_);
zstdThrowIfError(rc);
Cursor cursor(data);
auto result = IOBuf::createCombined(ZSTD_compressBound(cursor.totalLength()));
ZSTD_outBuffer out;
out.dst = result->writableTail();
out.size = result->capacity();
out.pos = 0;
for (auto buffer = cursor.peekBytes(); !buffer.empty();) {
ZSTD_inBuffer in;
in.src = buffer.data();
in.size = buffer.size();
for (in.pos = 0; in.pos != in.size;) {
rc = ZSTD_compressStream(zcs, &out, &in);
zstdThrowIfError(rc);
}
cursor.skip(in.size);
buffer = cursor.peekBytes();
}
rc = ZSTD_endStream(zcs, &out);
zstdThrowIfError(rc);
CHECK_EQ(rc, 0);
result->append(out.pos);
return result;
}
std::unique_ptr<IOBuf> ZlibStreamDecompressor::decompress(const IOBuf* in) {
auto out = IOBuf::create(decompressor_buffer_growth_);
auto appender = folly::io::Appender(out.get(), decompressor_buffer_growth_);
const IOBuf* crtBuf = in;
size_t offset = 0;
while (true) {
// Advance to the next IOBuf if necessary
DCHECK_GE(crtBuf->length(), offset);
if (crtBuf->length() == offset) {
crtBuf = crtBuf->next();
offset = 0;
if (crtBuf == in) {
// We hit the end of the IOBuf chain, and are done.
break;
}
}
if (status_ == Z_STREAM_END) {
// we convert this into a stream error
status_ = Z_STREAM_ERROR;
// we should probably bump up a counter here
LOG(ERROR) << "error uncompressing buffer: reached end of zlib data "
"before the end of the buffer";
return nullptr;
}
// Ensure there is space in the output IOBuf
appender.ensure(decompressor_buffer_minsize_);
DCHECK_GT(appender.length(), 0);
const size_t origAvailIn = crtBuf->length() - offset;
zlibStream_.next_in = const_cast<uint8_t*>(crtBuf->data() + offset);
zlibStream_.avail_in = origAvailIn;
zlibStream_.next_out = appender.writableData();
zlibStream_.avail_out = appender.length();
status_ = inflate(&zlibStream_, Z_PARTIAL_FLUSH);
if (status_ != Z_OK && status_ != Z_STREAM_END) {
LOG(INFO) << "error uncompressing buffer: r=" << status_;
return nullptr;
}
// Adjust the input offset ahead
auto inConsumed = origAvailIn - zlibStream_.avail_in;
offset += inConsumed;
// Move output buffer ahead
auto outMove = appender.length() - zlibStream_.avail_out;
appender.append(outMove);
}
return out;
}
IOBuf::IOBuf(CopyBufferOp /* op */,
const void* buf,
uint64_t size,
uint64_t headroom,
uint64_t minTailroom)
: IOBuf(CREATE, headroom + size + minTailroom) {
advance(headroom);
if (size > 0) {
assert(buf != nullptr);
memcpy(writableData(), buf, size);
append(size);
}
}
/**
* @details
* Gets the next chunk of data from the device (if it exists).
*/
void TestChunker::next(QIODevice* device)
{
// Read data off the device.
device->open(QIODevice::ReadOnly);
QByteArray array = device->readAll();
device->close();
if (array.size() == 0)
{
std::cout << "Nothing to read!" << std::endl;
return;
}
// Get writable data object.
WritableData writableData(0);
try
{
writableData = getDataStorage(_size);
}
catch (const QString& e)
{
std::cout << "Unexpected exception: " << e.toStdString() << std::endl;
return;
}
// If the writable data object is not valid, then return.
if (!writableData.isValid())
return;
// ... or, block until writable data is valid.
while (!writableData.isValid()) {
writableData = getDataStorage(_size);
usleep(1);
}
// Write some test data.
writableData.write((void*)array.data(), _size, 0);
}