/*
 *   Copyright (c) 1999, 2000, 2001, 2002, 2003, 2004, 2005, 2006, 2007
 *   NOVELL (All rights reserved)
 *
 *   Copyright (c) 2013
 *   Canonical Ltd. (All rights reserved)
 *
 *   This program is free software; you can redistribute it and/or
 *   modify it under the terms of version 2 of the GNU General Public
 *   License published by the Free Software Foundation.
 *
 *   This program is distributed in the hope that it will be useful,
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *   GNU General Public License for more details.
 *
 *   You should have received a copy of the GNU General Public License
 *   along with this program; if not, contact Novell, Inc.
 */

#include <assert.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <errno.h>
#include <fcntl.h>

#include <string>
#include <sstream>
#include <sys/apparmor.h>

#include "lib.h"
#include "parser.h"
#include "profile.h"
#include "libapparmor_re/apparmor_re.h"

#include <unistd.h>
#include <linux/unistd.h>


#define SD_CODE_SIZE (sizeof(u8))
#define SD_STR_LEN (sizeof(u16))

using namespace std;

int __sd_serialize_profile(int option, aa_kernel_interface *kernel_interface,
			   Profile *prof, int cache_fd);

static void print_error(int error)
{
	switch (error) {
	case -ESPIPE:
		PERROR(_("Bad write position\n"));
		break;
	case -ENOMEM:
		PERROR(_("Out of memory\n"));
		break;
	case -EFAULT:
		PERROR(_("Couldn't copy profile: Bad memory address\n"));
		break;
	case -EPROTO:
		PERROR(_("Profile doesn't conform to protocol\n"));
		break;
	case -EBADMSG:
		PERROR(_("Profile does not match signature\n"));
		break;
	case -EPROTONOSUPPORT:
		PERROR(_("Profile version not supported by Apparmor module\n"));
		break;
	case -EEXIST:
		PERROR(_("Profile already exists\n"));
		break;
	case -ENOENT:
		PERROR(_("Profile doesn't exist\n"));
		break;
	case -EPERM:
		PERROR(_("%s: Permission denied. You need policy admin privileges to manage profiles.\n\n"),
		       progname);
		break;
	case -EACCES:
		PERROR(_("%s: Access denied. You need policy admin privileges to manage profiles.\n\n"),
		       progname);
		break;
	default:
		PERROR(_("Unknown error (%d): %s\n"), -error, strerror(-error));
		break;
	}
}

int load_profile(int option, aa_kernel_interface *kernel_interface,
		 Profile *prof, int cache_fd)
{
	int retval = 0;
	int error = 0;

	PDEBUG("Serializing policy for %s.\n", prof->name);
	retval = __sd_serialize_profile(option, kernel_interface, prof, cache_fd);

	if (retval < 0) {
		error = retval;	/* yeah, we'll just report the last error */
		switch (option) {
		case OPTION_ADD:
			PERROR(_("%s: Unable to add \"%s\".  "),
			       progname, prof->name);
			print_error(error);
			break;
		case OPTION_REPLACE:
			PERROR(_("%s: Unable to replace \"%s\".  "),
			       progname, prof->name);
			print_error(error);
			break;
		case OPTION_REMOVE:
			PERROR(_("%s: Unable to remove \"%s\".  "),
			       progname, prof->name);
			print_error(error);
			break;
		case OPTION_STDOUT:
			PERROR(_("%s: Unable to write to stdout\n"),
			       progname);
			break;
		case OPTION_OFILE:
			PERROR(_("%s: Unable to write to output file\n"),
			       progname);
			break;
		default:
			PERROR(_("%s: ASSERT: Invalid option: %d\n"),
			       progname, option);
			exit(1);
			break;
		}

	} else if (conf_verbose) {
		switch (option) {
		case OPTION_ADD:
			printf(_("Addition succeeded for \"%s\".\n"),
			       prof->name);
			break;
		case OPTION_REPLACE:
			printf(_("Replacement succeeded for \"%s\".\n"),
			       prof->name);
			break;
		case OPTION_REMOVE:
			printf(_("Removal succeeded for \"%s\".\n"),
			       prof->name);
			break;
		case OPTION_STDOUT:
		case OPTION_OFILE:
			break;
		default:
			PERROR(_("%s: ASSERT: Invalid option: %d\n"),
			       progname, option);
			exit(1);
			break;
		}
	}

	return error;
}



