path: root/src/bin/pg_dump/compress_lz4.c

/*-------------------------------------------------------------------------
 *
 * compress_lz4.c
 *	 Routines for archivers to write a LZ4 compressed data stream.
 *
 * Portions Copyright (c) 1996-2025, PostgreSQL Global Development Group
 * Portions Copyright (c) 1994, Regents of the University of California
 *
 * IDENTIFICATION
 *	   src/bin/pg_dump/compress_lz4.c
 *
 *-------------------------------------------------------------------------
 */
#include "postgres_fe.h"
#include <unistd.h>

#include "compress_lz4.h"
#include "pg_backup_utils.h"

#ifdef USE_LZ4
#include <lz4frame.h>

/*
 * LZ4F_HEADER_SIZE_MAX first appeared in v1.7.5 of the library.
 * Redefine it for installations with a lesser version.
 */
#ifndef LZ4F_HEADER_SIZE_MAX
#define LZ4F_HEADER_SIZE_MAX	32
#endif

/*---------------------------------
 * Common to both compression APIs
 *---------------------------------
 */

/*
 * (de)compression state used by both the Compressor and Stream APIs.
 */
typedef struct LZ4State
{
	/*
	 * Used by the Stream API to keep track of the file stream.
	 */
	FILE	   *fp;

	LZ4F_preferences_t prefs;

	LZ4F_compressionContext_t ctx;
	LZ4F_decompressionContext_t dtx;

	/*
	 * Used by the Stream API's lazy initialization.
	 */
	bool		inited;

	/*
	 * Used by the Stream API to distinguish between compression and
	 * decompression operations.
	 */
	bool		compressing;

	/*
	 * I/O buffer area.
	 */
	char	   *buffer;			/* buffer for compressed data */
	size_t		buflen;			/* allocated size of buffer */
	size_t		bufdata;		/* amount of valid data currently in buffer */
	/* These fields are used only while decompressing: */
	size_t		bufnext;		/* next buffer position to decompress */
	char	   *outbuf;			/* buffer for decompressed data */
	size_t		outbuflen;		/* allocated size of outbuf */
	size_t		outbufdata;		/* amount of valid data currently in outbuf */
	size_t		outbufnext;		/* next outbuf position to return */

	/*
	 * Used by both APIs to keep track of error codes.
	 */
	size_t		errcode;
} LZ4State;

/*
 * LZ4State_compression_init
 *		Initialize the required LZ4State members for compression.
 *
 * Write the LZ4 frame header in a buffer keeping track of its length. Users of
 * this function can choose when and how to write the header to a file stream.
 *
 * Returns true on success. In case of a failure returns false, and stores the
 * error code in state->errcode.
 */
static bool
LZ4State_compression_init(LZ4State *state)
{
	size_t		status;

	/*
	 * Compute size needed for buffer, assuming we will present at most
	 * DEFAULT_IO_BUFFER_SIZE input bytes at a time.
	 */
	state->buflen = LZ4F_compressBound(DEFAULT_IO_BUFFER_SIZE, &state->prefs);

	/*
	 * Add some slop to ensure we're not forced to flush every time.
	 *
	 * The present slop factor of 50% is chosen so that the typical output
	 * block size is about 128K when DEFAULT_IO_BUFFER_SIZE = 128K.  We might
	 * need a different slop factor to maintain that equivalence if
	 * DEFAULT_IO_BUFFER_SIZE is changed dramatically.
	 */
	state->buflen += state->buflen / 2;

	/*
	 * LZ4F_compressBegin requires a buffer that is greater or equal to
	 * LZ4F_HEADER_SIZE_MAX. Verify that the requirement is met.
	 */
	if (state->buflen < LZ4F_HEADER_SIZE_MAX)
		state->buflen = LZ4F_HEADER_SIZE_MAX;

	status = LZ4F_createCompressionContext(&state->ctx, LZ4F_VERSION);
	if (LZ4F_isError(status))
	{
		state->errcode = status;
		return false;
	}

	state->buffer = pg_malloc(state->buflen);

	/*
	 * Insert LZ4 header into buffer.
	 */
	status = LZ4F_compressBegin(state->ctx,
								state->buffer, state->buflen,
								&state->prefs);
	if (LZ4F_isError(status))
	{
		state->errcode = status;
		return false;
	}

	state->bufdata = status;

	return true;
}

/*----------------------
 * Compressor API
 *----------------------
 */

/* Private routines that support LZ4 compressed data I/O */

