Using Import Address Table hooking for testing
For a procedure that is free of side effects, it is a relatively easy task to create a unit test that achieves sufficient code coverage by testing all (or at least all interesting) combinations of input data and verifying the computed results.
If, however, the procedure is not free of side effects, the state (global variables, external data, etc.) modified by the procedure has to be taken into account. A solid testcase has to test both the effects of the state on the correctness of the procedure and the correctness of the procedure’s modifications on the state. As a consequence, the testcase has to initialize the state to a ‘known state’ before running each test and validate the state after each test.
If the state can be queried and easily modified by the testcase’s initialization code to construct the individual testing scenarios, writing such a testcase may be laborious but does not pose a real problem. Things are a little different if the state is outside of the programmers control or at least hard to query and alter. This case is especially common when writing code that interfaces the operating system – a common example may be a library that uses the filesystem or the registry to query and modify data. While both filesystem and registry are accessibly to the testing code, several problems arise:
- The state is unknown on test case startup. Any affected files/keys have to be initialized to a defined state first.
- Parts of the data may be unaccessible for security reasons. As an example, the procedure under test may need to read settings from HKLM. The affected keys under HKLM thus belong to the state that affects the execution of the procedure and our test case needs to run the procedure using different initial states, i.e. using different values for the keys in HKLM. But, as the testcase should be able to run as normal user (maybe to ensure LUA compliance), modifying these keys is forbidden.
- The state may be altered concurrently by other programs running on the same system.
Such a test case may easily contain more code initializing and validating state (key and file creation/deletion etc) than actual test code. Needless to say, creating such test cases is a pain.
For some testcases, touching real files and keys may not only be troublesome but is actually of minor interest. For example, the objective of a testcase might be to check error handling code. In such situations, we are not really interested in modifying/querying the actual state of the system but rather in having the procedure under test see a specific state (e.g. error condition) or make specific changes to the state. The idea is thus to intercept the calls into the OS libraries by appropriate mocks or stubs: Instead of acting on the real state, we give the procedure under test the illusion of executing in a specific state and intercept its modifications. That way, initialization of state becomes trivial, the problem of concurrent access is mitigated and error conditions can be easily simulated.
The question of course is how to create such mocks or stubs given that the code under test will usually be statically linked to the appropriate OS libraries. One easy and rather nice solution is to use IAT hooking.
Import Address Table hooking
The basic idea of IAT hooking is to take a module’s Import Address Table (where the loader puts the function pointers of imported functions) and patch specific entries. Much has been written about IAT hooking, so I will skip the details. In contrast to other hooking techniques, IAT hooking has at least 2 interesting properties.
- Hooking only requires exchanging a single pointer, which can be done using an interlocked instruction. In comparison to other hooking/patching techniques, IAT hooking can thus be considered low-risk.
- Modifications affect a specific module only. If both module A and B import ReadFileW and module A’s IAT is patched to use FooReadFileW instead, module B remains unaffected.
Especially the second property makes IAT hooking interesting for use by test code, as the modifications can be scoped to affect the module under test only (assuming that test code and code under test are located in different modules).
The following listing shows a simple function that shows how to access a loaded PE image in order to hook a single entry in a module’s IAT. Note that the function is capable of hooking named imports only and does not provide any thread-safety (see my additional remark at the end of the post).
#define PtrFromRva( base, rva ) ( ( ( PBYTE ) base ) + rva )
/*++
Routine Description:
Replace the function pointer in a module's IAT.
Parameters:
Module - Module to use IAT from.
ImportedModuleName - Name of imported DLL from which
function is imported.
ImportedProcName - Name of imported function.
AlternateProc - Function to be written to IAT.
OldProc - Original function.
Return Value:
S_OK on success.
(any HRESULT) on failure.
--*/
HRESULT PatchIat(
__in HMODULE Module,
__in PSTR ImportedModuleName,
__in PSTR ImportedProcName,
__in PVOID AlternateProc,
__out_opt PVOID *OldProc
)
{
PIMAGE_DOS_HEADER DosHeader = ( PIMAGE_DOS_HEADER ) Module;
PIMAGE_NT_HEADERS NtHeader;
PIMAGE_IMPORT_DESCRIPTOR ImportDescriptor;
UINT Index;
_ASSERTE( Module );
_ASSERTE( ImportedModuleName );
_ASSERTE( ImportedProcName );
_ASSERTE( AlternateProc );
NtHeader = ( PIMAGE_NT_HEADERS )
PtrFromRva( DosHeader, DosHeader->e_lfanew );
if( IMAGE_NT_SIGNATURE != NtHeader->Signature )
{
return HRESULT_FROM_WIN32( ERROR_BAD_EXE_FORMAT );
}
ImportDescriptor = ( PIMAGE_IMPORT_DESCRIPTOR )
PtrFromRva( DosHeader,
NtHeader->OptionalHeader.DataDirectory
[ IMAGE_DIRECTORY_ENTRY_IMPORT ].VirtualAddress );
//
// Iterate over import descriptors/DLLs.
//
for ( Index = 0;
ImportDescriptor[ Index ].Characteristics != 0;
Index++ )
{
PSTR dllName = ( PSTR )
PtrFromRva( DosHeader, ImportDescriptor[ Index ].Name );
if ( 0 == _strcmpi( dllName, ImportedModuleName ) )
{
//
// This the DLL we are after.
//
PIMAGE_THUNK_DATA Thunk;
PIMAGE_THUNK_DATA OrigThunk;
if ( ! ImportDescriptor[ Index ].FirstThunk ||
! ImportDescriptor[ Index ].OriginalFirstThunk )
{
return E_INVALIDARG;
}
Thunk = ( PIMAGE_THUNK_DATA )
PtrFromRva( DosHeader,
ImportDescriptor[ Index ].FirstThunk );
OrigThunk = ( PIMAGE_THUNK_DATA )
PtrFromRva( DosHeader,
ImportDescriptor[ Index ].OriginalFirstThunk );
for ( ; OrigThunk->u1.Function != NULL;
OrigThunk++, Thunk++ )
{
if ( OrigThunk->u1.Ordinal & IMAGE_ORDINAL_FLAG )
{
//
// Ordinal import - we can handle named imports
// ony, so skip it.
//
continue;
}
PIMAGE_IMPORT_BY_NAME import = ( PIMAGE_IMPORT_BY_NAME )
PtrFromRva( DosHeader, OrigThunk->u1.AddressOfData );
if ( 0 == strcmp( ImportedProcName,
( char* ) import->Name ) )
{
//
// Proc found, patch it.
//
DWORD junk;
MEMORY_BASIC_INFORMATION thunkMemInfo;
//
// Make page writable.
//
VirtualQuery(
Thunk,
&thunkMemInfo,
sizeof( MEMORY_BASIC_INFORMATION ) );
if ( ! VirtualProtect(
thunkMemInfo.BaseAddress,
thunkMemInfo.RegionSize,
PAGE_EXECUTE_READWRITE,
&thunkMemInfo.Protect ) )
{
return HRESULT_FROM_WIN32( GetLastError() );
}
//
// Replace function pointers (non-atomically).
//
if ( OldProc )
{
*OldProc = ( PVOID ) ( DWORD_PTR )
Thunk->u1.Function;
}
#ifdef _WIN64
Thunk->u1.Function = ( ULONGLONG ) ( DWORD_PTR )
AlternateProc;
#else
Thunk->u1.Function = ( DWORD ) ( DWORD_PTR )
AlternateProc;
#endif
//
// Restore page protection.
//
if ( ! VirtualProtect(
thunkMemInfo.BaseAddress,
thunkMemInfo.RegionSize,
thunkMemInfo.Protect,
&junk ) )
{
return HRESULT_FROM_WIN32( GetLastError() );
}
return S_OK;
}
}
//
// Import not found.
//
return HRESULT_FROM_WIN32( ERROR_PROC_NOT_FOUND );
}
}
//
// DLL not found.
//
return HRESULT_FROM_WIN32( ERROR_MOD_NOT_FOUND );
}
Using IAT hooks to create stubs or mocks
Having IAT hooking at hands, it is now straightforwanrd to implement and install stubs or mocks. Given a procedure that, for example, uses some of the Reg* functions to query and modify the registry, all we have to do is implement stubs having the same signature as the corresponding Reg* functions and install them using the technique described above. The testcase will then, though being statically linked to advapi32, call our stubs instead of the real registry routines. Within the stub, we are free to delegate to the real registry routines as required, provided that the stubs are either located in a different module (s.t. the IAT hooks do not apply) or that these calls are made using the ‘original’ function pointers.
Various scenarios come to mind where such hooks can help testing, two shall now be discussed.
Scenario 1: Testing Error Checking
When interfacing the Win32 API, error handling code like the following is ubiquitous:
...
res = RegQueryValueEx(
key,
Name,
0,
&dataType,
Buffer,
&dataRead );
if ( ERROR_ACCESS_DENIED == res )
{
...
}
else if ( ERROR_SUCCESS != res )
{
...
}
else
{
...
}
In order to achieve full code coverarge for this code block, we have to implement at least 3 test cases with RegQueryValueEx returning ERROR_ACCESS_DENIED, ERROR_SUCCESS and some other error code, respectively. Under normal conditions, this would require the initialization and teardown code of each of these three testcases to modify the registry appropriately. Using IAT hooks, we can leave the registry untouched and instead use three different alternate implementations of RegQueryValueEx, each returning the appropriate error code and updating any out-parameters.
When using C++ rather than C, we can even save a considerable amount of typing by getting creative with templates.
As an example, consider a procedure that, depending on the parameter values passed, writes to either HKCU or HKLM. As writing to HKLM is permitted to admins only, it is vital to test that the procedure fails gracefully when access to certain keys is forbidden. In order to simulate the following conditions, we may choose to implement a templatized function.
- Simulate normal user – HKCU is allowed for write access, HKLM not
- Simulate admin – both HKCU and HKLM allowed for write access
- Simulate weird ACL settings – deny access to some or all keys in both HKCU and HKLM
At the beginning of each testcase, FailRegCreateKeyEx with appropriate arguments is then installed as a hook for RegCreateKeyEx.
template< BOOL FailOnHkcu, BOOL FailOnHklm >
static LONG FailRegCreateKeyEx (
__in HKEY hKey,
__in LPCWSTR lpSubKey,
__reserved DWORD Reserved,
__in_opt LPWSTR lpClass,
__in DWORD dwOptions,
__in REGSAM samDesired,
__in_opt LPSECURITY_ATTRIBUTES lpSecurityAttributes,
__out PHKEY phkResult,
__out_opt LPDWORD lpdwDisposition
)
{
//
// N.B. Use locals to avoid constant expression-warnings.
//
BOOL failOnHklm = FailOnHklm;
BOOL failOnHkcu = FailOnHkcu;
if ( hKey == HKEY_LOCAL_MACHINE && failOnHklm ||
hKey == HKEY_CURRENT_USER && failOnHkcu )
{
return ERROR_ACCESS_DENIED;
}
else
{
//
// Assume that we are in a different module, so using
// RegCreateKeyEx will not re-enter the hook.
//
return RegCreateKeyEx(
hKey,
lpSubKey,
Reserved,
lpClass,
dwOptions,
samDesired,
lpSecurityAttributes,
phkResult,
lpdwDisposition );
}
}
Scenario 2: Redirecting access
Rather than only saving us from typing complicated initialization and teardown code, IAT hooking also comes in handy for addressing problem 2. Again consider the scenario where we are to write a test case for an API that reads and writes to both HKLM and HKCU. HKLM may contain machine-wide settings which only administrators can modify using the code under test. HKCU may contain optional per-user settings, which (if present) override machine-wide settings.
In order to test writing machine-wide settings, the testcase process obviously needs administrative privileges, which are normally unavailable (assuming that you always run as limited user). Running the testcase as a different (adminstrative) user is both uncomfortable, makes debugging harder and is also risky – unlimately, a bug in the tested code could also harm your system.
A simple solution to this problem might be to explicitly grant access to the specific keys in HKLM s.t. they become writable. If this is not feasible or not flexible enough, IAT hooking can help again. In order to circumvent the problem of accessing keys in HKLM, we redirect all accesses to HKLM to some temporary key in HKCU.
The following code snippet illsutrates the basic idea (the code is not suitable for all usage scenarios of RegCreateKeyEx – but this usually is not required either). When asked to creare a key in HKLM, the hook instead creates a key in the location referred to by the global variable RedirectPathHklm:
static LONG VirtRegCreateKeyEx (
__in HKEY hKey,
__in LPCWSTR lpSubKey,
__reserved DWORD Reserved,
__in_opt LPWSTR lpClass,
__in DWORD dwOptions,
__in REGSAM samDesired,
__in_opt LPSECURITY_ATTRIBUTES lpSecurityAttributes,
__out PHKEY phkResult,
__out_opt LPDWORD lpdwDisposition
)
{
//
// Create virtualized key if neccessary
//
HKEY virtKey = NULL;
BOOL closeKey = FALSE;
LONG res = 0;
if ( hKey == HKEY_LOCAL_MACHINE )
{
res = RegCreateKeyEx(
HKEY_CURRENT_USER,
RedirectPathHklm,
0,
NULL,
0,
KEY_ALL_ACCESS,
NULL,
&virtKey,
NULL );
if ( NOERROR != res )
return res;
closeKey = TRUE;
}
else
{
virtKey = hKey;
}
res = RegCreateKeyEx(
virtKey,
lpSubKey,
Reserved,
lpClass,
dwOptions,
samDesired,
lpSecurityAttributes,
phkResult,
lpdwDisposition );
if ( closeKey )
RegCloseKey( virtKey );
return res;
}
Multithreading
Up to this point, the discussion has ignored the issue of multithreading. As the Import Address Table is a (module-) global resource, concurrent access by different threads is subject to appropriate synchronization. If the IAT is only adjusted once during initialization of our test suite and does not need to be touched it again, this might not be a problem. More likely, however, is the requirement to modify (and reset) the IAT for each individual test case, which, of course will lead to all sorts of nasty race conditions when multiple testcases are executed in parallel on different threads.
Fortunately, this is not an inherent limitation of this approach – rather, to create an implementation safe for use in a multithreaded environment, it is well conceivable to implement something like thread local IAT hooks – hooks that only apply to a specifiy thread – maybe a topic for another day.
Conclusion
Though a bit tempting, using IAT hooks in each and every testcase would of course thwart the whole idea of a testcase – if the real OS APIs are never called but always intercepted by hooks, the test suite can quickly become pretty much worthless. However, as I have shown in this post, used consciously, IAT hooking can indeed make writing specific testcases substantially easier. 