#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <signal.h>
#include <linux/unistd.h>
#include <asm/ldt.h>
#include <asm/ucontext.h>

unsigned char stack[65536];

void my_segv(int signum, siginfo_t *info, void *c)
{
  struct sigcontext *ctx = &((struct ucontext *)c)->uc_mcontext;
  printf("In sighandler: esp=%lx\n", ctx->esp);
  /* Skip HLT */
  ctx->eip++;
}

_syscall3(int, modify_ldt, int, func, void *, ptr, unsigned long, bytecount);

static int set_ldt_entry(int entry, unsigned long base, unsigned int limit,
	      int seg_32bit_flag, int contents, int read_only_flag,
	      int limit_in_pages_flag, int seg_not_present, int useable)
{
  struct modify_ldt_ldt_s ldt_info;
  ldt_info.entry_number = entry;
  ldt_info.base_addr = base;
  ldt_info.limit = limit;
  ldt_info.seg_32bit = seg_32bit_flag;
  ldt_info.contents = contents;
  ldt_info.read_exec_only = read_only_flag;
  ldt_info.limit_in_pages = limit_in_pages_flag;
  ldt_info.seg_not_present = seg_not_present;
  ldt_info.useable = useable;

  return modify_ldt(1, &ldt_info, sizeof(ldt_info));
}

int main(int argc, char *argv[])
{
  unsigned short _ss, new_ss;
  unsigned long _esp, new_esp;
  int is32 = 0;
  stack_t sig_stack;
  struct sigaction sa;

  if (argc > 1 && !strncmp(argv[1], "32", 3))
    is32 = 1;

  sig_stack.ss_sp = stack;
  sig_stack.ss_flags = 0;
  sig_stack.ss_size = sizeof(stack);
  if (sigaltstack(&sig_stack, NULL)) {
    perror("sigaltstack()");
    return 1;
  }

  sa.sa_sigaction = my_segv;
  sigemptyset(&sa.sa_mask);
  sa.sa_flags = SA_ONSTACK | SA_SIGINFO;
  sigaction(SIGSEGV, &sa, NULL);

  /* Get SS, ESP */
  asm volatile(
    "movw %%ss, %0\n"
    "movl %%esp, %1\n"
    :"=m"(_ss), "=m"(_esp)
  );
  printf("sp=%#lx, ss=%#hx\n", _esp, _ss);

  /* Force to LDT */
  new_ss = _ss | 4;
  /* Create the LDT entry */
  set_ldt_entry(new_ss >> 3, 0, 0xbffff, is32, MODIFY_LDT_CONTENTS_DATA, 0, 1, 0, 0);

  /* Do the trick... Switch stack and then switch context. */
  asm volatile(
    "movw %3, %%ss\n"		/* Load our LDT selector to SS */
    "hlt\n"			/* Force the context switch */
    "movw %1, %%ss\n"		/* Restore SS ASAP */
    "movl %%esp, %0\n"		/* Get new ESP */
    "movl %2, %%esp\n"		/* Restore ESP */
    :"=m"(new_esp)
    : "m"(_ss), "m"(_esp), "m"(new_ss)
  );
  printf("Now sp=%#lx, ss=%#hx\n", new_esp, new_ss);

  /* See what we've got... */
  if (new_esp > 0xc0000000) {
    printf("BUG!\n");
  } else if (new_esp != _esp) {
    printf("Esp changed, strange...\n");
  } else {
    printf("No bug here! What CPU is this?\n");
    if (!is32)
      system("cat /proc/cpuinfo");
  }
  return 0;
}