/**
 * Demonstrates the use of a signal handler to create data lazily
 *
 * Compile with -std=gnu++11
 *
 *
 * 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.
 */

#include <set>
// using std::set
#include <algorithm>
// using std::fill
#include <exception>
// using std::terminate, std::exception
#include <stdexcept>
// using std::logic_error
#include <system_error>
// using std::system_error, std::system_category
#include <cerrno>
// using errno
#include <cassert>
// using assert
#include <cstdio>
// using std::fprintf
#include <cstddef>
// using std::size_t
#include <signal.h>
// using sigaction
#include <sys/mman.h>
// using mprotect, mmap
#include <unistd.h>
// using sysconf


/**
 * Private C++ implementation follows
 */
namespace {

  /**
   * Returns bytes per memory page
   */
  std::size_t pagelen()
  {
    return sysconf(_SC_PAGESIZE);
  }

  /**
   * 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
   */
  class SegvHandler
  {
    typedef std::set<void*>::iterator iterator;

    /**
     * pointer to the start of the dedicated memory area or nullptr
     */
    void* base_addr;
    /**
     * pointer past the end of the dedicated memory area or nullptr
     */
    void* mapped_end;
    /**
     * Addresses of pages with read or read/write access permissions
     */
    std::set<void*> mapped_ro, mapped_rw;

    /**
     * Fills page with meaningful data
     *
     * TODO: Stub
     */
    void populate_page(void* page)
    {
      unsigned* typed = static_cast<unsigned*>(page);
      unsigned* page_end = typed + pagelen() / sizeof(*typed);
      std::fill(typed, page_end, 0xDEADBEEF);
    }

    /**
     * Flushes changed page to shared storage or whatever
     *
     * TODO: Stub
     */
    void commit_page(void* page)
    {
      char* byte_addr = static_cast<char*>(page);
      char* byte_base = static_cast<char*>(this->base_addr);
      std::size_t offset = byte_addr - byte_base;
      std::printf("Range [%zu, %zu) changed\n", offset, offset + pagelen());
    }

    /**
     * Returns page address of segfault
     */
    static void* get_page(const siginfo_t* siginfo)
    {
      std::size_t addr = reinterpret_cast<std::size_t>(siginfo->si_addr);
      addr &= ~(pagelen() - 1);
      return reinterpret_cast<void*>(addr);
    }

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

    /**
     * Throws an std::system_error constructed from errno
     */
    static void throw_sys_err(const char* what)
    {
      throw std::system_error(errno, std::system_category(), what);
    }

    /**
     * Invokes mprotect. Wraps errors in std::system_error
     */
    static void change_protection(void* page, int permissions,
				  std::size_t length = pagelen())
    {
      if(mprotect(page, length, permissions))
	throw_sys_err("mprotect");
    }

    /**
     * Permits write access to page. Marks page as dirty
     *
     * Precondition: page is removed from this->mapped_ro
     */
    void map_rw(void* page)
    {
      change_protection(page, PROT_READ | PROT_WRITE);
      this->mapped_rw.insert(page);
    }

    /**
     * Populates page and permits read access
     */
    void map_ro(void* page)
    {
      change_protection(page, PROT_READ | PROT_WRITE);
      this->populate_page(page);
      change_protection(page, PROT_READ);
      this->mapped_ro.insert(page);
    }

    /**
     * Returns true if this segfault can be avoided
     */
    bool is_magic_segfault(const siginfo_t* siginfo) const
    {
      void* addr = siginfo->si_addr;
      return siginfo->si_code == SEGV_ACCERR
	&& addr >= this->base_addr && addr < this->mapped_end
	&& ! this->mapped_rw.count(get_page(siginfo));
    }

    /**
     * Signal handler compatible with sigaction
     */
    static void signal_handler(int signum, siginfo_t* siginfo, void* ucontext)
    {
      try {
	assert(signum == SIGSEGV);
	SegvHandler* self = global_self();
	void* page = get_page(siginfo);
	if(! self->is_magic_segfault(siginfo))
	  real_segfault(siginfo);
	if(self->mapped_ro.erase(page))
	  self->map_rw(page);
	else
	  self->map_ro(page);
      } catch(std::exception& err) {
	std::fprintf(stderr, "segfault handler: %s\n", err.what());
	std::terminate();
      }
    }

  public:
    /**
     * Initializes empty, unmapped SegvHandler
     */
    SegvHandler()
      : base_addr(nullptr),
	mapped_end(nullptr)
    {}

