/*
 * Copyright (c) 2018      Andreas Schneider <asn@samba.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * 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, see <http://www.gnu.org/licenses/>.
 */

#include "config.h"

#include <errno.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <stdint.h>
#include <stdarg.h>
#include <stdbool.h>
#include <unistd.h>
#include <pthread.h>

#ifdef HAVE_SYS_PRCTL_H
#include <sys/prctl.h>
#endif
#ifdef HAVE_SYS_TIME_H
#include <sys/time.h>
#endif
#ifdef HAVE_SYS_RESOURCE_H
#include <sys/types.h>
#include <sys/resource.h>
#endif

enum privwrap_dbglvl_e {
	PRIVWRAP_LOG_ERROR = 0,
	PRIVWRAP_LOG_WARN,
	PRIVWRAP_LOG_DEBUG,
	PRIVWRAP_LOG_TRACE
};

/* GCC have printf type attribute check. */
#ifdef HAVE_FUNCTION_ATTRIBUTE_FORMAT
#define PRINTF_ATTRIBUTE(a,b) __attribute__ ((__format__ (__printf__, a, b)))
#else
#define PRINTF_ATTRIBUTE(a,b)
#endif /* HAVE_FUNCTION_ATTRIBUTE_FORMAT */

#ifdef HAVE_CONSTRUCTOR_ATTRIBUTE
#define CONSTRUCTOR_ATTRIBUTE __attribute__ ((constructor))
#else
#define CONSTRUCTOR_ATTRIBUTE
#endif /* HAVE_CONSTRUCTOR_ATTRIBUTE */

#ifdef HAVE_DESTRUCTOR_ATTRIBUTE
#define DESTRUCTOR_ATTRIBUTE __attribute__ ((destructor))
#else
#define DESTRUCTOR_ATTRIBUTE
#endif

#ifndef FALL_THROUGH
# ifdef HAVE_FALLTHROUGH_ATTRIBUTE
#  define FALL_THROUGH __attribute__ ((fallthrough))
# else /* HAVE_FALLTHROUGH_ATTRIBUTE */
#  define FALL_THROUGH ((void)0)
# endif /* HAVE_FALLTHROUGH_ATTRIBUTE */
#endif /* FALL_THROUGH */

#ifdef HAVE_ADDRESS_SANITIZER_ATTRIBUTE
#define DO_NOT_SANITIZE_ADDRESS_ATTRIBUTE __attribute__((no_sanitize_address))
#else
#define DO_NOT_SANITIZE_ADDRESS_ATTRIBUTE
#endif

#ifndef MIN
#define MIN(a,b) ((a)<(b)?(a):(b))
#endif

#ifndef ZERO_STRUCT
#define ZERO_STRUCT(x) memset((char *)&(x), 0, sizeof(x))
#endif

#ifndef ZERO_STRUCTP
#define ZERO_STRUCTP(x) do { \
		if ((x) != NULL) \
			memset((char *)(x), 0, sizeof(*(x))); \
	} while(0)
#endif

#ifndef SAFE_FREE
#define SAFE_FREE(x) do { if ((x) != NULL) {free(x); (x)=NULL;} } while(0)
#endif

#ifndef discard_const
#define discard_const(ptr) ((void *)((uintptr_t)(ptr)))
#endif

#ifndef discard_const_p
#define discard_const_p(type, ptr) ((type *)discard_const(ptr))
#endif

#define UNUSED(x) (void)(x)

bool priv_wrapper_enabled(void);

void privwrap_constructor(void) CONSTRUCTOR_ATTRIBUTE;
void privwrap_destructor(void) DESTRUCTOR_ATTRIBUTE;

static void privwrap_log(enum privwrap_dbglvl_e dbglvl, const char *func, const char *format, ...) PRINTF_ATTRIBUTE(3, 4);
# define PRIVWRAP_LOG(dbglvl, ...) privwrap_log((dbglvl), __func__, __VA_ARGS__)

/*
 * DO_NOT_SANITIZE_ADDRESS_ATTRIBUTE is used to avoid libasan to check all calls
 * used by setrlimit(). setrlimit() is called by every process too early,
 * already from _dl_init(), before call of __asan_init(), thus libasan is not
 * yet fully initialized. This causes SIGSEGV and process crash.
 */