enum sd_code {
	SD_U8,
	SD_U16,
	SD_U32,
	SD_U64,
	SD_NAME,		/* same as string except it is items name */
	SD_STRING,
	SD_BLOB,
	SD_STRUCT,
	SD_STRUCTEND,
	SD_LIST,
	SD_LISTEND,
	SD_ARRAY,
	SD_ARRAYEND,
	SD_OFFSET
};

const char *sd_code_names[] = {
	"SD_U8",
	"SD_U16",
	"SD_U32",
	"SD_U64",
	"SD_NAME",
	"SD_STRING",
	"SD_BLOB",
	"SD_STRUCT",
	"SD_STRUCTEND",
	"SD_LIST",
	"SD_LISTEND",
	"SD_ARRAY",
	"SD_ARRAYEND",
	"SD_OFFSET"
};


static inline void sd_write8(std::ostringstream &buf, u8 b)
{
	buf.write((const char *) &b, 1);
}

static inline void sd_write16(std::ostringstream &buf, u16 b)
{
	u16 tmp;
	tmp = cpu_to_le16(b);
	buf.write((const char *) &tmp, 2);
}

static inline void sd_write32(std::ostringstream &buf, u32 b)
{
	u32 tmp;
	tmp = cpu_to_le32(b);
	buf.write((const char *) &tmp, 4);
}

static inline void sd_write64(std::ostringstream &buf, u64 b)
{
	u64 tmp;
	tmp = cpu_to_le64(b);
	buf.write((const char *) &tmp, 8);
}

static inline void sd_write_uint8(std::ostringstream &buf, u8 b)
{
	sd_write8(buf, SD_U8);
	sd_write8(buf, b);
}

static inline void sd_write_uint16(std::ostringstream &buf, u16 b)
{
	sd_write8(buf, SD_U16);
	sd_write16(buf, b);
}

static inline void sd_write_uint32(std::ostringstream &buf, u32 b)
{
	sd_write8(buf, SD_U32);
	sd_write32(buf, b);
}

static inline void sd_write_uint64(std::ostringstream &buf, u64 b)
{
	sd_write8(buf, SD_U64);
	sd_write64(buf, b);
}

static inline void sd_write_name(std::ostringstream &buf, const char *name)
{
	PDEBUG("Writing name '%s'\n", name ? name : "(null)");
	if (name) {
		sd_write8(buf, SD_NAME);
		sd_write16(buf, strlen(name) + 1);
		buf.write(name, strlen(name) + 1);
	}
}

static inline void sd_write_blob(std::ostringstream &buf, const void *b, int buf_size, const char *name)
{
	sd_write_name(buf, name);
	sd_write8(buf, SD_BLOB);
	sd_write32(buf, buf_size);
	buf.write((const char *) b, buf_size);
}


static char zeros[64];
#define align64(X) (((X) + (typeof(X)) 7) & ~((typeof(X)) 7))
static inline void sd_write_aligned_blob(std::ostringstream &buf, void *b, int b_size,
				 const char *name)
{
	sd_write_name(buf, name);
	/* pad calculation MUST come after name is written */
	size_t pad = align64(buf.tellp() + ((std::streamoff) 5l)) - (buf.tellp() + ((std::streamoff) 5l));
	sd_write8(buf, SD_BLOB);
	sd_write32(buf, b_size + pad);
	buf.write(zeros, pad);
	buf.write((const char *) b, b_size);
}

static void sd_write_strn(std::ostringstream &buf, const char *b, int size, const char *name)
{
	sd_write_name(buf, name);
	sd_write8(buf, SD_STRING);
	sd_write16(buf, size);
	buf.write(b, size);
}

static inline void sd_write_string(std::ostringstream &buf, const char *b, const char *name)
{
	sd_write_strn(buf, b, strlen(b) + 1, name);
}

static inline void sd_write_struct(std::ostringstream &buf, const char *name)
{
	sd_write_name(buf, name);
	sd_write8(buf, SD_STRUCT);
}

static inline void sd_write_structend(std::ostringstream &buf)
{
	sd_write8(buf, SD_STRUCTEND);
}

