hook: Update capnhook source files from latest master

https://github.com/decafcode/capnhook

src/main/hook/Module.mk

@@ -1,10 +1,14 @@
 libs            += hook
 
 src_hook        := \
+    args.c \
     com-proxy.c \
     d3d9.c \
+    hr.c \
+    iobuf.c \
     iohook.c \
     pe.c \
     peb.c \
+    process.c \
     table.c \
 

src/main/hook/args.c

@@ -0,0 +1,137 @@
+#include <windows.h>
+
+#include <assert.h>
+#include <stdbool.h>
+#include <stdint.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include "hook/args.h"
+
+/* This does not handle escaped double quotes inside args correctly yet */
+
+static HRESULT args_push(
+        int *argc,
+        char ***argv,
+        const char *begin,
+        const char *end)
+{
+    int tmp_argc;
+    char **tmp_argv;
+    size_t nchars;
+    char *str;
+
+    nchars = end - begin;
+    str = malloc(nchars + 1);
+
+    if (str == NULL) {
+        goto fail;
+    }
+
+    memcpy(str, begin, nchars);
+    str[nchars] = '\0';
+
+    tmp_argc = *argc + 1;
+    tmp_argv = realloc(*argv, tmp_argc * sizeof(char **));
+
+    if (tmp_argv == NULL) {
+        goto fail;
+    }
+
+    tmp_argv[tmp_argc - 1] = str;
+
+    *argv = tmp_argv;
+    *argc = tmp_argc;
+
+    return S_OK;
+
+fail:
+    free(str);
+
+    return E_OUTOFMEMORY;
+}
+
+HRESULT args_recover(int *argc_out, char ***argv_out)
+{
+    int argc;
+    char **argv;
+    const char *begin;
+    const char *pos;
+    bool quote;
+    HRESULT hr;
+
+    assert(argc_out != NULL);
+    assert(argv_out != NULL);
+
+    *argc_out = 0;
+    *argv_out = NULL;
+
+    argc = 0;
+    argv = NULL;
+    quote = false;
+
+    for (begin = pos = GetCommandLine() ; *pos ; pos++) {
+        switch (*pos) {
+            case '"':
+                if (!quote) {
+                    quote = true;
+                    begin = pos + 1;
+                } else {
+                    hr = args_push(&argc, &argv, begin, pos);
+
+                    if (FAILED(hr)) {
+                        goto fail;
+                    }
+
+                    quote = false;
+                    begin = NULL;
+                }
+
+                break;
+
+            case ' ':
+                if (!quote && begin != NULL) {
+                    args_push(&argc, &argv, begin, pos);
+                    begin = NULL;
+                }
+
+                break;
+
+            default:
+                if (begin == NULL) {
+                    begin = pos;
+                }
+
+                break;
+        }
+    }
+
+    if (begin != NULL && !quote) {
+        hr = args_push(&argc, &argv, begin, pos);
+
+        if (FAILED(hr)) {
+            goto fail;
+        }
+    }
+
+    *argc_out = argc;
+    *argv_out = argv;
+
+    return S_OK;
+
+fail:
+    args_free(argc, argv);
+
+    return hr;
+}
+
+void args_free(int argc, char **argv)
+{
+    int i;
+
+    for (i = 0 ; i < argc ; i++) {
+        free(argv[i]);
+    }
+
+    free(argv);
+}

src/main/hook/args.h

@@ -0,0 +1,6 @@
+#pragma once
+
+#include <windows.h>
+
+HRESULT args_recover(int *argc, char ***argv);
+void args_free(int argc, char **argv);

src/main/hook/com-proxy.c

@@ -1,90 +1,177 @@
-#include <unknwn.h>
 #include <windows.h>
+#include <unknwn.h>
 
+#include <assert.h>
+#include <stdint.h>
 #include <stdlib.h>
 
 #include "hook/com-proxy.h"
 
-#include "util/defs.h"
+static void com_proxy_free(struct com_proxy *proxy);
-#include "util/mem.h"
+static HRESULT STDMETHODCALLTYPE com_proxy_query_interface(
+        IUnknown *unk,
+        REFIID iid,
+        void **iface);
+static ULONG STDMETHODCALLTYPE com_proxy_addref(IUnknown *unk);
+static ULONG STDMETHODCALLTYPE com_proxy_release(IUnknown *unk);
 
-#ifdef _WIN64
+#ifdef __amd64
 
-/***** 64-BIT TRAMPOLINE *****/
+    /***** 64-BIT TRAMPOLINE *****/
 
 #define SLOT_OFFSET 0x0A
 static const uint8_t com_proxy_tramp[] = {
-    /* mov rcx, [rcx+8]         ; Replace this with this->real */
+    /* mov rcx, [rcx+8]         ; Get this->real */
-    0x48,
+    0x48, 0x8B, 0x49, 0x08,
-    0x8B,
-    0x49,
-    0x08,
 
-    /* mov rax, [rcx]           ; Get real->vtbl */
+    /* mov rax, [rcx]           ; Get this->vtbl */
-    0x48,
+    0x48, 0x8B, 0x01,
-    0x8B,
-    0x01,
 
     /* mov rax, [rax+XX]        ; Get vtbl->slot_XX */
-    0x48,
+    0x48, 0x8B, 0x80, -1, -1, -1, -1,
-    0x8B,
-    0x80,
-    -1,
-    -1,
-    -1,
-    -1,
 
     /* jmp rax                  ; Continue to slot_XX */
-    0xFF,
+    0xFF, 0xE0,
-    0xE0,
 };
 
 #else
 
-/***** 32-BIT TRAMPOLINE *****/
+    /***** 32-BIT TRAMPOLINE *****/
 
 #define SLOT_OFFSET 0x0F
 static const uint8_t com_proxy_tramp[] = {
     /* mov eax, [esp+4]         ; Get this */
-    0x8B,
+    0x8B, 0x44, 0x24, 0x04,
-    0x44,
-    0x24,
-    0x04,
 
     /* mov eax, [eax+4]         ; Get this->real */
-    0x8B,
+    0x8B, 0x40, 0x04,
-    0x40,
-    0x04,
 
     /* mov [esp+4], eax         ; Replace this with this->real on stack */
-    0x89,
+    0x89, 0x44, 0x24, 0x04,
-    0x44,
-    0x24,
-    0x04,
 
-    /* mov ecx, [eax]           ; Get real->vtbl */
+    /* mov ecx, [eax]           ; Get this->vtbl */
-    0x8B,
+    0x8B, 0x08,
-    0x08,
 
     /* mov ecx, [ecx+XX]        ; Get vtbl->slot_XX */
-    0x8B,
+    0x8B, 0x89, -1, -1, -1, -1,
-    0x89,
-    -1,
-    -1,
-    -1,
-    -1,
 
     /* jmp ecx                  ; Continue to slot_XX */
-    0xFF,
+    0xFF, 0xE1
-    0xE1};
+};
 
 #endif
 