static void
ReadDataFromArchiveLZ4(ArchiveHandle *AH, CompressorState *cs)
{
	size_t		r;
	size_t		readbuflen;
	char	   *outbuf;
	char	   *readbuf;
	LZ4F_decompressionContext_t ctx = NULL;
	LZ4F_decompressOptions_t dec_opt;
	LZ4F_errorCode_t status;

	memset(&dec_opt, 0, sizeof(dec_opt));
	status = LZ4F_createDecompressionContext(&ctx, LZ4F_VERSION);
	if (LZ4F_isError(status))
		pg_fatal("could not create LZ4 decompression context: %s",
				 LZ4F_getErrorName(status));

	outbuf = pg_malloc(DEFAULT_IO_BUFFER_SIZE);
	readbuf = pg_malloc(DEFAULT_IO_BUFFER_SIZE);
	readbuflen = DEFAULT_IO_BUFFER_SIZE;
	while ((r = cs->readF(AH, &readbuf, &readbuflen)) > 0)
	{
		char	   *readp;
		char	   *readend;

		/* Process one chunk */
		readp = readbuf;
		readend = readbuf + r;
		while (readp < readend)
		{
			size_t		out_size = DEFAULT_IO_BUFFER_SIZE;
			size_t		read_size = readend - readp;

			status = LZ4F_decompress(ctx, outbuf, &out_size,
									 readp, &read_size, &dec_opt);
			if (LZ4F_isError(status))
				pg_fatal("could not decompress: %s",
						 LZ4F_getErrorName(status));

			ahwrite(outbuf, 1, out_size, AH);
			readp += read_size;
		}
	}

	pg_free(outbuf);
	pg_free(readbuf);

	status = LZ4F_freeDecompressionContext(ctx);
	if (LZ4F_isError(status))
		pg_fatal("could not free LZ4 decompression context: %s",
				 LZ4F_getErrorName(status));
}

static void
WriteDataToArchiveLZ4(ArchiveHandle *AH, CompressorState *cs,
					  const void *data, size_t dLen)
{
	LZ4State   *state = (LZ4State *) cs->private_data;
	size_t		remaining = dLen;

	while (remaining > 0)
	{
		size_t		chunk;
		size_t		required;
		size_t		status;

		/* We don't try to present more than DEFAULT_IO_BUFFER_SIZE bytes */
		chunk = Min(remaining, (size_t) DEFAULT_IO_BUFFER_SIZE);

		/* If not enough space, must flush buffer */
		required = LZ4F_compressBound(chunk, &state->prefs);
		if (required > state->buflen - state->bufdata)
		{
			cs->writeF(AH, state->buffer, state->bufdata);
			state->bufdata = 0;
		}

		status = LZ4F_compressUpdate(state->ctx,
									 state->buffer + state->bufdata,
									 state->buflen - state->bufdata,
									 data, chunk, NULL);

		if (LZ4F_isError(status))
			pg_fatal("could not compress data: %s",
					 LZ4F_getErrorName(status));

		state->bufdata += status;

		data = ((const char *) data) + chunk;
		remaining -= chunk;
	}
}

static void
EndCompressorLZ4(ArchiveHandle *AH, CompressorState *cs)
{
	LZ4State   *state = (LZ4State *) cs->private_data;
	size_t		required;
	size_t		status;

	/* Nothing needs to be done */
	if (!state)
		return;

	/* We might need to flush the buffer to make room for LZ4F_compressEnd */
	required = LZ4F_compressBound(0, &state->prefs);
	if (required > state->buflen - state->bufdata)
	{
		cs->writeF(AH, state->buffer, state->bufdata);
		state->bufdata = 0;
	}

	status = LZ4F_compressEnd(state->ctx,
							  state->buffer + state->bufdata,
							  state->buflen - state->bufdata,
							  NULL);
	if (LZ4F_isError(status))
		pg_fatal("could not end compression: %s",
				 LZ4F_getErrorName(status));
	state->bufdata += status;

	/* Write the final bufferload */
	cs->writeF(AH, state->buffer, state->bufdata);

	status = LZ4F_freeCompressionContext(state->ctx);
	if (LZ4F_isError(status))
		pg_fatal("could not end compression: %s",
				 LZ4F_getErrorName(status));

	pg_free(state->buffer);
	pg_free(state);

	cs->private_data = NULL;
}

/*
 * Public routines that support LZ4 compressed data I/O
 */
