python3_dll_redirect.cpp

//
// Having multiple Python environments in a single process is a problem.
// The code in this gist is not a complete solution but contains some
// ideas of how the problem might be worked around.
//
// Python extension modules link against python3.dll which redirects
// to the full python3x.dll module, but if multiple Python runtimes
// are loaded then when loading an extension module the python3.dll
// dependency resolves to the first instance of python3.dll that was
// loaded, which may not be the correct one.
//
// To attempt to solve this we can redirect the incorrect python3.dll
// to the correct one before doing any Python imports.
//
// The idea works as follows:
//
// 1. Get the Ldr data table entry for the python3.dll we want our
//    import to use.
//
//    This is loaded as part of initialising Python, or
//    we can find it from the path and load it using LoadLibraryEx with
//    an altered search path.
//
//    The function FindLoaderEntry does this for us once we have the
//    module handle.
//
// 2. Replace all other entries in the Ldr table where the dll base
//    name is the same as our target one (i.e. the other python3.dll).
//
//    When import calls LoadLibrary to load the c extension the DLL
//    resolution will use the updated entry and python3.dll will be
//    resolved to the one we want.
//
//    This is done by UpdateLoaderTable.
//
// 3. Import the module as normal.
//
// 4. Restore the Ldr table using RestoreLoaderTable.
//
//    This may not strictly be necessary if the other Python
//    environment is also patching the Ldr table, but we cannot
//    be sure of that so better to restore it.
//
//
// NOTES:
//
// This technique relies on various undocumented features of Windows
// and messing with the Ldr table may be detected by anti-malware
// or other end point security software.
//
// It doesn't solve the problem of other DLLs that may conflict
// between Python environments. For example, other pyd files
// named the same (same version of Python) that have already been
// loaded and resolved by the time the second interpreter is
// started.
//
#include <Python.h>
#include <winternl.h>

// These structs are not in the Windows headers
struct _PEB_LDR_DATA
{
    ULONG Length;
    BOOLEAN Initialized;
    PVOID SsHandle;
    LIST_ENTRY InLoadOrderModuleList;
    LIST_ENTRY InMemoryOrderModuleList;
    LIST_ENTRY InInitializationOrderModuleList;
};

struct _LDR_DATA_TABLE_ENTRY
{
    LIST_ENTRY InLoadOrderLinks;
    LIST_ENTRY InMemoryOrderLinks;
    LIST_ENTRY InInitializationOrderLinks;
    PVOID DllBase;
    PVOID EntryPoint;
    ULONG SizeOfImage;
    UNICODE_STRING FullDllName;
    UNICODE_STRING BaseDllName;
};

// Function pointers for the undocumented Loader Lock functions
HMODULE gNtDLL = NULL;
NTSTATUS (WINAPI *_LdrLockLoaderLock)(ULONG Flags, PULONG State, PULONG_PTR Cookie) = NULL;
NTSTATUS (WINAPI *_LdrUnlockLoaderLock)(ULONG Flags, ULONG_PTR Cookie) = NULL;


bool LdrLockInit()
{
    if (NULL == _LdrLockLoaderLock || NULL == _LdrUnlockLoaderLock)
    {
        gNtDLL = LoadLibraryA("ntdll.dll");
        if (NULL == gNtDLL)
            return false;

        *((FARPROC*)&_LdrLockLoaderLock) = GetProcAddress(gNtDLL, "LdrLockLoaderLock");
        *((FARPROC*)&_LdrUnlockLoaderLock) = GetProcAddress(gNtDLL, "LdrUnlockLoaderLock");
    }

    return NULL != _LdrLockLoaderLock && NULL != _LdrUnlockLoaderLock;
}


NTSTATUS LdrLockLoaderLock(ULONG Flags, PULONG State, PULONG_PTR Cookie)
{
    if (!LdrLockInit())
        return -1;
    return _LdrLockLoaderLock(Flags, State, Cookie);
}


NTSTATUS LdrUnlockLoaderLock(ULONG Flags, ULONG_PTR Cookie)
{
    if (!LdrLockInit())
        return -1;
    return _LdrUnlockLoaderLock(Flags, Cookie);
}