    /**
     * Returns singleton
     */
    static SegvHandler* global_self()
    {
      static SegvHandler* self = new SegvHandler();
      return self;
    }

    /**
     * Installs the global signal handler
     */
    static void install()
    {
      struct sigaction action;
      action.sa_sigaction = &SegvHandler::signal_handler;
      action.sa_mask = sigset_t();
      action.sa_flags = SA_SIGINFO;
      if(sigaction(SIGSEGV, &action, nullptr))
	throw_sys_err("sigaction");
    }

    /**
     * Allocates the dedicated memory area
     */
    void init_mapping()
    {
      if(this->base_addr)
	throw std::logic_error("segfault handler double initialization");
      const std::size_t mapping_len = pagelen() * 16; // TODO: placeholder
      void* mapped = mmap(nullptr, mapping_len, PROT_READ | PROT_WRITE,
			  MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
      if(mapped == MAP_FAILED)
	throw_sys_err("mmap");
      change_protection(mapped, PROT_NONE, mapping_len);
      this->base_addr = mapped;
      this->mapped_end = static_cast<char*>(mapped) + mapping_len;
    }

    /**
     * Commits all dirty pages and marks them as clean
     */
    void commit_changes()
    {
      iterator first = this->mapped_rw.begin();
      iterator last = this->mapped_rw.end();
      while(first != last) {
	void* page = *first;
	this->commit_page(page);
	change_protection(page, PROT_READ);
	this->mapped_rw.erase(first++);
	this->mapped_ro.insert(page);
      }
    }

    /**
     * Discards all dirty pages. Will re-populate them on demand
     */
    void discard_changes()
    {
      iterator first = this->mapped_rw.begin();
      iterator last = this->mapped_rw.end();
      while(first != last) {
	void* page = *first;
	change_protection(page, PROT_NONE);
	this->mapped_rw.erase(first++);
      }      
    }

    /**
     * Returns starting address of the dedicated memory area
     */
    void* get_base_addr() const
    { return this->base_addr; }
  };

  /**
   * Wraps a void functor so that is compatible to C
   *
   * Exceptions are converted to errno values.
   * Exceptions that indicate programming errors result in EINVAL
   * and output on stderr.
   *
   * \return 0 on success, -1 on exceptions
   */
  template<class Callable>
  int c_style_call(Callable&& function)
  {
    try {
      function();
      return 0;
    } catch(std::system_error& err) {
      errno = err.code().value();
    } catch(std::bad_alloc& err) {
      errno = ENOMEM;
    } catch(std::exception& err) {
      std::fprintf(stderr, "segfault handler: %s\n", err.what());
      errno = EINVAL;
    }
    return -1;
  }

} // namespace


/**
 * Public C interface follows
 */
extern "C" {

  /**
   * Installs and initializes the segmentation fault handler
   *
   * \return 0 on success, -1 otherwise. Sets errno
   */
  int sigsegv_install()
  {
    auto lambda = []() {
      SegvHandler::global_self()->init_mapping();
      SegvHandler::install();
    };
    return c_style_call(lambda);
  }

  /**
   * Commits all changes
   *
   * \return 0 on success, -1 otherwise. Sets errno
   */
  int sigsegv_commit()
  {
    auto lambda = []() {
      SegvHandler::global_self()->commit_changes();
    };
    return c_style_call(lambda);
  }

  /**
   * Discards all changes
   *
   * \return 0 on success, -1 otherwise. Sets errno
   */
  int sigsegv_discard()
  {
    auto lambda = []() {
      SegvHandler::global_self()->discard_changes();
    };
    return c_style_call(lambda);
  }

  /**
   * Returns base address of the dedicated memory area
   */
  void* sigsegv_baseptr()
  {
    void* base = nullptr;
    auto lambda = [&base]() {
      base = SegvHandler::global_self()->get_base_addr();
    };
    c_style_call(lambda);
    return base;
  }

} // extern "C"


/**
 * Some simple testing code
 *
 * Observe it with strace
 */
int main()
{
  sigsegv_install();
  unsigned* base = static_cast<unsigned*>(sigsegv_baseptr());
  std::printf("Accessing RO %p = 0x%x\n", base + 3, base[3]);
  base[3] = 0;
  std::puts("Committing");
  sigsegv_commit();
  std::size_t otherpage = pagelen()/sizeof(unsigned) + 3;
  std::puts("Making direct RW access");
  base[otherpage] = 0;
  std::puts("Discarding");
  sigsegv_discard();
  std::printf("Accessing RO %p = 0x%x\n", base + otherpage, base[otherpage]);
  return 0;
}