mirror of
https://github.com/pbatard/rufus.git
synced 2025-06-04 00:29:00 -04:00
[core] handle search improvements
* Fix memory leaks due to not releasing the heap * Also speed up lookups by not reopening the same process twice in a row * Also make the usb_debug variable global so we can use uuprintf everywhere
This commit is contained in:
Pete Batard
parent
fec0813073
commit
d1927ac1ef
7 changed files with 304 additions and 66 deletions
290
src/search.c
290
src/search.c
|
|
@ -34,9 +34,18 @@
|
|||
#include "missing.h"
|
||||
#include "msapi_utf8.h"
|
||||
|
||||
// Process Hacker does some filtering using Object Types, but this doesn't help us.
|
||||
// Keep this option, just in case.
|
||||
// #define USE_OBJECT_TYPES
|
||||
|
||||
PF_TYPE_DECL(NTAPI, PVOID, RtlCreateHeap, (ULONG, PVOID, SIZE_T, SIZE_T, PVOID, PRTL_HEAP_PARAMETERS));
|
||||
PF_TYPE_DECL(NTAPI, PVOID, RtlDestroyHeap, (PVOID));
|
||||
PF_TYPE_DECL(NTAPI, PVOID, RtlAllocateHeap, (PVOID, ULONG, SIZE_T));
|
||||
PF_TYPE_DECL(NTAPI, BOOLEAN, RtlFreeHeap, (PVOID, ULONG, PVOID));
|
||||
#ifdef USE_OBJECT_TYPES
|
||||
PF_TYPE_DECL(NTAPI, VOID, RtlInitUnicodeString, (PUNICODE_STRING, PCWSTR));
|
||||
PF_TYPE_DECL(NTAPI, BOOLEAN, RtlEqualUnicodeString, (PCUNICODE_STRING, PCUNICODE_STRING, BOOLEAN));
|
||||
#endif
|
||||
|
||||
PF_TYPE_DECL(NTAPI, NTSTATUS, NtQuerySystemInformation, (SYSTEM_INFORMATION_CLASS, PVOID, ULONG, PULONG));
|
||||
PF_TYPE_DECL(NTAPI, NTSTATUS, NtQueryObject, (HANDLE, OBJECT_INFORMATION_CLASS, PVOID, ULONG, PULONG));
|
||||
|
|
@ -60,10 +69,12 @@ static char* NtStatusError(NTSTATUS Status) {
|
|||
switch (Status) {
|
||||
case STATUS_UNSUCCESSFUL:
|
||||
return "Operation Failed";
|
||||
case STATUS_NOT_IMPLEMENTED:
|
||||
return "Not Implemented";
|
||||
case STATUS_BUFFER_OVERFLOW:
|
||||
return "Buffer Overflow";
|
||||
case STATUS_NOT_IMPLEMENTED:
|
||||
return "Not Implemented";
|
||||
case STATUS_INFO_LENGTH_MISMATCH:
|
||||
return "Info Length Mismatch";
|
||||
case STATUS_INVALID_HANDLE:
|
||||
return "Invalid Handle.";
|
||||
case STATUS_INVALID_PARAMETER:
|
||||
|
|
@ -77,9 +88,11 @@ static char* NtStatusError(NTSTATUS Status) {
|
|||
case STATUS_OBJECT_TYPE_MISMATCH:
|
||||
return "Wrong Type";
|
||||
case STATUS_OBJECT_NAME_INVALID:
|
||||
return "Object Name invalid";
|
||||
return "Object Name Invalid";
|
||||
case STATUS_OBJECT_NAME_NOT_FOUND:
|
||||
return "Object Name not found";
|
||||
case STATUS_OBJECT_PATH_INVALID:
|
||||
return "Object Path Invalid";
|
||||
case STATUS_SHARING_VIOLATION:
|
||||
return "Sharing Violation";
|
||||
case STATUS_INSUFFICIENT_RESOURCES:
|
||||
|
|
@ -92,6 +105,45 @@ static char* NtStatusError(NTSTATUS Status) {
|
|||
}
|
||||
}
|
||||
|
||||
|
||||
static NTSTATUS PhCreateHeap(VOID)
|
||||
{
|
||||
NTSTATUS status = STATUS_SUCCESS;
|
||||
|
||||
if (PhHeapHandle != NULL)
|
||||
return STATUS_ALREADY_COMPLETE;
|
||||
|
||||
PF_INIT_OR_SET_STATUS(RtlCreateHeap, Ntdll);
|
||||
|
||||
if (NT_SUCCESS(status)) {
|
||||
PhHeapHandle = pfRtlCreateHeap(HEAP_NO_SERIALIZE | HEAP_GROWABLE, NULL, 2 * MB, 1 * MB, NULL, NULL);
|
||||
if (PhHeapHandle == NULL)
|
||||
status = STATUS_UNSUCCESSFUL;
|
||||
}
|
||||
|
||||
return status;
|
||||
}
|
||||
|
||||
static NTSTATUS PhDestroyHeap(VOID)
|
||||
{
|
||||
NTSTATUS status = STATUS_SUCCESS;
|
||||
|
||||
if (PhHeapHandle == NULL)
|
||||
return STATUS_ALREADY_COMPLETE;
|
||||
|
||||
PF_INIT_OR_SET_STATUS(RtlDestroyHeap, Ntdll);
|
||||
|
||||
if (NT_SUCCESS(status)) {
|
||||
if (pfRtlDestroyHeap(PhHeapHandle) == NULL) {
|
||||
PhHeapHandle = NULL;
|
||||
} else {
|
||||
status = STATUS_UNSUCCESSFUL;
|
||||
}
|
||||
}
|
||||
|
||||
return status;
|
||||
}
|
||||
|
||||
/**
|
||||
* Allocates a block of memory.
|
||||
*
|
||||
|
|
@ -102,15 +154,11 @@ static char* NtStatusError(NTSTATUS Status) {
|
|||
*/
|
||||
static PVOID PhAllocate(SIZE_T Size)
|
||||
{
|
||||
PF_INIT_OR_OUT(RtlCreateHeap, Ntdll);
|
||||
PF_INIT_OR_OUT(RtlAllocateHeap, Ntdll);
|
||||
PF_INIT(RtlAllocateHeap, Ntdll);
|
||||
if (pfRtlAllocateHeap == NULL)
|
||||
return NULL;
|
||||
|
||||
if (PhHeapHandle == NULL) {
|
||||
PhHeapHandle = pfRtlCreateHeap(HEAP_GROWABLE, NULL, 2 * MB, 1 * MB, NULL, NULL);
|
||||
}
|
||||
return pfRtlAllocateHeap(PhHeapHandle, 0, Size);
|
||||
out:
|
||||
return NULL;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -122,6 +170,7 @@ out:
|
|||
static VOID PhFree(PVOID Memory)
|
||||
{
|
||||
PF_INIT(RtlFreeHeap, Ntdll);
|
||||
|
||||
if (pfRtlFreeHeap != NULL)
|
||||
pfRtlFreeHeap(PhHeapHandle, 0, Memory);
|
||||
}
|
||||
|
|
@ -137,13 +186,13 @@ static VOID PhFree(PVOID Memory)
|
|||
NTSTATUS PhEnumHandlesEx(PSYSTEM_HANDLE_INFORMATION_EX *Handles)
|
||||
{
|
||||
static ULONG initialBufferSize = 0x10000;
|
||||
NTSTATUS status;
|
||||
NTSTATUS status = STATUS_SUCCESS;
|
||||
PVOID buffer;
|
||||
ULONG bufferSize;
|
||||
|
||||
PF_INIT(NtQuerySystemInformation, Ntdll);
|
||||
if (pfNtQuerySystemInformation == NULL)
|
||||
return STATUS_NOT_IMPLEMENTED;
|
||||
PF_INIT_OR_SET_STATUS(NtQuerySystemInformation, Ntdll);
|
||||
if (!NT_SUCCESS(status))
|
||||
return status;
|
||||
|
||||
bufferSize = initialBufferSize;
|
||||
buffer = PhAllocate(bufferSize);
|
||||
|
|
@ -187,7 +236,7 @@ NTSTATUS PhEnumHandlesEx(PSYSTEM_HANDLE_INFORMATION_EX *Handles)
|
|||
*/
|
||||
NTSTATUS PhOpenProcess(PHANDLE ProcessHandle, ACCESS_MASK DesiredAccess, HANDLE ProcessId)
|
||||
{
|
||||
NTSTATUS status = STATUS_NOT_IMPLEMENTED;
|
||||
NTSTATUS status = STATUS_SUCCESS;
|
||||
OBJECT_ATTRIBUTES objectAttributes;
|
||||
CLIENT_ID clientId;
|
||||
|
||||
|
|
@ -196,7 +245,9 @@ NTSTATUS PhOpenProcess(PHANDLE ProcessHandle, ACCESS_MASK DesiredAccess, HANDLE
|
|||
return 0;
|
||||
}
|
||||
|
||||
PF_INIT_OR_OUT(NtOpenProcess, Ntdll);
|
||||
PF_INIT_OR_SET_STATUS(NtOpenProcess, Ntdll);
|
||||
if (!NT_SUCCESS(status))
|
||||
return status;
|
||||
|
||||
clientId.UniqueProcess = ProcessId;
|
||||
clientId.UniqueThread = NULL;
|
||||
|
|
@ -204,10 +255,85 @@ NTSTATUS PhOpenProcess(PHANDLE ProcessHandle, ACCESS_MASK DesiredAccess, HANDLE
|
|||
InitializeObjectAttributes(&objectAttributes, NULL, 0, NULL, NULL);
|
||||
status = pfNtOpenProcess(ProcessHandle, DesiredAccess, &objectAttributes, &clientId);
|
||||
|
||||
out:
|
||||
return status;
|
||||
}
|
||||
|
||||
#ifdef USE_OBJECT_TYPES
|
||||
NTSTATUS PhEnumObjectTypes(POBJECT_TYPES_INFORMATION *ObjectTypes)
|
||||
{
|
||||
NTSTATUS status = STATUS_SUCCESS;
|
||||
PVOID buffer;
|
||||
ULONG bufferSize;
|
||||
ULONG returnLength;
|
||||
|
||||
PF_INIT_OR_SET_STATUS(NtQueryObject, Ntdll);
|
||||
if (!NT_SUCCESS(status))
|
||||
return status;
|
||||
|
||||
bufferSize = 0x1000;
|
||||
buffer = PhAllocate(bufferSize);
|
||||
|
||||
while ((status = pfNtQueryObject(NULL, ObjectTypesInformation, buffer, bufferSize, &returnLength)) == STATUS_INFO_LENGTH_MISMATCH) {
|
||||
PhFree(buffer);
|
||||
bufferSize *= 2;
|
||||
|
||||
// Fail if we're resizing the buffer to something very large.
|
||||
if (bufferSize > PH_LARGE_BUFFER_SIZE)
|
||||
return STATUS_INSUFFICIENT_RESOURCES;
|
||||
|
||||
buffer = PhAllocate(bufferSize);
|
||||
}
|
||||
|
||||
if (!NT_SUCCESS(status)) {
|
||||
PhFree(buffer);
|
||||
return status;
|
||||
}
|
||||
|
||||
*ObjectTypes = (POBJECT_TYPES_INFORMATION)buffer;
|
||||
|
||||
return status;
|
||||
}
|
||||
|
||||
ULONG PhGetObjectTypeNumber(PUNICODE_STRING TypeName)
|
||||
{
|
||||
NTSTATUS status = STATUS_SUCCESS;
|
||||
POBJECT_TYPES_INFORMATION objectTypes;
|
||||
POBJECT_TYPE_INFORMATION objectType;
|
||||
ULONG objectIndex = -1;
|
||||
ULONG i;
|
||||
|
||||
PF_INIT_OR_SET_STATUS(RtlEqualUnicodeString, NtDll);
|
||||
if (!NT_SUCCESS(status))
|
||||
return -1;
|
||||
|
||||
status = PhEnumObjectTypes(&objectTypes);
|
||||
if (NT_SUCCESS(status)) {
|
||||
objectType = PH_FIRST_OBJECT_TYPE(objectTypes);
|
||||
|
||||
for (i = 0; i < objectTypes->NumberOfTypes; i++) {
|
||||
if (pfRtlEqualUnicodeString(&objectType->TypeName, TypeName, TRUE)) {
|
||||
if (nWindowsVersion >= WINDOWS_8_1) {
|
||||
objectIndex = objectType->TypeIndex;
|
||||
break;
|
||||
} else if (nWindowsVersion >= WINDOWS_7) {
|
||||
objectIndex = i + 2;
|
||||
break;
|
||||
} else {
|
||||
objectIndex = i + 1;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
objectType = PH_NEXT_OBJECT_TYPE(objectType);
|
||||
}
|
||||
|
||||
PhFree(objectTypes);
|
||||
}
|
||||
|
||||
return objectIndex;
|
||||
}
|
||||
#endif
|
||||
|
||||
/**
|
||||
* Search all the processes and list the ones that have a specific handle open.
|
||||
*
|
||||
|
|
@ -219,42 +345,63 @@ out:
|
|||
*/
|
||||
BOOL SearchProcess(char* HandleName, BOOL bPartialMatch, BOOL bIgnoreSelf)
|
||||
{
|
||||
NTSTATUS status;
|
||||
PSYSTEM_HANDLE_INFORMATION_EX handles;
|
||||
POBJECT_NAME_INFORMATION buffer;
|
||||
NTSTATUS status = STATUS_SUCCESS;
|
||||
PSYSTEM_HANDLE_INFORMATION_EX handles = NULL;
|
||||
POBJECT_NAME_INFORMATION buffer = NULL;
|
||||
#ifdef USE_OBJECT_TYPES
|
||||
UNICODE_STRING fileTypeName;
|
||||
ULONG fileObjectTypeIndex = -1;
|
||||
#endif
|
||||
ULONG_PTR i;
|
||||
ULONG_PTR pid[2];
|
||||
ULONG_PTR last_access_denied_pid = 0;
|
||||
ULONG bufferSize;
|
||||
USHORT wHandleNameLen;
|
||||
WCHAR *wHandleName;
|
||||
WCHAR *wHandleName = NULL;
|
||||
HANDLE dupHandle = NULL;
|
||||
HANDLE processHandle = NULL;
|
||||
BOOLEAN bFound = FALSE;
|
||||
char exe_path[2][MAX_PATH];
|
||||
int cur;
|
||||
char exe[2][MAX_PATH];
|
||||
int cur_exe, cur_pid;
|
||||
|
||||
status = STATUS_NOT_IMPLEMENTED;
|
||||
PF_INIT(NtQueryObject, Ntdll);
|
||||
PF_INIT(NtDuplicateObject, NtDll);
|
||||
PF_INIT(NtClose, NtDll);
|
||||
if ((pfNtQueryObject != NULL) && (pfNtClose != NULL) && (pfNtDuplicateObject != NULL))
|
||||
status = 0;
|
||||
PF_INIT_OR_SET_STATUS(NtQueryObject, Ntdll);
|
||||
PF_INIT_OR_SET_STATUS(NtDuplicateObject, NtDll);
|
||||
PF_INIT_OR_SET_STATUS(NtClose, NtDll);
|
||||
#ifdef USE_OBJECT_TYPES
|
||||
PF_INIT(RtlInitUnicodeString, NtDll);
|
||||
#endif
|
||||
|
||||
if (NT_SUCCESS(status))
|
||||
status = PhCreateHeap();
|
||||
|
||||
if (NT_SUCCESS(status))
|
||||
status = PhEnumHandlesEx(&handles);
|
||||
|
||||
if (!NT_SUCCESS(status)) {
|
||||
uprintf("Could not enumerate handles: %s", NtStatusError(status));
|
||||
return FALSE;
|
||||
uprintf("Warning: Could not enumerate process handles: %s", NtStatusError(status));
|
||||
goto out;
|
||||
}
|
||||
|
||||
|
||||
exe_path[0][0] = 0;
|
||||
cur = 1;
|
||||
pid[0] = (ULONG_PTR)NULL;
|
||||
cur_pid = 1;
|
||||
exe[0][0] = 0;
|
||||
cur_exe = 1;
|
||||
|
||||
wHandleName = utf8_to_wchar(HandleName);
|
||||
wHandleNameLen = (USHORT) wcslen(wHandleName);
|
||||
wHandleNameLen = (USHORT)wcslen(wHandleName);
|
||||
|
||||
bufferSize = 0x200;
|
||||
buffer = PhAllocate(bufferSize);
|
||||
if (buffer == NULL)
|
||||
goto out;
|
||||
|
||||
#ifdef USE_OBJECT_TYPES
|
||||
pfRtlInitUnicodeString(&fileTypeName, L"File");
|
||||
fileObjectTypeIndex = PhGetObjectTypeNumber(&fileTypeName);
|
||||
if (fileObjectTypeIndex < 0)
|
||||
uprintf("Warning: Could not get Object Index for file types");
|
||||
}
|
||||
#endif
|
||||
|
||||
for (i = 0; ; i++) {
|
||||
ULONG attempts = 8;
|
||||
|
|
@ -264,28 +411,53 @@ BOOL SearchProcess(char* HandleName, BOOL bPartialMatch, BOOL bIgnoreSelf)
|
|||
pfNtClose(dupHandle);
|
||||
dupHandle = NULL;
|
||||
}
|
||||
if (processHandle != NULL) {
|
||||
if (processHandle != NtCurrentProcess())
|
||||
pfNtClose(processHandle);
|
||||
processHandle = NULL;
|
||||
|
||||
#ifdef USE_OBJECT_TYPES
|
||||
// Only look for File objects type
|
||||
if ((fileObjectTypeIndex >= 0 ) && (handleInfo->ObjectTypeIndex != (USHORT)fileObjectTypeIndex))
|
||||
continue;
|
||||
#endif
|
||||
|
||||
// Update the current handle's process PID and compare against last
|
||||
pid[cur_pid] = handleInfo->UniqueProcessId;
|
||||
|
||||
if (pid[0] != pid[1]) {
|
||||
cur_pid = (cur_pid + 1) % 2;
|
||||
// Close the previous handle
|
||||
if (processHandle != NULL) {
|
||||
if (processHandle != NtCurrentProcess())
|
||||
pfNtClose(processHandle);
|
||||
processHandle = NULL;
|
||||
}
|
||||
}
|
||||
|
||||
CHECK_FOR_USER_CANCEL;
|
||||
|
||||
// Don't bother with processes we can't access
|
||||
if (handleInfo->UniqueProcessId == last_access_denied_pid)
|
||||
continue;
|
||||
|
||||
// Exit loop condition
|
||||
if (i >= handles->NumberOfHandles)
|
||||
break;
|
||||
|
||||
// Get the process that created the handle we are after
|
||||
// TODO: We probably should keep a list of the most recent processes and perform a lookup
|
||||
// instead of Opening the process every time
|
||||
status = PhOpenProcess(&processHandle, PROCESS_DUP_HANDLE | PROCESS_QUERY_INFORMATION,
|
||||
(HANDLE)handleInfo->UniqueProcessId);
|
||||
// There exists some processes we can't access
|
||||
if (!NT_SUCCESS(status))
|
||||
continue;
|
||||
// Open the process to which the handle we are after belongs, if not already opened
|
||||
if (pid[0] != pid[1]) {
|
||||
status = PhOpenProcess(&processHandle, PROCESS_DUP_HANDLE | PROCESS_QUERY_INFORMATION,
|
||||
(HANDLE)handleInfo->UniqueProcessId);
|
||||
// There exists some processes we can't access
|
||||
if (!NT_SUCCESS(status)) {
|
||||
uuprintf("SearchProcess: Could not open process %ld: %s",
|
||||
handleInfo->UniqueProcessId, NtStatusError(status));
|
||||
processHandle = NULL;
|
||||
if (status == STATUS_ACCESS_DENIED) {
|
||||
last_access_denied_pid = handleInfo->UniqueProcessId;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
// Must duplicate the handle onto our own process, before we can access its properties
|
||||
// Now duplicate this handle onto our own process, so that we can access its properties
|
||||
if (processHandle == NtCurrentProcess()) {
|
||||
if (bIgnoreSelf)
|
||||
continue;
|
||||
|
|
@ -293,7 +465,6 @@ BOOL SearchProcess(char* HandleName, BOOL bPartialMatch, BOOL bIgnoreSelf)
|
|||
} else {
|
||||
status = pfNtDuplicateObject(processHandle, (HANDLE)handleInfo->HandleValue,
|
||||
NtCurrentProcess(), &dupHandle, 0, 0, 0);
|
||||
// Why does it always work for Process Hacker and not me???
|
||||
if (!NT_SUCCESS(status))
|
||||
continue;
|
||||
}
|
||||
|
|
@ -304,18 +475,23 @@ BOOL SearchProcess(char* HandleName, BOOL bPartialMatch, BOOL bIgnoreSelf)
|
|||
|
||||
// A loop is needed because the I/O subsystem likes to give us the wrong return lengths...
|
||||
do {
|
||||
status = pfNtQueryObject(dupHandle, ObjectBasicInformation + 1,
|
||||
buffer, bufferSize, &bufferSize);
|
||||
ULONG returnSize;
|
||||
status = pfNtQueryObject(dupHandle, ObjectNameInformation, buffer, bufferSize, &returnSize);
|
||||
if (status == STATUS_BUFFER_OVERFLOW || status == STATUS_INFO_LENGTH_MISMATCH ||
|
||||
status == STATUS_BUFFER_TOO_SMALL) {
|
||||
uuprintf("SearchProcess: Realloc from %d to %d", bufferSize, returnSize);
|
||||
bufferSize = returnSize;
|
||||
PhFree(buffer);
|
||||
buffer = PhAllocate(bufferSize);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
} while (--attempts);
|
||||
if (!NT_SUCCESS(status))
|
||||
if (!NT_SUCCESS(status)) {
|
||||
uuprintf("SearchProcess: NtQueryObject failed for handle %X of process %ld: %s",
|
||||
handleInfo->HandleValue, handleInfo->UniqueProcessId, NtStatusError(status));
|
||||
continue;
|
||||
}
|
||||
|
||||
// Don't bother comparing if we are looking for full match and the length is different
|
||||
if ((!bPartialMatch) && (wHandleNameLen != buffer->Name.Length))
|
||||
|
|
@ -335,11 +511,11 @@ BOOL SearchProcess(char* HandleName, BOOL bPartialMatch, BOOL bIgnoreSelf)
|
|||
}
|
||||
|
||||
// TODO: only list processes with conflicting access rights (ignore "Read attributes" or "Synchronize")
|
||||
if (GetModuleFileNameExU(processHandle, 0, exe_path[cur], MAX_PATH - 1)) {
|
||||
if (GetModuleFileNameExU(processHandle, 0, exe[cur_exe], MAX_PATH - 1)) {
|
||||
// Avoid printing the same path repeatedly
|
||||
if (strcmp(exe_path[0], exe_path[1]) != 0) {
|
||||
uprintf("o %s", exe_path[cur]);
|
||||
cur = (cur + 1) % 2;
|
||||
if (strcmp(exe[0], exe[1]) != 0) {
|
||||
uprintf("o %s", exe[cur_exe]);
|
||||
cur_exe = (cur_exe + 1) % 2;
|
||||
}
|
||||
} else {
|
||||
uprintf("o Unknown (Process ID %d)", GetProcessId(processHandle));
|
||||
|
|
@ -354,5 +530,7 @@ out:
|
|||
|
||||
free(wHandleName);
|
||||
PhFree(buffer);
|
||||
PhFree(handles);
|
||||
PhDestroyHeap();
|
||||
return bFound;
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue