#include <err.h>
#include <stdio.h>
#include <stdint.h>
#include <signal.h>
#include <setjmp.h>
#include <string.h>
#include <xmmintrin.h>

#ifdef __linux__
#include <asm/prctl.h>
#include <sys/prctl.h>
extern int arch_prctl(int, unsigned long);
#endif

static void sethandler(int sig, void (*handler)(int, siginfo_t *, void *),
		       int flags)
{
	struct sigaction sa;
	memset(&sa, 0, sizeof(sa));
	sa.sa_sigaction = handler;
	sa.sa_flags = SA_SIGINFO | flags;
	sigemptyset(&sa.sa_mask);
	if (sigaction(sig, &sa, 0))
		err(1, "sigaction");
}

static void clearhandler(int sig)
{
	struct sigaction sa;
	memset(&sa, 0, sizeof(sa));
	sa.sa_handler = SIG_DFL;
	sigemptyset(&sa.sa_mask);
	if (sigaction(sig, &sa, 0))
		err(1, "sigaction");
}

static jmp_buf sigill_jmp;
static void sigill(int sig, siginfo_t *si, void *ctx_void)
{
	siglongjmp(sigill_jmp, 1);
}

/* returns GDT limit */
static uint16_t show_gdt_and_idt(void)
{
	struct {
		unsigned short limit;
		unsigned long base;
	} __attribute__((packed)) val;

	uint16_t ret;

	__asm__ ("sgdt %0" : "=m" (val));
	printf("GDT:    base = 0x%016lX limit = 0x%04hX\n", val.base, val.limit);

	ret = (val.limit + 1) / 8;

	__asm__ ("sidt %0" : "=m" (val));
	printf("IDT:    base = 0x%016lX limit = 0x%04hX\n", val.base, val.limit);

	return ret;
}

static void show_ldt(void)
{
	unsigned short ldt;
	asm ("sldt %0" : "=rm" (ldt));
	printf("LDT:    0x%04hX\n", ldt);
}

static void show_tr(void)
{
	unsigned short tr;
	asm ("str %0" : "=rm" (tr));
	printf("TR:     0x%04X\n", tr);
}

static void show_msw(void)
{
	static const struct msw_bit {
		int pos;
		const char *str;
	} bits[] = {
		{0, "PE"},
		{1, "MP"},
		{2, "EM"},
		{3, "TS"},	/* this is actually interesting */
		{4, "ET"},
		{5, "NE"},
	};

	unsigned short msw, remaining;
	asm ("smsww %0" : "=rm" (msw));
	printf("MSW:    0x%04X", msw);

	remaining = msw;
	for (int i = 0; i < sizeof(bits) / sizeof(bits[0]); i++) {
		if (msw & (1 << bits[i].pos)) {
			printf(" %s", bits[i].str);
			remaining &= ~(1 << bits[i].pos);
		}
	}
	if (remaining)
		printf(" unknown:0x%x", remaining);

	printf("\n");
}

static void show_flags(void)
{
	static char const * const flag_bits[] = {
		[0]  = "CF",
		[1]  = "FIXED",
		[2]  = "PF",
		[4]  = "AF",
		[6]  = "ZF",
		[7]  = "SF",
		[8]  = "TF",  /* Detects 'si 100000' in gdb */
		[9]  = "IF",
		[10] = "DF",
		[11] = "OF",
		[14] = "NT",
		[16] = "RF",  /* Invisible to pushf. */
		[17] = "VM",
		[18] = "AC",
		[19] = "VIF",
		[20] = "VIP",
		[21] = "ID",
	};

	unsigned long flags;
	asm ("sub $128, %%sp\n\t"
	     "xor %%ax,%%ax\n\t"  /* Make arithmetic flags be deterministic */
	     "add $1, %%ax\n\t"   /* Clear ZF and PF */
	     "pushf\n\t"
	     "pop %0\n\t"
	     "add $128, %%sp" : "=r" (flags) :
	     : "flags",
	       "eax");
	printf("FLAGS:  0x%016lX", flags);
	for (int i = 0; i < sizeof(flag_bits) / sizeof(flag_bits[0]); i++) {
		if (flag_bits[i] && (flags & (1 << i))) {
			printf(" %s", flag_bits[i]);
			flags &= ~(1 << i);
		}
	}
	printf(" IOPL=%lu", (flags >> 12) & 3);
	flags &= ~0x3000;
	if (flags)
		printf(" unknown:0x%lx", flags);
	printf("\n");
}

static void show_rdtscp(void)
{
	unsigned int cpu;

	sethandler(SIGILL, sigill, 0);
	if (sigsetjmp(sigill_jmp, 0)) {
		printf("RDTSCP: not supported\n");
	} else {
		__builtin_ia32_rdtscp(&cpu);
		printf("RDTSCP: cpu %d\n", cpu);
	}
	clearhandler(SIGILL);
}

static const char *user_seg_types[] = {
	"RO data",
	"RW data",
	"RO data, exp-down",
	"RW data, exp-down",
	"XO code",
	"XR code",
	"XO code, conforming",
	"XR code, confirming",
};

static void show_segment(uint16_t index, int ldt)
{
	uint32_t has_limit = 0, has_ar = 0, limit, ar;
	uint32_t selector = (index << 3) | (ldt << 2) | 3;

	asm ("lsl %[selector], %[limit]\n\t"
	     "jnz 1f\n\t"
	     "movl $1, %[has_limit]\n\t"
	     "1:"
	     : [limit] "=r" (limit), [has_limit] "+rm" (has_limit)
	     : [selector] "r" (selector));
	asm ("larl %[selector], %[ar]\n\t"
	     "jnz 1f\n\t"
	     "movl $1, %[has_ar]\n\t"
	     "1:"
	     : [ar] "=r" (ar), [has_ar] "+rm" (has_ar)
	     : [selector] "r" (selector));

	if (!has_limit && !has_ar)
		return;

	printf("%s entry %02hu: ", (ldt ? "LDT" : "GDT"), index);

	if (has_limit)
		printf(" limit = 0x%08X", limit);

	if (has_ar) {
#define ARBITS(low, high) ((ar >> low) & ((1 << (high - low + 1)) - 1))
#define ARSTR(bit, str) (ARBITS(bit,bit) ? " " str : "")
		printf(" access rights = 0x%08X\n"
		       "               type=%x (%s%s) DPL=%d%s%s%s%s%s%s",
		       ar, ARBITS(8,11), user_seg_types[ARBITS(9,11)],
		       (ARBITS(8,8) ? ", accessed" : ""),
		       ARBITS(13,14),
		       ARSTR(12, "S"), ARSTR(15, "P"),
		       ARSTR(20, "AVL"),  /* AVL is reserved for the OS */
		       ARSTR(21, "L"), ARSTR(22, "D/B"), ARSTR(23, "G"));
#undef ARSTR
#undef ARBITS
	}

	printf("\n");
}

int main()
{
#if defined(__linux__) && defined(__x86_64__)
	// Optional: gives some visibility into how TLS works.
	arch_prctl(ARCH_SET_GS, 500);
#endif

	uint16_t n_gdt_entries = show_gdt_and_idt();
	show_ldt();
	show_tr();
	show_msw();
	show_flags();
	show_rdtscp();
	for (int i = 0; i <= n_gdt_entries; i++)
		show_segment(i, 0);
	for (int i = 0; i < 10; i++)
		show_segment(i, 1);
}