dll loader

#include <windows.h>
#include <winnt.h>
#include <stdio.h>

typedef BOOL(WINAPI *DllMainFunc)(HINSTANCE, DWORD, LPVOID);

// New helpers
BOOL ProcessTLS(LPVOID imageBase);
BOOL ApplySectionProtections(LPVOID imageBase);
BOOL ProcessExceptionDirectory(LPVOID imageBase);

LPVOID LoadAndMapDLL(const char *dllPath);
BOOL ApplyRelocations(LPVOID imageBase, ULONGLONG delta);
BOOL ResolveImports(LPVOID imageBase);

int main(int argc, char *argv[])
{
    if (argc != 2)
    {
        printf("Usage: %s <path_to_dll>\n", argv[0]);
        return 1;
    }

    printf("Mapping\n");
    LPVOID imageBase = LoadAndMapDLL(argv[1]);
    if (!imageBase)
    {
        printf("[-] Failed to load DLL.\n");
        return 1;
    }

    PIMAGE_DOS_HEADER dos = (PIMAGE_DOS_HEADER)imageBase;
    PIMAGE_NT_HEADERS nt = (PIMAGE_NT_HEADERS)((BYTE *)imageBase + dos->e_lfanew);
    ULONGLONG delta = (ULONGLONG)imageBase - nt->OptionalHeader.ImageBase;

    printf("Relocations\n");
    if (!ApplyRelocations(imageBase, delta))
    {
        printf("[-] Failed to apply relocations.\n");
        VirtualFree(imageBase, 0, MEM_RELEASE);
        return 1;
    }

    printf("Imports\n");
    if (!ResolveImports(imageBase))
    {
        printf("[-] Failed to resolve imports.\n");
        VirtualFree(imageBase, 0, MEM_RELEASE);
        return 1;
    }

    printf("TLS\n");
    if (!ProcessTLS(imageBase))
    {
        printf("[-] Failed to process TLS.\n");
        VirtualFree(imageBase, 0, MEM_RELEASE);
        return 1;
    }

    printf("Sec protection\n");
    if (!ApplySectionProtections(imageBase))
    {
        printf("[-] Failed to apply section protections.\n");
        // Non‐fatal in many cases, so we continue
    }

    printf("Exception Dir\n");
    if (!ProcessExceptionDirectory(imageBase))
    {
        printf("[-] Warning: failed to register exception handlers.\n");
    }

    // Call DllMain
    printf("DLL Main\n");
    DllMainFunc entry = (DllMainFunc)((BYTE *)imageBase + nt->OptionalHeader.AddressOfEntryPoint);
    if (!entry((HINSTANCE)imageBase, DLL_PROCESS_ATTACH, NULL))
    {
        printf("[-] DllMain returned FALSE.\n");
        VirtualFree(imageBase, 0, MEM_RELEASE);
        return 1;
    }

    printf("[+] DLL loaded, TLS callbacks run, and DllMain called successfully.\n");
    return 0;
}

LPVOID LoadAndMapDLL(const char *dllPath)
{
    HANDLE hFile = CreateFileA(dllPath, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
    if (hFile == INVALID_HANDLE_VALUE)
        return NULL;

    HANDLE hMapping = CreateFileMappingA(hFile, NULL, PAGE_READONLY, 0, 0, NULL);
    if (!hMapping)
    {
        CloseHandle(hFile);
        return NULL;
    }

    BYTE *fileData = (BYTE *)MapViewOfFile(hMapping, FILE_MAP_READ, 0, 0, 0);
    if (!fileData)
    {
        CloseHandle(hMapping);
        CloseHandle(hFile);
        return NULL;
    }

    PIMAGE_DOS_HEADER dos = (PIMAGE_DOS_HEADER)fileData;
    if (dos->e_magic != IMAGE_DOS_SIGNATURE)
        goto cleanup;

    PIMAGE_NT_HEADERS nt = (PIMAGE_NT_HEADERS)(fileData + dos->e_lfanew);
    if (nt->Signature != IMAGE_NT_SIGNATURE)
        goto cleanup;

    SIZE_T imageSize = nt->OptionalHeader.SizeOfImage;
    LPVOID imageBase = VirtualAlloc(NULL, imageSize, MEM_RESERVE | MEM_COMMIT, PAGE_EXECUTE_READWRITE);
    if (!imageBase)
        goto cleanup;

    // Copy headers
    memcpy(imageBase, fileData, nt->OptionalHeader.SizeOfHeaders);

    // Copy sections
    PIMAGE_SECTION_HEADER section = IMAGE_FIRST_SECTION(nt);
    for (int i = 0; i < nt->FileHeader.NumberOfSections; i++, section++)
    {
        LPVOID dest = (BYTE *)imageBase + section->VirtualAddress;
        LPVOID src = fileData + section->PointerToRawData;
        memcpy(dest, src, section->SizeOfRawData);
    }

    UnmapViewOfFile(fileData);
    CloseHandle(hMapping);
    CloseHandle(hFile);
    return imageBase;

cleanup:
    UnmapViewOfFile(fileData);
    CloseHandle(hMapping);
    CloseHandle(hFile);
    return NULL;
}

BOOL ApplyRelocations(LPVOID imageBase, ULONGLONG delta)
{
    if (delta == 0)
        return TRUE;

    PIMAGE_DOS_HEADER dos = (PIMAGE_DOS_HEADER)imageBase;
    PIMAGE_NT_HEADERS nt = (PIMAGE_NT_HEADERS)((BYTE *)imageBase + dos->e_lfanew);

    DWORD relocRVA = nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress;
    DWORD relocSize = nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].Size;
    if (!relocRVA || !relocSize)
        return FALSE;

    PIMAGE_BASE_RELOCATION reloc = (PIMAGE_BASE_RELOCATION)((BYTE *)imageBase + relocRVA);
    DWORD processed = 0;

    while (processed < relocSize)
    {
        DWORD count = (reloc->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(WORD);
        WORD *list = (WORD *)((BYTE *)reloc + sizeof(IMAGE_BASE_RELOCATION));

        for (DWORD i = 0; i < count; i++)
        {
            WORD entry = list[i];
            if ((entry >> 12) == IMAGE_REL_BASED_DIR64)
            {
                ULONGLONG *patch = (ULONGLONG *)((BYTE *)imageBase + reloc->VirtualAddress + (entry & 0x0FFF));
                *patch += delta;
            }
        }
        processed += reloc->SizeOfBlock;
        reloc = (PIMAGE_BASE_RELOCATION)((BYTE *)reloc + reloc->SizeOfBlock);
    }
    return TRUE;
}

BOOL ResolveImports(LPVOID imageBase)
{
    PIMAGE_DOS_HEADER dos = (PIMAGE_DOS_HEADER)imageBase;
    PIMAGE_NT_HEADERS nt = (PIMAGE_NT_HEADERS)((BYTE *)imageBase + dos->e_lfanew);

    DWORD importRVA = nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
    if (!importRVA)
        return TRUE;

    PIMAGE_IMPORT_DESCRIPTOR desc = (PIMAGE_IMPORT_DESCRIPTOR)((BYTE *)imageBase + importRVA);
    while (desc->Name)
    {
        char *dllName = (char *)((BYTE *)imageBase + desc->Name);
        HMODULE hMod = LoadLibraryA(dllName);
        if (!hMod)
            return FALSE;

        PIMAGE_THUNK_DATA orig = (PIMAGE_THUNK_DATA)((BYTE *)imageBase + desc->OriginalFirstThunk);
        PIMAGE_THUNK_DATA thunk = (PIMAGE_THUNK_DATA)((BYTE *)imageBase + desc->FirstThunk);

        for (; orig->u1.AddressOfData; orig++, thunk++)
        {
            FARPROC fn;
            if (orig->u1.Ordinal & IMAGE_ORDINAL_FLAG)
            {
                fn = GetProcAddress(hMod, (LPCSTR)(orig->u1.Ordinal & 0xFFFF));
            }
            else
            {
                PIMAGE_IMPORT_BY_NAME ibn = (PIMAGE_IMPORT_BY_NAME)((BYTE *)imageBase + orig->u1.AddressOfData);
                fn = GetProcAddress(hMod, ibn->Name);
            }
            if (!fn)
                return FALSE;
            thunk->u1.Function = (ULONGLONG)fn;
        }
        desc++;
    }
    return TRUE;
}

// — TLS support —//
BOOL ProcessTLS(LPVOID imageBase)
{
    PIMAGE_DOS_HEADER dos = (PIMAGE_DOS_HEADER)imageBase;
    PIMAGE_NT_HEADERS nt = (PIMAGE_NT_HEADERS)((BYTE *)imageBase + dos->e_lfanew);
    IMAGE_DATA_DIRECTORY dir = nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_TLS];
    if (!dir.VirtualAddress)
        return TRUE; // no TLS

    PIMAGE_TLS_DIRECTORY64 tls = (PIMAGE_TLS_DIRECTORY64)((BYTE *)imageBase + dir.VirtualAddress);

    // Copy raw TLS template to the newly-allocated TLS area
    PVOID tlsDataStart = (BYTE *)imageBase + (SIZE_T)tls->StartAddressOfRawData - nt->OptionalHeader.ImageBase;
    SIZE_T tlsSize = tls->EndAddressOfRawData - tls->StartAddressOfRawData;
    memcpy((PVOID)tls->StartAddressOfRawData, tlsDataStart, tlsSize);

    // Run all callbacks with PROCESS_ATTACH
    PIMAGE_TLS_CALLBACK *cb = (PIMAGE_TLS_CALLBACK *)tls->AddressOfCallBacks;
    for (; cb && *cb; cb++)
    {
        (*cb)((PVOID)imageBase, DLL_PROCESS_ATTACH, NULL);
    }
    return TRUE;
}

// — Section memory protections —//
BOOL ApplySectionProtections(LPVOID imageBase)
{
    PIMAGE_DOS_HEADER dos = (PIMAGE_DOS_HEADER)imageBase;
    PIMAGE_NT_HEADERS nt = (PIMAGE_NT_HEADERS)((BYTE *)imageBase + dos->e_lfanew);
    PIMAGE_SECTION_HEADER sec = IMAGE_FIRST_SECTION(nt);

    for (int i = 0; i < nt->FileHeader.NumberOfSections; i++, sec++)
    {
        DWORD oldProt, newProt = 0;
        BOOL exec = (sec->Characteristics & IMAGE_SCN_MEM_EXECUTE) != 0;
        BOOL read = (sec->Characteristics & IMAGE_SCN_MEM_READ) != 0;
        BOOL write = (sec->Characteristics & IMAGE_SCN_MEM_WRITE) != 0;

        if (exec)
        {
            if (write)
                newProt = PAGE_EXECUTE_READWRITE;
            else if (read)
                newProt = PAGE_EXECUTE_READ;
            else
                newProt = PAGE_EXECUTE;
        }
        else
        {
            if (write)
                newProt = PAGE_READWRITE;
            else if (read)
                newProt = PAGE_READONLY;
            else
                newProt = PAGE_NOACCESS;
        }

        VirtualProtect((BYTE *)imageBase + sec->VirtualAddress,
                       sec->Misc.VirtualSize, newProt, &oldProt);
    }
    return TRUE;
}

// — Exception directory (64-bit only) —//
BOOL ProcessExceptionDirectory(LPVOID imageBase)
{
#ifdef _WIN64
    PIMAGE_DOS_HEADER dos = (PIMAGE_DOS_HEADER)imageBase;
    PIMAGE_NT_HEADERS64 nt = (PIMAGE_NT_HEADERS64)((BYTE *)imageBase + dos->e_lfanew);
    IMAGE_DATA_DIRECTORY dir = nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION];
    if (!dir.VirtualAddress || !dir.Size)
        return TRUE;

    // Register the runtime function table for SEH
    RUNTIME_FUNCTION *table = (RUNTIME_FUNCTION *)((BYTE *)imageBase + dir.VirtualAddress);
    ULONG count = dir.Size / sizeof(RUNTIME_FUNCTION);
    if (!RtlAddFunctionTable(table, count, (DWORD64)imageBase))
        return FALSE;
#endif
    return TRUE;
}