void
InitCompressorLZ4(CompressorState *cs, const pg_compress_specification compression_spec)
{
	LZ4State   *state;

	cs->readData = ReadDataFromArchiveLZ4;
	cs->writeData = WriteDataToArchiveLZ4;
	cs->end = EndCompressorLZ4;

	cs->compression_spec = compression_spec;

	/*
	 * Read operations have access to the whole input. No state needs to be
	 * carried between calls.
	 */
	if (cs->readF)
		return;

	state = pg_malloc0(sizeof(*state));
	if (cs->compression_spec.level >= 0)
		state->prefs.compressionLevel = cs->compression_spec.level;

	if (!LZ4State_compression_init(state))
		pg_fatal("could not initialize LZ4 compression: %s",
				 LZ4F_getErrorName(state->errcode));

	cs->private_data = state;
}

/*----------------------
 * Compress Stream API
 *----------------------
 */


/*
 * LZ4 equivalent to feof() or gzeof().  Return true iff there is no
 * more buffered data and the end of the input file has been reached.
 */
static bool
LZ4Stream_eof(CompressFileHandle *CFH)
{
	LZ4State   *state = (LZ4State *) CFH->private_data;

	return state->outbufnext >= state->outbufdata &&
		state->bufnext >= state->bufdata &&
		feof(state->fp);
}

static const char *
LZ4Stream_get_error(CompressFileHandle *CFH)
{
	LZ4State   *state = (LZ4State *) CFH->private_data;
	const char *errmsg;

	if (LZ4F_isError(state->errcode))
		errmsg = LZ4F_getErrorName(state->errcode);
	else
		errmsg = strerror(errno);

	return errmsg;
}

/*
 * Initialize an already alloc'ed LZ4State struct for subsequent calls.
 *
 * Creates the necessary contexts for either compression or decompression. When
 * compressing data (indicated by compressing=true), it additionally writes the
 * LZ4 header in the output buffer.
 *
 * It's expected that a not-yet-initialized LZ4State will be zero-filled.
 *
 * Returns true on success. In case of a failure returns false, and stores the
 * error code in state->errcode.
 */
static bool
LZ4Stream_init(LZ4State *state, bool compressing)
{
	size_t		status;

	if (state->inited)
		return true;

	state->compressing = compressing;

	if (state->compressing)
	{
		if (!LZ4State_compression_init(state))
			return false;
	}
	else
	{
		status = LZ4F_createDecompressionContext(&state->dtx, LZ4F_VERSION);
		if (LZ4F_isError(status))
		{
			state->errcode = status;
			return false;
		}

		state->buflen = DEFAULT_IO_BUFFER_SIZE;
		state->buffer = pg_malloc(state->buflen);
		state->outbuflen = DEFAULT_IO_BUFFER_SIZE;
		state->outbuf = pg_malloc(state->outbuflen);
	}

	state->inited = true;
	return true;
}

/*
 * The workhorse for reading decompressed content out of an LZ4 compressed
 * stream.
 *
 * It will read up to 'ptrsize' decompressed content, or up to the new line
 * char if one is found first when the eol_flag is set.
 *
 * Returns the number of bytes of decompressed data copied into the ptr
 * buffer, or -1 in case of error.
 */
static int
LZ4Stream_read_internal(LZ4State *state, void *ptr, int ptrsize, bool eol_flag)
{
	int			dsize = 0;
	int			remaining = ptrsize;

	/* Lazy init */
	if (!LZ4Stream_init(state, false /* decompressing */ ))
	{
		pg_log_error("unable to initialize LZ4 library: %s",
					 LZ4F_getErrorName(state->errcode));
		return -1;
	}

	/* Loop until postcondition is satisfied */
	while (remaining > 0)
	{
		/*
		 * If we already have some decompressed data, return that.
		 */
		if (state->outbufnext < state->outbufdata)
		{
			char	   *outptr = state->outbuf + state->outbufnext;
			size_t		readlen = state->outbufdata - state->outbufnext;
			bool		eol_found = false;

			if (readlen > remaining)
				readlen = remaining;
			/* If eol_flag is set, don't read beyond a newline */
			if (eol_flag)
			{
				char	   *eolptr = memchr(outptr, '\n', readlen);

				if (eolptr)
				{
					readlen = eolptr - outptr + 1;
					eol_found = true;
				}
			}
			memcpy(ptr, outptr, readlen);
			ptr = ((char *) ptr) + readlen;
			state->outbufnext += readlen;
			dsize += readlen;
			remaining -= readlen;
			if (eol_found || remaining == 0)
				break;
			/* We must have emptied outbuf */
			Assert(state->outbufnext >= state->outbufdata);
		}

		/*
		 * If we don't have any pending compressed data, load more into
		 * state->buffer.
		 */
		if (state->bufnext >= state->bufdata)
		{
			size_t		rsize;

			rsize = fread(state->buffer, 1, state->buflen, state->fp);
			if (rsize < state->buflen && !feof(state->fp))
			{
				pg_log_error("could not read from input file: %m");
				return -1;
			}
			if (rsize == 0)
				break;			/* must be EOF */
			state->bufdata = rsize;
			state->bufnext = 0;
		}

		/*
		 * Decompress some data into state->outbuf.
		 */
		{
			size_t		status;
			size_t		outlen = state->outbuflen;
			size_t		inlen = state->bufdata - state->bufnext;

			status = LZ4F_decompress(state->dtx,
									 state->outbuf, &outlen,
									 state->buffer + state->bufnext,
									 &inlen,
									 NULL);
			if (LZ4F_isError(status))
			{
				state->errcode = status;
				pg_log_error("could not read from input file: %s",
							 LZ4F_getErrorName(state->errcode));
				return -1;
			}
			state->bufnext += inlen;
			state->outbufdata = outlen;
			state->outbufnext = 0;
		}
	}

	return dsize;
}

