#if SEP_STACK_SEG

#define _GNU_SOURCE
#include "atomic.h"
#include "libc.h"
#include "syscall.h"
#include <ldt.h>
#include <malloc.h>
#include <stdint.h>
#include <string.h>

extern uintptr_t __stack_base;

static int modify_ldt(int func, void *ptr, unsigned long bytecount) {
	return syscall(SYS_modify_ldt, func, ptr, bytecount);
}

static void update_ldt(int idx, unsigned long addr, size_t len) {
	struct user_desc stack_desc;
	stack_desc.entry_number = idx;
	stack_desc.base_addr = (unsigned long)addr;
	stack_desc.contents = 0; /* data */
	stack_desc.limit = (int)((len - 1) >> 12);
	stack_desc.limit_in_pages = 1;
	stack_desc.read_exec_only = 0;
	stack_desc.seg_not_present = 0;
	stack_desc.seg_32bit = 1;
	stack_desc.useable = 1;

	if (modify_ldt(1, &stack_desc, sizeof(stack_desc)) == -1)
		a_crash();
}

static int find_available_ldt_slot(void) {
	static const int SEG_DESC_P_BIT_OFF = 47;

	uint64_t ldt[1];

	/* read the current LDT */
	int ldt_len = modify_ldt(0, ldt, sizeof(ldt));
	if (ldt_len == -1)
		a_crash();

	if (ldt_len == 0)
		/* LDT is currently empty */
		return 0;

	for (int i = 0; i < (ldt_len/sizeof(ldt[0])); i++)
		if (((ldt[i] >> SEG_DESC_P_BIT_OFF) & 1) == 0)
			return i;

	/* crash if no available slot was found */
	a_crash();

	/* unreachable; just for suppressing compiler warning */
	return 0;
}

#define SEG_SEL_LDT 4
#define SEG_SEL_CPL3 3
#define SEG_SEL_TO_IDX(sel) ((sel) >> 3)
#define SEG_IDX_TO_LDT_SEL(idx) ((idx) | SEG_SEL_LDT | SEG_SEL_CPL3)

__attribute__((__visibility__("hidden")))
void __restrict_segments(void)
{
	uintptr_t limit, stack_base = __stack_base;
	int data_seg_sel, data_ldt_sel;

	__asm__ __volatile__ ("mov %%ds, %0" : "=r"(data_seg_sel));

	if ((data_seg_sel & SEG_SEL_LDT) == SEG_SEL_LDT) {
		data_ldt_sel = data_seg_sel;

		/* Read the current limit from the segment register rather than
		 * relying on __stack_base, since __stack_base is in the
		 * default data segment and could potentially be subject to
		 * memory corruption. */
		__asm__ __volatile__ ("lsl %1, %0" : "=r"(limit) : "m"(data_seg_sel));

		if (limit < stack_base)
			return;
	} else
		data_ldt_sel = SEG_IDX_TO_LDT_SEL(find_available_ldt_slot());

	update_ldt(SEG_SEL_TO_IDX(data_ldt_sel), 0, stack_base);

	/* Reload the DS and ES segment registers from the new or updated LDT
	 * entry. */
	__asm__ __volatile__ (
	  "mov %0, %%ds\n\t"
	  "mov %0, %%es\n\t"
	  ::
	  "r"(data_ldt_sel)
	);
}

extern char **__environ;

/* Programs and much of the libc code expect to be able to access the arguments,
 * environment, and auxv in DS, but they are initially located on the stack.  This
 * function moves them to the heap. This uses strdup to copy data from the stack,
 * so it must run before segment limits are restricted.
 */
__attribute__((__visibility__("hidden")))
void __safestack_dup_argv_env_aux(int argc, char ***argvp, char ***envpp)
{
	char **argv = *argvp;
	char **envp = *envpp;
	char **environ_end = envp;
	size_t *auxv, *auxv_end;
	char **new_argv = 0;

	while (*environ_end) environ_end++;

	auxv_end = (size_t *)environ_end + 1;
	while (*auxv_end) auxv_end++;
	auxv_end++;

	new_argv = malloc((uintptr_t)auxv_end - (uintptr_t)argvp);
	if (!new_argv)
		a_crash();

	*new_argv = (char *)argc;
	new_argv++;

	*argvp = new_argv;

	for (int i = 0; i < argc; i++)
		new_argv[i] = strdup(argv[i]);
	new_argv += argc;
	*new_argv = NULL;
	new_argv++;

	*envpp = __environ = new_argv;
	while (envp != environ_end) {
		*new_argv = strdup(*envp);
		envp++;
		new_argv++;
	}
	*new_argv = NULL;
	envp++;
	new_argv++;

	libc.auxv = (size_t *)new_argv;
	memcpy(new_argv, envp, (uintptr_t)auxv_end - (uintptr_t)envp);

	__restrict_segments();
}

#endif /*SEP_STACK_SEG*/