[FileLocksmith]Move hanging operations to a distinct thread (#22806)

src/modules/FileLocksmith/FileLocksmithLibInterop/NtdllExtensions.cpp

@@ -1,6 +1,8 @@
 #include "pch.h"
 
 #include "NtdllExtensions.h"
+#include <thread>
+#include <atomic>
 
 #define STATUS_INFO_LENGTH_MISMATCH ((LONG)0xC0000004)
 
@@ -165,66 +167,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;
+        previous_i = i;
-        DWORD pid = handle_info->ProcessId;
 
-        HANDLE process_handle = NULL;
+        // The system calls we use in this block were reported to hang on some machines.
-        auto iter = pid_to_handle.find(pid);
+        // We need to offload the cycle to another thread and keep track of progress to terminate and resume when needed.
-        if (iter != pid_to_handle.end())
+        // Unfortunately, there are no alternative APIs to what we're using that accept timeouts. (NtQueryObject and GetFileType)
-        {
+        auto offload_function = std::thread([&] {
-            process_handle = iter->second;
+            for (; i < handle_count; i++)
-        }
-        else
-        {
-            process_handle = OpenProcess(PROCESS_DUP_HANDLE, FALSE, pid);
-            if (!process_handle)
             {
-                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:
+        offload_function.detach();
-        // https://stackoverflow.com/questions/46384048/enumerate-handles
+        do
-        // 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.
+            Sleep(200); // Timeout in milliseconds for detecting that the system hang on getting information for a handle.
-            continue;
+            if (i >= handle_count)
-        }
+            {
+                // We're done.
+                break;
+            }
 
-        ULONG return_length;
+            if (previous_i >= i)
-        auto status = NtQueryObject(handle_copy, ObjectTypeInformation, object_info_buffer.data(), (ULONG)object_info_buffer.size(), &return_length);
+            {
-        if (NT_ERROR(status))
+                // The thread looks like it's hanging on some handle. Let's kill it and resume.
-        {
-            // Ignore this handle.
-            CloseHandle(handle_copy);
-            continue;
-        }
 
-        auto object_type_info = (OBJECT_TYPE_INFORMATION*)object_info_buffer.data();
+                // 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.
-        auto type_name = unicode_to_str(object_type_info->Name);
+                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)

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<std::wstring> modules;

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<char*>(&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))
                 {