/*
 * Compress size bytes from ptr and write them to the stream.
 */
static void
LZ4Stream_write(const void *ptr, size_t size, CompressFileHandle *CFH)
{
	LZ4State   *state = (LZ4State *) CFH->private_data;
	size_t		remaining = size;

	/* Lazy init */
	if (!LZ4Stream_init(state, true))
		pg_fatal("unable to initialize LZ4 library: %s",
				 LZ4F_getErrorName(state->errcode));

	while (remaining > 0)
	{
		size_t		chunk;
		size_t		required;
		size_t		status;

		/* We don't try to present more than DEFAULT_IO_BUFFER_SIZE bytes */
		chunk = Min(remaining, (size_t) DEFAULT_IO_BUFFER_SIZE);

		/* If not enough space, must flush buffer */
		required = LZ4F_compressBound(chunk, &state->prefs);
		if (required > state->buflen - state->bufdata)
		{
			errno = 0;
			if (fwrite(state->buffer, 1, state->bufdata, state->fp) != state->bufdata)
			{
				errno = (errno) ? errno : ENOSPC;
				pg_fatal("error during writing: %m");
			}
			state->bufdata = 0;
		}

		status = LZ4F_compressUpdate(state->ctx,
									 state->buffer + state->bufdata,
									 state->buflen - state->bufdata,
									 ptr, chunk, NULL);
		if (LZ4F_isError(status))
			pg_fatal("error during writing: %s", LZ4F_getErrorName(status));
		state->bufdata += status;

		ptr = ((const char *) ptr) + chunk;
		remaining -= chunk;
	}
}

/*
 * fread() equivalent implementation for LZ4 compressed files.
 */
static size_t
LZ4Stream_read(void *ptr, size_t size, CompressFileHandle *CFH)
{
	LZ4State   *state = (LZ4State *) CFH->private_data;
	int			ret;

	if ((ret = LZ4Stream_read_internal(state, ptr, size, false)) < 0)
		pg_fatal("could not read from input file: %s", LZ4Stream_get_error(CFH));

	return (size_t) ret;
}

/*
 * fgetc() equivalent implementation for LZ4 compressed files.
 */
static int
LZ4Stream_getc(CompressFileHandle *CFH)
{
	LZ4State   *state = (LZ4State *) CFH->private_data;
	unsigned char c;

	if (LZ4Stream_read_internal(state, &c, 1, false) <= 0)
	{
		if (!LZ4Stream_eof(CFH))
			pg_fatal("could not read from input file: %s", LZ4Stream_get_error(CFH));
		else
			pg_fatal("could not read from input file: end of file");
	}

	return c;
}

/*
 * fgets() equivalent implementation for LZ4 compressed files.
 */
static char *
LZ4Stream_gets(char *ptr, int size, CompressFileHandle *CFH)
{
	LZ4State   *state = (LZ4State *) CFH->private_data;
	int			ret;

	ret = LZ4Stream_read_internal(state, ptr, size - 1, true);

	/*
	 * LZ4Stream_read_internal returning 0 or -1 means that it was either an
	 * EOF or an error, but gets_func is defined to return NULL in either case
	 * so we can treat both the same here.
	 */
	if (ret <= 0)
		return NULL;

	/*
	 * Our caller expects the return string to be NULL terminated and we know
	 * that ret is greater than zero.
	 */
	ptr[ret - 1] = '\0';

	return ptr;
}

