//
// Die Klasse MemoryProtect hat die Fähigkeit den Speicher
// aus dem ein Objekt das von ihr abgleitet hat gegen Schreiben
// zu schützen. Nützlich nur zu debug-Zwecken
//

#include "stdafx.h"
#include <new>          // std::bad_alloc
#include <unordered_set>
#include <windows.h>
#include <mutex>        // std::mutex

using namespace std;

DWORD getPageSize() {
  static DWORD pagesize = [] {
    SYSTEM_INFO si;
    GetSystemInfo(&si);
    return si.dwPageSize;
  }();
  return pagesize;
}

template <class T>
class MemoryProtect {
  protected:
    typedef T cageclass_t;

  public:
    MemoryProtect() {}

    /* static */ void *operator new(size_t size);
    /* static */ void *operator new[](size_t size);
    /* static */ void operator delete(void *p);
    /* static */ void operator delete[](void *p);

    bool protectMemoryCage();
    bool unprotectMemoryCage();

    static bool checkPointer(const T *);
    static void g_beValidOrCrash(const T *);
    void beValidOrCrash() const;
  private:
    static unordered_set<const T*> c_validpointerlist;
    static mutex c_mtx;
    size_t m_actualsize;
};

template <class T>
/* static */
unordered_set<const T*> MemoryProtect<T>::c_validpointerlist;

template <class T>
/* static */
mutex MemoryProtect<T>::c_mtx;

template <class T>
/* static */
bool MemoryProtect<T>::checkPointer(const T *obj) {
  lock_guard<mutex> lck (c_mtx);
  return c_validpointerlist.end() != c_validpointerlist.find(obj);
}

template <class T>
/* static */
void MemoryProtect<T>::g_beValidOrCrash(const T *obj) {
  if (!checkPointer(obj)) {
    // Purposefully crash and spawn the Crashreport...
/*    void *p = 0;
    *p = 0;*/
  }
}

template <class T>
void MemoryProtect<T>::beValidOrCrash() const {
  g_beValidOrCrash(static_cast<const T *>(this));
}

template <class T>
bool MemoryProtect<T>::protectMemoryCage() {
  DWORD oldprotect;
  return (TRUE == VirtualProtect(this, m_actualsize, PAGE_READONLY, &oldprotect));
}

template <class T>
bool MemoryProtect<T>::unprotectMemoryCage() {
  DWORD oldprotect;
  return (TRUE == VirtualProtect(this, m_actualsize, PAGE_READWRITE, &oldprotect));
}

template <class T>
// /* static */
void *MemoryProtect<T>::operator new(size_t size)
{
  // we dont use VirtualAlloc to be similar to the
  // original malloc behaviour
  // we need to use memory that spans over cpu-pages
  // for VirtualProtect
  DWORD pagesize = getPageSize();
  size_t actualsize = size_t((size + (pagesize - 1)) / pagesize) * pagesize;
  void *p = _aligned_malloc(actualsize, pagesize);
  if (!p)
    throw bad_alloc();
  ((T *)p)->m_actualsize = actualsize;
  lock_guard<mutex> lck (c_mtx);
  MemoryProtect<T>::c_validpointerlist.insert(((T *)p));
  return p;
}

template <class T>
void *MemoryProtect<T>::operator new[](size_t size)
{
  // we dont use VirtualAlloc to be similar to the
  // original malloc behaviour
  // we need to use memory that spans over cpu-pages
  // for VirtualProtect
  DWORD pagesize = getPageSize();
  size_t actualsize = size_t(size + (pagesize - 1) / pagesize) * pagesize;
  void *p = _aligned_malloc(actualsize, pagesize);
  if (!p)
    throw bad_alloc();
  ((T *)p)->m_actualsize = actualsize;
  lock_guard<mutex> lck (c_mtx);
  MemoryProtect<T>::c_validpointerlist.insert(((T *)p));
  return p;
}

template <class T>
void MemoryProtect<T>::operator delete(void *p)
{
  if (nullptr == p)
    throw std::bad_alloc();
  ((T *)p)->unprotectMemoryCage();
  lock_guard<mutex> lck (c_mtx);
  if (0 == MemoryProtect<T>::c_validpointerlist.erase(((T*)p)))
    throw std::bad_alloc();
  _aligned_free(p);
}

template <class T>
void MemoryProtect<T>::operator delete[](void *p)
{
  if (nullptr == p)
    throw bad_alloc();
  (T *)p->unprotectMemoryCage();
  lock_guard<mutex> lck (c_mtx);
  if (0 == c_validpointerlist.erase(p))
    throw bad_alloc();
  _aligned_free(p);
}

//

class A : public MemoryProtect<A> {
  int blubber;
};

int _tmain(int argc, _TCHAR* argv[])
{
  A *blah = new A();
  blah->protectMemoryCage();
  if (!A::checkPointer(blah)) {
    int damn = 4;
  }
  if (A::checkPointer(blah+1)) {
    int damn = 5;
  }
  blah->beValidOrCrash();
  //((unsigned int *)blah)[0] = 9;
  delete blah;
    return 0;
}