ByteArray DecompressFile(const FilePath& path)
{
BinaryReader reader(path);
if (!reader)
{
return ByteArray();
}
const size_t inputBufferSize = ZSTD_DStreamInSize();
const auto pInputBuffer = std::make_unique<Byte[]>(inputBufferSize);
const size_t outputBufferSize = ZSTD_DStreamOutSize();
const auto pOutputBuffer = std::make_unique<Byte[]>(outputBufferSize);
ZSTD_DStream* const dStream = ZSTD_createDStream();
if (!dStream)
{
return ByteArray();
}
const size_t initResult = ZSTD_initDStream(dStream);
if (ZSTD_isError(initResult))
{
ZSTD_freeDStream(dStream);
return ByteArray();
}
size_t toRead = initResult;
Array<Byte> buffer;
while (const size_t read = static_cast<size_t>(reader.read(pInputBuffer.get(), toRead)))
{
ZSTD_inBuffer input = { pInputBuffer.get(), read, 0 };
while (input.pos < input.size)
{
ZSTD_outBuffer output = { pOutputBuffer.get(), outputBufferSize, 0 };
toRead = ZSTD_decompressStream(dStream, &output, &input);
if (ZSTD_isError(toRead))
{
ZSTD_freeDStream(dStream);
return ByteArray();
}
buffer.insert(buffer.end(), pOutputBuffer.get(), pOutputBuffer.get() + output.pos);
}
}
ZSTD_freeDStream(dStream);
return ByteArray(std::move(buffer));
}
int LLVMFuzzerTestOneInput(const uint8_t *src, size_t size)
{
seed = FUZZ_seed(&src, &size);
/* Allocate all buffers and contexts if not already allocated */
if (!buf) {
buf = malloc(kBufSize);
FUZZ_ASSERT(buf);
}
if (!dstream) {
dstream = ZSTD_createDStream();
FUZZ_ASSERT(dstream);
FUZZ_ASSERT(!ZSTD_isError(ZSTD_initDStream(dstream)));
} else {
FUZZ_ASSERT(!ZSTD_isError(ZSTD_resetDStream(dstream)));
}
while (size > 0) {
ZSTD_inBuffer in = makeInBuffer(&src, &size);
while (in.pos != in.size) {
ZSTD_outBuffer out = makeOutBuffer();
size_t const rc = ZSTD_decompressStream(dstream, &out, &in);
if (ZSTD_isError(rc)) goto error;
if (rc == 0) FUZZ_ASSERT(!ZSTD_isError(ZSTD_resetDStream(dstream)));
}
}
error:
#ifndef STATEFUL_FUZZING
ZSTD_freeDStream(dstream); dstream = NULL;
#endif
return 0;
}
int testStreamingAPI(void)
{
size_t const outBuffSize = ZSTD_DStreamOutSize();
char* const outBuff = malloc(outBuffSize);
ZSTD_DStream* const stream = ZSTD_createDStream();
ZSTD_inBuffer input = { COMPRESSED, COMPRESSED_SIZE, 0 };
size_t outputPos = 0;
int needsInit = 1;
if (outBuff == NULL) {
DISPLAY("ERROR: Could not allocate memory\n");
return 1;
}
if (stream == NULL) {
DISPLAY("ERROR: Could not create dstream\n");
return 1;
}
while (1) {
ZSTD_outBuffer output = {outBuff, outBuffSize, 0};
if (needsInit) {
size_t const ret = ZSTD_initDStream(stream);
if (ZSTD_isError(ret)) {
DISPLAY("ERROR: %s\n", ZSTD_getErrorName(ret));
return 1;
}
}
{
size_t const ret = ZSTD_decompressStream(stream, &output, &input);
if (ZSTD_isError(ret)) {
DISPLAY("ERROR: %s\n", ZSTD_getErrorName(ret));
return 1;
}
if (ret == 0) {
needsInit = 1;
}
}
if (memcmp(outBuff, EXPECTED + outputPos, output.pos) != 0) {
DISPLAY("ERROR: Wrong decoded output produced\n");
return 1;
}
outputPos += output.pos;
if (input.pos == input.size && output.pos < output.size) {
break;
}
}
free(outBuff);
ZSTD_freeDStream(stream);
DISPLAY("Streaming API OK\n");
return 0;
}
/*
* Clean up the decompressor.
*/
static int
zstd_filter_close(struct archive_read_filter *self)
{
struct private_data *state;
state = (struct private_data *)self->data;
ZSTD_freeDStream(state->dstream);
free(state->out_block);
free(state);
return (ARCHIVE_OK);
}
static int
pkg_archive_cookie_close(void *p)
{
struct pkg_archive_stdio_cookie *cookie = p;
if (cookie->zstream) {
ZSTD_freeDStream(cookie->zstream);
free(cookie->rdbuf);
}
fclose(cookie->fs);
free(cookie);
return (0);
}
static int decompress_zstd(const char *inbuf, char *outbuf, u64 compress_len,
u64 decompress_len)
{
#if !BTRFSRESTORE_ZSTD
error("btrfs not compiled with zstd support");
return -1;
#else
ZSTD_DStream *strm;
size_t zret;
int ret = 0;
ZSTD_inBuffer in = {inbuf, compress_len, 0};
ZSTD_outBuffer out = {outbuf, decompress_len, 0};
strm = ZSTD_createDStream();
if (!strm) {
error("zstd create failed");
return -1;
}
zret = ZSTD_initDStream(strm);
if (ZSTD_isError(zret)) {
error("zstd init failed: %s", ZSTD_getErrorName(zret));
ret = -1;
goto out;
}
zret = ZSTD_decompressStream(strm, &out, &in);
if (ZSTD_isError(zret)) {
error("zstd decompress failed %s\n", ZSTD_getErrorName(zret));
ret = -1;
goto out;
}
if (zret != 0) {
error("zstd frame incomplete");
ret = -1;
goto out;
}
out:
ZSTD_freeDStream(strm);
return ret;
#endif
}
static void read_zstd(int f, int fd, const char *arg)
{
ZSTD_inBuffer zin;
ZSTD_outBuffer zout;
size_t const inbufsz = ZSTD_DStreamInSize();
zin.src = malloc(inbufsz);
zout.size = ZSTD_DStreamOutSize();
zout.dst = malloc(zout.size);
if (!zin.src || !zout.dst)
goto zstd_r_no_stream;
ZSTD_DStream* const stream = ZSTD_createDStream();
if (!stream)
goto zstd_r_no_stream;
if (ZSTD_isError(ZSTD_initDStream(stream)))
goto zstd_r_error;
size_t s;
while ((s = read(f, (void*)zin.src, inbufsz)) > 0)
{
zin.size = s;
zin.pos = 0;
while (zin.pos < zin.size)
{
zout.pos = 0;
size_t w = ZSTD_decompressStream(stream, &zout, &zin);
if (ZSTD_isError(w))
goto zstd_r_error;
if (write(fd, zout.dst, zout.pos) != (ssize_t)zout.pos)
goto zstd_r_error;
}
}
zstd_r_error:
ZSTD_freeDStream(stream);
zstd_r_no_stream:
free((void*)zin.src);
free(zout.dst);
close(f);
close(fd);
}
/*
* Initialize the filter object
*/
static int
zstd_bidder_init(struct archive_read_filter *self)
{
struct private_data *state;
const size_t out_block_size = ZSTD_DStreamOutSize();
void *out_block;
ZSTD_DStream *dstream;
self->code = ARCHIVE_FILTER_ZSTD;
self->name = "zstd";
state = (struct private_data *)calloc(sizeof(*state), 1);
out_block = (unsigned char *)malloc(out_block_size);
dstream = ZSTD_createDStream();
if (state == NULL || out_block == NULL || dstream == NULL) {
free(out_block);
free(state);
ZSTD_freeDStream(dstream); /* supports free on NULL */
archive_set_error(&self->archive->archive, ENOMEM,
"Can't allocate data for zstd decompression");
return (ARCHIVE_FATAL);
}
self->data = state;
state->out_block_size = out_block_size;
state->out_block = out_block;
state->dstream = dstream;
self->read = zstd_filter_read;
self->skip = NULL; /* not supported */
self->close = zstd_filter_close;
state->eof = 0;
state->in_frame = 0;
return (ARCHIVE_OK);
}
bool DecompressFileToFile(const FilePath& inputPath, const FilePath& outputPath)
{
BinaryReader reader(inputPath);
if (!reader)
{
return false;
}
const size_t inputBufferSize = ZSTD_DStreamInSize();
const auto pInputBuffer = std::make_unique<Byte[]>(inputBufferSize);
const size_t outputBufferSize = ZSTD_DStreamOutSize();
const auto pOutputBuffer = std::make_unique<Byte[]>(outputBufferSize);
ZSTD_DStream* const dStream = ZSTD_createDStream();
if (!dStream)
{
return false;
}
const size_t initResult = ZSTD_initDStream(dStream);
if (ZSTD_isError(initResult))
{
ZSTD_freeDStream(dStream);
return false;
}
size_t toRead = initResult;
BinaryWriter writer(outputPath);
if (!writer)
{
ZSTD_freeDStream(dStream);
return false;
}
while (const size_t read = static_cast<size_t>(reader.read(pInputBuffer.get(), toRead)))
{
ZSTD_inBuffer input = { pInputBuffer.get(), read, 0 };
while (input.pos < input.size)
{
ZSTD_outBuffer output = { pOutputBuffer.get(), outputBufferSize, 0 };
toRead = ZSTD_decompressStream(dStream, &output, &input);
if (ZSTD_isError(toRead))
{
writer.clear();
ZSTD_freeDStream(dStream);
return false;
}
writer.write(pOutputBuffer.get(), output.pos);
}
}
ZSTD_freeDStream(dStream);
return true;
}
bool DecompressToFile(const ByteArrayView view, const FilePath& outputPath)
{
const size_t inputBufferSize = ZSTD_DStreamInSize();
const auto pInputBuffer = std::make_unique<Byte[]>(inputBufferSize);
const size_t outputBufferSize = ZSTD_DStreamOutSize();
const auto pOutputBuffer = std::make_unique<Byte[]>(outputBufferSize);
ZSTD_DStream* const dStream = ZSTD_createDStream();
if (!dStream)
{
return false;
}
const size_t initResult = ZSTD_initDStream(dStream);
if (ZSTD_isError(initResult))
{
ZSTD_freeDStream(dStream);
return false;
}
size_t toRead = initResult;
BinaryWriter writer(outputPath);
if (!writer)
{
ZSTD_freeDStream(dStream);
return false;
}
ReaderView reader(view.data(), view.size());
for (;;)
{
const size_t read = std::min<size_t>(toRead, view.size() - static_cast<size_t>(reader.getPos()));
if (read == 0)
{
break;
}
reader.read(pInputBuffer.get(), read);
ZSTD_inBuffer input = { pInputBuffer.get(), read, 0 };
while (input.pos < input.size)
{
ZSTD_outBuffer output = { pOutputBuffer.get(), outputBufferSize, 0 };
toRead = ZSTD_decompressStream(dStream, &output, &input);
if (ZSTD_isError(toRead))
{
writer.clear();
ZSTD_freeDStream(dStream);
return false;
}
writer.write(pOutputBuffer.get(), output.pos);
}
}
ZSTD_freeDStream(dStream);
return true;
}
std::unique_ptr<IOBuf> ZSTDCodec::doUncompress(
const IOBuf* data,
uint64_t uncompressedLength) {
auto zds = ZSTD_createDStream();
SCOPE_EXIT {
ZSTD_freeDStream(zds);
};
auto rc = ZSTD_initDStream(zds);
zstdThrowIfError(rc);
ZSTD_outBuffer out{};
ZSTD_inBuffer in{};
auto outputSize = ZSTD_DStreamOutSize();
if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH) {
outputSize = uncompressedLength;
} else {
auto decompressedSize =
ZSTD_getDecompressedSize(data->data(), data->length());
if (decompressedSize != 0 && decompressedSize < outputSize) {
outputSize = decompressedSize;
}
}
IOBufQueue queue(IOBufQueue::cacheChainLength());
Cursor cursor(data);
for (rc = 0;;) {
if (in.pos == in.size) {
auto buffer = cursor.peekBytes();
in.src = buffer.data();
in.size = buffer.size();
in.pos = 0;
cursor.skip(in.size);
if (rc > 1 && in.size == 0) {
throw std::runtime_error(to<std::string>("ZSTD: incomplete input"));
}
}
if (out.pos == out.size) {
if (out.pos != 0) {
queue.postallocate(out.pos);
}
auto buffer = queue.preallocate(outputSize, outputSize);
out.dst = buffer.first;
out.size = buffer.second;
out.pos = 0;
outputSize = ZSTD_DStreamOutSize();
}
rc = ZSTD_decompressStream(zds, &out, &in);
zstdThrowIfError(rc);
if (rc == 0) {
break;
}
}
if (out.pos != 0) {
queue.postallocate(out.pos);
}
if (in.pos != in.size || !cursor.isAtEnd()) {
throw std::runtime_error("ZSTD: junk after end of data");
}
if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
queue.chainLength() != uncompressedLength) {
throw std::runtime_error("ZSTD: invalid uncompressed length");
}
return queue.move();
}
static void Uninit(ZSTD_DStream *stream)
{
ZSTD_freeDStream(stream);
}
static gint
rspamd_client_finish_handler (struct rspamd_http_connection *conn,
struct rspamd_http_message *msg)
{
struct rspamd_client_request *req =
(struct rspamd_client_request *)conn->ud;
struct rspamd_client_connection *c;
struct ucl_parser *parser;
GError *err;
const rspamd_ftok_t *tok;
c = req->conn;
if (!c->req_sent) {
c->req_sent = TRUE;
rspamd_http_connection_reset (c->http_conn);
rspamd_http_connection_read_message (c->http_conn,
c->req,
c->fd,
&c->timeout,
c->ev_base);
return 0;
}
else {
if (rspamd_http_message_get_body (msg, NULL) == NULL || msg->code != 200) {
err = g_error_new (RCLIENT_ERROR, msg->code, "HTTP error: %d, %.*s",
msg->code,
(gint)msg->status->len, msg->status->str);
req->cb (c, msg, c->server_name->str, NULL, req->input, req->ud, err);
g_error_free (err);
return 0;
}
tok = rspamd_http_message_find_header (msg, "compression");
if (tok) {
/* Need to uncompress */
rspamd_ftok_t t;
t.begin = "zstd";
t.len = 4;
if (rspamd_ftok_casecmp (tok, &t) == 0) {
ZSTD_DStream *zstream;
ZSTD_inBuffer zin;
ZSTD_outBuffer zout;
guchar *out;
gsize outlen, r;
zstream = ZSTD_createDStream ();
ZSTD_initDStream (zstream);
zin.pos = 0;
zin.src = msg->body_buf.begin;
zin.size = msg->body_buf.len;
if ((outlen = ZSTD_getDecompressedSize (zin.src, zin.size)) == 0) {
outlen = ZSTD_DStreamOutSize ();
}
out = g_malloc (outlen);
zout.dst = out;
zout.pos = 0;
zout.size = outlen;
while (zin.pos < zin.size) {
r = ZSTD_decompressStream (zstream, &zout, &zin);
if (ZSTD_isError (r)) {
err = g_error_new (RCLIENT_ERROR, 500,
"Decompression error: %s",
ZSTD_getErrorName (r));
req->cb (c, msg, c->server_name->str, NULL,
req->input, req->ud, err);
g_error_free (err);
ZSTD_freeDStream (zstream);
g_free (out);
return 0;
}
if (zout.pos == zout.size) {
/* We need to extend output buffer */
zout.size = zout.size * 1.5 + 1.0;
zout.dst = g_realloc (zout.dst, zout.size);
}
}
ZSTD_freeDStream (zstream);
parser = ucl_parser_new (0);
if (!ucl_parser_add_chunk (parser, zout.dst, zout.pos)) {
err = g_error_new (RCLIENT_ERROR, msg->code, "Cannot parse UCL: %s",
ucl_parser_get_error (parser));
ucl_parser_free (parser);
req->cb (c, msg, c->server_name->str, NULL, req->input, req->ud, err);
g_error_free (err);
g_free (zout.dst);
return 0;
}
g_free (zout.dst);
}
else {
err = g_error_new (RCLIENT_ERROR, 500,
"Invalid compression method");
req->cb (c, msg, c->server_name->str, NULL,
req->input, req->ud, err);
g_error_free (err);
return 0;
}
}
else {
parser = ucl_parser_new (0);
if (!ucl_parser_add_chunk (parser, msg->body_buf.begin, msg->body_buf.len)) {
err = g_error_new (RCLIENT_ERROR, msg->code, "Cannot parse UCL: %s",
ucl_parser_get_error (parser));
ucl_parser_free (parser);
req->cb (c, msg, c->server_name->str, NULL, req->input, req->ud, err);
g_error_free (err);
return 0;
}
}
req->cb (c, msg, c->server_name->str, ucl_parser_get_object (
parser), req->input, req->ud, NULL);
ucl_parser_free (parser);
}
return 0;
}
static int
pkg_section_maybe_uncompress(FILE *fs, struct pkg_section_hdr *hdr, void **uncompressed,
size_t *sz)
{
ZSTD_DStream *zstream;
ZSTD_inBuffer zin;
ZSTD_outBuffer zout;
unsigned char *out, *in;
size_t outlen, r;
if (hdr->flags & PKG_FORMAT_FLAGS_ZSTD) {
outlen = hdr->additional;
if (outlen > UINT32_MAX) {
/* Some garbadge instead of size */
pkg_emit_error("invalid uncompressed size: %zu", outlen);
return (EPKG_FATAL);
}
/* Read input */
in = xmalloc(hdr->size);
if (fread(in, 1, hdr->size, fs) != hdr->size) {
return (EPKG_FATAL);
}
zstream = ZSTD_createDStream();
ZSTD_initDStream(zstream);
zin.pos = 0;
zin.src = in;
zin.size = hdr->size;
if (outlen == 0 &&
(outlen = ZSTD_getDecompressedSize(zin.src, zin.size)) == 0) {
outlen = ZSTD_DStreamOutSize();
}
out = xmalloc(outlen);
zout.dst = out;
zout.pos = 0;
zout.size = outlen;
while (zin.pos < zin.size) {
r = ZSTD_decompressStream(zstream, &zout, &zin);
if (ZSTD_isError(r)) {
pkg_emit_error("cannot decompress data: %s",
ZSTD_getErrorName(r));
ZSTD_freeDStream(zstream);
free(out);
free(in);
return (EPKG_FATAL);
}
if (zout.pos == zout.size) {
/* We need to extend output buffer */
zout.size = zout.size * 1.5 + 1.0;
out = xrealloc(zout.dst, zout.size);
zout.dst = out;
}
}
ZSTD_freeDStream(zstream);
free(in);
*uncompressed = out;
*sz = zout.pos;
}
else {
*sz = hdr->size;
out = xmalloc(hdr->size);
if (fread(out, 1, hdr->size, fs) != hdr->size) {
free(out);
return (EPKG_FATAL);
}
*uncompressed = out;
}
return (EPKG_OK);
}
