#define _GNU_SOURCE

#include <stdio.h>
#include <sys/prctl.h>
#include <asm/prctl.h>
#include <sys/ucontext.h>
#include <sys/syscall.h>
#include <unistd.h>
#include <errno.h>
#include <signal.h>
#include <string.h>
#include <err.h>

static volatile sig_atomic_t got_sigill = 0;

extern int arch_prctl(int, unsigned long);

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");
}

#ifdef __x86_64__
# define REG_IP REG_RIP
#else
# define REG_IP REG_EIP
#endif

static void sigill(int sig, siginfo_t *si, void *ctx_void)
{
	ucontext_t *ctx = (ucontext_t*)ctx_void;

	got_sigill = 1;
	ctx->uc_mcontext.gregs[REG_IP] += 2;
}

static int nerrs = 0;

static void test_syscall(long nr)
{
	extern const char syscall_rip[];
	unsigned long rcx = 1;
	unsigned long orig_rcx = rcx;
	unsigned short orig_ss, ss;
	unsigned long orig_flags, flags;
	long ret;
	int local_nerrs = 0;

#ifdef __x86_64__
	arch_prctl(ARCH_SET_GS, 500);  // Enable segment 13
	asm volatile ("mov %0,%%ss" : : "r" ((13 << 3) | 3));
#endif

	got_sigill = 0;

	asm volatile ("mov %%ss,%[orig_ss]\n\t"
		      "pushf\n\t"
		      "pop %[orig_flags]\n\t"
		      "syscall\n\t"
		      "syscall_rip:"
		      "mov %%ss,%[ss]\n\t"
		      "pushf\n\t"
		      "pop %[flags]"
		      : "+c" (rcx), [orig_ss] "=rm" (orig_ss), [ss] "=rm" (ss),
			"=a" (ret), [flags] "=rm" (flags),
			[orig_flags] "=rm" (orig_flags)
		      : "a" (nr) :
#ifdef __x86_64__
			"r11"
#endif
		);

	if (got_sigill) {
		printf("[OK]\tsyscall %lx: SIGILL\n", nr);
	} else {
		const char *outcome;
		if (rcx != (unsigned long)syscall_rip) {
			local_nerrs++;
			outcome = "[FAIL]";
		} else {
			outcome = "[OK]";
		}
		printf("%s\tsyscall %lx: RCX = %lX  RIP = %lX\n",
		       outcome, nr, rcx, (unsigned long)syscall_rip);

		if (ret != (long)-ENOSYS && nr != SYS_fork) {
			local_nerrs++;
			outcome = "[FAIL]";
		} else {
			outcome = "[OK]";
		}
		printf("%s\tsyscall %lx: AX = %lx\n",
		       outcome, nr, ret);

		printf("[NOTE]\tsyscall %lx: orig RCX = %lx  ss = %hx  orig_ss = %hx  flags = %lx  orig_flags = %lx\n",
		       nr, orig_rcx, ss, orig_ss, flags, orig_flags);
	}

	nerrs += local_nerrs;
}

static void test_int80(long nr)
{
	long ret;
	int local_nerrs = 0;

	got_sigill = 0;

	asm volatile ("int $0x80"
		      : "=a" (ret)
		      : "a" (nr)
		      :
		);

	if (got_sigill) {
		printf("[OK]\tint80 %lx: SIGILL\n", nr);
	} else {
		const char *outcome;

		if (ret != (long)-ENOSYS) {
			local_nerrs++;
			outcome = "[FAIL]";
		} else {
			outcome = "[OK]";
		}
		printf("%s\tint80 %lx: AX = %lx\n",
		       outcome, nr, ret);
	}

	nerrs += local_nerrs;
}

static void test_syscallfn(long nr)
{
	long ret;
	int local_nerrs = 0;

	ret = syscall(nr);

	const char *outcome;

	if (ret != -1 || errno != ENOSYS) {
		local_nerrs++;
		outcome = "[FAIL]";
	} else {
		outcome = "[OK]";
	}
	printf("%s\tsyscall(%lx): ret = %ld, errno = %d\n",
	       outcome, nr, ret, errno);

	nerrs += local_nerrs;
}

int main()
{
	sethandler(SIGILL, sigill, 0);

	test_int80(0xFFFF);
	test_int80(0x4000FFFF);

	test_syscall(0xFFFF);
	test_syscall(0x4000FFFF);

	test_syscallfn(0xFFFF);

	test_syscall(SYS_fork);

	return nerrs ? 1 : 0;
}