#define _GNU_SOURCE  // RTLD_NEXT
#include <stdio.h>
#include <stdbool.h>
#include <stdlib.h>
#include <stdarg.h>
#include <string.h>
#include <limits.h>
#include <unistd.h>
#include <errno.h>
#include <fcntl.h>
#include <dlfcn.h>
#include <spawn.h>  // posix_spawn
#include <sys/socket.h>
#include <sys/un.h>
#include "libintercept.h"


// Returns whether successful
static bool sendall(int sock, const char *buffer, size_t length) {
	size_t cursor = 0;
	while (cursor < length) {
		ssize_t nw = send(sock, buffer + cursor, length - cursor, 0);
		if (nw <= 0) return false;
		cursor += nw;
	}
	return true;
}

static void try_transmit_invocation(const char *pathname, char *const argv[]) {
	const char *socketpath = getenv(COMM_SOCKET_ENVVAR);
	if (socketpath == NULL) return;

	// The SOCK_CLOEXEC option is technically unnecessary, but let's be careful.
	int sock = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0);
	if (sock < 0) return;

	struct sockaddr_un addr;
	memset(&addr, 0, sizeof addr);
	addr.sun_family = AF_UNIX;
	strncpy(addr.sun_path, socketpath, sizeof(addr.sun_path) - 1);

	int ret = connect(sock, (const struct sockaddr*)&addr, sizeof addr);
	if (ret < 0) return;

	size_t nargs = 0;
	while (argv[nargs] != NULL) nargs++;

	if (!sendall(sock, (const char*)&nargs, 8)) goto cleanup;
	if (!sendall(sock, pathname, strlen(pathname) + 1)) goto cleanup;
	for (size_t i = 0; i < nargs; i++) {
		if (!sendall(sock, argv[i], strlen(argv[i]) + 1)) goto cleanup;
	}

cleanup:
	close(sock);
}

// Abridged from glibc posix/execl.c, LGPL copyright FSF
int execl(const char *pathname, const char *arg, ...) {
	size_t argc;
	va_list ap;
	va_start(ap, arg);
	for (argc = 1; va_arg(ap, const char*) && argc < INT_MAX; argc++) {}
	va_end(ap);
	if (argc >= INT_MAX) {
		errno = E2BIG;
		return -1;
	}

	char *argv[argc + 1];
	argv[0] = (char*)arg;
	va_start(ap, arg);
	for (size_t i = 1; i < argc; i++) argv[i] = va_arg(ap, char*);
	argv[argc] = NULL;
	va_end(ap);

	try_transmit_invocation(pathname, argv);

	int (*real_execv)(const char*, char *const[]) = dlsym(RTLD_NEXT, "execv");
	return real_execv(pathname, argv);
}

int execlp(const char *file, const char *arg, ...) {
	size_t argc;
	va_list ap;
	va_start(ap, arg);
	for (argc = 1; va_arg(ap, const char*) && argc < INT_MAX; argc++) {}
	va_end(ap);
	if (argc >= INT_MAX) {
		errno = E2BIG;
		return -1;
	}

	char *argv[argc + 1];
	argv[0] = (char*)arg;
	va_start(ap, arg);
	for (size_t i = 1; i < argc; i++) argv[i] = va_arg(ap, char*);
	argv[argc] = NULL;
	va_end(ap);

	try_transmit_invocation(file, argv);

	int (*real_execvp)(const char*, char *const[]) = dlsym(RTLD_NEXT, "execvp");
	return real_execvp(file, argv);
}

int execle(const char *pathname, const char *arg, ...) {
	size_t argc;
	va_list ap;
	va_start(ap, arg);
	for (argc = 1; va_arg(ap, const char*) && argc < INT_MAX; argc++) {}
	char *const *envp = va_arg(ap, char *const *);
	va_end(ap);
	if (argc >= INT_MAX) {
		errno = E2BIG;
		return -1;
	}

	char *argv[argc + 1];
	argv[0] = (char*)arg;
	va_start(ap, arg);
	for (size_t i = 1; i < argc; i++) argv[i] = va_arg(ap, char*);
	argv[argc] = NULL;
	va_end(ap);

	try_transmit_invocation(pathname, argv);

	int (*real_execvpe)(const char*, char *const[], char *const[]) = dlsym(RTLD_NEXT, "execvpe");
	return real_execvpe(pathname, argv, envp);
}

int execv(const char *pathname, char *const argv[]) {
	try_transmit_invocation(pathname, argv);

	int (*real_execv)(const char*, char *const[]) = dlsym(RTLD_NEXT, "execv");
	return real_execv(pathname, argv);
}

int execvp(const char *file, char *const argv[]) {
	try_transmit_invocation(file, argv);

	int (*real_execvp)(const char*, char *const[]) = dlsym(RTLD_NEXT, "execvp");
	return real_execvp(file, argv);
}

int execve(const char *pathname, char *const argv[], char *const envp[]) {
	try_transmit_invocation(pathname, argv);

	int (*real_execve)(const char*, char *const[], char *const[]) = dlsym(RTLD_NEXT, "execve");
	return real_execve(pathname, argv, envp);
}

int execvpe(const char *file, char *const argv[], char *const envp[]) {
	try_transmit_invocation(file, argv);

	int (*real_execvpe)(const char*, char *const[], char *const[]) = dlsym(RTLD_NEXT, "execvpe");
	return real_execvpe(file, argv, envp);
}

int posix_spawn(
	pid_t *pid, const char *path,
	const posix_spawn_file_actions_t *file_actions,
	const posix_spawnattr_t *attrp,
	char *const argv[], char *const envp[]
) {
	try_transmit_invocation(path, argv);

	int (*real_posix_spawn)(
		pid_t*, const char*,
		const posix_spawn_file_actions_t*,
		const posix_spawnattr_t*,
		char *const[], char *const[]
	) = dlsym(RTLD_NEXT, "posix_spawn");
	return real_posix_spawn(pid, path, file_actions, attrp, argv, envp);
}