static inline void sd_write_array(std::ostringstream &buf, const char *name, int size)
{
	sd_write_name(buf, name);
	sd_write8(buf, SD_ARRAY);
	sd_write16(buf, size);
}

static inline void sd_write_arrayend(std::ostringstream &buf)
{
	sd_write8(buf, SD_ARRAYEND);
}

static inline void sd_write_list(std::ostringstream &buf, const char *name)
{
	sd_write_name(buf, name);
	sd_write8(buf, SD_LIST);
}

static inline void sd_write_listend(std::ostringstream &buf)
{
	sd_write8(buf, SD_LISTEND);
}

void sd_serialize_identity(std::ostringstream &buf, struct value_list* l)
{
	int i;
	int list_sz = list_len(l);

	sd_write_array(buf, "identity", list_sz);
	for (i = 0; i < list_sz; i++) {
		sd_write_string(buf, l->value, NULL);
		l = l->next;
	}
	sd_write_arrayend(buf);
}

void sd_serialize_perm(std::ostringstream &buf, aa_perms &perms)
{
	sd_write_uint32(buf, 0);	/* reserved */
	sd_write_uint32(buf, perms.allow);
	sd_write_uint32(buf, perms.deny);
	sd_write_uint32(buf, perms.subtree);
	sd_write_uint32(buf, perms.cond);
	sd_write_uint32(buf, perms.kill);
	sd_write_uint32(buf, perms.complain);
	sd_write_uint32(buf, perms.prompt);
	sd_write_uint32(buf, perms.audit);
	sd_write_uint32(buf, perms.quiet);
	sd_write_uint32(buf, perms.hide);
	sd_write_uint32(buf, perms.xindex);
	sd_write_uint32(buf, perms.tag);
	sd_write_uint32(buf, perms.label);
}

void sd_serialize_permstable(std::ostringstream &buf, vector <aa_perms> &perms_table, bool filedfa)
{
	sd_write_struct(buf, "perms");
	sd_write_name(buf, "version");
	sd_write_uint32(buf, 1);
	sd_write_array(buf, NULL, perms_table.size());
	for (size_t i = 0; i < perms_table.size(); i++) {
		sd_serialize_perm(buf, perms_table[i]);
	}
	sd_write_arrayend(buf);
	sd_write_structend(buf);
	/* currently only the file dfa makes use of owner conditional accept */
	if (filedfa && perms_table.size() > 0) {
		sd_write_name(buf, "permsv");
		sd_write_uint32(buf, 3);
	}
}

void sd_serialize_dfa(std::ostringstream &buf, void *dfa, size_t size,
		      vector <aa_perms> &perms_table, bool filedfa)
{
	if (dfa) {
		if (kernel_supports_permstable32 && perms_table.size() > 0) {
			//fprintf(stderr, "writing perms table %d\n", size);
			sd_serialize_permstable(buf, perms_table, filedfa);
		} else {
			//fprintf(stderr, "skipping permtable32 %d, size %d\n", kernel_supports_permstable32, perms_table.size());
		}
		sd_write_aligned_blob(buf, dfa, size, "aadfa");
	}
}

void sd_serialize_rlimits(std::ostringstream &buf, struct aa_rlimits *limits)
{
	if (!limits->specified)
		return;

	sd_write_struct(buf, "rlimits");
	sd_write_uint32(buf, limits->specified);
	sd_write_array(buf, NULL, RLIM_NLIMITS);
	for (int i = 0; i < RLIM_NLIMITS; i++) {
		sd_write_uint64(buf, limits->limits[i]);
	}
	sd_write_arrayend(buf);
	sd_write_structend(buf);
}

