/**
 * Demonstrates the use of a signal handler to create data lazily
 *
 *
 * Copyright 2014 Florian Philipp
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#define _GNU_SOURCE
/* using tdestroy */

#include <stddef.h>
/* using size_t */
#include <stdbool.h>
/* using bool */
#include <stdio.h>
/* using printf, fprintf, perror, fputs */
#include <stdlib.h>
/* using exit */
#include <assert.h>
/* using assert */
#include <errno.h>
/* using errno */
#include <unistd.h>
/* using sysconf */
#include <sys/mman.h>
/* using mprotect, mmap, munmap */
#include <search.h>
/* using tsearch, tfind, tdestroy */
#include <signal.h>
/* using sigaction */


/**
 * Segmentation fault handler
 *
 * Uses a dedicated memory area that is filled with data on demand.
 * A singleton in production code, but not necessarily in testing code
 *
 * TODO: Make thread-safe
 *
 * BUG: Not async-signal-safe
 */
struct _SegvHandler
{
  /**
   * pointer to the start of the dedicated memory area or NULL
   */
  void* base_addr;
  /**
   * pointer past the end of the dedicated memory area or NULL
   */
  void* mapped_end;
  /**
   * search.h tree of addresses of pages with read permissions
   */
  void* mapped_ro;
  /**
   * search.h tree of addresses of pages with write permissions
   */
  void* mapped_rw;
  int async_err;
};


/**
 * Signal handler singleton
 *
 * Default initialization with 0 is valid for all attributes
 */
static struct _SegvHandler _segv_global_self;


/**
 * Comparison of memory addresses. Used with search.h
 */
static int _segv_page_cmp(const void* first, const void* second)
{
  if(first < second)
    return -1;
  else if(first == second)
    return 0;
  else
    return 1;
}

/**
 * Sets async_err from errno unless an error is already set
 */
static void _segv_store_errno(struct _SegvHandler* self)
{
  if(! self->async_err)
    self->async_err = errno;
}

/**
 * Returns bytes per memory page
 */
static size_t _segv_pagelen()
{
  return sysconf(_SC_PAGESIZE);
}

/**
 * Fills page with meaningful data
 *
 * TODO: Stub
 */
static void _segv_populate_page(void* page)
{
  unsigned* typed = page;
  unsigned* end = page + _segv_pagelen() / sizeof(*typed);
  for(; typed != end; ++typed)
    *typed = 0xDEADBEEF;
}

/**
 * Flushes changed page to shared storage or whatever
 *
 * TODO: Stub
 */
static void _segv_commit_page(struct _SegvHandler* self, void* page)
{
  char* byte_addr = page;
  char* byte_base = self->base_addr;
  size_t offset = byte_addr - byte_base;
  printf("Range [%zu, %zu) changed\n", offset, offset + _segv_pagelen());
}

/**
 * Returns page address of segfault
 */
static void* _segv_get_page(const siginfo_t* siginfo)
{
  size_t addr = (size_t) siginfo->si_addr;
  addr &= ~(_segv_pagelen() - 1);
  return (void*) addr;
}

/**
 * Invokes handler for segmentation faults that cannot be handled otherwise
 *
 * TODO: Replace with original segfault handler
 */
static void _segv_real_segfault(const siginfo_t* siginfo)
{
  fprintf(stderr, "SEGMENTATION FAULT %p\n", siginfo->si_addr);
  exit(EXIT_FAILURE);
}

/**
 * Permits write access to page. Marks page as dirty
 *
 * Precondition: page is removed from self->mapped_ro
 * Postcondition on error: Page is still write-protected
 */
static int _segv_map_rw(struct _SegvHandler* self, void* page)
{
  size_t pagelen = _segv_pagelen();
  if(mprotect(page, pagelen, PROT_READ | PROT_WRITE))
    goto err_rtrn;
  if(tsearch(page, &self->mapped_rw, _segv_page_cmp) == NULL)
    goto err_nomem;
  return 0;
 err_nomem:
  mprotect(page, pagelen, PROT_READ);
  errno = ENOMEM;
 err_rtrn:
  return -1;
}

/**
 * Populates page and permits read access
 */
static int _segv_map_ro(struct _SegvHandler* self, void* page)
{
  size_t pagelen = _segv_pagelen();
  if(mprotect(page, pagelen, PROT_READ | PROT_WRITE))
    goto err_rtrn;
  _segv_populate_page(page);
  if(mprotect(page, pagelen, PROT_READ))
    goto err_rtrn;
  if(! tsearch(page, &self->mapped_ro, _segv_page_cmp))
    goto err_nomem;
  return 0;
 err_nomem:
  mprotect(page, pagelen, PROT_NONE);
  errno = ENOMEM;
 err_rtrn:
  return -1;
}

/**
 * Returns true if this segfault can be avoided
 */
static bool _segv_is_magic_segfault(const struct _SegvHandler* self,
				    const siginfo_t* siginfo)
{
  void* addr = siginfo->si_addr;
  if(siginfo->si_code != SEGV_ACCERR)
    return false;
  if(addr < self->base_addr || addr >= self->mapped_end)
    return false;
  void* page = _segv_get_page(siginfo);
  if(tfind(page, &self->mapped_rw, _segv_page_cmp))
    return false;
  return true;
}

