penter.c

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <windows.h>

#define BE(_0,_1,_2,_3,_4,_5,_6,_7) ( \
    (0x ## _7 ## ull) << 56 | \
    (0x ## _6 ## ull) << 48 | \
    (0x ## _5 ## ull) << 40 | \
    (0x ## _4 ## ull) << 32 | \
    (0x ## _3 ## ull) << 24 | \
    (0x ## _2 ## ull) << 16 | \
    (0x ## _1 ## ull) <<  8 | \
    (0x ## _0 ## ull) <<  0   \
)

static void my_func(void*);
static ULONG my_custom_tls_index = 0xFFFFFFFF;

static void lib_init(void) {
    my_custom_tls_index = TlsAlloc();
    TlsSetValue(my_custom_tls_index, NULL);
}

// TODO(NeGate): we could avoid using rcx to cut down on registers and eventually
// move all the pushes (except for RAX) into deeper pieces of the code which aren't run
// as often:
//   mov ecx, [rax]
//   mov rax, gs:[0x30]
//   lea rcx,[rax+rcx*8+0x1480]
//
//  vs
//
//   mov eax, [rax]
//   lea rax, [rax*8+0x1480]
//   add rax, gs:[0x30]
#pragma code_seg(".text")
__declspec(allocate(".text")) uint64_t _penter[] = {
    BE(
        90,90,90,90,90,   // 5 heckin nops
        50,               // push rax
        48,B8             // mov rax, my_custom_tls_index
    ),
    (uint64_t) &my_custom_tls_index, // abs64 relocation lmao
    BE(
        83,38,FF,         // cmp rax, 0xFFFFFFFF
        75,0B,            // jne DONT_EXIT
        0F,1F,00          // 2-byte nop
    ),
    BE(
        58,               // pop rax
        C3,               // ret
        66,0F,1F,44,00,00 // 6-byte nop
    ),
    // DONT_EXIT:
    BE(
        51,               // push rcx
        52,               // push rdx
        49,50,            // push r8
        49,51,            // push r9
        49,52             // push r10
    ),
    BE(
        49,53,            // push r11
        48,83,EC,20,      // sub rsp, 0x20
        8B,08             // mov ecx, [rax]
    ),
    // check if we're already inside of penter, if so then don't call again
    BE(
        65,48,8B,04,25,30,00,00),BE(00, // mov rax, gs:[0x30]
        0F,1F,80,00,00,00,00
    ),
    BE(
        48,8D,8C,C8,80,14,00,00 // lea rcx,[rax+rcx*8+0x1480]
    ),
    BE(
        83,39,00,         // cmp dword [rcx], 0
        75,31,            // jne RET (aka 49 bytes ahead)
        0F,1F,00          // nop
    ),
    BE(
        C7,01,01,00,00,00,// mov dword [rcx], 1
        48,B8             // mov rax, my_func
    ),
    (uint64_t) my_func,   // abs64 relocation lmao
    BE(
        51,               // push rcx
        66,90,            // nop

        // return address
        48,8B,4C,24,60    // mov rcx, [rsp+0x58]
    ),
    BE(
        48,83,EC,08,      // sub rsp, 0x8
        FF,D0,            // call rax
        66,90             // nop
    ),
    BE(
        48,83,C4,08,      // add rsp, 0x8
        59,               // pop rcx
        0F,1F,00          // nop
    ),
    BE(
        C7,01,00,00,00,00,// mov dword [rcx], 0
        // RET:
        66,90             // nop
    ),
    BE(
        48,83,C4,20,      // add rsp, 0x20
        49,5B,            // pop r11
        49,5A             // pop r10
    ),
    BE(
        49,59,            // pop r9
        49,58,            // pop r8
        5A,               // pop rdx
        59,               // pop rcx
        58,               // pop rax
        C3                // ret
    ),
};

__declspec(allocate(".text")) uint64_t _pexit[] = {
    BE(
        90,90,90,90,90,   // 5 heckin nops
        50,               // push rax
        48,B8             // mov rax, my_custom_tls_index
    ),
    (uint64_t) &my_custom_tls_index, // abs64 relocation lmao
    BE(
        83,38,FF,         // cmp rax, 0xFFFFFFFF
        75,0B,            // jne DONT_EXIT
        0F,1F,00          // 2-byte nop
    ),
    BE(
        58,               // pop rax
        C3,               // ret
        66,0F,1F,44,00,00 // 6-byte nop
    ),
    // DONT_EXIT:
    BE(
        51,               // push rcx
        52,               // push rdx
        49,50,            // push r8
        49,51,            // push r9
        49,52             // push r10
    ),
    BE(
        49,53,            // push r11
        48,83,EC,20,      // sub rsp, 0x20
        8B,08             // mov ecx, [rax]
    ),
    // check if we're already inside of penter, if so then don't call again
    BE(
        65,48,8B,04,25,30,00,00),BE(00, // mov rax, gs:[0x30]
        0F,1F,80,00,00,00,00
    ),
    BE(
        48,8D,8C,C8,80,14,00,00 // lea rcx,[rax+rcx*8+0x1480]
    ),
    BE(
        83,39,00,         // cmp dword [rcx], 0
        75,31,            // jne RET (aka 49 bytes ahead)
        0F,1F,00          // nop
    ),
    BE(
        C7,01,01,00,00,00,// mov dword [rcx], 1
        48,B8             // mov rax, my_func
    ),
    (uint64_t) my_func,   // abs64 relocation lmao
    BE(
        51,               // push rcx
        31,C9,            // xor ecx,ecx
        0F,1F,44,00,00    // nop
    ),
    BE(
        48,83,EC,08,      // sub rsp, 0x8
        FF,D0,            // call rax
        66,90             // nop
    ),
    BE(
        48,83,C4,08,      // add rsp, 0x8
        59,               // pop rcx
        0F,1F,00          // nop
    ),
    BE(
        C7,01,00,00,00,00,// mov dword [rcx], 0
        // RET:
        66,90             // nop
    ),
    BE(
        48,83,C4,20,      // add rsp, 0x20
        49,5B,            // pop r11
        49,5A             // pop r10
    ),
    BE(
        49,59,            // pop r9
        49,58,            // pop r8
        5A,               // pop rdx
        59,               // pop rcx
        58,               // pop rax
        C3                // ret
    ),
};

static void my_func(void* ret_addr) {
    if (ret_addr) {
        printf("Hello, World! %p\n", ret_addr);
    } else {
        printf("Bye, World!\n");
    }

    ((void(*)()) _penter)();
    ((void(*)()) _penter)();
    ((void(*)()) _penter)();
    ((void(*)()) _penter)();
}

int main(void) {
    int a = GetCurrentThreadId();

    ((void(*)()) _penter)();
    lib_init();
    ((void(*)()) _penter)();
    ((void(*)()) _pexit)();
    return 0;
}