DO_NOT_SANITIZE_ADDRESS_ATTRIBUTE
static void privwrap_log(enum privwrap_dbglvl_e dbglvl,
		      const char *func,
		      const char *format, ...)
{
	char buffer[1024];
	va_list va;
	const char *d;
	unsigned int lvl = 0;
	const char *prefix = "PRIVWRAP";

	d = getenv("PRIV_WRAPPER_DEBUGLEVEL");
	if (d != NULL) {
		lvl = atoi(d);
	}

	if (lvl < dbglvl) {
		return;
	}

	va_start(va, format);
	vsnprintf(buffer, sizeof(buffer), format, va);
	va_end(va);

	switch (dbglvl) {
		case PRIVWRAP_LOG_ERROR:
			prefix = "PRIVWRAP_ERROR";
			break;
		case PRIVWRAP_LOG_WARN:
			prefix = "PRIVWRAP_WARN";
			break;
		case PRIVWRAP_LOG_DEBUG:
			prefix = "PRIVWRAP_DEBUG";
			break;
		case PRIVWRAP_LOG_TRACE:
			prefix = "PRIVWRAP_TRACE";
			break;
	}

	fprintf(stderr,
		"%s(%d) - %s: %s\n",
		prefix, (int)getpid(), func, buffer);
}

/*********************************************************
 * PRIVWRAP LOADING LIBC FUNCTIONS
 *********************************************************/

#include <dlfcn.h>

typedef int (*__libc_chroot)(const char *path);
#ifdef HAVE_PRCTL
typedef int (*__libc_prctl)(int option, ...);
#endif
#ifdef HAVE_PLEDGE
typedef int (*__libc_pledge)(const char *promises,
			     const char *execpromises);
#endif
#ifdef HAVE_SETRLIMIT
typedef int (*__libc_setrlimit)(int resource,
				const struct rlimit *rlp);
#endif

#define PRIVWRAP_SYMBOL_ENTRY(i) \
	union { \
		__libc_##i f; \
		void *obj; \
	} _libc_##i

struct privwrap_libc_symbols {
#ifdef HAVE_PRCTL
	PRIVWRAP_SYMBOL_ENTRY(prctl);
#endif
#ifdef HAVE_SETRLIMIT
	PRIVWRAP_SYMBOL_ENTRY(setrlimit);
#endif
	PRIVWRAP_SYMBOL_ENTRY(chroot);
#ifdef HAVE_PLEDGE
	PRIVWRAP_SYMBOL_ENTRY(pledge);
#endif
};

struct privwrap {
	bool enabled;
	struct {
		void *handle;
		struct privwrap_libc_symbols symbols;
	} libc;
};

static struct privwrap privwrap;

#define LIBC_NAME "libc.so"

enum privwrap_lib {
    PRIVWRAP_LIBC,
};

static const char *privwrap_str_lib(enum privwrap_lib lib)
{
	switch (lib) {
	case PRIVWRAP_LIBC:
		return "libc";
	}

	/* Compiler would warn us about unhandled enum value if we get here */
	return "unknown";
}

DO_NOT_SANITIZE_ADDRESS_ATTRIBUTE
static void *privwrap_load_lib_handle(enum privwrap_lib lib)
{
	int flags = RTLD_LAZY;
	void *handle = NULL;
	int i;

#ifdef RTLD_DEEPBIND
	const char *env_preload = getenv("LD_PRELOAD");
	const char *env_deepbind = getenv("PRIV_WRAPPER_DISABLE_DEEPBIND");
	bool enable_deepbind = true;

	/* Don't do a deepbind if we run with libasan */
	if (env_preload != NULL && strlen(env_preload) < 1024) {
		const char *p = strstr(env_preload, "libasan.so");
		if (p != NULL) {
			enable_deepbind = false;
		}
	}

	if (env_deepbind != NULL && strlen(env_deepbind) >= 1) {
		enable_deepbind = false;
	}

	if (enable_deepbind) {
		flags |= RTLD_DEEPBIND;
	}
#endif

	switch (lib) {
	case PRIVWRAP_LIBC:
		handle = privwrap.libc.handle;
#ifdef LIBC_SO
		if (handle == NULL) {
			handle = dlopen(LIBC_SO, flags);

			privwrap.libc.handle = handle;
		}
#endif
		if (handle == NULL) {
			for (i = 10; i >= 0; i--) {
				char soname[256] = {0};

				snprintf(soname, sizeof(soname), "libc.so.%d", i);
				handle = dlopen(soname, flags);
				if (handle != NULL) {
					break;
				}
			}

			privwrap.libc.handle = handle;
		}
		break;
	}

	if (handle == NULL) {
#ifdef RTLD_NEXT
		handle = privwrap.libc.handle = RTLD_NEXT;
#else
		PRIVWRAP_LOG(PRIVWRAP_LOG_ERROR,
			  "Failed to dlopen library: %s\n",
			  dlerror());
		exit(-1);
#endif
	}

	return handle;
}