-static HRESULT STDCALL
+HRESULT com_proxy_wrap(
-com_proxy_query_interface(IUnknown *ptr, REFIID iid, void **iface)
+        struct com_proxy **out,
+        void *real,
+        size_t vtbl_size)
 {
-    struct com_proxy *self = (struct com_proxy *) ptr;
+    struct com_proxy *proxy;
-    IUnknown *obj = self->real; /* Not necessarily the real IUnknown* */
+    void **vtbl;
+    uint8_t *cur_tramp;
+    size_t nslots;
+    size_t i;
+    HRESULT hr;
+
+    assert(out != NULL);
+    assert(real != NULL);
+
+    *out = NULL;
+
+    proxy = calloc(1, sizeof(*proxy));
+
+    if (proxy == NULL) {
+        hr = E_OUTOFMEMORY;
+
+        goto end;
+    }
+
+    proxy->vptr = malloc(vtbl_size);
+
+    if (proxy->vptr == NULL) {
+        hr = E_OUTOFMEMORY;
+
+        goto end;
+    }
+
+    nslots = vtbl_size / sizeof(void *);
+
+    proxy->tramps = VirtualAlloc(
+            NULL,
+            sizeof(com_proxy_tramp) * nslots,
+            MEM_RESERVE | MEM_COMMIT,
+            PAGE_EXECUTE_READWRITE);
+
+    if (proxy->tramps == NULL) {
+        hr = E_OUTOFMEMORY;
+
+        goto end;
+    }
+
+    proxy->real = real;
+
+    /* Set up proxied IUnknown impl */
+
+    vtbl = proxy->vptr;
+    vtbl[0] = com_proxy_query_interface;
+    vtbl[1] = com_proxy_addref;
+    vtbl[2] = com_proxy_release;
+
+    /* Populate trampoline code for remaining vtbl entries */
+
+    for (i = 3 /* Skip IUnknown */ ; i < nslots ; i++) {
+        cur_tramp = proxy->tramps + i * sizeof(com_proxy_tramp);
+
+        /* Copy template */
+        memcpy(cur_tramp, com_proxy_tramp, sizeof(com_proxy_tramp));
+
+        /* Patch XX into vtbl lookup (see definition of tramp) */
+        *((uint32_t *) (cur_tramp + SLOT_OFFSET)) = i * sizeof(void *);
+
+        /* Set vtable entry */
+        vtbl[i] = cur_tramp;
+    }
+
+    *out = proxy;
+    proxy = NULL;
+    hr = S_OK;
+
+end:
+    com_proxy_free(proxy);
+
+    return hr;
+}
+
+static void com_proxy_free(struct com_proxy *proxy)
+{
+    if (proxy == NULL) {
+        return;
+    }
+
+    if (proxy->cleanup_ctx != NULL) {
+        proxy->cleanup_ctx(proxy->ctx);
+    }
+
+    if (proxy->tramps != NULL) {
+        VirtualFree(proxy->tramps, 0, MEM_RELEASE);
+    }
+
+    free(proxy->vptr);
+    free(proxy);
+}
+
+static HRESULT STDMETHODCALLTYPE com_proxy_query_interface(
+        IUnknown *unk,
+        REFIID iid,
+        void **iface)
+{
+    struct com_proxy *proxy;
+    IUnknown *obj;
+
+    assert(unk != NULL);
+
+    proxy = (struct com_proxy *) unk;
+    obj = proxy->real; /* Not necessarily this object's canonical IUnknown */
 
     /* To some extent, COM is designed to support shennanigans like these.
        We can safely pass the call straight through to the underlying
@@ -104,70 +191,35 @@ com_proxy_query_interface(IUnknown *ptr, REFIID iid, void **iface)
     return IUnknown_QueryInterface(obj, iid, iface);
 }
 
-static ULONG STDCALL com_proxy_addref(IUnknown *ptr)
+static ULONG STDMETHODCALLTYPE com_proxy_addref(IUnknown *unk)
 {
-    struct com_proxy *self = (struct com_proxy *) ptr;
+    struct com_proxy *proxy;
-    IUnknown *obj = self->real;
+    IUnknown *obj;
+
+    assert(unk != NULL);
+
+    proxy = (struct com_proxy *) unk;
+    obj = proxy->real;
 
     return IUnknown_AddRef(obj);
 }
 
-static ULONG STDCALL com_proxy_release(IUnknown *ptr)
+static ULONG STDMETHODCALLTYPE com_proxy_release(IUnknown *unk)
 {
-    struct com_proxy *self = (struct com_proxy *) ptr;
+    struct com_proxy *proxy;
-    IUnknown *obj = self->real;
+    IUnknown *real;
     ULONG result;
 
-    result = IUnknown_Release(obj);
+    assert(unk != NULL);
+
+    proxy = (struct com_proxy *) unk;
+    real = proxy->real;
+    result = IUnknown_Release(real);
 
     if (!result) {
         /* Last ref to underlying object released */
-        VirtualFree(self->tramps, 0, MEM_RELEASE);
+        com_proxy_free(proxy);
-        free(self->vptr);
-        free(self);
     }
 
     return result;
 }
-
-struct com_proxy *com_proxy_wrap(void *iface, size_t vtbl_size)
-{
-    struct com_proxy *self;
-    void **vtbl;
-    uint8_t *cur_tramp;
-    uint32_t nslots;
-    uint32_t i;
-
-    nslots = vtbl_size / sizeof(void *);
-
-    self = xmalloc(sizeof(*self));
-    self->vptr = xmalloc(vtbl_size);
-    self->real = iface;
-    self->tramps = VirtualAlloc(
-        NULL,
-        sizeof(com_proxy_tramp) * nslots,
-        MEM_RESERVE | MEM_COMMIT,
-        PAGE_EXECUTE_READWRITE);
-
-    /* Set up proxied IUnknown impl */
-    vtbl = self->vptr;
-    vtbl[0] = com_proxy_query_interface;
-    vtbl[1] = com_proxy_addref;
-    vtbl[2] = com_proxy_release;
-
-    /* Populate trampoline code for remaining vtbl entries */
-    for (i = 3 /* Skip IUnknown */; i < nslots; i++) {
-        cur_tramp = self->tramps + i * sizeof(com_proxy_tramp);
-
-        /* Copy template */
-        memcpy(cur_tramp, com_proxy_tramp, sizeof(com_proxy_tramp));
-
-        /* Patch XX into vtbl lookup (see definition of tramp) */
-        *((uint32_t *) (cur_tramp + SLOT_OFFSET)) = i * sizeof(void *);
-
-        /* Set vtable entry */
-        vtbl[i] = cur_tramp;
-    }
-
-    return self;
-}

src/main/hook/com-proxy.h

@@ -1,14 +1,9 @@
-#ifndef HOOK_COM_PROXY_H
+#pragma once
-#define HOOK_COM_PROXY_H
 
 #include <stddef.h>
 #include <stdint.h>
 
-/* N.B. Here be dragons. You really ought to be fairly familiar with COM
+#define com_proxy_downcast(self) ((struct com_proxy *) self)
-   before using this, otherwise you risk ignoring the subtler issues at your
-   peril. */
-
-#define COM_PROXY_UNWRAP(self) (((struct com_proxy *) self)->real)
 
 struct com_proxy {
     /* Pointer to vtable filled with trampolines. Edit these as you please.
@@ -18,13 +13,20 @@ struct com_proxy {
     /* Interface pointer wrapped by this proxy. */
     void *real;
 
+    /* Context pointer for use by hook code. */
+    void *ctx;
+
+    /* Optional cleanup callback, will be called during com_proxy deallocation
+       to clean up *ctx. */
+    void (*cleanup_ctx)(void *ctx);
+
     /* Dynamically generated x86 trampoline code. The initial vtable entries
        all point into code located here. */
     uint8_t *tramps;
 };
 
 /* Wrap a COM interface pointer in a proxy. This is an object that acts just
-   like the object that it wraps, but has a freely editable vtable, which you
+   like the object that it wraps but has a freely editable vtable, which you
    can modify in order to intercept a subset of the interface's method calls.
 
    By default, all the vtable slots contain dynamically generated trampolines
@@ -34,9 +36,8 @@ struct com_proxy {
    NOTE! This does not AddRef the underlying interface.
 
    NOTE! This function wraps COM POINTERS, not COM OBJECTS (since the latter
-   is, in general, impossible). If you're insufficiently versed in COM to
+   is, in general, impossible). Consequences of this distinction include the
-   understand the difference... well, you really should be, but the following
+   following:
-   observations are a start:
 
    1. Do not wrap IUnknown pointers with this function. This will break the
       IUnknown::QueryInterface contract. This refers to _the_ unique
@@ -44,11 +45,11 @@ struct com_proxy {
       extends IUnknown). Wrapping the unique IUnknown* for an object will
       cause it to no longer be unique.
 
-   2. Callers can "jailbreak" your wrapper using IUnknown::QueryInterface.
+   2. Callers can "jailbreak" your wrapper using IUnknown::QueryInterface
+      unless you provide a custom QueryInterface implementation to prevent them
+      from doing so. */
 
