// SPDX-License-Identifier: GPL 2.0+ OR BSD-3-Clause
/*
 * Copyright 2015 Google Inc.
 */

#include <common.h>
#include <compiler.h>
#include <linux/kernel.h>
#include <linux/types.h>

#include "lz4_wrapper.h"

#define LZ4F_MAGIC 0x184D2204

#define LZ4_MAX_BLOCK_SIZE_64K 4
#define LZ4_MAX_BLOCK_SIZE_256K 5
#define LZ4_MAX_BLOCK_SIZE_1M 6
#define LZ4_MAX_BLOCK_SIZE_4M 7

struct lz4_block_header {
	union {
		u32 raw;
		struct {
			u32 size : 31;
			u32 not_compressed : 1;
		};
	};
	/* + size bytes of data */
	/* + u32 block_checksum iff has_block_checksum is set */
} __packed;

static u16 LZ4_readLE16(const void *src)
{
	return le16_to_cpu(*(u16 *)src);
}
static void LZ4_copy4(void *dst, const void *src)
{
	*(u32 *)dst = *(u32 *)src;
}
static void LZ4_copy8(void *dst, const void *src)
{
	*(u64 *)dst = *(u64 *)src;
}

/* Unaltered (except removing unrelated code) from github.com/Cyan4973/lz4. */
#include "lz4.c" /* #include for inlining, do not link! */

bool is_data_lz4_compressed(const void *src, size_t srcn)
{
	if (srcn < sizeof(struct lz4_frame_header)) {
		return false;
	}

	struct lz4_frame_header h;
	memcpy(&h, src, sizeof(h));
	return le32_to_cpu(h.magic) == LZ4F_MAGIC;
}

int check_and_extract_lz4_frame_header(const void *src, size_t srcn,
				       struct lz4_frame_header *out,
				       size_t *data_offset)
{
	const void *in = src;
	const struct lz4_frame_header *h = in;

	if (srcn < sizeof(*h) + sizeof(u64) + sizeof(u8))
		return -EINVAL; /* input overrun */

	/* We assume there's always only a single, standard frame. */
	if (le32_to_cpu(h->magic) != LZ4F_MAGIC || h->version != 1) {
		fprintf(stderr, "invalid magic %x, %d\n",
			le32_to_cpu(h->magic) != LZ4F_MAGIC, h->version);
		return -EPROTONOSUPPORT; /* unknown format */
	}
	if (h->reserved0 || h->reserved1 || h->reserved2)
		return -EINVAL; /* reserved must be zero */
	if (!h->independent_blocks) {
		fprintf(stderr, "no support for independent block\n");
		return -EPROTONOSUPPORT; /* we can't support this yet */
	}

	in += sizeof(*h);
	if (h->has_content_size)
		in += sizeof(u64);
	in += sizeof(u8);
	*data_offset = in - src;
	*out = *h;
	return 0;
}

/**
 * Decompress one block of data from lz4 data.
 *
 * @param src	The lz4 compressed data.
 * @param srcn	Size of |src|.
 * @param dst	Output pointer for storing the decompress block.
 * @param dstn	Size of |dst|.
 * @param block_size	Output pointer for size of current compressed block.
 * @param decompressed_size Output pointer for the amount of data decompressed.
 */
int decompress_block(const void *src, size_t srcn, void *dst, int dstn,
		     size_t *block_size, size_t *decompressed_size)
{
	const void *in = src;
	void *out = dst;
	const void *end = out + dstn;
	*decompressed_size = 0;
	struct lz4_block_header b;
	b.raw = le32_to_cpu(*(u32 *)in);
	*block_size = sizeof(struct lz4_block_header) + b.size;
	in += sizeof(struct lz4_block_header);
	if (in - src + b.size > srcn) {
		fprintf(stderr, "input overrun\n");
		return -EINVAL; /* input overrun */
	}

	if (!b.size) {
		return 0; /* decompression successful */
	}

	if (b.not_compressed) {
		size_t size = min((ptrdiff_t)b.size, end - out);
		memcpy(out, in, size);
		out += size;
		if (size < b.size) {
			fprintf(stderr, "output overrun\n");
			return -ENOBUFS; /* output overrun */
		}
	} else {
		/* constant folding essential, do not touch params! */
		int ret = LZ4_decompress_generic(in, out, b.size, end - out,
						 endOnInputSize, full, 0,
						 noDict, out, NULL, 0);
		if (ret < 0) {
			fprintf(stderr, "LZ4_decompress_generic error %d\n",
				ret);
			return -EPROTO; /* decompression error */
		}
		out += ret;
	}

	*decompressed_size = out - dst;
	return 0;
}

int ulz4fn(const void *src, size_t srcn, void *dst, size_t *dstn)
{
	const void *src_end = src + srcn;
	const void *dst_end = dst + *dstn;
	struct lz4_frame_header h;
	size_t data_offset = 0;
	int ret =
		check_and_extract_lz4_frame_header(src, srcn, &h, &data_offset);
	if (ret) {
		return ret;
	}
	const void *data = (uint8_t *)src + data_offset;
	void *out = dst;
	size_t decompressed_size = 0, block_size;
	do {
		if ((ret = decompress_block(data, src_end - data, out,
					   dst_end - out, &block_size,
					   &decompressed_size))) {
			return ret;
		}
		out += decompressed_size;
		data += block_size + (h.has_block_checksum ? sizeof(u32) : 0);
	} while (decompressed_size);

	*dstn = out - dst;
	return 0;
}

int get_lz4_max_block_size(const struct lz4_frame_header *h, size_t *out)
{
	switch (h->max_block_size) {
	case LZ4_MAX_BLOCK_SIZE_64K:
		*out = 64 * 1024;
		break;
	case LZ4_MAX_BLOCK_SIZE_256K:
		*out = 256 * 1024;
		break;
	case LZ4_MAX_BLOCK_SIZE_1M:
		*out = 1024 * 1024;
		break;
	case LZ4_MAX_BLOCK_SIZE_4M:
		*out = 4 * 1024 * 1024;
		break;
	default:
		printf("Invalid block max size code");
		return -1;
	}
	return 0;
}

int lz4_decompress(const struct lz4_frame_header *h, const uint8_t *src,
		   size_t *srcn, uint8_t *dst, size_t *dstn,
		   bool *has_more_data)
{
	int ret = 0;
	size_t max_block_size = 0;
	if ((ret = get_lz4_max_block_size(h, &max_block_size))) {
		return ret;
	}

	const uint8_t *in_curr = src;
	const uint8_t *in_end = src + *srcn;
	uint8_t *out_curr = dst;
	uint8_t *out_end = dst + *dstn;
	*dstn = 0;
	*srcn = 0;
	*has_more_data = true;
	while (out_curr + max_block_size <= out_end) {
		size_t compressed_size = 0, decompressed_size = 0;
		if ((ret = decompress_block(in_curr, in_end - in_curr, out_curr,
					    out_end - out_curr,
					    &compressed_size,
					    &decompressed_size))) {
			return ret;
		}
		in_curr += compressed_size +
			   (h->has_block_checksum ? sizeof(u32) : 0);
		out_curr += decompressed_size;
		// If no more data can be decompressed, exit the loop
		if (!decompressed_size) {
			*has_more_data = false;
			break;
		}
	}

	*dstn = out_curr - dst;
	*srcn = in_curr - src;
	return 0;
}