/*-------------------------------------------------------------------------
 *
 * slowio.c
 *    LD_PRELOAD library to slow down destination writes during cp.
 *
 * Environment variables:
 *    SLOWIO_TARGET     - if set, only slow writes to fds whose path contains
 *                        this substring (e.g. "archives")
 *
 * Delays are applied only when the process name is "cp", so the library can
 * be inherited broadly without slowing unrelated children.
 *
 * Example:
 *    LD_PRELOAD=./slowio.so SLOWIO_TARGET=archives/000000010000000000000002 \
 *        cp src dst
 *
 *-------------------------------------------------------------------------
 */

#define _GNU_SOURCE

#include <dlfcn.h>
#include <errno.h>
#include <fcntl.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <time.h>

#ifdef __linux__
#include <sys/sendfile.h>
#endif

static ssize_t (*next_write)(int fd, const void *buf, size_t count);
static ssize_t (*next_pwrite)(int fd, const void *buf, size_t count, off_t offset);
#ifdef __linux__
static ssize_t (*next_copy_file_range)(int fd_in, loff_t *off_in,
									   int fd_out, loff_t *off_out,
									   size_t len, unsigned int flags);
static ssize_t (*next_sendfile)(int out_fd, int in_fd, off_t *offset, size_t count);
#endif

static bool symbols_initialized;
static int	is_cp_process = -1;		/* tri-state cache for process_is_cp() */

static void
init_symbols(void)
{
	if (symbols_initialized)
		return;

	next_write = (ssize_t (*)(int, const void *, size_t))
		dlsym(RTLD_NEXT, "write");
	next_pwrite = (ssize_t (*)(int, const void *, size_t, off_t))
		dlsym(RTLD_NEXT, "pwrite");
#ifdef __linux__
	next_copy_file_range = (ssize_t (*)(int, loff_t *, int, loff_t *, size_t, unsigned int))
		dlsym(RTLD_NEXT, "copy_file_range");
	next_sendfile = (ssize_t (*)(int, int, off_t *, size_t))
		dlsym(RTLD_NEXT, "sendfile");
#endif

	symbols_initialized = true;

	srand(time(NULL));
}

static bool
name_is_cp(const char *name)
{
	const char *base;

	if (name == NULL || name[0] == '\0')
		return false;

	base = strrchr(name, '/');
	base = base ? base + 1 : name;

	return strcmp(base, "cp") == 0;
}

static bool
process_is_cp(void)
{
	if (is_cp_process >= 0)
		return is_cp_process;

#if defined(__APPLE__)
	is_cp_process = name_is_cp(getprogname());
#elif defined(__linux__)
	{
		char		comm[32];
		FILE	   *file;

		is_cp_process = 0;
		file = fopen("/proc/self/comm", "r");
		if (file != NULL)
		{
			if (fgets(comm, sizeof(comm), file) != NULL)
			{
				char	   *nl = strchr(comm, '\n');

				if (nl != NULL)
					*nl = '\0';
				is_cp_process = name_is_cp(comm);
			}
			fclose(file);
		}

		if (!is_cp_process)
		{
			char		cmdline[256];
			ssize_t		n;
			int			fd;

			fd = open("/proc/self/cmdline", O_RDONLY);
			if (fd >= 0)
			{
				n = read(fd, cmdline, sizeof(cmdline) - 1);
				close(fd);
				if (n > 0)
				{
					cmdline[n] = '\0';
					is_cp_process = name_is_cp(cmdline);
				}
			}
		}
	}
#else
	is_cp_process = 0;
#endif

	return is_cp_process;
}

static bool
fd_is_target(int fd)
{
	const char *needle = getenv("SLOWIO_TARGET");
	char		path[4096];

	if (needle == NULL || needle[0] == '\0')
		return true;

#if defined(__APPLE__)
	if (fcntl(fd, F_GETPATH, path) == -1)
		return false;
#elif defined(__linux__)
	{
		char		linkpath[64];
		ssize_t		n;

		snprintf(linkpath, sizeof(linkpath), "/proc/self/fd/%d", fd);
		n = readlink(linkpath, path, sizeof(path) - 1);
		if (n < 0)
			return false;
		path[n] = '\0';
	}
#else
	return true;
#endif

	return strstr(path, needle) != NULL;
}

static void
maybe_delay(int fd)
{
	if (process_is_cp() && fd_is_target(fd))
	{
		if (rand() % 10 == 0) usleep(50000);
	}
}

ssize_t
write(int fd, const void *buf, size_t count)
{
	ssize_t		rc;

	init_symbols();
	if (next_write == NULL)
	{
		errno = ENOSYS;
		return -1;
	}

	rc = next_write(fd, buf, count);
	if (rc > 0)
		maybe_delay(fd);
	return rc;
}

ssize_t
pwrite(int fd, const void *buf, size_t count, off_t offset)
{
	ssize_t		rc;

	init_symbols();
	if (next_pwrite == NULL)
	{
		errno = ENOSYS;
		return -1;
	}

	rc = next_pwrite(fd, buf, count, offset);
	if (rc > 0)
		maybe_delay(fd);
	return rc;
}

#ifdef __linux__
ssize_t
copy_file_range(int fd_in, loff_t *off_in,
				int fd_out, loff_t *off_out,
				size_t len, unsigned int flags)
{
	size_t		chunk = 65536;
	ssize_t		total = 0;

	init_symbols();
	if (next_copy_file_range == NULL)
	{
		errno = ENOSYS;
		return -1;
	}

	while (len > 0)
	{
		size_t		n = len < chunk ? len : chunk;
		ssize_t		rc;

		rc = next_copy_file_range(fd_in, off_in, fd_out, off_out, n, flags);
		if (rc < 0)
			return total > 0 ? total : -1;
		if (rc == 0)
			break;

		total += rc;
		len -= (size_t) rc;
		maybe_delay(fd_out);
	}

	return total;
}

ssize_t
sendfile(int out_fd, int in_fd, off_t *offset, size_t count)
{
	size_t		chunk = 65536;
	ssize_t		total = 0;

	init_symbols();
	if (next_sendfile == NULL)
	{
		errno = ENOSYS;
		return -1;
	}

	while (count > 0)
	{
		size_t		n = count < chunk ? count : chunk;
		ssize_t		rc;

		rc = next_sendfile(out_fd, in_fd, offset, n);
		if (rc < 0)
			return total > 0 ? total : -1;
		if (rc == 0)
			break;

		total += rc;
		count -= (size_t) rc;
		maybe_delay(out_fd);
	}

	return total;
}
#endif