// Find an entry in the DLL loader table for a module
_LDR_DATA_TABLE_ENTRY* FindLoaderEntry(HMODULE module)
{
    // Release the GIL before acquiring the loader lock
    Py_BEGIN_ALLOW_THREADS

    // Acquire the loader lock before iterating over the loader data
    ULONG_PTR cookie(0);
    LdrLockLoaderLock(0, NULL, &cookie);

#ifdef _WIN64
    PPEB peb = reinterpret_cast<PPEB>(__readgsqword(0x60));
#else
    PPEB peb = reinterpret_cast<PPEB>(__readfsdword(0x30));
#endif

    // Look through the loaded modules to find the entry relating to our specific python3.dll
    _PEB_LDR_DATA* data = reinterpret_cast<_PEB_LDR_DATA*>(peb->Ldr);
    _LDR_DATA_TABLE_ENTRY *found = NULL;
    LIST_ENTRY* first = data->InLoadOrderModuleList.Flink, *link = first;
    do
    {
        _LDR_DATA_TABLE_ENTRY* entry = reinterpret_cast<_LDR_DATA_TABLE_ENTRY*>(link);
        if (entry->DllBase == module)
        {
            found = entry;
            break;
        }

        link = link->Flink;
    }
    while (link != first);

    // Release the loader lock
    if (cookie)
        LdrUnlockLoaderLock(0, cookie);

    Py_END_ALLOW_THREADS

    return found;
}


// Update the loader table to redirect to our module
std::map<_LDR_DATA_TABLE_ENTRY*, _LDR_DATA_TABLE_ENTRY>
UpdateLoaderTable(_LDR_DATA_TABLE_ENTRY* entry)
{
    std::map<_LDR_DATA_TABLE_ENTRY*, _LDR_DATA_TABLE_ENTRY> updates;

    if (NULL == entry || !entry->EntryPoint)
        return updates;

    // Release the GIL before acquiring the loader lock
    Py_BEGIN_ALLOW_THREADS

    // Acquire the loader lock before updating the loader data
    ULONG_PTR cookie(0);
    LdrLockLoaderLock(0, NULL, &cookie);

#ifdef _WIN64
    PPEB peb = reinterpret_cast<PPEB>(__readgsqword(0x60));
#else
    PPEB peb = reinterpret_cast<PPEB>(__readfsdword(0x30));
#endif

    // Redirect all DLLs with the same base name to our entry
    _PEB_LDR_DATA* data = reinterpret_cast<_PEB_LDR_DATA*>(peb->Ldr);
    LIST_ENTRY* first = data->InLoadOrderModuleList.Flink, *link = first;
    do
    {
        _LDR_DATA_TABLE_ENTRY* current = reinterpret_cast<_LDR_DATA_TABLE_ENTRY*>(link);
        if (current->BaseDllName.Length == entry->BaseDllName.Length
            && current->DllBase != entry->DllBase
            && 0 == wcsnicmp(current->BaseDllName.Buffer,
                             entry->BaseDllName.Buffer,
                             entry->BaseDllName.Length))
        {
            // Add this entry to the list of updates before modifying it
            updates.insert(std::make_pair(current, *current));

            // Update this one to match the entry we want to use
            current->DllBase = entry->DllBase;
            current->EntryPoint = entry->EntryPoint;
            current->SizeOfImage = entry->SizeOfImage;
        }

        link = link->Flink;
    }
    while (link != first);

    // Release the loader lock
    if (cookie)
        LdrUnlockLoaderLock(0, cookie);

    Py_END_ALLOW_THREADS

    return updates;
}


void RestoreLoaderTable(const std::map<_LDR_DATA_TABLE_ENTRY*, _LDR_DATA_TABLE_ENTRY>& updates)
{
    if (updates.empty())
        return;

    // Release the GIL before acquiring the loader lock
    Py_BEGIN_ALLOW_THREADS

    // Acquire the loader lock before updating the loader data
    ULONG_PTR cookie(0);
    LdrLockLoaderLock(0, NULL, &cookie);

#ifdef _WIN64
    PPEB peb = reinterpret_cast<PPEB>(__readgsqword(0x60));
#else
    PPEB peb = reinterpret_cast<PPEB>(__readfsdword(0x30));
#endif

    _PEB_LDR_DATA* data = reinterpret_cast<_PEB_LDR_DATA*>(peb->Ldr);
    LIST_ENTRY* first = data->InLoadOrderModuleList.Flink, *link = first;
    do
    {
        _LDR_DATA_TABLE_ENTRY* entry = reinterpret_cast<_LDR_DATA_TABLE_ENTRY*>(link);
        auto found = updates.find(entry);
        if (found != updates.end())
        {
            // Restore this entry back to how it was
            entry->DllBase = found->second.DllBase;
            entry->EntryPoint = found->second.EntryPoint;
            entry->SizeOfImage = found->second.SizeOfImage;
        }
        link = link->Flink;
    }
    while (link != first);

    // Release the loader lock
    if (cookie)
        LdrUnlockLoaderLock(0, cookie);

    Py_END_ALLOW_THREADS
}