void sd_serialize_xtable(std::ostringstream &buf, const char * const *table)
{
	size_t count;

	if (!table[4])
		return;
	sd_write_struct(buf, "xtable");
	count = 0;
	for (int i = 4; i < AA_EXEC_COUNT; i++) {
		if (table[i])
			count++;
	}

	sd_write_array(buf, NULL, count);
	for (size_t i = 4; i < AA_EXEC_COUNT; i++) {
		if (!table[i]) {
			continue;
		}
		size_t len = strlen(table[i]) + 1;

		/* if its a namespace make sure the second : is overwritten
		 * with 0, so that the namespace and name are \0 separated
		 */
		if (*table[i] == ':') {
			/* don't change table */
			char *ent = strdup(table[i]);
			if (!ent) {
				PERROR("Failed to allocate memory: %s\n", strerror(errno));
				/* exiting here because there's no
				 * error handling for serialize functions */
				exit(errno);
			}
			char *tmp = ent + 1;
			strsep(&tmp, ":");
			sd_write_strn(buf, ent, len, NULL);
			free(ent);
		}
		else
			sd_write_strn(buf, table[i], len, NULL);
	}

	sd_write_arrayend(buf);
	sd_write_structend(buf);
}

void sd_serialize_xattrs(std::ostringstream &buf, struct cond_entry_list xattrs)
{
	int count;
	struct cond_entry *entry;

	if (!(xattrs.list))
		return;

	count = 0;
	for (entry = xattrs.list; entry; entry = entry->next) {
		count++;
	}

	sd_write_struct(buf, "xattrs");
	sd_write_array(buf, NULL, count);
	for (entry = xattrs.list; entry; entry = entry->next) {
		sd_write_string(buf, entry->name, NULL);
	}
	sd_write_arrayend(buf);
	sd_write_structend(buf);
}

void sd_serialize_profile(std::ostringstream &buf, Profile *profile,
			 int flattened)
{
	uint64_t allowed_caps;

	sd_write_struct(buf, "profile");
	if (flattened) {
		assert(profile->parent);
		sd_write_string(buf, profile->get_name(false).c_str(), NULL);
	} else {
		sd_write_string(buf, profile->name, NULL);
	}

	/* only emit this if current kernel at least supports "create" */
	if (perms_create) {
		if (profile->xmatch) {
		  sd_serialize_dfa(buf, profile->xmatch, profile->xmatch_size, profile->xmatch_perms_table, false);
			sd_write_uint32(buf, profile->xmatch_len);
		}
	}

	/* added in 4.13, unfortunately there is no features flag */
	if (profile->flags.disconnected_path) {
		sd_write_string(buf, profile->flags.disconnected_path,
				"disconnected");
	}

	if (profile->flags.disconnected_ipc && features_supports_flag_disconnected_ipc) {
		sd_write_string(buf, profile->flags.disconnected_ipc,
				"disconnected_ipc");
	}

	if (profile->flags.signal && features_supports_flag_signal) {
		sd_write_name(buf, "kill");
		sd_write_uint32(buf, profile->flags.signal);
	}

	if (profile->flags.error && features_supports_flag_error) {
		sd_write_name(buf, "error");
		sd_write_uint32(buf, profile->flags.error);
	}

	sd_write_struct(buf, "flags");
	/* used to be flags.debug, but that's no longer supported */
	sd_write_uint32(buf, profile->flags.flags);
	sd_write_uint32(buf, profile_mode_packed(profile->flags.mode));
	sd_write_uint32(buf, profile->flags.audit);
	sd_write_structend(buf);
	if (profile->flags.path) {
		int flags = 0;
		if (profile->flags.path & PATH_CHROOT_REL)
			flags |= 0x8;
		if (profile->flags.path & PATH_MEDIATE_DELETED)
			flags |= 0x10000;
		if (profile->flags.path & PATH_ATTACH)
			flags |= 0x4;
		if (profile->flags.path & PATH_CHROOT_NSATTACH)
			flags |= 0x10;
		if (profile->flags.path & PATH_IPC_ATTACH)
			flags |= 0x20;

		sd_write_name(buf, "path_flags");
		sd_write_uint32(buf, flags);
	}

#define low_caps(X) ((u32) ((X) & 0xffffffff))
#define high_caps(X) ((u32) (((X) >> 32) & 0xffffffff))
	allowed_caps = (profile->caps.allow) & ~profile->caps.deny;
	sd_write_uint32(buf, low_caps(allowed_caps));
	sd_write_uint32(buf, low_caps(allowed_caps & profile->caps.audit));
	sd_write_uint32(buf, low_caps(profile->caps.deny & profile->caps.quiet));
	sd_write_uint32(buf, 0);

	sd_write_struct(buf, "caps64");
	sd_write_uint32(buf, high_caps(allowed_caps));
	sd_write_uint32(buf, high_caps(allowed_caps & profile->caps.audit));
	sd_write_uint32(buf, high_caps(profile->caps.deny & profile->caps.quiet));
	sd_write_uint32(buf, 0);
	sd_write_structend(buf);

	if (features_supports_identity_names)
		sd_serialize_identity(buf, profile->identities);

	sd_serialize_xattrs(buf, profile->xattrs);

	sd_serialize_rlimits(buf, &profile->rlimits);

	/* choice to support / downgrade needs to already have been made */
	if (features_supports_networkv8) {
		/* nothing - encoded in policydb */
	} else if (profile->net.allow && features_supports_network) {
		size_t i;
		sd_write_array(buf, "net_allowed_af", get_af_max());
		for (i = 0; i < get_af_max(); i++) {
			u16 allowed = profile->net.allow[i] &
				~profile->net.deny[i];
			sd_write_uint16(buf, allowed);
			sd_write_uint16(buf, allowed & profile->net.audit[i]);
			sd_write_uint16(buf, profile->net.deny[i] & profile->net.quiet[i]);
		}
		sd_write_arrayend(buf);
	} else if (profile->net.allow)
		pwarn(WARN_RULE_NOT_ENFORCED, _("profile %s network rules not enforced\n"), profile->name);

	if (profile->policy.dfa) {
		sd_write_struct(buf, "policydb");
		sd_serialize_dfa(buf, profile->policy.dfa, profile->policy.size,
				 profile->policy.perms_table, false);
		if (kernel_supports_permstable32) {
			sd_serialize_xtable(buf, profile->exec_table);

		}
		sd_write_structend(buf);
	}

	sd_serialize_dfa(buf, profile->dfa.dfa, profile->dfa.size,
			 profile->dfa.perms_table, true);
	if (profile->dfa.dfa) {
		// fprintf(stderr, "profile %s: dfa xtable\n", profile->name);
		sd_serialize_xtable(buf, profile->exec_table);
	}
	sd_write_structend(buf);
}

size_t recompress_policy_zstd(const char *compressed_data, size_t compressed_size,  char **recompressed_buffer)
{
	unsigned long long const fcs = ZSTD_getFrameContentSize(compressed_data, compressed_size);
	if (fcs == ZSTD_CONTENTSIZE_UNKNOWN || fcs == ZSTD_CONTENTSIZE_ERROR) {
		PERROR(_("Recompression error: could not determine frame content size\n"));
		return 0;
	}
	size_t raw_size = (size_t)fcs;

	if (raw_size == 0) {
		PERROR(_("Recompression error: raw_size is zero\n"));
		return 0;
	}
	const char *zstd_payload = compressed_data;
	size_t zstd_payload_size = compressed_size;

	/* Allocate and decompress */
	char *raw_buf = (char *)malloc(raw_size);
	if (!raw_buf) {
		PERROR(_("Recompression error: out of memory\n"));
		return 0;
	}

	size_t dec_sz = ZSTD_decompress(raw_buf, raw_size, zstd_payload, zstd_payload_size);
	if (ZSTD_isError(dec_sz)) {
		free(raw_buf);
		PERROR(_("Decompression error: %s\n"), ZSTD_getErrorName(dec_sz));
		return 0;
	}

	if (dec_sz != raw_size) {
		free(raw_buf);
		PERROR(_("Decompression size mismatch: expected %zu, got %zu\n"),
			   raw_size, dec_sz);
		return 0;
	}

	/* Recompress at current compress_level */
	size_t out_sz = compress_policy_zstd(raw_buf, raw_size, recompressed_buffer);
	free(raw_buf);
	return out_sz;
}

size_t compress_policy_zstd(const char* raw_data_str, size_t raw_data_size, char** compressed_buffer)
{
	size_t bound = ZSTD_compressBound(raw_data_size);
	*compressed_buffer = (char*) malloc(bound);
	if (!*compressed_buffer) {
		PERROR(_("Compression error: out of memory\n"));
		return 0;
	}

	size_t compressed_size = ZSTD_compress(*compressed_buffer, bound, raw_data_str, raw_data_size, zstd_compress_level);

	if (ZSTD_isError(compressed_size)) {
		free(*compressed_buffer);
		*compressed_buffer = NULL;
		PERROR(_("Compression error: %s\n"),  ZSTD_getErrorName(compressed_size));
		return 0;
	}
	return compressed_size;
}

