[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,12 +167,29 @@ 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)
     {
+        previous_i = i;
+
+        // 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++)
+            {
+                process_handle = NULL;
+                handle_copy = NULL;
+
                 auto handle_info = info_ptr->Handles + i;
                 DWORD pid = handle_info->ProcessId;
 
-        HANDLE process_handle = NULL;
                 auto iter = pid_to_handle.find(pid);
                 if (iter != pid_to_handle.end())
                 {
@@ -195,14 +214,14 @@ std::vector<NtdllExtensions::HandleInfo> NtdllExtensions::handles() noexcept
                 //     continue;
                 // }
 
-        HANDLE handle_copy;
+                HANDLE local_handle_copy;
-
+                auto dh_result = DuplicateHandle(process_handle, (HANDLE)handle_info->Handle, GetCurrentProcess(), &local_handle_copy, 0, 0, DUPLICATE_SAME_ACCESS);
-        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;
                 }
+                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);
@@ -210,21 +229,54 @@ std::vector<NtdllExtensions::HandleInfo> NtdllExtensions::handles() noexcept
                 {
                     // 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 = file_handle_to_kernel_name(handle_copy, object_info_buffer);
+                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 });
                 }
 
-        result.push_back(HandleInfo{ pid, handle_info->Handle, type_name, file_name });
                 CloseHandle(handle_copy);
+                handle_copy = NULL;
+            }
+        });
+
+        offload_function.detach();
+        do
+        {
+            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 (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);
+
     }
 
     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))
                 {