ares/rust/ares_guard/c-src/guard.c

#include <assert.h>
#include <errno.h>
#include <setjmp.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <unistd.h>

#include "guard.h"

#define GD_PAGE_BITS              14ULL
#define GD_PAGE_SIZE              (1ULL << GD_PAGE_BITS)    // 16 KB
#define GD_PAGE_MASK              (GD_PAGE_SIZE - 1)
#define GD_PAGE_ROUND_DOWN(foo)   (foo & (~GD_PAGE_MASK))

typedef struct GD_state GD_state;
struct GD_state {
  uintptr_t         guard_p;
  const uintptr_t  *stack_pp;
  const uintptr_t  *alloc_pp;
  GD_buflistnode   *buffer_list;
  struct sigaction  prev_sigsegv_sa;
  struct sigaction  prev_sigbus_sa;
};

static GD_state gd = { 
  .guard_p = 0,
  .stack_pp = NULL,
  .alloc_pp = NULL,
  .buffer_list = NULL,
  .prev_sigsegv_sa = { .sa_sigaction = NULL, .sa_flags = 0 },
  .prev_sigbus_sa = { .sa_sigaction = NULL, .sa_flags = 0 },
};

static guard_result
_protect_page(void *address, int prot)
{
  if (mprotect(address, GD_PAGE_SIZE, prot)) {
    fprintf(stderr, "guard: prot: mprotect error %d\r\n", errno);
    fprintf(stderr, "%s\r\n", strerror(errno));
    return guard_mprotect | errno;
  }

  return guard_success;
}

// Center the guard page.
static guard_result
_focus_guard()
{
  uintptr_t stack_p = *gd.stack_pp;
  uintptr_t alloc_p = *gd.alloc_pp;
  uintptr_t old_guard_p = gd.guard_p;
  uintptr_t new_guard_p;
  guard_result   err = guard_success;

  if (stack_p == 0 || alloc_p == 0) {
    fprintf(stderr, "guard: focus: stack or alloc pointer is null\r\n");
    return guard_null;
  } else if (stack_p == alloc_p) {
    return guard_oom;
  }

  // Compute new guard page.
  new_guard_p = GD_PAGE_ROUND_DOWN((stack_p + alloc_p) / 2);
  if (new_guard_p == old_guard_p) {
    return guard_oom;
  }

  // Mark new guard page.
  if ((err = _protect_page((void *)new_guard_p, PROT_NONE))) {
    fprintf(stderr, "guard: focus: mark error\r\n");
    return err;
  }

  // Update guard page tracker.
  gd.guard_p = new_guard_p;

  // Unmark the old guard page if there is one.
  if (old_guard_p) {
    if ((err = _protect_page((void *)old_guard_p, PROT_READ | PROT_WRITE))) {
      fprintf(stderr, "guard: focus: unmark error\r\n");
      return err;
    }
  }

  return guard_success;
}

static void
_signal_handler(int sig, siginfo_t *si, void *unused)
{
  uintptr_t sig_addr;
  guard_result err = guard_success;

  assert(gd.guard_p);

  if (sig != SIGSEGV && sig != SIGBUS) {
    fprintf(stderr, "guard: handler: invalid signal: %d\r\n", sig);
    assert(0);
  }

  sig_addr = (uintptr_t)si->si_addr;

  if (sig_addr >= gd.guard_p &&
      sig_addr <  gd.guard_p + GD_PAGE_SIZE)
  {
    err = _focus_guard();
    if (err) {
      siglongjmp(gd.buffer_list->buffer, err);
    }
  }
  else {
    switch (sig) {
      case SIGSEGV: {
        if (gd.prev_sigsegv_sa.sa_sigaction != NULL) {
          gd.prev_sigsegv_sa.sa_sigaction(sig, si, unused);
        } else if (gd.prev_sigsegv_sa.sa_handler != NULL) {
          gd.prev_sigsegv_sa.sa_handler(sig);
        } else {
          assert(0);
        }
        break;
      }
      case SIGBUS: {
        if (gd.prev_sigbus_sa.sa_sigaction != NULL) {
          gd.prev_sigbus_sa.sa_sigaction(sig, si, unused);
        } else if (gd.prev_sigbus_sa.sa_handler != NULL) {
          gd.prev_sigbus_sa.sa_handler(sig);
        } else {
          assert(0);
        }
      }
    }
  }
}

