Skip to content

Commit

Permalink
Bug，GetProcAddress时发生死递归
Browse files Browse the repository at this point in the history
  • Loading branch information
@mingkuang-Chuyu
mingkuang-Chuyu committed Jun 14, 2024
1 parent fc95f8e commit 98f0213
Show file tree
Hide file tree
Showing 4 changed files with 202 additions and 98 deletions.
170 changes: 170 additions & 0 deletions src/Thunks/YY_Thunks.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -638,6 +638,28 @@ namespace YY::Thunks::internal
}


/// <summary>
/// 计算字符串长度，某些场景我们特意不依赖wcslen之类的，防止发生死锁。
/// </summary>
/// <typeparam name="Char"></typeparam>
/// <param name="_szString"></param>
/// <param name="_cchMaxLength"></param>
/// <returns></returns>
template<typename Char>
constexpr size_t StringLength(_In_z_ const Char* _szString, size_t _cchMaxLength = -1)
{
if (!_szString)
return 0;

size_t _cchString = 0;
for (;_cchMaxLength && *_szString;--_cchMaxLength, ++_szString)
{
++_cchString;
}

return _cchString;
}

template<typename Char>
class StringBuffer
{
Expand Down Expand Up @@ -927,6 +949,21 @@ namespace YY::Thunks::internal
return memcmp(_Left.Buffer, _Right.Buffer, _Left.Length) == 0;
}

static constexpr UNICODE_STRING __fastcall MakeNtString(_In_z_ const wchar_t* _szString)
{
const auto _cbString = StringLength(_szString) * sizeof(_szString[0]);
UNICODE_STRING _Result = { (USHORT)max(UINT16_MAX, _cbString), (USHORT)max(UINT16_MAX, _cbString + sizeof(_szString[0])), const_cast<PWSTR>(_szString) };
return _Result;
}

static constexpr ANSI_STRING __fastcall MakeNtString(_In_z_ const char* _szString)
{
const auto _cbString = StringLength(_szString) * sizeof(_szString[0]);

ANSI_STRING _Result = { (USHORT)max(UINT16_MAX, _cbString), (USHORT)max(UINT16_MAX, _cbString + sizeof(_szString[0])), const_cast<PSTR>(_szString)};
return _Result;
}