static void *_privwrap_bind_symbol(enum privwrap_lib lib, const char *fn_name)
{
	void *handle;
	void *func;

	handle = privwrap_load_lib_handle(lib);

	func = dlsym(handle, fn_name);
	if (func == NULL) {
		PRIVWRAP_LOG(PRIVWRAP_LOG_ERROR,
			  "Failed to find %s: %s\n",
			  fn_name,
			  dlerror());
		exit(-1);
	}

	PRIVWRAP_LOG(PRIVWRAP_LOG_TRACE,
		  "Loaded %s from %s",
		  fn_name,
		  privwrap_str_lib(lib));

	return func;
}

/*
 * These macros have a thread race condition on purpose!
 *
 * This is an optimization to avoid locking each time we check if the symbol is
 * bound.
 */
#define privwrap_bind_symbol_libc(sym_name) \
	if (privwrap.libc.symbols._libc_##sym_name.obj == NULL) { \
		privwrap.libc.symbols._libc_##sym_name.obj = \
			_privwrap_bind_symbol(PRIVWRAP_LIBC, #sym_name); \
	}

DO_NOT_SANITIZE_ADDRESS_ATTRIBUTE
static void __pwrap_bind_symbol_all_once(void)
{
#ifdef HAVE_SETRLIMIT
	privwrap_bind_symbol_libc(setrlimit);
#endif
	privwrap_bind_symbol_libc(chroot);
#ifdef HAVE_PRCTL
	privwrap_bind_symbol_libc(prctl);
#endif
#ifdef HAVE_PLEDGE
	privwrap_bind_symbol_libc(pledge);
#endif
}

DO_NOT_SANITIZE_ADDRESS_ATTRIBUTE
static void pwrap_bind_symbol_all(void)
{
	static pthread_once_t all_symbol_binding_once = PTHREAD_ONCE_INIT;

	pthread_once(&all_symbol_binding_once, __pwrap_bind_symbol_all_once);
}

/****************************************************************************
 *                               IMPORTANT
 ****************************************************************************
 *
 * Functions especially from libc need to be loaded individually, you can't
 * load all at once or gdb will segfault at startup. The same applies to
 * valgrind and has probably something todo with with the linker.  So we need
 * load each function at the point it is called the first time.
 *
 ****************************************************************************/

static int libc_chroot(const char *path)
{
	pwrap_bind_symbol_all();

	return privwrap.libc.symbols._libc_chroot.f(path);
}

#ifdef HAVE_PRCTL
static int libc_vprctl(int option,
		       va_list ap)
{
	unsigned long arg2;
	unsigned long arg3;
	unsigned long arg4;
	unsigned long arg5;

	pwrap_bind_symbol_all();

	arg2 = va_arg(ap, unsigned long);
	arg3 = va_arg(ap, unsigned long);
	arg4 = va_arg(ap, unsigned long);
	arg5 = va_arg(ap, unsigned long);

	return privwrap.libc.symbols._libc_prctl.f(option, arg2, arg3, arg4, arg5);
}
#endif /* HAVE_PRCTL */

#ifdef HAVE_SETRLIMIT
DO_NOT_SANITIZE_ADDRESS_ATTRIBUTE
static int libc_setrlimit(int resource,
			  const struct rlimit *rlp)
{
	pwrap_bind_symbol_all();

	return privwrap.libc.symbols._libc_setrlimit.f(resource, rlp);
}
#endif /* HAVE_SETRLIMIT */

/*********************************************************
 * PRIVWRAP HELPER FUNCTIONS
 *********************************************************/

DO_NOT_SANITIZE_ADDRESS_ATTRIBUTE
bool priv_wrapper_enabled(void)
{
	const char *env;

	privwrap.enabled = false;

	env = getenv("PRIV_WRAPPER");
	if (env != NULL && env[0] == '1') {
		privwrap.enabled = true;
	}

	return privwrap.enabled;
}

/****************************
 * CHROOT
 ***************************/

/* Allows also a normal (non root) user to call chroot() */
static int privwrap_chroot(const char *path)
{
	bool privwrap_on = priv_wrapper_enabled();
	const char *env = getenv("PRIV_WRAPPER_CHROOT_DISABLE");

	if (privwrap_on && env != NULL && env[0] == '1') {
		return 0;
	}

	return libc_chroot(path);
}

int chroot(const char *path)
{
	return privwrap_chroot(path);
}

/****************************
 * PLEDGE
 ***************************/

#ifdef HAVE_PLEDGE
static int libc_pledge(const char *promises, const char *execpromises)
{
	pwrap_bind_symbol_all();

	return privwrap.libc.symbols._libc_pledge.f(promises, execpromises);
}

static int privwrap_pledge(const char *promises, const char *execpromises)
{
	bool privwrap_on = priv_wrapper_enabled();
	const char *env = getenv("PRIV_WRAPPER_PLEDGE_DISABLE");

	if (privwrap_on && env != NULL && env[0] == '1') {
		return 0;
	}

	return libc_pledge(promises, execpromises);
}

int pledge(const char *promises, const char *execpromises)
{
	return privwrap_pledge(promises, execpromises);
}
#endif /* HAVE_PLEDGE */

/****************************
 * PRCTL
 ***************************/

#ifdef HAVE_PRCTL
static bool privwrap_prctl_disabled(const char *option)
{
	const char *env = getenv("PRIV_WRAPPER_PRCTL_DISABLE");
	const char *p;
	int cmp;

	if (env == NULL || strlen(env) == 0) {
		return false;
	}

	cmp = strcasecmp(env, "ALL");
	if (cmp == 0) {
		return true;
	}

	p = strstr(env, option);
	if (p == NULL) {
		return false;
	}

	return true;
}


/*
 * We provide fine grained conrol for some (those seen in openssh) prctl options
 * - PR_SET_SECCOMP  ... works only in kernel with CONFIG_SECCOMP_FILTER
 * - PR_SET_NO_NEW_PRIVS ... once set, cannot be reset
 * - PR_SET_DUMPABLE  ... always succeeds if called with valid arguments
 */
static int privwrap_vprctl(int option, va_list va)
{
	bool privwrap_on = priv_wrapper_enabled();
	bool option_disabled = false;
	va_list ap;
	int rc = 0;

	if (!privwrap_on) {
		return libc_vprctl(option, va);
	}

	va_copy(ap, va);

	/* PRIV_WRAPPER_PRCTL_DISABLE=PR_SET_SECCOMP|PR_SET_NO_NEW_PRIVS */
	switch(option) {
#ifdef PR_SET_SECCOMP
	case PR_SET_SECCOMP:
		option_disabled = privwrap_prctl_disabled("PR_SET_SECCOMP");
		break;
#endif
#ifdef PR_SET_NO_NEW_PRIVS
	case PR_SET_NO_NEW_PRIVS:
		option_disabled = privwrap_prctl_disabled("PR_SET_NO_NEW_PRIVS");
		break;
#endif
#ifdef PR_SET_DUMPABLE
	case PR_SET_DUMPABLE:
		option_disabled = privwrap_prctl_disabled("PR_SET_DUMPABLE");
		break;
#endif
	}

	if (!option_disabled) {
		rc = libc_vprctl(option, va);
	}

	va_end(ap);

	return rc;
}

int prctl(int option, ...)
{
	va_list va;
	int rc;

	va_start(va, option);

	rc = privwrap_vprctl(option, va);

	va_end(va);

	return rc;
}
#endif /* HAVE_PRCTL */

#ifdef HAVE_SETRLIMIT
static bool privwrap_setrlimit_disabled(const char *resource)
{
	const char *env = getenv("PRIV_WRAPPER_SETRLIMIT_DISABLE");
	const char *p = NULL;
	int cmp;

	if (env == NULL || strlen(env) == 0) {
		return false;
	}

	cmp = strcasecmp(env, "ALL");
	if (cmp == 0) {
		return true;
	}

	p = strstr(env, resource);
	if (p == NULL) {
		return false;
	}

	return true;
}

DO_NOT_SANITIZE_ADDRESS_ATTRIBUTE
static int privwrap_setrlimit(int resource, const struct rlimit *rlp)
{
	bool setrlimit_disabled = false;
	int rc = 0;

	if (!priv_wrapper_enabled()) {
		return libc_setrlimit(resource, rlp);
	}

	switch(resource) {
#ifdef RLIMIT_CPU
	case RLIMIT_CPU:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_CPU");
		break;
#endif /* RLIMIT_CPU */
#ifdef RLIMIT_FSIZE
	case RLIMIT_FSIZE:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_FSIZE");
		break;
#endif /* RLIMIT_FSIZE */
#ifdef RLIMIT_DATA
	case RLIMIT_DATA:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_DATA");
		break;
#endif /* RLIMIT_DATA */
#ifdef RLIMIT_STACK
	case RLIMIT_STACK:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_STACK");
		break;
#endif /* RLIMIT_STACK */
#ifdef RLIMIT_CORE
	case RLIMIT_CORE:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_CORE");
		break;
#endif /* RLIMIT_CORE */
#ifdef RLIMIT_RSS
	case RLIMIT_RSS:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_RSS");
		break;
#endif /* RLIMIT_RSS */
#ifdef RLIMIT_NOFILE
	case RLIMIT_NOFILE:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_NOFILE");
		break;
#endif /* RLIMIT_NOFILE */
#ifdef RLIMIT_AS
	case RLIMIT_AS:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_AS");
		break;
#endif /* RLIMIT_AS */
#ifdef RLIMIT_NPROC
	case RLIMIT_NPROC:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_NPROC");
		break;
#endif /* RLIMIT_NPROC */
#ifdef RLIMIT_MEMLOCK
	case RLIMIT_MEMLOCK:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_MEMLOCK");
		break;
#endif /* RLIMIT_MEMLOCK */
#ifdef RLIMIT_LOCKS
	case RLIMIT_LOCKS:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_LOCKS");
		break;
#endif /* RLIMIT_LOCKS */
#ifdef RLIMIT_SIGPENDING
	case RLIMIT_SIGPENDING:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_SIGPENDING");
		break;
#endif /* RLIMIT_SIGPENDING */
#ifdef RLIMIT_MSGQUEUE
	case RLIMIT_MSGQUEUE:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_MSGQUEUE");
		break;
#endif /* RLIMIT_MSGQUEUE */
#ifdef RLIMIT_NICE
	case RLIMIT_NICE:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_NICE");
		break;
#endif /* RLIMIT_NICE */
#ifdef RLIMIT_RTPRIO
	case RLIMIT_RTPRIO:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_RTPRIO");
		break;
#endif /* RLIMIT_RTPRIO */
#ifdef RLIMIT_RTTIME
	case RLIMIT_RTTIME:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_RTTIME");
		break;
#endif /* RLIMIT_RTTIME */
#ifdef RLIMIT_NLIMMITS
	case RLIMIT_NLIMITS:
		setrlimit_disabled = privwrap_setrlimit_disabled("RLIMIT_NLIMMITS");
		break;
#endif /* RLIMIT_NLIMMITS */
	}

	if (!setrlimit_disabled) {
		rc = libc_setrlimit(resource, rlp);
	}

	return rc;
}

DO_NOT_SANITIZE_ADDRESS_ATTRIBUTE
#ifdef HAVE_SETRLIMIT_RLIMIT_RESOURCE_T
int setrlimit(__rlimit_resource_t resource, const struct rlimit *rlp)
#else
int setrlimit(int resource, const struct rlimit *rlp)
#endif
{
	return privwrap_setrlimit(resource, rlp);
}
#endif /* HAVE_SETRLIMIT */

/****************************
 * CONSTRUCTOR
 ***************************/
void privwrap_constructor(void)
{
	PRIVWRAP_LOG(PRIVWRAP_LOG_WARN, "priv_wrapper loaded");
}

/****************************
 * DESTRUCTOR
 ***************************/

void privwrap_destructor(void)
{
	if (privwrap.libc.handle != NULL
#ifdef RTLD_NEXT
	    && privwrap.libc.handle != RTLD_NEXT
#endif
			) {
		dlclose(privwrap.libc.handle);
	}
}
