Add experimental support for interposing vmethods of known classes.

The hairiest bit is the abuse of compiler-specific pointer-to-member
internals in order to provide more or less transparent API.

CMakeLists.txt

@@ -111,6 +111,9 @@ IF(UNIX)
   SET(CMAKE_CXX_FLAGS_RELWITHDEBINFO "-g -Wall -Wno-unused-variable")
   SET(CMAKE_CXX_FLAGS "-fvisibility=hidden -m32 -march=i686 -mtune=generic -std=c++0x")
   SET(CMAKE_C_FLAGS "-fvisibility=hidden -m32 -march=i686 -mtune=generic")
+ELSEIF(MSVC)
+  # for msvc, tell it to always use 8-byte pointers to member functions to avoid confusion
+  SET(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /vmg /vmm")
 ENDIF()
 
 # use shared libraries for protobuf

library/CMakeLists.txt

@@ -27,6 +27,7 @@ include/Core.h
 include/ColorText.h
 include/DataDefs.h
 include/DataIdentity.h
+include/VTableInterpose.h
 include/LuaWrapper.h
 include/LuaTools.h
 include/Error.h
@@ -53,6 +54,7 @@ SET(MAIN_SOURCES
 Core.cpp
 ColorText.cpp
 DataDefs.cpp
+VTableInterpose.cpp
 LuaWrapper.cpp
 LuaTypes.cpp
 LuaTools.cpp

library/DataDefs.cpp

@@ -35,6 +35,7 @@ distribution.
 // must be last due to MS stupidity
 #include "DataDefs.h"
 #include "DataIdentity.h"
+#include "VTableInterpose.h"
 
 #include "MiscUtils.h"
 
@@ -214,6 +215,13 @@ virtual_identity::virtual_identity(size_t size, TAllocateFn alloc,
 {
 }
 
+virtual_identity::~virtual_identity()
+{
+    // Remove interpose entries, so that they don't try accessing this object later
+    for (int i = interpose_list.size()-1; i >= 0; i--)
+        interpose_list[i]->remove();
+}
+
 /* Vtable name to identity lookup. */
 static std::map<std::string, virtual_identity*> name_lookup;
 

library/VTableInterpose.cpp

@@ -0,0 +1,238 @@
+/*
+https://github.com/peterix/dfhack
+Copyright (c) 2009-2011 Petr Mrázek (peterix@gmail.com)
+
+This software is provided 'as-is', without any express or implied
+warranty. In no event will the authors be held liable for any
+damages arising from the use of this software.
+
+Permission is granted to anyone to use this software for any
+purpose, including commercial applications, and to alter it and
+redistribute it freely, subject to the following restrictions:
+
+1. The origin of this software must not be misrepresented; you must
+not claim that you wrote the original software. If you use this
+software in a product, an acknowledgment in the product documentation
+would be appreciated but is not required.
+
+2. Altered source versions must be plainly marked as such, and
+must not be misrepresented as being the original software.
+
+3. This notice may not be removed or altered from any source
+distribution.
+*/
+
+#include "Internal.h"
+
+#include <string>
+#include <vector>
+#include <map>
+
+#include "MemAccess.h"
+#include "Core.h"
+#include "VersionInfo.h"
+#include "VTableInterpose.h"
+
+#include "MiscUtils.h"
+
+using namespace DFHack;
+
+/*
+ *  Code for accessing method pointers directly. Very compiler-specific.
+ */
+
+#if defined(_MSC_VER)
+
+struct MSVC_MPTR {
+    void *method;
+    intptr_t this_shift;
+};
+
+bool DFHack::is_vmethod_pointer_(void *pptr)
+{
+    auto pobj = (MSVC_MPTR*)pptr;
+    if (!pobj->method) return false;
+
+    // MSVC implements pointers to vmethods via thunks.
+    // This expects that they all follow a very specific pattern.
+    auto pval = (unsigned*)pobj->method;
+    switch (pval[0]) {
+    case 0x60FF018BU: // mov eax, [ecx]; jmp [eax+0x??]
+    case 0xA0FF018BU: // mov eax, [ecx]; jmp [eax+0x????????]
+        return true;
+    default:
+        return false;
+    }
+}
+
+int DFHack::vmethod_pointer_to_idx_(void *pptr)
+{
+    auto pobj = (MSVC_MPTR*)pptr;
+    if (!pobj->method || pobj->this_shift != 0) return -1;
+
+    auto pval = (unsigned*)pobj->method;
+    switch (pval[0]) {
+    case 0x60FF018BU: // mov eax, [ecx]; jmp [eax+0x??]
+        return ((int8_t)pval[1])/sizeof(void*);
+    case 0xA0FF018BU: // mov eax, [ecx]; jmp [eax+0x????????]
+        return ((int32_t)pval[1])/sizeof(void*);
+    default:
+        return -1;
+    }
+}
+
+void* DFHack::method_pointer_to_addr_(void *pptr)
+{
+    if (is_vmethod_pointer_(pptr)) return NULL;
+    auto pobj = (MSVC_MPTR*)pptr;
+    return pobj->method;
+}
+
+void DFHack::addr_to_method_pointer_(void *pptr, void *addr)
+{
+    auto pobj = (MSVC_MPTR*)pptr;
+    pobj->method = addr;
+    pobj->this_shift = 0;
+}
+
+#elif defined(__GXX_ABI_VERSION)
+
+struct GCC_MPTR {
+    intptr_t method;
+    intptr_t this_shift;
+};
+
+bool DFHack::is_vmethod_pointer_(void *pptr)
+{
+    auto pobj = (GCC_MPTR*)pptr;
+    return (pobj->method & 1) != 0;
+}
+
+int DFHack::vmethod_pointer_to_idx_(void *pptr)
+{
+    auto pobj = (GCC_MPTR*)pptr;
+    if ((pobj->method & 1) == 0 || pobj->this_shift != 0)
+        return -1;
+    return (pobj->method-1)/sizeof(void*);
+}
+
+void* DFHack::method_pointer_to_addr_(void *pptr)
+{
+    auto pobj = (GCC_MPTR*)pptr;
+    if ((pobj->method & 1) != 0 || pobj->this_shift != 0)
+        return NULL;
+    return (void*)pobj->method;
+}
+
+void DFHack::addr_to_method_pointer_(void *pptr, void *addr)
+{
+    auto pobj = (GCC_MPTR*)pptr;
+    pobj->method = (intptr_t)addr;
+    pobj->this_shift = 0;
+}
+
+#else
+#error Unknown compiler type
+#endif
+
+void *virtual_identity::get_vmethod_ptr(int idx)
+{
+    assert(idx >= 0);
+    void **vtable = (void**)vtable_ptr;
+    if (!vtable) return NULL;
+    return vtable[idx];
+}
+
+bool virtual_identity::set_vmethod_ptr(int idx, void *ptr)
+{
+    assert(idx >= 0);
+    void **vtable = (void**)vtable_ptr;
+    if (!vtable) return NULL;
+    return Core::getInstance().p->patchMemory(&vtable[idx], &ptr, sizeof(void*));
+}
+
+void VMethodInterposeLinkBase::set_chain(void *chain)
+{
+    saved_chain = chain;
+    addr_to_method_pointer_(chain_mptr, chain);
+}
+
+VMethodInterposeLinkBase::VMethodInterposeLinkBase(virtual_identity *host, int vmethod_idx, void *interpose_method, void *chain_mptr)
+    : host(host), vmethod_idx(vmethod_idx), interpose_method(interpose_method), chain_mptr(chain_mptr),
+      saved_chain(NULL), next(NULL), prev(NULL)
+{
+}
+
+VMethodInterposeLinkBase::~VMethodInterposeLinkBase()
+{
+    if (is_applied())
+        remove();
+}
+
+bool VMethodInterposeLinkBase::apply()
+{
+    if (is_applied())
+        return true;
+    if (!host->vtable_ptr)
+        return false;
+
+    // Retrieve the current vtable entry
+    void *old_ptr = host->get_vmethod_ptr(vmethod_idx);
+    assert(old_ptr != NULL);
+
+    // Check if there are other interpose entries for the same slot
+    VMethodInterposeLinkBase *old_link = NULL;
+
+    for (int i = host->interpose_list.size()-1; i >= 0; i--)
+    {
+        if (host->interpose_list[i]->vmethod_idx != vmethod_idx)
+            continue;
+
+        old_link = host->interpose_list[i];
+        assert(old_link->next == NULL && old_ptr == old_link->interpose_method);
+        break;
+    }
+
+    // Apply the new method ptr
+    if (!host->set_vmethod_ptr(vmethod_idx, interpose_method))
+        return false;
+
+    set_chain(old_ptr);
+
+    // Link into the chain if any
+    if (old_link)
+    {
+        old_link->next = this;
+        prev = old_link;
+    }
+
+    return true;
+}
+
+void VMethodInterposeLinkBase::remove()
+{
+    if (!is_applied())
+        return;
+
+    // Remove from the list in the identity
+    for (int i = host->interpose_list.size()-1; i >= 0; i--)
+        if (host->interpose_list[i] == this)
+            vector_erase_at(host->interpose_list, i);
+
+    // Remove from the chain
+    if (prev)
+        prev->next = next;
+
+    if (next)
+    {
+        next->set_chain(saved_chain);
+        next->prev = prev;
+    }
+    else
+    {
+        host->set_vmethod_ptr(vmethod_idx, saved_chain);
+    }
+
+    prev = next = NULL;
+    set_chain(NULL);
+}

library/include/DataDefs.h

@@ -292,6 +292,8 @@ namespace DFHack
     typedef virtual_class *virtual_ptr;
 #endif
 
+    class DFHACK_EXPORT VMethodInterposeLinkBase;
+
     class DFHACK_EXPORT virtual_identity : public struct_identity {
         static std::map<void*, virtual_identity*> known;
 
@@ -299,6 +301,9 @@ namespace DFHack
 
         void *vtable_ptr;
 
+        friend class VMethodInterposeLinkBase;
+        std::vector<VMethodInterposeLinkBase*> interpose_list;
+
     protected:
         virtual void doInit(Core *core);
 
@@ -306,10 +311,14 @@ namespace DFHack
 
         bool can_allocate() { return struct_identity::can_allocate() && (vtable_ptr != NULL); }
 
+        void *get_vmethod_ptr(int index);
+        bool set_vmethod_ptr(int index, void *ptr);
+
     public:
         virtual_identity(size_t size, TAllocateFn alloc,
                          const char *dfhack_name, const char *original_name,
                          virtual_identity *parent, const struct_field_info *fields);
+        ~virtual_identity();
 
         virtual identity_type type() { return IDTYPE_CLASS; }
 
@@ -337,6 +346,8 @@ namespace DFHack
                               : (this == get(instance_ptr));
         }
 
+        template<class P> static P get_vmethod_ptr(P selector);
+
     public:
         bool can_instantiate() { return can_allocate(); }
         virtual_ptr instantiate() { return can_instantiate() ? (virtual_ptr)do_allocate() : NULL; }

library/include/DataFuncs.h

@@ -75,28 +75,31 @@ namespace df {
     cur_lua_ostream_argument name(state);
 
 #define INSTANTIATE_RETURN_TYPE(FArgs) \
-    template<FW_TARGSC class RT> struct return_type<RT (*) FArgs> { typedef RT type; }; \
+    template<FW_TARGSC class RT> struct return_type<RT (*) FArgs> { \
-    template<FW_TARGSC class RT, class CT> struct return_type<RT (CT::*) FArgs> { typedef RT type; };
+        typedef RT type; \
+        static const bool is_method = false; \
+    }; \
+    template<FW_TARGSC class RT, class CT> struct return_type<RT (CT::*) FArgs> { \
+        typedef RT type; \
+        typedef CT class_type; \
+        static const bool is_method = true; \
+    };
 
 #define INSTANTIATE_WRAPPERS(Count, FArgs, Args, Loads) \
     template<FW_TARGS> struct function_wrapper<void (*) FArgs, true> { \
-        static const bool is_method = false; \
         static const int num_args = Count; \
         static void execute(lua_State *state, int base, void (*cb) FArgs) { Loads; INVOKE_VOID(cb Args); } \
     }; \
     template<FW_TARGSC class RT> struct function_wrapper<RT (*) FArgs, false> { \
-        static const bool is_method = false; \
         static const int num_args = Count; \
         static void execute(lua_State *state, int base, RT (*cb) FArgs) { Loads; INVOKE_RV(cb Args); } \
     }; \
     template<FW_TARGSC class CT> struct function_wrapper<void (CT::*) FArgs, true> { \
-        static const bool is_method = true; \
         static const int num_args = Count+1; \
         static void execute(lua_State *state, int base, void (CT::*cb) FArgs) { \
             LOAD_CLASS() Loads; INVOKE_VOID((self->*cb) Args); } \
     }; \
     template<FW_TARGSC class RT, class CT> struct function_wrapper<RT (CT::*) FArgs, false> { \
-        static const bool is_method = true; \
         static const int num_args = Count+1; \
         static void execute(lua_State *state, int base, RT (CT::*cb) FArgs) { \
             LOAD_CLASS(); Loads; INVOKE_RV((self->*cb) Args); } \

library/include/VTableInterpose.h

@@ -0,0 +1,158 @@
+/*
+https://github.com/peterix/dfhack
+Copyright (c) 2009-2011 Petr Mrázek (peterix@gmail.com)
+
+This software is provided 'as-is', without any express or implied
+warranty. In no event will the authors be held liable for any
+damages arising from the use of this software.
+
+Permission is granted to anyone to use this software for any
+purpose, including commercial applications, and to alter it and
+redistribute it freely, subject to the following restrictions:
+
+1. The origin of this software must not be misrepresented; you must
+not claim that you wrote the original software. If you use this
+software in a product, an acknowledgment in the product documentation
+would be appreciated but is not required.
+
+2. Altered source versions must be plainly marked as such, and
+must not be misrepresented as being the original software.
+
+3. This notice may not be removed or altered from any source
+distribution.
+*/
+
+#pragma once
+
+#include "DataFuncs.h"
+
+namespace DFHack
+{
+    template<bool> struct StaticAssert;
+    template<> struct StaticAssert<true> {};
+
+#define STATIC_ASSERT(condition) { StaticAssert<(condition)>(); }
+
+    /* Wrapping around compiler-specific representation of pointers to methods. */
+
+#if defined(_MSC_VER)
+#define METHOD_POINTER_SIZE (sizeof(void*)*2)
+#elif defined(__GXX_ABI_VERSION)
+#define METHOD_POINTER_SIZE (sizeof(void*)*2)
+#else
+#error Unknown compiler type
+#endif
+
+#define ASSERT_METHOD_POINTER(type) \
+    STATIC_ASSERT(df::return_type<type>::is_method && sizeof(type)==METHOD_POINTER_SIZE);
+
+    DFHACK_EXPORT bool is_vmethod_pointer_(void*);
+    DFHACK_EXPORT int vmethod_pointer_to_idx_(void*);
+    DFHACK_EXPORT void* method_pointer_to_addr_(void*);
+    DFHACK_EXPORT void addr_to_method_pointer_(void*,void*);
+
+    template<class T> bool is_vmethod_pointer(T ptr) {
+        ASSERT_METHOD_POINTER(T);
+        return is_vmethod_pointer_(&ptr);
+    }
+    template<class T> int vmethod_pointer_to_idx(T ptr) {
+        ASSERT_METHOD_POINTER(T);
+        return vmethod_pointer_to_idx_(&ptr);
+    }
+    template<class T> void *method_pointer_to_addr(T ptr) {
+        ASSERT_METHOD_POINTER(T);
+        return method_pointer_to_addr_(&ptr);
+    }
+    template<class T> T addr_to_method_pointer(void *addr) {
+        ASSERT_METHOD_POINTER(T);
+        T rv;
+        addr_to_method_pointer_(&rv, addr);
+        return rv;
+    }
+
+    /* Access to vmethod pointers from the vtable. */
+
+    template<class P>
+    P virtual_identity::get_vmethod_ptr(P selector)
+    {
+        typedef typename df::return_type<P>::class_type host_class;
+        virtual_identity &identity = host_class::_identity;
+        int idx = vmethod_pointer_to_idx(selector);
+        return addr_to_method_pointer<P>(identity.get_vmethod_ptr(idx));
+    }
+
+    /* VMethod interpose API
+
+       Usage:
+
+       struct my_hack : public someclass {
+           typedef someclass interpose_base;
+
+           DEFINE_VMETHOD_INTERPOSE(void, foo, (int arg)) {
+               ...
+               INTERPOSE_NEXT(foo)(arg) // call the original
+               ...
+           }
+       };
+
+       IMPLEMENT_VMETHOD_INTERPOSE(my_hack, foo);
+
+       void init() {
+           my_hack::interpose_foo.apply()
+       }
+     */
+
+#define DEFINE_VMETHOD_INTERPOSE(rtype, name, args) \
+    typedef rtype (interpose_base::*interpose_ptr_##name)args; \
+    static DFHack::VMethodInterposeLink<interpose_base,interpose_ptr_##name> interpose_##name; \
+    rtype interpose_fn_##name args
+
+#define IMPLEMENT_VMETHOD_INTERPOSE(class,name) \
+    DFHack::VMethodInterposeLink<class::interpose_base,class::interpose_ptr_##name> \
+        class::interpose_##name(&class::interpose_base::name, &class::interpose_fn_##name);
+
+#define INTERPOSE_NEXT(name) (this->*interpose_##name.chain)
+
+    class DFHACK_EXPORT VMethodInterposeLinkBase {
+        /*
+          These link objects try to:
+          1) Allow multiple hooks into the same vmethod
+          2) Auto-remove hooks when a plugin is unloaded.
+        */
+
+        virtual_identity *host; // Class with the vtable
+        int vmethod_idx;
+        void *interpose_method; // Pointer to the code of the interposing method
+        void *chain_mptr;       // Pointer to the chain field below
+
+        void *saved_chain;      // Previous pointer to the code
+        VMethodInterposeLinkBase *next, *prev; // Other hooks for the same method
+
+        void set_chain(void *chain);
+    public:
+        VMethodInterposeLinkBase(virtual_identity *host, int vmethod_idx, void *interpose_method, void *chain_mptr);
+        ~VMethodInterposeLinkBase();
+
+        bool is_applied() { return saved_chain != NULL; }
+        bool apply();
+        void remove();
+    };
+
+    template<class Base, class Ptr>
+    class VMethodInterposeLink : public VMethodInterposeLinkBase {
+    public:
+        Ptr chain;
+
+        operator Ptr () { return chain; }
+
+        template<class Ptr2>
+        VMethodInterposeLink(Ptr target, Ptr2 src)
+            : VMethodInterposeLinkBase(
+                &Base::_identity,
+                vmethod_pointer_to_idx(target),
+                method_pointer_to_addr(src),
+                &chain
+              )
+        { src = target; /* check compatibility */ }
+    };
+}

plugins/devel/CMakeLists.txt

@@ -17,3 +17,4 @@ DFHACK_PLUGIN(stockcheck stockcheck.cpp)
 DFHACK_PLUGIN(stripcaged stripcaged.cpp)
 DFHACK_PLUGIN(rprobe rprobe.cpp)
 DFHACK_PLUGIN(nestboxes nestboxes.cpp)
+DFHACK_PLUGIN(vshook vshook.cpp)

plugins/devel/vshook.cpp

@@ -0,0 +1,61 @@
+#include "Core.h"
+#include <Console.h>
+#include <Export.h>
+#include <PluginManager.h>
+#include <modules/Gui.h>
+#include <vector>
+#include <cstdio>
+#include <stack>
+#include <string>
+#include <cmath>
+
+#include <VTableInterpose.h>
+#include "df/graphic.h"
+#include "df/viewscreen_titlest.h"
+
+using std::vector;
+using std::string;
+using std::stack;
+using namespace DFHack;
+
+using df::global::gps;
+
+DFHACK_PLUGIN("vshook");
+
+struct title_hook : df::viewscreen_titlest {
+    typedef df::viewscreen_titlest interpose_base;
+
+    DEFINE_VMETHOD_INTERPOSE(void, render, ())
+    {
+        INTERPOSE_NEXT(render)();
+
+        if (gps) {
+            uint8_t *buf = gps->screen;
+            int32_t *stp = gps->screentexpos;
+
+            for (const char *p = "DFHack " DFHACK_VERSION; *p; p++) {
+                *buf++ = *p; *buf++ = 7; *buf++ = 0; *buf++ = 1;
+                *stp++ = 0;
+            }
+        }
+    }
+};
+
+IMPLEMENT_VMETHOD_INTERPOSE(title_hook, render);
+
+DFhackCExport command_result plugin_init ( color_ostream &out, std::vector <PluginCommand> &commands)
+{
+    if (gps)
+    {
+        if (!title_hook::interpose_render.apply())
+            out.printerr("Could not interpose viewscreen_titlest::render\n");
+    }
+
+    return CR_OK;
+}
+
+DFhackCExport command_result plugin_shutdown ( color_ostream &out )
+{
+    title_hook::interpose_render.remove();
+    return CR_OK;
+}