From 0b86941a3bb75e74f9bcdab9225b0205f2747df4 Mon Sep 17 00:00:00 2001
From: Bartosz Sosnowski <bartosz@janeasystems.com>
Date: Thu, 19 Sep 2019 15:33:28 +0200
Subject: [PATCH] Make FancyZone distinguish different modern apps.

---
 src/modules/fancyzones/lib/FancyZones.cpp | 14 +++---
 src/modules/fancyzones/lib/ZoneWindow.cpp |  7 ++-
 src/modules/fancyzones/lib/util.h         | 57 ++++++++++++++++++-----
 3 files changed, 54 insertions(+), 24 deletions(-)

diff --git a/src/modules/fancyzones/lib/FancyZones.cpp b/src/modules/fancyzones/lib/FancyZones.cpp
index 2dc732cecf..efeb224528 100644
--- a/src/modules/fancyzones/lib/FancyZones.cpp
+++ b/src/modules/fancyzones/lib/FancyZones.cpp
@@ -168,12 +168,11 @@ IFACEMETHODIMP_(void) FancyZones::WindowCreated(HWND window) noexcept
 {
     if (m_settings->GetSettings().appLastZone_moveWindows)
     {
-        wchar_t processPath[MAX_PATH] = { 0 };
-        DWORD modulePathSize = GetProcessPath(window, processPath, static_cast<DWORD>(MAX_PATH));
-        if (modulePathSize > 0) 
+        auto processPath = GetProcessPath(window);
+        if (!processPath.empty()) 
         {
             INT zoneIndex = -1;
-            LRESULT res = RegistryHelpers::GetAppLastZone(window, processPath, &zoneIndex);
+            LRESULT res = RegistryHelpers::GetAppLastZone(window, processPath.data(), &zoneIndex);
             if ((res == ERROR_SUCCESS) && (zoneIndex != -1))
             {
                 MoveWindowIntoZoneByIndex(window, zoneIndex);
@@ -670,11 +669,10 @@ void FancyZones::MoveSizeEndInternal(HWND window, POINT const& ptScreen, require
     {
         ::RemoveProp(window, ZONE_STAMP);
 
-        wchar_t processPath[MAX_PATH]{};
-        DWORD processPathSize = GetProcessPath(window, processPath, static_cast<DWORD>(MAX_PATH));
-        if (processPathSize > 0)
+        auto processPath = GetProcessPath(window);
+        if (!processPath.empty())
         {
-            RegistryHelpers::SaveAppLastZone(window, processPath, -1);
+            RegistryHelpers::SaveAppLastZone(window, processPath.data(), -1);
         }
     }
 }
diff --git a/src/modules/fancyzones/lib/ZoneWindow.cpp b/src/modules/fancyzones/lib/ZoneWindow.cpp
index d7b1666c7e..837a5adb0a 100644
--- a/src/modules/fancyzones/lib/ZoneWindow.cpp
+++ b/src/modules/fancyzones/lib/ZoneWindow.cpp
@@ -1252,14 +1252,13 @@ int ZoneWindow::GetSwitchButtonIndexFromPoint(POINT ptClient) noexcept
 
 IFACEMETHODIMP_(void) ZoneWindow::SaveWindowProcessToZoneIndex(HWND window) noexcept
 {
-    wchar_t processPath[MAX_PATH] = { 0 };
-    DWORD processPathSize = GetProcessPath(window, processPath, static_cast<DWORD>(MAX_PATH));
-    if (processPathSize > 0)
+    auto processPath = GetProcessPath(window);
+    if (!processPath.empty())
     {
         DWORD zoneIndex = static_cast<DWORD>(m_activeZoneSet->GetZoneIndexFromWindow(window));
         if (zoneIndex != -1)
         {
-            RegistryHelpers::SaveAppLastZone(window, processPath, zoneIndex);
+            RegistryHelpers::SaveAppLastZone(window, processPath.data(), zoneIndex);
         }
     }
 }
diff --git a/src/modules/fancyzones/lib/util.h b/src/modules/fancyzones/lib/util.h
index 2204036375..532ae64f3e 100644
--- a/src/modules/fancyzones/lib/util.h
+++ b/src/modules/fancyzones/lib/util.h
@@ -134,19 +134,52 @@ inline void ParseDeviceId(PCWSTR deviceId, PWSTR parsedId, size_t size)
     }
 }
 
-inline DWORD GetProcessPath(HWND window, LPWSTR processPath, DWORD processPathMaxSize) noexcept
+inline std::wstring GetProcessPathByPID(DWORD pid)
+{
+    wil::unique_handle process(OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, TRUE, pid));
+    std::wstring name;
+    if (process && process.get() != INVALID_HANDLE_VALUE)
+    {
+        name.resize(MAX_PATH);
+        DWORD name_length = static_cast<DWORD>(name.length());
+        QueryFullProcessImageNameW(process.get(), 0, (LPWSTR)name.data(), &name_length);
+        name.resize(name_length);
+    }
+    return name;
+}
+
+inline std::wstring GetProcessPath(HWND window) noexcept
 {
+    const static std::wstring app_frame_host = L"ApplicationFrameHost.exe";
     DWORD pid{};
     GetWindowThreadProcessId(window, &pid);
-
-    DWORD numCopiedChars = 0;
-    wil::unique_handle windowProcessHandle(OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, TRUE, pid));
-    if (windowProcessHandle && (windowProcessHandle.get() != INVALID_HANDLE_VALUE))
+    auto name = GetProcessPathByPID(pid);
+    if (name.length() >= app_frame_host.length() &&
+        name.compare(name.length() - app_frame_host.length(), app_frame_host.length(), app_frame_host) == 0)
     {
-        // numCopiedChars first holds the size of processPath[], will then hold amount of characters returned by QueryFullProcessImageNameW
-        // if QueryFullProcessImageNameW fails, numCopiedChars will be zero.
-        numCopiedChars = processPathMaxSize;
-        QueryFullProcessImageNameW(windowProcessHandle.get(), 0, processPath, &numCopiedChars);
-    }
-    return numCopiedChars;
-}
\ No newline at end of file
+        // It is a UWP app. We will enumarate the windows and look for one created
+        // by something with a different PID
+        DWORD new_pid = pid;
+        EnumChildWindows(window, [](HWND hwnd, LPARAM param) -> BOOL
+        {
+            auto new_pid_ptr = reinterpret_cast<DWORD*>(param);
+            DWORD pid;
+            GetWindowThreadProcessId(hwnd, &pid);
+            if (pid != *new_pid_ptr)
+            {
+                *new_pid_ptr = pid;
+                return FALSE;
+            }
+            else
+            {
+                return TRUE;
+            }
+        }, reinterpret_cast<LPARAM>(&new_pid));
+        // If we have a new pid, get the new name.
+        if (new_pid != pid)
+        {
+            return GetProcessPathByPID(new_pid);
+        }
+    }
+    return name;
+}