diff --git a/src/modules/FileLocksmith/FileLocksmithLibInterop/NtdllExtensions.cpp b/src/modules/FileLocksmith/FileLocksmithLibInterop/NtdllExtensions.cpp index 7bab774922..cf9510bfb4 100644 --- a/src/modules/FileLocksmith/FileLocksmithLibInterop/NtdllExtensions.cpp +++ b/src/modules/FileLocksmith/FileLocksmithLibInterop/NtdllExtensions.cpp @@ -1,6 +1,8 @@ #include "pch.h" #include "NtdllExtensions.h" +#include +#include #define STATUS_INFO_LENGTH_MISMATCH ((LONG)0xC0000004) @@ -165,66 +167,116 @@ std::vector NtdllExtensions::handles() noexcept std::vector object_info_buffer(DefaultResultBufferSize); - for (ULONG i = 0; i < info_ptr->HandleCount; i++) + std::atomic i = 0; + std::atomic handle_count = info_ptr->HandleCount; + std::atomic process_handle = NULL; + std::atomic 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) + offload_function.detach(); + do { - // Ignore this handle. - continue; - } + 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; + } - 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); - continue; - } + if (previous_i >= i) + { + // The thread looks like it's hanging on some handle. Let's kill it and resume. - auto object_type_info = (OBJECT_TYPE_INFORMATION*)object_info_buffer.data(); - auto type_name = unicode_to_str(object_type_info->Name); + // 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); - std::wstring file_name = file_handle_to_kernel_name(handle_copy, object_info_buffer); + // Close Handles that might be lingering. + if (handle_copy!=NULL) + { + CloseHandle(handle_copy); + } + i++; + break; + } + previous_i = i; + } while (1); - 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); } for (auto [pid, handle] : pid_to_handle) diff --git a/src/modules/FileLocksmith/FileLocksmithLibInterop/NtdllExtensions.h b/src/modules/FileLocksmith/FileLocksmithLibInterop/NtdllExtensions.h index 57227ac37d..64a09299a8 100644 --- a/src/modules/FileLocksmith/FileLocksmithLibInterop/NtdllExtensions.h +++ b/src/modules/FileLocksmith/FileLocksmithLibInterop/NtdllExtensions.h @@ -27,7 +27,7 @@ private: public: struct ProcessInfo { - DWORD pid; + DWORD pid = 0; std::wstring name; std::wstring user; std::vector modules; diff --git a/src/modules/FileLocksmith/FileLocksmithLibInterop/interop.cpp b/src/modules/FileLocksmith/FileLocksmithLibInterop/interop.cpp index f4406b8406..3968348997 100644 --- a/src/modules/FileLocksmith/FileLocksmithLibInterop/interop.cpp +++ b/src/modules/FileLocksmith/FileLocksmithLibInterop/interop.cpp @@ -35,6 +35,7 @@ namespace FileLocksmith::Interop std::wstring paths_file() { +#pragma warning(suppress : 4691) // Weird warning about System::String from referenced library not being the one expected (?!) std::wstring path = from_system_string(interop::Constants::AppDataPath()); path += L"\\"; path += constants::nonlocalizable::PowerToyName; @@ -102,7 +103,7 @@ namespace FileLocksmith::Interop while (!finished) { - WCHAR ch; + WCHAR ch{}; // We have to read data like this if (!stream.read(reinterpret_cast(&ch), 2)) { @@ -157,7 +158,7 @@ namespace FileLocksmith::Interop auto exec_path = executable_path(); - SHELLEXECUTEINFOW exec_info; + SHELLEXECUTEINFOW exec_info{}; exec_info.cbSize = sizeof(exec_info); exec_info.fMask = SEE_MASK_NOCLOSEPROCESS; exec_info.hwnd = NULL; @@ -181,7 +182,7 @@ namespace FileLocksmith::Interop static System::Boolean SetDebugPrivilege() { HANDLE hToken; - TOKEN_PRIVILEGES tp; + TOKEN_PRIVILEGES tp{}; LUID luid; if (OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hToken) != 0) @@ -230,7 +231,7 @@ namespace FileLocksmith::Interop bool elevated = false; if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &token)) { - TOKEN_ELEVATION elevation; + TOKEN_ELEVATION elevation{}; DWORD size; if (GetTokenInformation(token, TokenElevation, &elevation, sizeof(elevation), &size)) {