/*
 * Finalize (de)compression of a stream. When compressing it will write any
 * remaining content and/or generated footer from the LZ4 API.
 */
static bool
LZ4Stream_close(CompressFileHandle *CFH)
{
	FILE	   *fp;
	LZ4State   *state = (LZ4State *) CFH->private_data;
	size_t		required;
	size_t		status;
	int			ret;

	fp = state->fp;
	if (state->inited)
	{
		if (state->compressing)
		{
			/* We might need to flush the buffer to make room */
			required = LZ4F_compressBound(0, &state->prefs);
			if (required > state->buflen - state->bufdata)
			{
				errno = 0;
				if (fwrite(state->buffer, 1, state->bufdata, state->fp) != state->bufdata)
				{
					errno = (errno) ? errno : ENOSPC;
					pg_log_error("could not write to output file: %m");
				}
				state->bufdata = 0;
			}

			status = LZ4F_compressEnd(state->ctx,
									  state->buffer + state->bufdata,
									  state->buflen - state->bufdata,
									  NULL);
			if (LZ4F_isError(status))
			{
				pg_log_error("could not end compression: %s",
							 LZ4F_getErrorName(status));
			}
			else
				state->bufdata += status;

			errno = 0;
			if (fwrite(state->buffer, 1, state->bufdata, state->fp) != state->bufdata)
			{
				errno = (errno) ? errno : ENOSPC;
				pg_log_error("could not write to output file: %m");
			}

			status = LZ4F_freeCompressionContext(state->ctx);
			if (LZ4F_isError(status))
				pg_log_error("could not end compression: %s",
							 LZ4F_getErrorName(status));
		}
		else
		{
			status = LZ4F_freeDecompressionContext(state->dtx);
			if (LZ4F_isError(status))
				pg_log_error("could not end decompression: %s",
							 LZ4F_getErrorName(status));
			pg_free(state->outbuf);
		}

		pg_free(state->buffer);
	}

	pg_free(state);
	CFH->private_data = NULL;

	errno = 0;
	ret = fclose(fp);
	if (ret != 0)
	{
		pg_log_error("could not close file: %m");
		return false;
	}

	return true;
}

static bool
LZ4Stream_open(const char *path, int fd, const char *mode,
			   CompressFileHandle *CFH)
{
	LZ4State   *state = (LZ4State *) CFH->private_data;

	if (fd >= 0)
		state->fp = fdopen(dup(fd), mode);
	else
		state->fp = fopen(path, mode);
	if (state->fp == NULL)
	{
		state->errcode = errno;
		return false;
	}

	return true;
}

static bool
LZ4Stream_open_write(const char *path, const char *mode, CompressFileHandle *CFH)
{
	char	   *fname;
	int			save_errno;
	bool		ret;

	fname = psprintf("%s.lz4", path);
	ret = CFH->open_func(fname, -1, mode, CFH);

	save_errno = errno;
	pg_free(fname);
	errno = save_errno;

	return ret;
}

/*
 * Public routines
 */
void
InitCompressFileHandleLZ4(CompressFileHandle *CFH,
						  const pg_compress_specification compression_spec)
{
	LZ4State   *state;

	CFH->open_func = LZ4Stream_open;
	CFH->open_write_func = LZ4Stream_open_write;
	CFH->read_func = LZ4Stream_read;
	CFH->write_func = LZ4Stream_write;
	CFH->gets_func = LZ4Stream_gets;
	CFH->getc_func = LZ4Stream_getc;
	CFH->eof_func = LZ4Stream_eof;
	CFH->close_func = LZ4Stream_close;
	CFH->get_error_func = LZ4Stream_get_error;

	CFH->compression_spec = compression_spec;
	state = pg_malloc0(sizeof(*state));
	if (CFH->compression_spec.level >= 0)
		state->prefs.compressionLevel = CFH->compression_spec.level;

	CFH->private_data = state;
}
#else							/* USE_LZ4 */
void
InitCompressorLZ4(CompressorState *cs,
				  const pg_compress_specification compression_spec)
{
	pg_fatal("this build does not support compression with %s", "LZ4");
}

void
InitCompressFileHandleLZ4(CompressFileHandle *CFH,
						  const pg_compress_specification compression_spec)
{
	pg_fatal("this build does not support compression with %s", "LZ4");
}
#endif							/* USE_LZ4 */