-   Generally this isn't an issue for DirectX objects, since nobody ever seems
+HRESULT com_proxy_wrap(
-   to use QueryInterface with them. */
+        struct com_proxy **out,
-
+        void *real,
-struct com_proxy *com_proxy_wrap(void *iface, size_t vtbl_size);
+        size_t vtbl_size);
-
-#endif

src/main/hook/hr.c

@@ -0,0 +1,33 @@
+#include <windows.h>
+
+#include <stdint.h>
+
+#include "hook/hr.h"
+
+uint32_t hr_to_win32_error(HRESULT hr)
+{
+    if (SUCCEEDED(hr)) {
+        return ERROR_SUCCESS;
+    } else if (HRESULT_FACILITY(hr) == FACILITY_WIN32) {
+        return HRESULT_CODE(hr);
+    } else {
+        switch (hr) {
+        case E_ABORT:           return ERROR_OPERATION_ABORTED;
+        case E_ACCESSDENIED:    return ERROR_ACCESS_DENIED;
+        case E_FAIL:            return ERROR_GEN_FAILURE;
+        case E_HANDLE:          return ERROR_INVALID_HANDLE;
+        case E_INVALIDARG:      return ERROR_INVALID_PARAMETER;
+        case E_NOINTERFACE:     return ERROR_INVALID_FUNCTION;
+        case E_NOTIMPL:         return ERROR_NOT_SUPPORTED;
+        case E_OUTOFMEMORY:     return ERROR_OUTOFMEMORY;
+        case E_POINTER:         return ERROR_INVALID_ADDRESS;
+        case E_UNEXPECTED:      return ERROR_INTERNAL_ERROR;
+        default:                return ERROR_INTERNAL_ERROR;
+        }
+    }
+}
+
+void hr_propagate_win32_(HRESULT hr)
+{
+    SetLastError(hr_to_win32_error(hr));
+}

src/main/hook/hr.h

@@ -0,0 +1,10 @@
+#pragma once
+
+#include <windows.h>
+
+#include <stdint.h>
+
+#define hr_propagate_win32(hr, r) (hr_propagate_win32_(hr), r)
+
+uint32_t hr_to_win32_error(HRESULT hr);
+void hr_propagate_win32_(HRESULT hr);

src/main/hook/iobuf.c