void sd_serialize_top_profile(std::ostringstream &buf, Profile *profile)
{
	uint32_t version;

	version = ENCODE_VERSION(force_complain, policy_version,
				 parser_abi_version, kernel_abi_version);

	sd_write_name(buf, "version");
	sd_write_uint32(buf, version);

	if (profile->ns) {
		sd_write_string(buf, profile->ns, "namespace");
	}

	sd_serialize_profile(buf, profile, profile->parent ? 1 : 0);
}

int __sd_serialize_profile(int option, aa_kernel_interface *kernel_interface,
			   Profile *prof, int cache_fd)
{
	autoclose int fd = -1;
	int error, policy_size, wsize;
	std::ostringstream work_area;

	switch (option) {
	case OPTION_ADD:
	case OPTION_REPLACE:
	case OPTION_REMOVE:
		break;
	case OPTION_STDOUT:
		fd = dup(1);
		if (fd < 0) {
			error = -errno;
			PERROR(_("Unable to open stdout - %s\n"),
			       strerror(errno));
			goto exit;
		}
		break;
	case OPTION_OFILE:
		fd = dup(fileno(ofile));
		if (fd < 0) {
			error = -errno;
			PERROR(_("Unable to open output file - %s\n"),
			       strerror(errno));
			goto exit;
		}
		break;
	default:
		error = -EINVAL;
		goto exit;
		break;
	}

	error = 0;

	if (option == OPTION_REMOVE) {
		if (kernel_load) {
			if (aa_kernel_interface_remove_policy(kernel_interface,
							      prof->fqname().c_str()) == -1)
				error = -errno;
		}
	} else {
		std::string raw_policy_string;
		const char *policy = NULL;
		char *compressed_buffer = NULL;
		int raw_policy_size;

		sd_serialize_top_profile(work_area, prof);
		raw_policy_string = work_area.str();
		raw_policy_size = raw_policy_string.size();

		if (zstd_compress_policy == ZSTD_COMPRESS_POLICY) {
			policy_size = compress_policy_zstd(raw_policy_string.c_str(), raw_policy_size, &compressed_buffer);
			if (policy_size == 0) {
				error = -1;
				goto exit;
			}
			policy = compressed_buffer;
		} else {
			policy = raw_policy_string.c_str();
			policy_size = raw_policy_size;
		}
		if (kernel_load) {
			if (option == OPTION_ADD &&
			    aa_kernel_interface_load_policy(kernel_interface,
							    policy, policy_size) == -1) {
				error = -errno;
			} else if (option == OPTION_REPLACE &&
				   aa_kernel_interface_replace_policy(kernel_interface,
								      policy, policy_size) == -1) {
				error = -errno;
			}
		} else if (option == OPTION_STDOUT || option == OPTION_OFILE) {
			if (zstd_compress_policy == ZSTD_COMPRESS_POLICY) {
				if (write_compressed_with_user_header(fd, policy, policy_size,
								      (uint8_t)zstd_compress_level) < 0)
					error = -errno;
			} else if (aa_kernel_interface_write_policy(fd, policy, policy_size) == -1) {
				error = -errno;
			}
		}

		if (cache_fd != -1) {
			/* Write uncompressed policy to the cache FD. It will be compressed later in parser loop if enabled. */
			wsize = write(cache_fd, raw_policy_string.c_str(), raw_policy_size);
			if (wsize < 0) {
				error = -errno;
			} else if (wsize < raw_policy_size) {
				PERROR(_("%s: Unable to write entire profile entry to cache\n"),
				       progname);
				error = -EIO;
			}
		}

		free(compressed_buffer);
	}

	if (!prof->hat_table.empty() && option != OPTION_REMOVE) {
		if (load_flattened_hats(prof, option, kernel_interface, cache_fd) == 0)
			return 0;
	}


exit:
	return error;
}

#ifdef UNIT_TEST

#include "unit_test.h"

/*
 * Generate data to compress with 4 bits / byte of entropy
 */
static void generate_test_data(char *buffer, size_t size)
{
	for (size_t i = 0; i < size; i++)
		buffer[i] = (char)(rand() & 0x0F);
}