template<size_t kLength>
static constexpr UNICODE_STRING __fastcall MakeStaticUnicodeString(const wchar_t (&_Right)[kLength])
{
Expand Down Expand Up @@ -976,3 +1013,136 @@ __if_exists(YY::Thunks::Fallback::_CRT_CONCATENATE(try_get_, _FUNCTION))

#undef __DEFINE_THUNK
#undef YY_Thunks_Implemented

static HMODULE __fastcall try_get_module(volatile HMODULE* pModule, const wchar_t* module_name, int Flags) noexcept
{
// First check to see if we've cached the module handle:
if (HMODULE const cached_handle = __crt_interlocked_read_pointer(pModule))
{
if (cached_handle == INVALID_HANDLE_VALUE)
{
return nullptr;
}

return cached_handle;
}

// If we haven't yet cached the module handle, try to load the library. If
// this fails, cache the sentinel handle value INVALID_HANDLE_VALUE so that
// we don't attempt to load the module again:
HMODULE new_handle = NULL;
if (Flags & USING_GET_MODULE_HANDLE)
{
new_handle = GetModuleHandleW(module_name);
}
else
{
// 我们不能直接使用 LoadLibraryExW，因为它可能被Thunk。
__if_exists(YY::Thunks::try_get_LoadLibraryExW)
{
const auto LoadLibraryExW = YY::Thunks::try_get_LoadLibraryExW();
if (!LoadLibraryExW)
return nullptr;
}

if (Flags & LOAD_AS_DATA_FILE)
{
new_handle = LoadLibraryExW(module_name, NULL, LOAD_LIBRARY_AS_DATAFILE);
}
else if (Flags & USING_UNSAFE_LOAD)
{
new_handle = LoadLibraryExW(module_name, nullptr, 0);
}
else
{
// 使用DLL安全加载
#if (YY_Thunks_Support_Version < NTDDI_WIN8)
if (!try_get_AddDllDirectory())
{
#if !defined(__USING_NTDLL_LIB)
const auto LdrLoadDll = try_get_LdrLoadDll();
if (!LdrLoadDll)
return nullptr;
#endif
wchar_t szFilePathBuffer[MAX_PATH] = {};
const auto _cchSystemPath = GetSystemDirectoryW(szFilePathBuffer, _countof(szFilePathBuffer));
if (_cchSystemPath == 0 || _cchSystemPath >= _countof(szFilePathBuffer))
{
// 回落普通加载，按理说 GetSystemDirectoryW不应该发生这样的失败。
new_handle = LoadLibraryExW(module_name, nullptr, 0);
}
else
{
auto _sModuleName = YY::Thunks::internal::MakeNtString(module_name);
LdrLoadDll(szFilePathBuffer, nullptr, &_sModuleName, &new_handle);
}
}
else
#endif
{
new_handle = LoadLibraryExW(module_name, nullptr, LOAD_LIBRARY_SEARCH_SYSTEM32);
}
}
}

if (!new_handle)
{
if (HMODULE const cached_handle = __crt_interlocked_exchange_pointer(pModule, INVALID_HANDLE_VALUE))
{
_ASSERTE(cached_handle == INVALID_HANDLE_VALUE);
}

return nullptr;
}

// Swap the new handle into the cache. If the cache no longer contained a
// null handle, then some other thread loaded the module and cached the
// handle while we were doing the same. In that case, we free the handle
// once to maintain the reference count:
if (HMODULE const cached_handle = __crt_interlocked_exchange_pointer(pModule, new_handle))
{
_ASSERTE(cached_handle == new_handle);
FreeLibrary(new_handle);
}

return new_handle;
}

static __forceinline void* __fastcall try_get_proc_address_from_first_available_module(
const ProcInfo& _ProcInfo
) noexcept
{
if (_ProcInfo.pfnCustomGetProcAddress)
{
return _ProcInfo.pfnCustomGetProcAddress(_ProcInfo);
}

HMODULE const module_handle = _ProcInfo.pfnGetModule();
if (!module_handle)
{
return nullptr;
}

// 无法直接调用GetProcAddress，因为GetProcAddress可能被Thunk
// 我们需要严格判断，避免发生死锁。
#if defined(__USING_NTDLL_LIB)
void* _pProc = nullptr;
if (uintptr_t(_ProcInfo.szProcName) > UINT16_MAX)
{
ANSI_STRING _sFunctionName = YY::Thunks::internal::MakeNtString(_ProcInfo.szProcName);
LdrGetProcedureAddress(module_handle, &_sFunctionName, 0, &_pProc);
}
else
{
LdrGetProcedureAddress(module_handle, nullptr, (WORD)uintptr_t(_ProcInfo.szProcName), &_pProc);
}
return _pProc;
#else // !defined(__USING_NTDLL_LIB)
__if_exists(YY::Thunks::try_get_GetProcAddress)
{
const auto GetProcAddress = YY::Thunks::try_get_GetProcAddress();
}

return reinterpret_cast<void*>(GetProcAddress(module_handle, _ProcInfo.szProcName));
#endif // defined(__USING_NTDLL_LIB)
}
74 changes: 2 additions & 72 deletions src/Thunks/YY_Thunks.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -303,62 +303,7 @@ static __forceinline T* __fastcall __crt_interlocked_read_pointer(T* const volat
// 该加载模式存在劫持风险，使用前请确认。
#define USING_UNSAFE_LOAD 0x00000004

static HMODULE __fastcall try_get_module(volatile HMODULE* pModule, const wchar_t* module_name, int Flags) noexcept
{
// First check to see if we've cached the module handle:
if (HMODULE const cached_handle = __crt_interlocked_read_pointer(pModule))
{
if (cached_handle == INVALID_HANDLE_VALUE)
{
return nullptr;
}

return cached_handle;
}

// If we haven't yet cached the module handle, try to load the library. If
// this fails, cache the sentinel handle value INVALID_HANDLE_VALUE so that
// we don't attempt to load the module again:
HMODULE new_handle = NULL;
if (Flags & LOAD_AS_DATA_FILE)
{
new_handle = LoadLibraryExW(module_name, NULL, LOAD_LIBRARY_AS_DATAFILE);
}
else if (Flags & USING_GET_MODULE_HANDLE)
{
new_handle = GetModuleHandleW(module_name);
}
else if (Flags & USING_UNSAFE_LOAD)
{
new_handle = LoadLibraryExW(module_name, nullptr, 0);
}
else
{
new_handle = LoadLibraryExW(module_name, nullptr, LOAD_LIBRARY_SEARCH_SYSTEM32);
}

if (!new_handle)
{
if (HMODULE const cached_handle = __crt_interlocked_exchange_pointer(pModule, INVALID_HANDLE_VALUE))
{
_ASSERTE(cached_handle == INVALID_HANDLE_VALUE);
}

return nullptr;
}

// Swap the new handle into the cache. If the cache no longer contained a
// null handle, then some other thread loaded the module and cached the
// handle while we were doing the same. In that case, we free the handle
// once to maintain the reference count:
if (HMODULE const cached_handle = __crt_interlocked_exchange_pointer(pModule, new_handle))
{
_ASSERTE(cached_handle == new_handle);
FreeLibrary(new_handle);
}

return new_handle;
}
static HMODULE __fastcall try_get_module(volatile HMODULE* pModule, const wchar_t* module_name, int Flags) noexcept;

#define _APPLY(_MODULE, _NAME, _FLAGS) \
static HMODULE __fastcall _CRT_CONCATENATE(try_get_module_, _MODULE)() noexcept \
Expand All @@ -382,22 +327,7 @@ struct ProcInfo

static __forceinline void* __fastcall try_get_proc_address_from_first_available_module(
const ProcInfo& _ProcInfo
) noexcept
{
if (_ProcInfo.pfnCustomGetProcAddress)
{
return _ProcInfo.pfnCustomGetProcAddress(_ProcInfo);
}

HMODULE const module_handle = _ProcInfo.pfnGetModule();
if (!module_handle)
{
return nullptr;
}

return reinterpret_cast<void*>(GetProcAddress(module_handle, _ProcInfo.szProcName));
}

) noexcept;

static __forceinline void* __cdecl invalid_function_sentinel() noexcept
{
Expand Down
52 changes: 28 additions & 24 deletions src/Thunks/api-ms-win-core-libraryloader.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -83,40 +83,44 @@ namespace YY::Thunks::Fallback

#if defined(__USING_NTDLL_LIB)
void* _pProc = nullptr;
ANSI_STRING _sFunctionName;
_sFunctionName.Length = strlen(_ProcInfo.szProcName);
_sFunctionName.MaximumLength = _sFunctionName.Length + 1;
_sFunctionName.Buffer = const_cast<PSTR>(_ProcInfo.szProcName);
ANSI_STRING _sFunctionName = YY::Thunks::internal::MakeNtString(_ProcInfo.szProcName);
const LONG _Status = LdrGetProcedureAddress(_hModule, &_sFunctionName, 0, &_pProc);
return _Status >= 0 ? _pProc : nullptr;
#else // !defined(__USING_NTDLL_LIB)
ULONG _uSize;
auto _pExport = (_IMAGE_EXPORT_DIRECTORY*)YY_ImageDirectoryEntryToData(_hModule, IMAGE_DIRECTORY_ENTRY_EXPORT, &_uSize);
if (!_pExport)
return nullptr;

if (_pExport->AddressOfNames == 0 || _pExport->AddressOfFunctions == 0 || _pExport->AddressOfNameOrdinals == 0)
return nullptr;

auto _puFunctions = (const DWORD*)(PBYTE(_hModule) + _pExport->AddressOfFunctions);
auto _puNames = (const DWORD*)(PBYTE(_hModule) + _pExport->AddressOfNames);
auto _puNameOrdinals = (const WORD*)(PBYTE(_hModule) + _pExport->AddressOfNameOrdinals);

for (DWORD i = 0; i != _pExport->NumberOfNames; ++i)
// 防止遇到畸形的DLL，抓个异常
__try
{
auto _uName = _puNames[i];
if (_uName == 0)
continue;
ULONG _uSize;
auto _pExport = (_IMAGE_EXPORT_DIRECTORY*)YY_ImageDirectoryEntryToData(_hModule, IMAGE_DIRECTORY_ENTRY_EXPORT, &_uSize);
if (!_pExport)
return nullptr;

if (_pExport->AddressOfNames == 0 || _pExport->AddressOfFunctions == 0 || _pExport->AddressOfNameOrdinals == 0)
return nullptr;

auto _szName = (char*)(PBYTE(_hModule) + _uName);
auto _puFunctions = (const DWORD*)(PBYTE(_hModule) + _pExport->AddressOfFunctions);
auto _puNames = (const DWORD*)(PBYTE(_hModule) + _pExport->AddressOfNames);
auto _puNameOrdinals = (const WORD*)(PBYTE(_hModule) + _pExport->AddressOfNameOrdinals);

if (StringCompare(_szName, _ProcInfo.szProcName, -1) == 0)
for (DWORD i = 0; i != _pExport->NumberOfNames; ++i)
{
auto _pfn = PBYTE(_hModule) + _puFunctions[_puNameOrdinals[i]];
return _pfn;
auto _uName = _puNames[i];
if (_uName == 0)
continue;

auto _szName = (char*)(PBYTE(_hModule) + _uName);

if (StringCompare(_szName, _ProcInfo.szProcName, -1) == 0)
{
auto _pfn = PBYTE(_hModule) + _puFunctions[_puNameOrdinals[i]];
return _pfn;
}
}
}

__except (EXCEPTION_EXECUTE_HANDLER)
{
}
return nullptr;
#endif // defined(__USING_NTDLL_LIB)
}
Expand Down
4 changes: 2 additions & 2 deletions src/YY-Thunks.UnitTest/YY-Thunks.UnitTest.vcxproj
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@
<FunctionLevelLinking>true</FunctionLevelLinking>
<IntrinsicFunctions>true</IntrinsicFunctions>
<AdditionalIncludeDirectories>$(VCInstallDir)UnitTest\include;%(AdditionalIncludeDirectories)</AdditionalIncludeDirectories>
<PreprocessorDefinitions>SymbolBuildTestPath=LR"($(ProjectDir)TestProject\)";YY_ThunksRootPath=LR"($(ProjectDir)..\..\)";MSBuildBinPath=LR"($(MSBuildBinPath)\)";__YY_Thunks_Unit_Test;WIN32;NDEBUG;YY_Thunks_Support_Version=NTDDI_WIN2K;__APPLY_CHROMIUM_WORKAROUNDS;%(PreprocessorDefinitions)</PreprocessorDefinitions>
<PreprocessorDefinitions>SymbolBuildTestPath=LR"($(ProjectDir)TestProject\)";YY_ThunksRootPath=LR"($(ProjectDir)..\..\)";MSBuildBinPath=LR"($(MSBuildBinPath)\)";__YY_Thunks_Unit_Test;WIN32;NDEBUG;YY_Thunks_Support_Version=NTDDI_WIN2K;__APPLY_CHROMIUM_WORKAROUNDS;__USING_NTDLL_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
<UseFullPaths>true</UseFullPaths>
<PrecompiledHeaderFile>pch.h</PrecompiledHeaderFile>
<PrecompiledHeader>Use</PrecompiledHeader>
Expand All @@ -117,7 +117,7 @@
<ClCompile>
<WarningLevel>Level3</WarningLevel>
<AdditionalIncludeDirectories>$(VCInstallDir)UnitTest\include;%(AdditionalIncludeDirectories)</AdditionalIncludeDirectories>
<PreprocessorDefinitions>SymbolBuildTestPath=LR"($(ProjectDir)TestProject\)";YY_ThunksRootPath=LR"($(ProjectDir)..\..\)";MSBuildBinPath=LR"($(MSBuildBinPath)\)";__YY_Thunks_Unit_Test;WIN32;_DEBUG;YY_Thunks_Support_Version=NTDDI_WIN2K;__APPLY_CHROMIUM_WORKAROUNDS;%(PreprocessorDefinitions)</PreprocessorDefinitions>
<PreprocessorDefinitions>SymbolBuildTestPath=LR"($(ProjectDir)TestProject\)";YY_ThunksRootPath=LR"($(ProjectDir)..\..\)";MSBuildBinPath=LR"($(MSBuildBinPath)\)";__YY_Thunks_Unit_Test;WIN32;_DEBUG;YY_Thunks_Support_Version=NTDDI_WIN2K;__APPLY_CHROMIUM_WORKAROUNDS;__USING_NTDLL_LIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
<UseFullPaths>true</UseFullPaths>
<PrecompiledHeaderFile>pch.h</PrecompiledHeaderFile>
<PrecompiledHeader>Use</PrecompiledHeader>
Expand Down

0 comments on commit 98f0213

Please sign in to comment.