@@ -0,0 +1,343 @@
+#include <assert.h>
+#include <string.h>
+
+#include "hook/iobuf.h"
+
+void iobuf_flip(struct const_iobuf *child, struct iobuf *parent)
+{
+    assert(child != NULL);
+    assert(parent != NULL);
+
+    child->bytes = parent->bytes;
+    child->pos = 0;
+    child->nbytes = parent->pos;
+}
+
+size_t iobuf_move(struct iobuf *dest, struct const_iobuf *src)
+{
+    size_t dest_avail;
+    size_t src_avail;
+    size_t chunksz;
+
+    assert(dest != NULL);
+    assert(dest->bytes != NULL || dest->nbytes == 0);
+    assert(dest->pos <= dest->nbytes);
+
+    assert(src != NULL);
+    assert(src->bytes != NULL || src->nbytes == 0);
+    assert(src->pos <= src->nbytes);
+
+    dest_avail = dest->nbytes - dest->pos;
+    src_avail = src->nbytes - src->pos;
+    chunksz = dest_avail < src_avail ? dest_avail : src_avail;
+
+    memcpy(&dest->bytes[dest->pos], &src->bytes[src->pos], chunksz);
+
+    dest->pos += chunksz;
+    src->pos += chunksz;
+
+    return chunksz;
+}
+
+size_t iobuf_shift(struct iobuf *dest, struct iobuf *src)
+{
+    struct const_iobuf span;
+
+    assert(dest != NULL);
+    assert(src != NULL);
+
+    iobuf_flip(&span, src);
+    iobuf_move(dest, &span);
+
+    memmove(src->bytes, &src->bytes[span.pos], span.nbytes - span.pos);
+    src->pos -= span.pos;
+
+    return span.pos;
+}
+
+HRESULT iobuf_read(struct const_iobuf *src, void *bytes, size_t nbytes)
+{
+    assert(src != NULL);
+    assert(bytes != NULL || nbytes == 0);
+
+    if (src->pos + nbytes > src->nbytes) {
+        return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
+    }
+
+    memcpy(bytes, &src->bytes[src->pos], nbytes);
+    src->pos += nbytes;
+
+    return S_OK;
+}
+
+HRESULT iobuf_read_8(struct const_iobuf *src, uint8_t *out)
+{
+    assert(src != NULL);
+    assert(out != NULL);
+
+    if (src->pos + sizeof(uint8_t) > src->nbytes) {
+        return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
+    }
+
+    *out = src->bytes[src->pos++];
+
+    return S_OK;
+}
+
+HRESULT iobuf_read_be16(struct const_iobuf *src, uint16_t *out)
+{
+    uint16_t value;
+
+    assert(src != NULL);
+    assert(out != NULL);
+
+    if (src->pos + sizeof(uint16_t) > src->nbytes) {
+        return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
+    }
+
+    value  = src->bytes[src->pos++] << 8;
+    value |= src->bytes[src->pos++];
+
+    *out = value;
+
+    return S_OK;
+}
+
+HRESULT iobuf_read_be32(struct const_iobuf *src, uint32_t *out)
+{
+    uint32_t value;
+
+    assert(src != NULL);
+    assert(out != NULL);
+
+    if (src->pos + sizeof(uint32_t) > src->nbytes) {
+        return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
+    }
+
+    value  = src->bytes[src->pos++] << 24;
+    value |= src->bytes[src->pos++] << 16;
+    value |= src->bytes[src->pos++] << 8;
+    value |= src->bytes[src->pos++];
+
+    *out = value;
+
+    return S_OK;
+}
+
+HRESULT iobuf_read_be64(struct const_iobuf *src, uint64_t *out)
+{
+    uint64_t value;
+
+    assert(src != NULL);
+    assert(out != NULL);
+
+    if (src->pos + sizeof(uint64_t) > src->nbytes) {
+        return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
+    }
+
+    value  = ((uint64_t) src->bytes[src->pos++]) << 56;
+    value |= ((uint64_t) src->bytes[src->pos++]) << 48;
+    value |= ((uint64_t) src->bytes[src->pos++]) << 40;
+    value |= ((uint64_t) src->bytes[src->pos++]) << 32;
+    value |= ((uint64_t) src->bytes[src->pos++]) << 24;
+    value |= ((uint64_t) src->bytes[src->pos++]) << 16;
+    value |= ((uint64_t) src->bytes[src->pos++]) << 8;
+    value |= ((uint64_t) src->bytes[src->pos++]);
+
+    *out = value;
+
+    return S_OK;
+}
+
+HRESULT iobuf_read_le16(struct const_iobuf *src, uint16_t *out)
+{
+    uint16_t value;
+
+    assert(src != NULL);
+    assert(out != NULL);
+
+    if (src->pos + sizeof(uint16_t) > src->nbytes) {
+        return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
+    }
+
+    value  = src->bytes[src->pos++];
+    value |= src->bytes[src->pos++] << 8;
+
+    *out = value;
+
+    return S_OK;
+}
+
+HRESULT iobuf_read_le32(struct const_iobuf *src, uint32_t *out)
+{
+    uint32_t value;
+
+    assert(src != NULL);
+    assert(out != NULL);
+
+    if (src->pos + sizeof(uint32_t) > src->nbytes) {
+        return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
+    }
+
+    value  = src->bytes[src->pos++];
+    value |= src->bytes[src->pos++] << 8;
+    value |= src->bytes[src->pos++] << 16;
+    value |= src->bytes[src->pos++] << 24;
+
+    *out = value;
+
+    return S_OK;
+}
+
+HRESULT iobuf_read_le64(struct const_iobuf *src, uint64_t *out)
+{
+    uint64_t value;
+
+    assert(src != NULL);
+    assert(out != NULL);
+
+    if (src->pos + sizeof(uint64_t) > src->nbytes) {
+        return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
+    }
+
+    value  = ((uint64_t) src->bytes[src->pos++]);
+    value |= ((uint64_t) src->bytes[src->pos++]) << 8;
+    value |= ((uint64_t) src->bytes[src->pos++]) << 16;
+    value |= ((uint64_t) src->bytes[src->pos++]) << 24;
+    value |= ((uint64_t) src->bytes[src->pos++]) << 32;
+    value |= ((uint64_t) src->bytes[src->pos++]) << 40;
+    value |= ((uint64_t) src->bytes[src->pos++]) << 48;
+    value |= ((uint64_t) src->bytes[src->pos++]) << 56;
+
+    *out = value;
+
+    return S_OK;
+}
+
+HRESULT iobuf_write(struct iobuf *dest, const void *bytes, size_t nbytes)
+{
+    assert(dest != NULL);
+    assert(bytes != NULL || nbytes == 0);
+
+    if (dest->pos + nbytes > dest->nbytes) {
+        return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
+    }
+
+    memcpy(&dest->bytes[dest->pos], bytes, nbytes);
+    dest->pos += nbytes;
+
+    return S_OK;
+}
+
+HRESULT iobuf_write_8(struct iobuf *dest, uint8_t value)
+{
+    assert(dest != NULL);
+
+    if (dest->pos + sizeof(uint8_t) > dest->nbytes) {
+        return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
+    }
+
+    dest->bytes[dest->pos++] = value;
+
+    return S_OK;
+}
+
+HRESULT iobuf_write_be16(struct iobuf *dest, uint16_t value)
+{
+    assert(dest != NULL);
+
+    if (dest->pos + sizeof(uint16_t) > dest->nbytes) {
+        return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
+    }
+
+    dest->bytes[dest->pos++] = value >> 8;
+    dest->bytes[dest->pos++] = value;
+
+    return S_OK;
+}
+
+HRESULT iobuf_write_be32(struct iobuf *dest, uint32_t value)
+{
+    assert(dest != NULL);
+
+    if (dest->pos + sizeof(uint32_t) > dest->nbytes) {
+        return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
+    }
+
+    dest->bytes[dest->pos++] = value >> 24;
+    dest->bytes[dest->pos++] = value >> 16;
+    dest->bytes[dest->pos++] = value >> 8;
+    dest->bytes[dest->pos++] = value;
+
+    return S_OK;
+}
+
+HRESULT iobuf_write_be64(struct iobuf *dest, uint64_t value)
+{
+    assert(dest != NULL);
+
+    if (dest->pos + sizeof(uint64_t) > dest->nbytes) {
+        return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
+    }
+
+    dest->bytes[dest->pos++] = value >> 56;
+    dest->bytes[dest->pos++] = value >> 48;
+    dest->bytes[dest->pos++] = value >> 40;
+    dest->bytes[dest->pos++] = value >> 32;
+    dest->bytes[dest->pos++] = value >> 24;
+    dest->bytes[dest->pos++] = value >> 16;
+    dest->bytes[dest->pos++] = value >> 8;
+    dest->bytes[dest->pos++] = value;
+
+    return S_OK;
+}
+
+HRESULT iobuf_write_le16(struct iobuf *dest, uint16_t value)
+{
+    assert(dest != NULL);
+
+    if (dest->pos + sizeof(uint16_t) > dest->nbytes) {
+        return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
+    }
+
+    dest->bytes[dest->pos++] = value;
+    dest->bytes[dest->pos++] = value >> 8;
+
+    return S_OK;
+}
+
+HRESULT iobuf_write_le32(struct iobuf *dest, uint32_t value)
+{
+    assert(dest != NULL);
+
+    if (dest->pos + sizeof(uint32_t) > dest->nbytes) {
+        return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
+    }
+
+    dest->bytes[dest->pos++] = value;
+    dest->bytes[dest->pos++] = value >> 8;
+    dest->bytes[dest->pos++] = value >> 16;
+    dest->bytes[dest->pos++] = value >> 24;
+
+    return S_OK;
+}
+
+HRESULT iobuf_write_le64(struct iobuf *dest, uint64_t value)
+{
+    assert(dest != NULL);
+
+    if (dest->pos + sizeof(uint64_t) > dest->nbytes) {
+        return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
+    }
+
+    dest->bytes[dest->pos++] = value;
+    dest->bytes[dest->pos++] = value >> 8;
+    dest->bytes[dest->pos++] = value >> 16;
+    dest->bytes[dest->pos++] = value >> 24;
+    dest->bytes[dest->pos++] = value >> 32;
+    dest->bytes[dest->pos++] = value >> 40;
+    dest->bytes[dest->pos++] = value >> 48;
+    dest->bytes[dest->pos++] = value >> 56;
+
+    return S_OK;
+}

src/main/hook/iobuf.h

@@ -0,0 +1,40 @@
+#pragma once
+
+#include <windows.h>
+
+#include <stddef.h>
+#include <stdint.h>
+
+struct iobuf {
+    uint8_t *bytes;
+    size_t nbytes;
+    size_t pos;
+};
+
+struct const_iobuf {
+    const uint8_t *bytes;
+    size_t nbytes;
+    size_t pos;
+};
+
+void iobuf_flip(struct const_iobuf *child, struct iobuf *parent);
+size_t iobuf_move(struct iobuf *dest, struct const_iobuf *src);
+size_t iobuf_shift(struct iobuf *dest, struct iobuf *src);
+
+HRESULT iobuf_read(struct const_iobuf *src, void *bytes, size_t nbytes);
+HRESULT iobuf_read_8(struct const_iobuf *src, uint8_t *value);
+HRESULT iobuf_read_be16(struct const_iobuf *src, uint16_t *value);
+HRESULT iobuf_read_be32(struct const_iobuf *src, uint32_t *value);
+HRESULT iobuf_read_be64(struct const_iobuf *src, uint64_t *value);
+HRESULT iobuf_read_le16(struct const_iobuf *src, uint16_t *value);
+HRESULT iobuf_read_le32(struct const_iobuf *src, uint32_t *value);
+HRESULT iobuf_read_le64(struct const_iobuf *src, uint64_t *value);
+
+HRESULT iobuf_write(struct iobuf *dest, const void *bytes, size_t nbytes);
+HRESULT iobuf_write_8(struct iobuf *dest, uint8_t value);
+HRESULT iobuf_write_be16(struct iobuf *dest, uint16_t value);
+HRESULT iobuf_write_be32(struct iobuf *dest, uint32_t value);
+HRESULT iobuf_write_be64(struct iobuf *dest, uint64_t value);
+HRESULT iobuf_write_le16(struct iobuf *dest, uint16_t value);
+HRESULT iobuf_write_le32(struct iobuf *dest, uint32_t value);
+HRESULT iobuf_write_le64(struct iobuf *dest, uint64_t value);

src/main/hook/iohook.h

@@ -1,12 +1,11 @@
-#ifndef HOOK_IOHOOK_H
+#pragma once
-#define HOOK_IOHOOK_H
 
 #include <windows.h>
 
 #include <stddef.h>
 #include <stdint.h>
 