/**
 * Signal handler compatible with sigaction
 */
static void _segv_signal_handler(int signum, siginfo_t* siginfo,
				 void* ucontext)
{
  assert(signum == SIGSEGV);
  struct _SegvHandler* self = &_segv_global_self;
  void* page = _segv_get_page(siginfo);
  if(! _segv_is_magic_segfault(self, siginfo))
    _segv_real_segfault(siginfo);
  if(tdelete(page, &self->mapped_ro, _segv_page_cmp)) {
    if(_segv_map_rw(self, page))
      goto err;
  }
  else {
    if(_segv_map_ro(self, page))
      goto err;
  }
  return;
 err:
  perror("segfault handler");
  exit(EXIT_FAILURE);
}

/**
 * Installs the global signal handler
 */
static int _segv_install()
{
  struct sigaction action = {
    .sa_sigaction = _segv_signal_handler,
    .sa_mask = {{0}},
    .sa_flags = SA_SIGINFO,
  };
  return sigaction(SIGSEGV, &action, NULL);
}

/**
 * Allocates the dedicated memory area
 */
static int _segv_init_mapping(struct _SegvHandler* self)
{
  int err;
  if(self->base_addr)
    goto err_double;
  size_t mapping_len = _segv_pagelen() * 16;
  void* mapped = mmap(NULL, mapping_len, PROT_READ | PROT_WRITE,
		      MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
  if(mapped == MAP_FAILED)
    goto err_rtrn;
  if(mprotect(mapped, mapping_len, PROT_NONE))
    goto err_unmap;
  self->base_addr = mapped;
  self->mapped_end = ((char*) mapped) + mapping_len;
  return 0;
 err_unmap:
  err = errno;
  munmap(mapped, mapping_len);
  errno = err;
  goto err_rtrn;
 err_double:
  fputs("segfault handler double initialization\n", stderr);
  errno = EINVAL;
 err_rtrn:
  return -1;
}

/**
 * Commits a dirty pages and marks it as clean
 *
 * Callback for tdestroy
 */
static void _segv_commit_action(void* page)
{
  struct _SegvHandler* self = &_segv_global_self;
  _segv_commit_page(self, page);
  size_t pagelen = _segv_pagelen();
  if(mprotect(page, pagelen, PROT_READ))
    goto err_store;
  if(! tsearch(page, &self->mapped_ro, _segv_page_cmp))
    goto err_mem;
  return;
 err_mem:
  mprotect(page, pagelen, PROT_NONE);
  errno = ENOMEM;
 err_store:
  _segv_store_errno(self);
}

/**
 * Discards a dirty page. Will re-populate it on demand
 *
 * Callback for tdestroy
 */
static void _segv_discard_action(void* page)
{
  if(mprotect(page, _segv_pagelen(), PROT_NONE))
    _segv_store_errno(&_segv_global_self);
}

/**
 * Removes all entries from the global mapped_rw set
 *
 * \param a callback called on every page
 * \return async_err
 */
static int _segv_clear_rw(void (*action)(void*))
{
  struct _SegvHandler* self = &_segv_global_self;
  self->async_err = 0;
  tdestroy(self->mapped_rw, action);
  self->mapped_rw = NULL;
  if(self->async_err)
    errno = self->async_err;
  return self->async_err;
}

/**
 * Installs and initializes the segmentation fault handler
 *
 * \return 0 on success, -1 otherwise. Sets errno
 */
int sigsegv_install()
{
  if(_segv_init_mapping(&_segv_global_self))
    return -1;
  return _segv_install();
}

/**
 * Commits all changes
 *
 * \return 0 on success, -1 otherwise. Sets errno
 */
int sigsegv_commit()
{
  return _segv_clear_rw(_segv_commit_action);
}

/**
 * Discards all changes
 *
 * \return 0 on success, -1 otherwise. Sets errno
 */
int sigsegv_discard()
{
  return _segv_clear_rw(_segv_discard_action);
}

/**
 * Returns starting address of the dedicated memory area
 */
void* sigsegv_baseptr()
{
  const struct _SegvHandler* self = &_segv_global_self;
  return self->base_addr;
}


/**
 * Some simple testing code
 *
 * Observe it with strace
 */
int main()
{
  if(sigsegv_install())
    goto err;
  unsigned* base = sigsegv_baseptr();
  printf("Accessing RO %p = 0x%x\n", base + 3, base[3]);
  base[3] = 0;
  puts("Committing");
  if(sigsegv_commit())
    goto err;
  size_t otherpage = _segv_pagelen() / sizeof(unsigned) + 3;
  puts("Making direct RW access");
  base[otherpage] = 0;
  puts("Discarding");
  if(sigsegv_discard())
    goto err;
  printf("Accessing RO %p = 0x%x\n", base + otherpage, base[otherpage]);
  return EXIT_SUCCESS;
 err:
  perror("sigsegv");
  return EXIT_FAILURE;
}