FancyZones: optimize elevation detection logic (#2103)

microsoft · Apr 13, 2020 · 86704ef · 86704ef
1 parent 6bb0f18
commit 86704ef
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 21 deletions.
diff --git a/src/common/common.cpp b/src/common/common.cpp
@@ -364,27 +364,31 @@ WindowState get_window_state(HWND hwnd)
     return RESTORED;
 }
 
-bool is_process_elevated()
+bool is_process_elevated(const bool use_cached_value)
 {
-    HANDLE token = nullptr;
-    bool elevated = false;
+    auto detection_func = []() {
+        HANDLE token = nullptr;
+        bool elevated = false;
 
-    if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &token))
-    {
-        TOKEN_ELEVATION elevation;
-        DWORD size;
-        if (GetTokenInformation(token, TokenElevation, &elevation, sizeof(elevation), &size))
+        if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &token))
         {
-            elevated = (elevation.TokenIsElevated != 0);
+            TOKEN_ELEVATION elevation;
+            DWORD size;
+            if (GetTokenInformation(token, TokenElevation, &elevation, sizeof(elevation), &size))
+            {
+                elevated = (elevation.TokenIsElevated != 0);
+            }
         }
-    }
 
-    if (token)
-    {
-        CloseHandle(token);
-    }
+        if (token)
+        {
+            CloseHandle(token);
+        }
 
-    return elevated;
+        return elevated;
+    };
+    static const bool cached_value = detection_func();
+    return use_cached_value ? cached_value : detection_func();
 }
 
 bool drop_elevated_privileges()

diff --git a/src/common/common.h b/src/common/common.h
@@ -61,7 +61,7 @@ enum WindowState
 WindowState get_window_state(HWND hwnd);
 
 // Returns true if the current process is running with elevated privileges
-bool is_process_elevated();
+bool is_process_elevated(const bool use_cached_value = true);
 
 // Drops the elevated privilages if present
 bool drop_elevated_privileges();
@@ -78,7 +78,7 @@ bool run_same_elevation(const std::wstring& file, const std::wstring& params);
 // Returns true if the current process is running from administrator account
 bool check_user_is_admin();
 
-//Returns true when one or more strings from vector found in string 
+// Returns true when one or more strings from vector found in string
 bool find_app_name_in_path(const std::wstring& where, const std::vector<std::wstring>& what);
 
 // Get the executable path or module name for modern apps

diff --git a/src/modules/fancyzones/lib/FancyZones.cpp b/src/modules/fancyzones/lib/FancyZones.cpp
@@ -867,10 +867,8 @@ void FancyZones::UpdateDragState(HWND window, require_write_lock) noexcept
         m_dragEnabled = !(shift | mouse);
     }
 
-    const bool windowElevated = IsProcessOfWindowElevated(window);
-    static const bool meElevated = is_process_elevated();
     static bool warning_shown = false;
-    if (windowElevated && !meElevated)
+    if (!is_process_elevated() && IsProcessOfWindowElevated(window))
     {
         m_dragEnabled = false;
         if (!warning_shown && !is_cant_drag_elevated_warning_disabled())

diff --git a/src/settings/main.cpp b/src/settings/main.cpp
@@ -567,7 +567,7 @@ int WINAPI WinMain(_In_ HINSTANCE hInstance, _In_opt_ HINSTANCE hPrevInstance, _
     Trace::RegisterProvider();
     CoInitialize(nullptr);
 
-    const bool should_try_drop_privileges = !initialize_com_security_policy_for_webview() && is_process_elevated();
+    const bool should_try_drop_privileges = !initialize_com_security_policy_for_webview() && is_process_elevated(false);
 
     if (should_try_drop_privileges)
     {