/* SPDX-License-Identifier: BSD-2-Clause */
/* Minimal DEFLATE decompressor for EROFS FreeBSD */
#include <sys/param.h>
#include <sys/systm.h>
#include <sys/kernel.h>
#include <contrib/zlib/zlib.h>

#include "compress.h"

struct z_erofs_deflate_ctx {
	struct erofs_stream_ctx pool;
	z_stream stream;
	bool initialized;
};

_Static_assert(sizeof(struct z_erofs_deflate_ctx) <=
    EROFS_STREAM_CTX_WRAPPER_SIZE, "Deflate stream wrapper exceeds UMA item");

static voidpf
z_erofs_deflate_alloc(voidpf opaque, uInt items, uInt size)
{
	struct erofs_stream_ctx *pool;
	size_t bytes;

	pool = opaque;
	if (__builtin_mul_overflow((size_t)items, (size_t)size, &bytes)) {
		pool->allocation_failed = true;
		return (NULL);
	}
	return (z_erofs_stream_ctx_alloc(pool, bytes));
}

static void
z_erofs_deflate_free(voidpf opaque, voidpf address)
{

	z_erofs_stream_ctx_free(opaque, address);
}

static int
z_erofs_load_deflate_config(struct erofs_sb_info *sbi,
    const struct erofs_super_block *dsb, const void *data, size_t size)
{
	const struct z_erofs_deflate_cfgs *deflate;

	(void)dsb;
	if (size < sizeof(*deflate))
		return (EINTEGRITY);
	deflate = data;
	if (deflate->windowbits < 8 || deflate->windowbits > 15)
		return (EOPNOTSUPP);
	sbi->deflate_windowbits = deflate->windowbits;
	return (0);
}

static int
z_erofs_deflate_error(int ret)
{

	switch (ret) {
	case Z_MEM_ERROR:
		return (ENOMEM);
	case Z_VERSION_ERROR:
		return (EOPNOTSUPP);
	case Z_NEED_DICT:
	case Z_DATA_ERROR:
	case Z_BUF_ERROR:
		return (EINTEGRITY);
	default:
		return (EIO);
	}
}

static int
z_erofs_deflate_ctx_init(struct erofs_stream_ctx *pool)
{
	struct z_erofs_deflate_ctx *ctx;
	int ret;

	ctx = (struct z_erofs_deflate_ctx *)pool;
	bzero(&ctx->stream, sizeof(ctx->stream));
	ctx->stream.zalloc = z_erofs_deflate_alloc;
	ctx->stream.zfree = z_erofs_deflate_free;
	ctx->stream.opaque = pool;
	ret = inflateInit2(&ctx->stream, -pool->sbi->deflate_windowbits);
	if (ret != Z_OK)
		return (z_erofs_deflate_error(ret));
	ctx->initialized = true;
	return (0);
}

static void
z_erofs_deflate_ctx_fini(struct erofs_stream_ctx *pool)
{
	struct z_erofs_deflate_ctx *ctx;

	ctx = (struct z_erofs_deflate_ctx *)pool;
	if (ctx->initialized) {
		(void)inflateEnd(&ctx->stream);
		ctx->initialized = false;
	}
}

static int
z_erofs_deflate_finish(const struct z_erofs_decompress_req *rq, int ret,
    uInt avail_in)
{

	if (rq->partial_decoding)
		return (0);
	if (ret != Z_STREAM_END || avail_in != 0)
		return (EINTEGRITY);
	return (0);
}

static int
z_erofs_deflate_decompress(const struct z_erofs_decompress_req *rq)
{
	struct erofs_stream_ctx *pool;
	struct z_erofs_deflate_ctx *ctx;
	z_stream *strm;
	uInt in_before, out_before;
	int error, ret;

	if (rq->sbi->deflate_windowbits < 8 ||
	    rq->sbi->deflate_windowbits > MAX_WBITS)
		return (EOPNOTSUPP);
	if (rq->inputsize > (size_t)(uInt)-1 ||
	    rq->outputsize > (size_t)(uInt)-1)
		return (EOVERFLOW);
	if (rq->outputsize == 0)
		return (EINTEGRITY);

	error = z_erofs_stream_ctx_get(rq->sbi, Z_EROFS_COMPRESSION_DEFLATE,
	    sizeof(*ctx), z_erofs_deflate_ctx_init, z_erofs_deflate_ctx_fini,
	    &pool);
	if (error != 0)
		return (error);
	ctx = (struct z_erofs_deflate_ctx *)pool;
	strm = &ctx->stream;
	pool->allocation_failed = false;
	ret = inflateReset2(strm, -rq->sbi->deflate_windowbits);
	if (ret != Z_OK) {
		error = z_erofs_deflate_error(ret);
		z_erofs_stream_ctx_put(pool, false);
		return (error);
	}
	strm->next_in = __DECONST(void *, rq->in);
	strm->avail_in = rq->inputsize;
	strm->next_out = rq->out;
	strm->avail_out = rq->outputsize;

	error = 0;
	ret = Z_OK;
	while (strm->avail_out != 0) {
		in_before = strm->avail_in;
		out_before = strm->avail_out;
		ret = inflate(strm, Z_SYNC_FLUSH);
		if (ret == Z_STREAM_END)
			break;
		if (ret != Z_OK) {
			error = z_erofs_deflate_error(ret);
			break;
		}
		if (strm->avail_in == in_before && strm->avail_out == out_before) {
			error = EINTEGRITY;
			break;
		}
	}
	if (error == 0 && strm->avail_out != 0)
		error = EINTEGRITY;
	else if (error == 0)
		error = z_erofs_deflate_finish(rq, ret, strm->avail_in);
	if (pool->allocation_failed)
		error = ENOMEM;
	z_erofs_stream_ctx_put(pool, error == 0);
	return (error);
}

const struct z_erofs_decompressor z_erofs_deflate_decomp = {
	.config = z_erofs_load_deflate_config,
	.decompress = z_erofs_deflate_decompress,
	.supports_subextent = 1,
	.name = "deflate",
};