Skip to content

Commit

Permalink
Actual offload thread instead of starting many threads
Browse files Browse the repository at this point in the history
  • Loading branch information
@jaimecbernardo
jaimecbernardo committed Dec 14, 2022
1 parent fc7e261 commit d758eee
Show file tree
Hide file tree
Showing 2 changed files with 104 additions and 70 deletions.
19 changes: 1 addition & 18 deletions src/modules/FileLocksmith/FileLocksmithLibInterop/NtdllBase.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,22 +63,5 @@ NTSTATUS Ntdll::NtQueryObject(
ULONG ObjectInformationLength,
PULONG ReturnLength)
{
// On very specific cases, NtQueryObject seems to hang.
// "file object opened by a kernel-mode driver with FILE_SYNCHRONOUS_IO_NONALERT" and handles opened with certain permissions were examples I could find being reported on other sources.
// Given this, lets do this call in a thread we can cancel.
std::packaged_task<NTSTATUS()> task([=]() {
return m_NtQueryObject(ObjectHandle, ObjectInformationClass, ObjectInformation, ObjectInformationLength, ReturnLength);
});
std::future<NTSTATUS> future = task.get_future();
std::thread myThread(std::move(task));
myThread.detach();

std::future_status status = future.wait_for(std::chrono::milliseconds(100));
if (status != std::future_status::ready)
{
return STATUS_UNSUCCESSFUL;
TerminateThread(myThread.native_handle(), 1);
}
return future.get();

return m_NtQueryObject(ObjectHandle, ObjectInformationClass, ObjectInformation, ObjectInformationLength, ReturnLength);
}
155 changes: 103 additions & 52 deletions src/modules/FileLocksmith/FileLocksmithLibInterop/NtdllExtensions.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
#include <future>
#include <thread>
#include <chrono>
#include <atomic>

#define STATUS_INFO_LENGTH_MISMATCH ((LONG)0xC0000004)

Expand Down Expand Up @@ -169,66 +170,116 @@ std::vector<NtdllExtensions::HandleInfo> NtdllExtensions::handles() noexcept

std::vector<BYTE> object_info_buffer(DefaultResultBufferSize);

for (ULONG i = 0; i < info_ptr->HandleCount; i++)
std::atomic<ULONG> i = 0;
std::atomic<ULONG> handle_count = info_ptr->HandleCount;
std::atomic<HANDLE> process_handle = NULL;
std::atomic<HANDLE> handle_copy = NULL;
ULONG previous_i;


while (i < handle_count)
{
auto handle_info = info_ptr->Handles + i;
DWORD pid = handle_info->ProcessId;
previous_i = i;

HANDLE process_handle = NULL;
auto iter = pid_to_handle.find(pid);
if (iter != pid_to_handle.end())
{
process_handle = iter->second;
}
else
{
process_handle = OpenProcess(PROCESS_DUP_HANDLE, FALSE, pid);
if (!process_handle)
// The system calls we use in this block were reported to hang on some machines.
// We need to offload the cycle to another thread and keep track of progress to terminate and resume when needed.
// Unfortunately, there are no alternative APIs to what we're using that accept timeouts. (NtQueryObject and GetFileType)
auto offload_function = std::thread([&] {
for (; i < handle_count; i++)
{
continue;
process_handle = NULL;
handle_copy = NULL;

auto handle_info = info_ptr->Handles + i;
DWORD pid = handle_info->ProcessId;

auto iter = pid_to_handle.find(pid);
if (iter != pid_to_handle.end())
{
process_handle = iter->second;
}
else
{
process_handle = OpenProcess(PROCESS_DUP_HANDLE, FALSE, pid);
if (!process_handle)
{
continue;
}
pid_to_handle[pid] = process_handle;
}

// According to this:
// https://stackoverflow.com/questions/46384048/enumerate-handles
// NtQueryObject could hang

// TODO uncomment and investigate
// if (handle_info->GrantedAccess == 0x0012019f) {
// continue;
// }

HANDLE local_handle_copy;
auto dh_result = DuplicateHandle(process_handle, (HANDLE)handle_info->Handle, GetCurrentProcess(), &local_handle_copy, 0, 0, DUPLICATE_SAME_ACCESS);
if (dh_result == 0)
{
// Ignore this handle.
continue;
}
handle_copy = local_handle_copy;

ULONG return_length;
auto status = NtQueryObject(handle_copy, ObjectTypeInformation, object_info_buffer.data(), (ULONG)object_info_buffer.size(), &return_length);
if (NT_ERROR(status))
{
// Ignore this handle.
CloseHandle(handle_copy);
handle_copy = NULL;
continue;
}

auto object_type_info = (OBJECT_TYPE_INFORMATION*)object_info_buffer.data();
auto type_name = unicode_to_str(object_type_info->Name);

std::wstring file_name;

if (type_name == L"File")
{
file_name = file_handle_to_kernel_name(handle_copy, object_info_buffer);
result.push_back(HandleInfo{ pid, handle_info->Handle, type_name, file_name });
}

CloseHandle(handle_copy);
handle_copy = NULL;
}
pid_to_handle[pid] = process_handle;
}

// According to this:
// https://stackoverflow.com/questions/46384048/enumerate-handles
// NtQueryObject could hang

// TODO uncomment and investigate
// if (handle_info->GrantedAccess == 0x0012019f) {
// continue;
// }

HANDLE handle_copy;

auto dh_result = DuplicateHandle(process_handle, (HANDLE)handle_info->Handle, GetCurrentProcess(), &handle_copy, 0, 0, DUPLICATE_SAME_ACCESS);
if (dh_result == 0)
{
// Ignore this handle.
continue;
}
});

ULONG return_length;
auto status = NtQueryObject(handle_copy, ObjectTypeInformation, object_info_buffer.data(), (ULONG)object_info_buffer.size(), &return_length);
if (NT_ERROR(status))
offload_function.detach();
do
{
// Ignore this handle.
CloseHandle(handle_copy);
continue;
}

auto object_type_info = (OBJECT_TYPE_INFORMATION*)object_info_buffer.data();
auto type_name = unicode_to_str(object_type_info->Name);

std::wstring file_name;
Sleep(200); // Timeout in milliseconds for detecting that the system hang on getting information for a handle.
if (i >= handle_count)
{
// We're done.
break;
}

if (type_name == L"File")
{
file_name = file_handle_to_kernel_name(handle_copy, object_info_buffer);
result.push_back(HandleInfo{ pid, handle_info->Handle, type_name, file_name });
}
if (previous_i >= i)
{
// The thread looks like it's hanging on some handle. Let's kill it and resume.

// HACK: This is unsafe and may leak something, but looks like there's no way to properly clean up a thread when it's hanging on a system call.
TerminateThread(offload_function.native_handle(), 1);

// Close Handles that might be lingering.
if (handle_copy!=NULL)
{
CloseHandle(handle_copy);
}
i++;
break;
}
previous_i = i;
} while (1);

CloseHandle(handle_copy);
}

for (auto [pid, handle] : pid_to_handle)
Expand Down

0 comments on commit d758eee

Please sign in to comment.