Last active
February 25, 2022 10:40
-
-
Save tonyroberts/74888762f0063238d4f7fd7c7d36f0f0 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
// | |
// 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 | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment