From 86704efcec94ed8c3de84b502dd8677a1206bd0e Mon Sep 17 00:00:00 2001
From: Andrey Nekrasov <yuyoyuppe@users.noreply.github.com>
Date: Mon, 13 Apr 2020 18:22:37 +0300
Subject: [PATCH] FancyZones: optimize elevation detection logic (#2103)

---
 src/common/common.cpp                     | 34 +++++++++++++----------
 src/common/common.h                       |  4 +--
 src/modules/fancyzones/lib/FancyZones.cpp |  4 +--
 src/settings/main.cpp                     |  2 +-
 4 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/src/common/common.cpp b/src/common/common.cpp
index 7e379b2976..c64c7178cb 100644
--- 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
index e99327a021..17cd475993 100644
--- 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
index ce05da6464..e932f4a57e 100644
--- 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
index 68411bd187..7f24e26b1f 100644
--- 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)
     {