/* libstackless prototype
 * Copyright (C) 2007  Michael FIG
 *
 * Michael FIG <michael@fig.org>, 2007-03-19
 */
#define _GNU_SOURCE 1
#include <signal.h>

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#include <string.h>
#include <errno.h>

#include <sys/mman.h>
#include <limits.h> /* PAGESIZE */
#ifndef PAGESIZE
# define PAGESIZE 4096
#endif

#define STACKLESS_MAX_SITES 256

/* This is a user-supplied function to start the application. */
int stackless_main(int argc, char **argv);

/* Put breakpoints on the return instructions, and intercept all the
   call or jump instructions in FUNC.  Return the maximum size of
   FUNC's activation record. */
size_t stackless_infiltrate(void *func);

/* Mess with page protections to allow ADDR to be modified. */
void stackless_reprotect(void *addr);

/* This function does nothing more than provide a call site for
   bootstrapping the activation trapper. */
int stackless_premain(int argc, char **argv);

static char *unprotected_page;
static char *site_address[STACKLESS_MAX_SITES];
static char site_contents[STACKLESS_MAX_SITES];
static int num_sites = 0;


/* Intercept a call to a child function, or a return to a parent
   function. */
void
stackless_trap_activation(int signum, siginfo_t *si, void *uctx)
{
  int i;
  ucontext_t *uc = uctx;
  int *eip = &uc->uc_mcontext.gregs[REG_EIP];
  printf("from breakpoint at 0x%08x\n", *eip);

  /* Restore the current sites to their defaults. */
  for (i = 0; i < num_sites; i ++) {
    stackless_reprotect(site_address[i]);
    *site_address[i] = site_contents[i];
  }

  /* Try to protect ourselves the most again. */
  stackless_reprotect(NULL);

  /* Resume where we left off. */
  (*eip) --;
}


/* Put breakpoints on all the call and return instructions in FUNC.
   Return the maximum size of FUNC's activation record. */
size_t
stackless_infiltrate(void *func)
{
  /* Do not infiltrate any return instructions for our premain. */
  int do_rets = (func != stackless_premain);

  /* FIXME: Use the disassembler. */
  site_address[num_sites] = func;

  site_contents[num_sites] = *site_address[num_sites];
  stackless_reprotect(site_address[num_sites]);
  *site_address[num_sites] = 0xcc; /* Instruction for debug trap. */

  /* Advance to the next site. */
  num_sites ++;

  /* Try to reprotect our page. */
  stackless_reprotect(NULL);

  /* FIXME: Return the number of bytes we need. */
  return 0;
}


/* Mess with page protections to allow ADDR to be modified. */
void
stackless_reprotect(void *addr)
{
  char *to_unprotect = ((char *)addr) - ((long) addr) % PAGESIZE;
  if (unprotected_page == to_unprotect) {
    return;
  }

  if (unprotected_page &&
      mprotect(unprotected_page, PAGESIZE, PROT_READ|PROT_EXEC)) {
    fprintf(stderr, "cannot reprotect %p: %s\n",
	    unprotected_page, strerror(errno));
  }
  if (to_unprotect &&
      mprotect(to_unprotect, PAGESIZE, PROT_READ|PROT_WRITE|PROT_EXEC)) {
    fprintf(stderr, "cannot unprotect %p: %s\n",
	    to_unprotect, strerror(errno));
  }
  else {
    unprotected_page = to_unprotect;
  }
}

/* This function does nothing more than provide a call site for
   bootstrapping the activation trapper. */
int
stackless_premain(int argc, char **argv)
{
  return stackless_main(argc, argv);
}


/* Stackless takes control over the main function; you must define
   your application's main in a function called "stackless_main". */
int
main(int argc, char **argv)
{
  /* The exit code from the user's stackless_main. */
  int ret;
  struct sigaction act, oldact;
  stack_t altstack;

  /* Use a temporary stack for our signal handler. */
  memset(&altstack, 0, sizeof(altstack));
  altstack.ss_size = SIGSTKSZ;
  altstack.ss_sp = malloc(altstack.ss_size);
  if (sigaltstack(&altstack, NULL)) {
    fprintf(stderr, "%s: cannot set up Stackless signal stack: %s\n",
	    argv[0], strerror(errno));
    exit(1);
  }

  /* Install our call/return trapper. */
  memset(&act, 0, sizeof(act));
  act.sa_sigaction = stackless_trap_activation;
  act.sa_flags = SA_SIGINFO | SA_ONSTACK;
  if (sigaction(SIGTRAP, &act, &oldact)) {
    fprintf(stderr, "%s: cannot install Stackless signal handler: %s\n",
	    argv[0], strerror(errno));
    exit(1);
  }

  /* Select stackless_premain as the first callsite. */
  stackless_infiltrate(stackless_premain);

  /* Run the user's main. */
  ret = stackless_premain(argc, argv);

  /* Restore the state of the call trapper. */
  sigaction(SIGTRAP, &oldact, NULL);
  return ret;
}


int
stackless_main(int argc, char **argv)
{
  // FIXME: Do something useful.
  printf("Hello, world!\n");
  return 0;
}

/*
 * Local variables:
 * compile-command:"gcc -Wall -g -O2 t.c -o t && ./t"
 * End:
 */