static int test_compress_roundtrip(void)
{
	int rc = 0;
	char *test_data = NULL;
	char *compressed = NULL;
	size_t test_size = 64 * 1024; /* 64KB */
	size_t comp_size;

	test_data = (char *)malloc(test_size);
	if (!test_data)
		return 1;

	generate_test_data(test_data, test_size);

	/* Set a valid compression level for testing */
	zstd_compress_level = 10;

	comp_size = compress_policy_zstd(test_data, test_size, &compressed);

	printf("roundtrip: input=%zu compressed=%zu ratio=%.2f%%\n",
	       test_size, comp_size, 100.0 * comp_size / test_size);

	MY_TEST(comp_size > 0, "compression produced output");
	MY_TEST(compressed != NULL, "compression allocated buffer");
	MY_TEST(comp_size < test_size, "compression reduced size");



	free(test_data);
	free(compressed);
	return rc;
}

static int test_compress_empty_input(void)
{
	int rc = 0;
	char *compressed = NULL;
	size_t comp_size;

	zstd_compress_level = 10;

	/* Empty input should still work (edge case) */
	comp_size = compress_policy_zstd("", 0, &compressed);
	MY_TEST(comp_size > 0, "empty input produces zstd frame");
	MY_TEST(compressed != NULL, "empty input allocated buffer");

	free(compressed);
	return rc;
}

static int test_compress_various_sizes(void)
{
	int rc = 0;
	char *test_data = NULL;
	char *compressed = NULL;
	size_t comp_size;
	/* Test sizes: 1KB, 64KB, 512KB */
	size_t sizes[] = {1024, 64 * 1024, 512 * 1024};

	zstd_compress_level = 10;

	for (size_t i = 0; i < sizeof(sizes)/sizeof(sizes[0]); i++) {
		size_t sz = sizes[i];
		test_data = (char *)malloc(sz);
		if (!test_data) {
			rc = 1;
			continue;
		}

		generate_test_data(test_data, sz);
		comp_size = compress_policy_zstd(test_data, sz, &compressed);

		char msg[128];
		snprintf(msg, sizeof(msg), "size %zuKB compressed (ratio=%.1f%%)",
		         sz / 1024, 100.0 * comp_size / sz);
		MY_TEST(comp_size > 0 && comp_size < sz, msg);

		free(test_data);
		free(compressed);
		test_data = NULL;
		compressed = NULL;
	}

	return rc;
}

static int test_recompress_buffer(void)
{
	int rc = 0;
	char *test_data = NULL;
	char *compressed = NULL;
	char *recompressed = NULL;
	size_t test_size = 100 * 1024; /* 100KB */
	size_t comp_size, recomp_size;

	test_data = (char *)malloc(test_size);
	if (!test_data)
		return 1;

	generate_test_data(test_data, test_size);
	zstd_compress_level = 10;

	/* 1. Compress initial data */
	comp_size = compress_policy_zstd(test_data, test_size, &compressed);
	MY_TEST(comp_size > 0, "initial compression success");

	/* 2. Recompress */
	/* recompress_policy_zstd takes (compressed_data, compressed_size, &out) */
	recomp_size = recompress_policy_zstd(compressed, comp_size, &recompressed);

	MY_TEST(recomp_size > 0, "recompression produced output");
	MY_TEST(recompressed != NULL, "recompression allocated buffer");

	/* 3. Verify consistency (optional: check if size is similar) */
	/* Zstd is deterministic for same level/params, so sizes should match exactly if using same level */
	MY_TEST(recomp_size == comp_size, "recompressed size matches initial compressed size");

	free(test_data);
	free(compressed);
	free(recompressed);
	return rc;
}

int main(void)
{
	int rc = 0;
	int retval;

	srand(time(NULL));
	progname = __FILE__;

	retval = test_compress_roundtrip();
	if (retval != 0)
		rc = retval;

	retval = test_compress_empty_input();
	if (retval != 0)
		rc = retval;

	retval = test_compress_various_sizes();
	if (retval != 0)
		rc = retval;

	retval = test_recompress_buffer();
	if (retval != 0)
		rc = retval;


	if (rc == 0)
		printf("All compression tests passed.\n");

	return rc;
}

#endif /* UNIT_TEST */