-#include "util/iobuf.h"
+#include "hook/iobuf.h"
 
 enum irp_op {
     IRP_OP_OPEN,
@@ -38,10 +37,14 @@ struct irp {
     uint64_t seek_pos;
 };
 
-typedef HRESULT (*irp_handler_t)(struct irp *irp);
+typedef HRESULT (*iohook_fn_t)(struct irp *irp);
-
-void iohook_init(const irp_handler_t *handlers, size_t nhandlers);
-HANDLE iohook_open_dummy_fd(void);
-HRESULT irp_invoke_next(struct irp *irp);
 
+HANDLE iohook_open_dummy_fd(void)
+#ifdef __GNUC__
+__attribute__((deprecated("Use iohook_open_nul_fd instead")))
 #endif
+;
+
+HRESULT iohook_open_nul_fd(HANDLE *fd);
+HRESULT iohook_push_handler(iohook_fn_t fn);
+HRESULT iohook_invoke_next(struct irp *irp);

src/main/hook/pe.c

@@ -1,20 +1,17 @@
 #include <windows.h>
 
+#include <assert.h>
 #include <stdbool.h>
 #include <stddef.h>
 #include <stdint.h>
+#include <stdlib.h>
+#include <string.h>
 
 #include "hook/pe.h"
 
-#include "util/log.h"
-
-typedef BOOL(WINAPI *dll_main_t)(HMODULE, uint32_t, void *);
-
-static const IMAGE_NT_HEADERS *pe_get_nt_header(HMODULE pe);
-static uint32_t
-pe_get_virtual_size(const IMAGE_SECTION_HEADER *sh, int nsections);
 static void *pe_offset(void *ptr, size_t off);
 static const void *pe_offsetc(const void *ptr, size_t off);
+static const IMAGE_NT_HEADERS *pe_get_nt_header(HMODULE pe);
 
 static void *pe_offset(void *ptr, size_t off)
 {
@@ -42,26 +39,6 @@ static const void *pe_offsetc(const void *ptr, size_t off)
     return base + off;
 }
 
-static uint32_t
-pe_get_virtual_size(const IMAGE_SECTION_HEADER *sh, int nsections)
-{
-    uint32_t sec_end;
-    uint32_t size;
-    int i;
-
-    size = 0;
-
-    for (i = 0; i < nsections; i++) {
-        sec_end = sh[i].VirtualAddress + sh[i].Misc.VirtualSize;
-
-        if (size < sec_end) {
-            size = sec_end;
-        }
-    }
-
-    return size;
-}
-
 static const IMAGE_NT_HEADERS *pe_get_nt_header(HMODULE pe)
 {
     const IMAGE_DOS_HEADER *dh;
@@ -76,13 +53,14 @@ static const IMAGE_NT_HEADERS *pe_get_nt_header(HMODULE pe)
 const pe_iid_t *pe_iid_get_first(HMODULE pe)
 {
     const IMAGE_NT_HEADERS *nth;
+    const IMAGE_DATA_DIRECTORY *idd;
     const IMAGE_IMPORT_DESCRIPTOR *iid;
 
+    assert(pe != NULL);
+
     nth = pe_get_nt_header(pe);
-    iid = pe_offsetc(
+    idd = &nth->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
-        pe,
+    iid = pe_offsetc(pe, idd->VirtualAddress);
-        nth->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT]
-            .VirtualAddress);
 
     if (iid == NULL || iid->Name == 0) {
         return NULL;
@@ -93,6 +71,9 @@ const pe_iid_t *pe_iid_get_first(HMODULE pe)
 
 const char *pe_iid_get_name(HMODULE pe, const pe_iid_t *iid)
 {
+    assert(pe != NULL);
+    assert(iid != NULL);
+
     return pe_offsetc(pe, iid->Name);
 }
 
@@ -100,6 +81,9 @@ const pe_iid_t *pe_iid_get_next(HMODULE pe, const pe_iid_t *iid)
 {
     const IMAGE_IMPORT_DESCRIPTOR *iid_next;
 
+    assert(pe != NULL);
+    assert(iid != NULL);
+
     iid_next = iid + 1;
 
     if (iid_next->Name != 0) {
@@ -109,27 +93,30 @@ const pe_iid_t *pe_iid_get_next(HMODULE pe, const pe_iid_t *iid)
     }
 }
 
-bool pe_iid_get_iat_entry(
+HRESULT pe_iid_get_iat_entry(
-    HMODULE pe, const pe_iid_t *iid, size_t n, struct pe_iat_entry *entry)
+        HMODULE pe,
+        const pe_iid_t *iid,
+        size_t n,
+        struct pe_iat_entry *entry)
 {
     const IMAGE_IMPORT_BY_NAME *import;
-    uintptr_t *import_rvas;
+    intptr_t *import_rvas;
     void **pointers;
 
-    if (iid->OriginalFirstThunk == 0) {
+    assert(pe != NULL);
-        log_warning("OriginalFirstThunk == 0");
+    assert(iid != NULL);
-    }
+    assert(entry != NULL);
 
     import_rvas = pe_offset(pe, iid->OriginalFirstThunk);
 
     if (import_rvas[n] == 0) {
         /* End of imports */
-        entry->name = NULL;
+        memset(entry, 0, sizeof(*entry));
-        entry->ordinal = 0;
-        entry->ppointer = NULL;
 
-        return false;
+        return S_FALSE;
-    } else if (import_rvas[n] & INTPTR_MIN) {
+    }
+
+    if (import_rvas[n] & INTPTR_MIN) {
         /* Ordinal import */
         entry->name = NULL;
         entry->ordinal = (uint16_t) import_rvas[n];
@@ -143,143 +130,43 @@ bool pe_iid_get_iat_entry(
     pointers = pe_offset(pe, iid->FirstThunk);
     entry->ppointer = &pointers[n];
 
-    return true;
+    return S_OK;
-}
-
-void pe_patch_pointer(void **ppointer, void *new_value)
-{
-    DWORD old_protect;
-
-    VirtualProtect(
-        ppointer, sizeof(void *), PAGE_EXECUTE_READWRITE, &old_protect);
-    *ppointer = new_value;
-    VirtualProtect(ppointer, sizeof(void *), old_protect, &old_protect);
-}
-
-HMODULE
-pe_explode(const uint8_t *bytes, uint32_t nbytes)
-{
-    HMODULE base;
-    const IMAGE_DOS_HEADER *dh;
-    const IMAGE_NT_HEADERS *nth;
-    const IMAGE_SECTION_HEADER *sh;
-    uint32_t virtual_size;
-    uint32_t vflags;
-    int i;
-
-    dh = (IMAGE_DOS_HEADER *) bytes;
-    nth = pe_offsetc(bytes, dh->e_lfanew);
-    sh = pe_offsetc(bytes, dh->e_lfanew + sizeof(*nth));
-
-    virtual_size = pe_get_virtual_size(sh, nth->FileHeader.NumberOfSections);
-    base = (HMODULE) VirtualAlloc(
-        (void *) nth->OptionalHeader.ImageBase,
-        virtual_size,
-        MEM_RESERVE,
-        PAGE_NOACCESS);
-
-    if (base == NULL) {
-        /* Try again, allowing any base address */
-        base = (HMODULE) VirtualAlloc(
-            NULL, virtual_size, MEM_RESERVE, PAGE_NOACCESS);
-
-        if (base == NULL) {
-            /* Aargh */
-            log_fatal(
-                "Failed to VirtualAlloc %#x bytes of address space",
-                virtual_size);
-        }
-    }
-
-    log_misc(
-        "Exploding PE, base %p actual %p",
-        (void *) nth->OptionalHeader.ImageBase,
-        base);
-
-    /* Commit header region */
-    VirtualAlloc(
-        (void *) base,
-        nth->OptionalHeader.SizeOfHeaders,
-        MEM_COMMIT,
-        PAGE_READWRITE);
-
-    memcpy(base, dh, sizeof(*dh));
-    memcpy(pe_offset(base, dh->e_lfanew), nth, sizeof(*nth));
-    memcpy(
-        pe_offset(base, dh->e_lfanew + sizeof(*nth)),
-        sh,
-        sizeof(*sh) * nth->FileHeader.NumberOfSections);
-
-    for (i = 0; i < nth->FileHeader.NumberOfSections; i++) {
-        vflags = sh[i].Characteristics & 0x20000000 ? PAGE_EXECUTE_READWRITE :
-                                                      PAGE_READWRITE;
-
-        VirtualAlloc(
-            pe_offset(base, sh[i].VirtualAddress),
-            sh[i].Misc.VirtualSize,
-            MEM_COMMIT,
-            vflags);
-
-        memcpy(
-            pe_offset(base, sh[i].VirtualAddress),
-            pe_offsetc(bytes, sh[i].PointerToRawData),
-            sh[i].SizeOfRawData);
-    }
-
-    return base;
-}
-
-void pe_relocate(HMODULE pe)
-{
-    const IMAGE_NT_HEADERS *nth;
-    const IMAGE_DATA_DIRECTORY *dde;
-    const IMAGE_BASE_RELOCATION *chunk;
-    intptr_t delta_va;
-    const uint16_t *reloc;
-    uintptr_t *addr_ptr;
-
-    nth = pe_get_nt_header(pe);
-    delta_va = (intptr_t) pe - nth->OptionalHeader.ImageBase;
-    dde = nth->OptionalHeader.DataDirectory + IMAGE_DIRECTORY_ENTRY_BASERELOC;
-
-    for (chunk = pe_offsetc(pe, dde->VirtualAddress);
-         (void *) chunk < pe_offsetc(pe, dde->VirtualAddress + dde->Size);
-         chunk = pe_offsetc(chunk, chunk->SizeOfBlock)) {
-        for (reloc = (uint16_t *) (chunk + 1);
-             (void *) reloc < pe_offsetc(chunk, chunk->SizeOfBlock);
-             reloc++) {
-            if (*reloc >> 12 == IMAGE_REL_BASED_HIGHLOW) {
-                addr_ptr =
-                    pe_offset(pe, chunk->VirtualAddress + (*reloc & 0x0FFF));
-                *addr_ptr += delta_va;
-            }
-        }
-    }
 }
 
 void *pe_get_export(HMODULE pe, const char *name, uint16_t ord)
 {
     const IMAGE_NT_HEADERS *nth;
+    const IMAGE_DATA_DIRECTORY *idd;
     const IMAGE_EXPORT_DIRECTORY *ied;
     const uint32_t *name_rvas;
     const uint32_t *target_rvas;
+    const char *name_va;
     DWORD i;
 
+    assert(pe != NULL);
+
     nth = pe_get_nt_header(pe);
-    ied = pe_offsetc(
+    idd = &nth->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
-        pe,
+    ied = pe_offsetc(pe, idd->VirtualAddress);
-        nth->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT]
-            .VirtualAddress);
 
     name_rvas = pe_offsetc(pe, ied->AddressOfNames);
     target_rvas = pe_offsetc(pe, ied->AddressOfFunctions);
 
     if (name != NULL) {
-        for (i = 0; i < ied->NumberOfNames; i++) {
+        for (i = 0 ; i < ied->NumberOfNames ; i++) {
-            if (name_rvas[i] != 0 &&
+            if (name_rvas[i] == 0) {
-                strcmp(pe_offsetc(pe, name_rvas[i]), name) == 0) {
+                /* Ordinal-only export, cannot match against this */
-                return pe_offset(pe, target_rvas[i]);
+                continue;
             }
+
+            name_va = pe_offsetc(pe, name_rvas[i]);
+
+            if (strcmp(name_va, name) != 0) {
+                /* Name did not match */
+                continue;
+            }
+
+            return pe_offset(pe, target_rvas[i]);
         }
 
         return NULL;
@@ -290,13 +177,46 @@ void *pe_get_export(HMODULE pe, const char *name, uint16_t ord)
     }
 }
 
-BOOL pe_call_dll_main(HMODULE pe, uint32_t reason, void *ctx)
+void *pe_get_entry_point(HMODULE pe)
 {
     const IMAGE_NT_HEADERS *nth;
-    dll_main_t dll_main;
+
+    assert(pe != NULL);
 
     nth = pe_get_nt_header(pe);
-    dll_main = pe_offset(pe, nth->OptionalHeader.AddressOfEntryPoint);
 
-    return dll_main(pe, reason, ctx);
+    return pe_offset(pe, nth->OptionalHeader.AddressOfEntryPoint);
+}
+
+HRESULT pe_patch(void *dest, const void *src, size_t nbytes)
+{
+    DWORD old_protect;
+    BOOL ok;
+
+    assert(dest != NULL);
+    assert(src != NULL);
+
+    ok = VirtualProtect(
+            dest,
+            nbytes,
+            PAGE_EXECUTE_READWRITE,
+            &old_protect);
+
+    if (!ok) {
+        return HRESULT_FROM_WIN32(GetLastError());
+    }
+
+    memcpy(dest, src, nbytes);
+
+    ok = VirtualProtect(
+            dest,
+            nbytes,
+            old_protect,
+            &old_protect);
+
+    if (!ok) {
+        return HRESULT_FROM_WIN32(GetLastError());
+    }
+
+    return S_OK;
 }

src/main/hook/pe.h

@@ -1,9 +1,9 @@
-#ifndef HOOK_PE_H
+#pragma once
-#define HOOK_PE_H
 
 #include <windows.h>
 
 #include <stdbool.h>
+#include <stddef.h>
 #include <stdint.h>
 
 typedef IMAGE_IMPORT_DESCRIPTOR pe_iid_t;
@@ -17,14 +17,11 @@ struct pe_iat_entry {
 const pe_iid_t *pe_iid_get_first(HMODULE pe);
 const char *pe_iid_get_name(HMODULE pe, const pe_iid_t *iid);
 const pe_iid_t *pe_iid_get_next(HMODULE pe, const pe_iid_t *iid);
-bool pe_iid_get_iat_entry(
+HRESULT pe_iid_get_iat_entry(
-    HMODULE pe, const pe_iid_t *iid, size_t n, struct pe_iat_entry *entry);
+        HMODULE pe,
+        const pe_iid_t *iid,
+        size_t n,
+        struct pe_iat_entry *entry);
 void *pe_get_export(HMODULE pe, const char *name, uint16_t ord);
-BOOL pe_call_dll_main(HMODULE pe, uint32_t reason, void *ctx);
+void *pe_get_entry_point(HMODULE pe);
-
+HRESULT pe_patch(void *dest, const void *src, size_t nbytes);
-void pe_patch_pointer(void **ppointer, void *new_value);
-
-HMODULE pe_explode(const uint8_t *bytes, uint32_t nbytes);
-void pe_relocate(HMODULE pe);
-
-#endif

src/main/hook/peb.c

@@ -1,14 +1,13 @@
 #include <windows.h>
 #include <winternl.h>
 
-#include "hook/peb.h"
+#include <assert.h>
 
-#include "util/defs.h"
+#include "hook/peb.h"
-#include "util/str.h"
 
 static const PEB *peb_get(void)
 {
-#ifdef __amd64
+#ifdef _M_AMD64
     return (const PEB *) __readgsqword(0x60);
 #else
     return (const PEB *) __readfsdword(0x30);
@@ -23,7 +22,7 @@ const peb_dll_t *peb_dll_get_first(void)
     peb = peb_get();
     node = peb->Ldr->InMemoryOrderModuleList.Flink;
 
-    return containerof(node, LDR_DATA_TABLE_ENTRY, InMemoryOrderLinks);
+    return CONTAINING_RECORD(node, LDR_DATA_TABLE_ENTRY, InMemoryOrderLinks);
 }
 
 const peb_dll_t *peb_dll_get_next(const peb_dll_t *dll)
@@ -31,6 +30,8 @@ const peb_dll_t *peb_dll_get_next(const peb_dll_t *dll)
     const PEB *peb;
     const LIST_ENTRY *node;
 
+    assert(dll != NULL);
+
     peb = peb_get();
     node = dll->InMemoryOrderLinks.Flink;
 
@@ -38,30 +39,12 @@ const peb_dll_t *peb_dll_get_next(const peb_dll_t *dll)
         return NULL;
     }
 
-    return containerof(node, LDR_DATA_TABLE_ENTRY, InMemoryOrderLinks);
+    return CONTAINING_RECORD(node, LDR_DATA_TABLE_ENTRY, InMemoryOrderLinks);
 }
 
-HMODULE
+HMODULE peb_dll_get_base(const peb_dll_t *dll)
-peb_dll_get_base(const peb_dll_t *dll)
 {
+    assert(dll != NULL);
+
     return dll->DllBase;
 }
-
-char *peb_dll_dup_name(const peb_dll_t *dll)
-{
-    const UNICODE_STRING *wstr;
-    char *name;
-    size_t i;
-
-    wstr = &dll->FullDllName;
-
-    for (i = wstr->Length / 2 - 1; i > 0; i--) {
-        if (wstr->Buffer[i] == L'\\') {
-            wstr_narrow(&wstr->Buffer[i + 1], &name);
-
-            return name;
-        }
-    }
-
-    return NULL;
-}

src/main/hook/peb.h

@@ -1,5 +1,4 @@
-#ifndef HOOK_PEB_H
+#pragma once
-#define HOOK_PEB_H
 
 #include <windows.h>
 #include <winternl.h>
@@ -10,5 +9,3 @@ const peb_dll_t *peb_dll_get_first(void);
 const peb_dll_t *peb_dll_get_next(const peb_dll_t *dll);
 HMODULE peb_dll_get_base(const peb_dll_t *dll);
 char *peb_dll_dup_name(const peb_dll_t *dll);
-
-#endif

src/main/hook/process.c

@@ -0,0 +1,196 @@
+#include <windows.h>
+#include <tlhelp32.h>
+
+#include <assert.h>
+#include <stdbool.h>
+#include <stdlib.h>
+
+#include "hook/hr.h"
+#include "hook/pe.h"
+#include "hook/process.h"
+
+static bool thread_match_startup(
+        const CONTEXT *ctx,
+        void *ntstart,
+        void *exe_entry)
+{
+#ifdef _M_AMD64
+    return  ctx->Rip == (DWORD64) ntstart &&
+            ctx->Rcx == (DWORD64) exe_entry;
+#else
+    return  ctx->Eip == (DWORD) ntstart &&
+            ctx->Eax == (DWORD) exe_entry;
+#endif
+}
+
+static void thread_patch_startup(
+        process_entry_t new_entry,
+        process_entry_t *orig_entry,
+        CONTEXT *ctx)
+{
+#ifdef _M_AMD64
+    *orig_entry = (void *) ctx->Rcx;
+    ctx->Rcx = (DWORD64) new_entry;
+#else
+    *orig_entry = (void *) ctx->Eax;
+    ctx->Eax = (DWORD) new_entry;
+#endif
+}
+
+static HRESULT process_hijack_try_thread(
+        process_entry_t new_entry,
+        process_entry_t *orig_entry,
+        DWORD thread_id)
+{
+    CONTEXT ctx;
+    HMODULE exe;
+    HMODULE ntdll;
+    void *exe_entry;
+    void *ntstart;
+    HANDLE thread;
+    HRESULT hr;
+    BOOL ok;
+
+    thread = NULL;
+
+    exe = GetModuleHandleW(NULL);
+
+    if (exe == NULL) {
+        /* uhhhh... */
+        hr = E_UNEXPECTED;
+
+        goto end;
+    }
+
+    ntdll = GetModuleHandleW(L"ntdll.dll");
+
+    if (ntdll == NULL) {
+        /* Another 2 + 2 = 5 situation */
+        hr = E_UNEXPECTED;
+
+        goto end;
+    }
+
+    exe_entry = pe_get_entry_point(exe);
+    ntstart = GetProcAddress(ntdll, "RtlUserThreadStart");
+
+    if (ntstart == NULL) {
+        /* TODO Deal with WinXP, for the poor souls still stuck on that OS.
+           XP starts threads at LdrInitializeThunk instead (I think) */
+        hr = E_NOTIMPL;
+
+        goto end;
+    }
+
+    thread = OpenThread(
+            THREAD_GET_CONTEXT | THREAD_SET_CONTEXT,
+            FALSE,
+            thread_id);
+
+    if (thread == NULL) {
+        hr = HRESULT_FROM_WIN32(GetLastError());
+
+        goto end;
+    }
+
+    memset(&ctx, 0, sizeof(ctx));
+#ifdef _M_AMD64
+    ctx.ContextFlags = CONTEXT_AMD64 | CONTEXT_FULL;
+#else
+    ctx.ContextFlags = CONTEXT_i386 | CONTEXT_FULL;
+#endif
+
+    ok = GetThreadContext(thread, &ctx);
+
+    if (!ok) {
+        hr = HRESULT_FROM_WIN32(GetLastError());
+
+        goto end;
+    }
+
+    if (thread_match_startup(&ctx, ntstart, exe_entry)) {
+        thread_patch_startup(new_entry, orig_entry, &ctx);
+        ok = SetThreadContext(thread, &ctx);
+
+        if (!ok) {
+            hr = HRESULT_FROM_WIN32(GetLastError());
+
+            goto end;
+        }
+
+        return S_OK;
+    } else {
+        return S_FALSE;
+    }
+
+end:
+    if (thread != NULL) {
+        CloseHandle(thread);
+    }
+
+    return hr;
+}
+
+HRESULT process_hijack_startup(
+        process_entry_t new_entry,
+        process_entry_t *orig_entry)
+{
+    THREADENTRY32 thread;
+    HANDLE snap;
+    DWORD pid;
+    HRESULT fault;
+    HRESULT hr;
+    BOOL ok;
+
+    assert(new_entry != NULL);
+    assert(orig_entry != NULL);
+
+    pid = GetCurrentProcessId();
+    snap = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
+
+    if (snap == INVALID_HANDLE_VALUE) {
+        hr = HRESULT_FROM_WIN32(GetLastError());
+
+        goto end;
+    }
+
+    thread.dwSize = sizeof(thread);
+    ok = Thread32First(snap, &thread);
+
+    if (!ok) {
+        hr = HRESULT_FROM_WIN32(GetLastError());
+
+        goto end;
+    }
+
+    /* Return this if we don't find anything suitable */
+    fault = E_FAIL;
+
+    do {
+        if (thread.th32OwnerProcessID != pid) {
+            continue;
+        }
+
+        hr = process_hijack_try_thread(
+                new_entry,
+                orig_entry,
+                thread.th32ThreadID);
+
+        if (hr == S_OK) {
+            /* Main thread successfully hijacked, finish up */
+            goto end;
+        } else if (FAILED(hr)) {
+            /* Latch this error code, but don't abort, keep trying. */
+            fault = hr;
+        }
+    } while (Thread32Next(snap, &thread));
+
+    hr = fault;
+
+end:
+    if (snap != INVALID_HANDLE_VALUE) {
+        CloseHandle(snap);
+    }
+
+    return hr;
+}

src/main/hook/process.h

@@ -0,0 +1,9 @@
+#pragma once
+
+#include <windows.h>
+
+typedef DWORD (CALLBACK *process_entry_t)(void);
+
+HRESULT process_hijack_startup(
+        process_entry_t new_entry,
+        process_entry_t *orig_entry);

src/main/hook/table.c

@@ -1,5 +1,6 @@
 #include <windows.h>
 
+#include <assert.h>
 #include <stddef.h>
 #include <stdlib.h>
 #include <string.h>
@@ -8,32 +9,44 @@
 #include "hook/peb.h"
 #include "hook/table.h"
 
-#include "util/mem.h"
+static const char apiset_prefix[] = "api-ms-win-core-";
+static const size_t apiset_prefix_len = sizeof(apiset_prefix) - 1;
 
 static void hook_table_apply_to_all(
-    const char *depname, const struct hook_symbol *syms, size_t nsyms);
+        const char *depname,
+        const struct hook_symbol *syms,
+        size_t nsyms);
 
 static void hook_table_apply_to_iid(
-    HMODULE target,
+        HMODULE target,
-    const pe_iid_t *iid,
+        const pe_iid_t *iid,
-    const struct hook_symbol *syms,
+        const struct hook_symbol *syms,
-    size_t nsyms);
+        size_t nsyms);
+
+static bool hook_table_match_module(
+        HMODULE target,
+        const char *iid_name,
+        const char *depname);
 
 static bool hook_table_match_proc(
-    const struct pe_iat_entry *iate, const struct hook_symbol *sym);
+        const struct pe_iat_entry *iate,
+        const struct hook_symbol *sym);
 
 static void hook_table_apply_to_all(
-    const char *depname, const struct hook_symbol *syms, size_t nsyms)
+        const char *depname,
+        const struct hook_symbol *syms,
+        size_t nsyms)
 {
     const peb_dll_t *dll;
     HMODULE pe;
 
-    for (dll = peb_dll_get_first(); dll != NULL; dll = peb_dll_get_next(dll)) {
+    for (   dll = peb_dll_get_first() ;
+            dll != NULL ;
+            dll = peb_dll_get_next(dll)) {
         pe = peb_dll_get_base(dll);
 
         if (pe == NULL) {
-            /* wtf? */
+            continue; /* ?? Happens sometimes. */
-            continue;
         }
 
         hook_table_apply(pe, depname, syms, nsyms);
@@ -41,25 +54,29 @@ static void hook_table_apply_to_all(
 }
 
 void hook_table_apply(
-    HMODULE target,
+        HMODULE target,
-    const char *depname,
+        const char *depname,
-    const struct hook_symbol *syms,
+        const struct hook_symbol *syms,
-    size_t nsyms)
+        size_t nsyms)
 {
     const pe_iid_t *iid;
     const char *iid_name;
 
+    assert(depname != NULL);
+    assert(syms != NULL || nsyms == 0);
+
     if (target == NULL) {
         /*  Call out, which will then call us back repeatedly. Awkward, but
             viewed from the outside it's good for usability. */
 
         hook_table_apply_to_all(depname, syms, nsyms);
     } else {
-        for (iid = pe_iid_get_first(target); iid != NULL;
+        for (   iid = pe_iid_get_first(target) ;
-             iid = pe_iid_get_next(target, iid)) {
+                iid != NULL ;
+                iid = pe_iid_get_next(target, iid)) {
             iid_name = pe_iid_get_name(target, iid);
 
-            if (_stricmp(iid_name, depname) == 0) {
+            if (hook_table_match_module(target, iid_name, depname)) {
                 hook_table_apply_to_iid(target, iid, syms, nsyms);
             }
         }
@@ -67,10 +84,10 @@ void hook_table_apply(
 }
 
 static void hook_table_apply_to_iid(
-    HMODULE target,
+        HMODULE target,
-    const pe_iid_t *iid,
+        const pe_iid_t *iid,
-    const struct hook_symbol *syms,
+        const struct hook_symbol *syms,
-    size_t nsyms)
+        size_t nsyms)
 {
     struct pe_iat_entry iate;
     size_t i;
@@ -79,8 +96,8 @@ static void hook_table_apply_to_iid(
 
     i = 0;
 
-    while (pe_iid_get_iat_entry(target, iid, i++, &iate)) {
+    while (pe_iid_get_iat_entry(target, iid, i++, &iate) == S_OK) {
-        for (j = 0; j < nsyms; j++) {
+        for (j = 0 ; j < nsyms ; j++) {
             sym = &syms[j];
 
             if (hook_table_match_proc(&iate, sym)) {
@@ -88,17 +105,70 @@ static void hook_table_apply_to_iid(
                     *sym->link = *iate.ppointer;
                 }
 
-                pe_patch_pointer(iate.ppointer, sym->patch);
+                pe_patch(iate.ppointer, &sym->patch, sizeof(sym->patch));
             }
         }
     }
 }
 
-static bool hook_table_match_proc(
+static bool hook_table_match_module(
-    const struct pe_iat_entry *iate, const struct hook_symbol *sym)
+        HMODULE target,
+        const char *iid_name,
+        const char *depname)
 {
-    if (sym->name != NULL && iate->name != NULL &&
+    HMODULE kernel32;
-        strcmp(sym->name, iate->name) == 0) {
+    int result;
+
+    /* OK, first do a straightforward match on the imported DLL name versus
+       the hook table DLL name. If it succeeds then we're done. */
+
+    result = _stricmp(iid_name, depname);
+
+    if (result == 0) {
+        return true;
+    }
+
+    /* If it failed then we have to check if this hook table targets kernel32.
+       We have to do some special processing around API sets in that case, so
+       stop here if kernel32 is not the subject of this hook table. */
+
+    if (_stricmp(depname, "kernel32.dll") != 0) {
+        return false;
+    }
+
+    /* There isn't really any good test for whether a DLL import wants a
+       concrete DLL or an abstract DLL providing a particular Windows API-set,
+       so we use a hacky check against the prefix. If the imported DLL name
+       looks like an apiset then we'll allow kernel32 hook tables to apply. */
+
+    result = _strnicmp(iid_name, apiset_prefix, apiset_prefix_len);
+
+    if (result != 0) {
+        return false;
+    }
+
+    /* ... EXCEPT for the case where we're hooking all DLLs loaded into the
+       process and we are currently examining kernel32 itself. In that case
+       there's some weird reference loops issues I don't entirely understand
+       right now. To avoid those, we just don't apply kernel32 hook tables to
+       kernel32 itself. */
+
+    kernel32 = GetModuleHandleW(L"kernel32.dll");
+
+    if (target == kernel32) {
+        return false;
+    }
+
+    return true;
+}
+
+static bool hook_table_match_proc(
+        const struct pe_iat_entry *iate,
+        const struct hook_symbol *sym)
+{
+    if (    sym->name != NULL &&
+            iate->name != NULL &&
+            strcmp(sym->name, iate->name) == 0) {
         return true;
     }
 

src/main/hook/table.h

@@ -1,5 +1,4 @@
-#ifndef HOOK_TABLE_H
+#pragma once
-#define HOOK_TABLE_H
 
 #include <windows.h>
 
@@ -14,9 +13,13 @@ struct hook_symbol {
 };
 
 void hook_table_apply(
-    HMODULE target,
+        HMODULE target,
-    const char *depname,
+        const char *depname,
-    const struct hook_symbol *syms,
+        const struct hook_symbol *syms,
-    size_t nsyms);
+        size_t nsyms);
 
-#endif
+void hook_table_revert(
+        HMODULE target,
+        const char *depname,
+        const struct hook_symbol *syms,
+        size_t nsyms);