// Registers the same handler function for SIGSEGV and SIGBUS.
static guard_result
_register_handlers()
{
  struct sigaction sa;
  sa.sa_flags = SA_SIGINFO;
  sa.sa_sigaction = _signal_handler;

  if (sigaction(SIGSEGV, &sa, &gd.prev_sigsegv_sa)) {
    fprintf(stderr, "guard: register: sigaction error\r\n");
    fprintf(stderr, "%s\r\n", strerror(errno));
    return guard_sigaction | errno;
  }

  if (sigaction(SIGBUS, &sa, &gd.prev_sigbus_sa)) {
    fprintf(stderr, "guard: register: sigaction error\r\n");
    fprintf(stderr, "%s\r\n", strerror(errno));
    return guard_sigaction | errno;
  }

  return guard_success;
}

guard_result
guard(
  void *(*f)(void *),
  void *closure,
  const uintptr_t *const s_pp,
  const uintptr_t *const a_pp,
  void **ret
) {
  GD_buflistnode  *new_buffer;
  guard_result err = guard_success;
  guard_result td_err = guard_success;

  if (gd.guard_p == 0) {
    assert(gd.buffer_list == NULL);

    gd.stack_pp = s_pp;
    gd.alloc_pp = a_pp;

    // Initialize the guard page.
    if ((err = _focus_guard())) {
      fprintf(stderr, "guard: initial focus error\r\n");
      goto exit;
    }

    // Register guard page signal handler.
    if ((err = _register_handlers())) {
      fprintf(stderr, "guard: registration error\r\n");
      goto clean;
    }
  } else {
    assert(gd.buffer_list != NULL);
  }

  // Setup new longjmp buffer.
  new_buffer = (GD_buflistnode *)malloc(sizeof(GD_buflistnode));
  if (new_buffer == NULL) {
    fprintf(stderr, "guard: malloc error\r\n");
    fprintf(stderr, "%s\r\n", strerror(errno));
    err = guard_malloc | errno;
    goto skip;
  }
  new_buffer->next = gd.buffer_list;
  gd.buffer_list = new_buffer;

  // Run given closure.
  if (!(err = sigsetjmp(gd.buffer_list->buffer, 1))) {
    *ret = f(closure);
  }

  // Restore previous longjmp buffer.
  gd.buffer_list = gd.buffer_list->next;
  free((void *)new_buffer);

skip:
  if (gd.buffer_list == NULL) {
    if (sigaction(SIGSEGV, &gd.prev_sigsegv_sa, NULL)) {
      fprintf(stderr, "guard: error replacing sigsegv handler\r\n");
      fprintf(stderr, "%s\r\n", strerror(errno));
      td_err = guard_sigaction | errno;

      if (!err) {
        err = td_err;
      }
    }
    if (sigaction(SIGBUS, &gd.prev_sigbus_sa, NULL)) {
      fprintf(stderr, "guard: error replacing sigbus handler\r\n");
      fprintf(stderr, "%s\r\n", strerror(errno));
      td_err = guard_sigaction | errno;

      if (!err) {
        err = td_err;
      }
    }

clean:
    // Unmark guard page.
    assert(gd.guard_p != 0);
    td_err = _protect_page((void *)gd.guard_p, PROT_READ | PROT_WRITE);
    if (td_err) {
      fprintf(stderr, "guard: unmark error\r\n");
      fprintf(stderr, "%s\r\n", strerror(errno));
      if (!err) {
        err = td_err;
      }
    }
    gd.guard_p = 0;
  }

exit:
  return err;
}