/*
 * Copyright (c) 2014-2015 Andy Lutomirski
 * GPL v2
 */

#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <time.h>
#include <err.h>
#include <asm/ldt.h>
#include <sys/syscall.h>

static unsigned short GDT3(int idx)
{
	return (idx << 3) | 3;
}

static unsigned short LDT3(int idx)
{
	return (idx << 3) | 7;
}

static int create_tls(int idx, unsigned int base)
{
	struct user_desc desc = {
		.entry_number    = idx,
		.base_addr       = base,
		.limit           = 0xfffff,
		.seg_32bit       = 1,
		.contents        = 0, /* Data, grow-up */
		.read_exec_only  = 0,
		.limit_in_pages  = 1,
		.seg_not_present = 0,
		.useable         = 0,
	};

	if (syscall(SYS_set_thread_area, &desc) != 0)
		err(1, "set_thread_area");

	return desc.entry_number;
}

static void do_gs_test(int idx, int np)
{
	unsigned short orig_gs, new_gs;
	int ax;

	/*
	 * Install a valid LDT entry (set_thread_area's hardening measures
	 * defeat the #NP part of this test).
	 */
	struct user_desc desc;
	memset(&desc, 0, sizeof(desc));
	desc.entry_number = 0;
	desc.limit = 0xffff;
	desc.seg_32bit = 1;
	desc.contents = 0; /* Data, grow-up */
	desc.read_exec_only  = 0;
	desc.limit_in_pages  = 1;
	desc.seg_not_present = 0;
	if (syscall(SYS_modify_ldt, 1, &desc, sizeof(desc)) != 0)
		err(1, "modify_ldt");

	if (np) {
		/* valid but not present */
		desc.seg_not_present = 1;
	} else {
		/* "empty" user_desc -> destroy the segment. */
		desc.limit = 0;
		desc.seg_32bit = 0;
		desc.limit_in_pages = 0;
		desc.read_exec_only = 1;
		desc.seg_not_present = 1;
	}

	struct timespec req = {
		.tv_sec = 0,
		.tv_nsec = 100000,
	};

	printf("[RUN]\tGS %s\n", (np ? "not present" : "deleted"));
	asm volatile ("mov %%gs,%0" : "=rm" (orig_gs));
	asm volatile ("mov %0,%%gs" : : "rm" (LDT3(0)));

	asm volatile ("int $0x80"
		      : "=a" (ax)
		      : "a" (SYS_modify_ldt), "b" (1), "c" (&desc),
			"d" (sizeof(desc)));
	if (ax != 0)
		err(1, "set_thread_area");

	/*
	 * Force rescheduling.  On 32-bit kernels, fast syscalls
	 * destroy DS and ES, so force int 80.
	 */
	asm volatile ("int $0x80"
		      : "=a" (ax)
		      : "a" (SYS_nanosleep), "b" (&req),
			"c" (0));

	asm volatile ("mov %%gs,%0" : "=rm" (new_gs));
	asm volatile ("mov %0,%%gs" : : "rm" (orig_gs));
	printf("[OK]\tGS changed from %x to %x\n", (unsigned)GDT3(idx),
	       (unsigned)new_gs);
}

int main()
{
	int ret;
	int idx = create_tls(-1, 0);
	printf("Allocated GDT index %d\n", idx);

	unsigned short orig_es;
	asm volatile ("mov %%es,%0" : "=rm" (orig_es));

	int errors = 0;
	int total = 1000;
	for (int i = 0; i < total; i++) {
		struct timespec req = {
			.tv_sec = 0,
			.tv_nsec = 100000,
		};
		int ret;

		asm volatile ("mov %0,%%es" : : "rm" (GDT3(idx)));

		/*
		 * Force rescheduling.  On 32-bit kernels, fast syscalls
		 * destroy DS and ES, so force int 80.
		 */
		asm volatile ("int $0x80"
			      : "=a" (ret)
			      : "a" (SYS_nanosleep), "b" (&req),
				"c" (0));

		unsigned short es;
		asm volatile ("mov %%es,%0" : "=rm" (es));
		asm volatile ("mov %0,%%es" : : "rm" (orig_es));
		if (es != GDT3(idx)) {
			if (errors == 0)
				printf("[FAIL]\tES changed from 0x%hx to 0x%hx\n",
				       GDT3(idx), es);
			errors++;
		}
	}

	if (errors) {
		printf("[FAIL]\tES was corrupted %d/%d times\n", errors, total);
		ret = 1;
	} else {
		printf("[OK]\tES was preserved\n");
		ret = 0;
	}

	do_gs_test(idx, 1);
	create_tls(idx, 0);
	do_gs_test(idx, 0);
	return ret;
}