From 33a81416fc1fd518a99c0b726164696f621e536b Mon Sep 17 00:00:00 2001
From: Den Delimarsky <1389609+dend@users.noreply.github.com>
Date: Mon, 3 May 2021 08:04:53 -0700
Subject: [PATCH] Proper handling for PID binding Now Espresso can be connected
 directly to the PowerToys PID.

---
 .../Espresso.Shell/Espresso.Shell.csproj      |  4 +++
 .../espresso/Espresso.Shell/Program.cs        | 30 +++++++++++++++++--
 src/modules/espresso/Espresso/dllmain.cpp     |  4 +--
 3 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/src/modules/espresso/Espresso.Shell/Espresso.Shell.csproj b/src/modules/espresso/Espresso.Shell/Espresso.Shell.csproj
index 4e2bd2fb4c..ea61957bbd 100644
--- a/src/modules/espresso/Espresso.Shell/Espresso.Shell.csproj
+++ b/src/modules/espresso/Espresso.Shell/Espresso.Shell.csproj
@@ -43,6 +43,10 @@
     <PackageReference Include="System.Runtime.Caching" Version="6.0.0-preview.1.21102.12" />
   </ItemGroup>
 
+  <ItemGroup>
+    <ProjectReference Include="..\..\..\common\ManagedCommon\ManagedCommon.csproj" />
+  </ItemGroup>
+
   <ItemGroup>
     <Compile Update="Program.cs">
       <CopyToOutputDirectory>Never</CopyToOutputDirectory>
diff --git a/src/modules/espresso/Espresso.Shell/Program.cs b/src/modules/espresso/Espresso.Shell/Program.cs
index 91393ea80e..4e16cd1135 100644
--- a/src/modules/espresso/Espresso.Shell/Program.cs
+++ b/src/modules/espresso/Espresso.Shell/Program.cs
@@ -4,6 +4,7 @@
 
 using Espresso.Shell.Core;
 using Espresso.Shell.Models;
+using ManagedCommon;
 using Newtonsoft.Json;
 using NLog;
 using System;
@@ -89,16 +90,30 @@ namespace Espresso.Shell
 
             timeOption.Required = false;
 
+            var powerToysPidOption = new Option<int>(
+                    aliases: new[] { "--ptpid", "-p" },
+                    getDefaultValue: () => 0,
+                    description: "Reference to PowerToys PID, when executed under the application context.")
+            {
+                Argument = new Argument<int>(() => 0)
+                {
+                    Arity = ArgumentArity.ZeroOrOne,
+                },
+            };
+
+            powerToysPidOption.Required = false;
+
             var rootCommand = new RootCommand
             {
                 configOption,
                 displayOption,
-                timeOption
+                timeOption,
+                powerToysPidOption
             };
 
             rootCommand.Description = appName;
 
-            rootCommand.Handler = CommandHandler.Create<string, bool, long>(HandleCommandLineArguments);
+            rootCommand.Handler = CommandHandler.Create<string, bool, long, int>(HandleCommandLineArguments);
 
             return rootCommand.InvokeAsync(args).Result;
         }
@@ -111,11 +126,12 @@ namespace Espresso.Shell
             Environment.Exit(exitCode);
         }
 
-        private static void HandleCommandLineArguments(string config, bool displayOn, long timeLimit)
+        private static void HandleCommandLineArguments(string config, bool displayOn, long timeLimit, int ptpid)
         {
             log.Info($"The value for --config is: {config}");
             log.Info($"The value for --display-on is: {displayOn}");
             log.Info($"The value for --time-limit is: {timeLimit}");
+            log.Info($"The value for --ptpid is: {ptpid}");
 
             if (!string.IsNullOrWhiteSpace(config))
             {
@@ -150,6 +166,14 @@ namespace Espresso.Shell
                     log.Info(errorString);
                     log.Debug(errorString);
                 }
+
+                if (ptpid != 0)
+                {
+                    RunnerHelper.WaitForPowerToysRunner(ptpid, () =>
+                    {
+                        Environment.Exit(0);
+                    });
+                }
             }
             else
             {
diff --git a/src/modules/espresso/Espresso/dllmain.cpp b/src/modules/espresso/Espresso/dllmain.cpp
index 6b7b897ceb..1b877910a3 100644
--- a/src/modules/espresso/Espresso/dllmain.cpp
+++ b/src/modules/espresso/Espresso/dllmain.cpp
@@ -81,8 +81,8 @@ private:
         // Get the configuration file that will be passed to the process.
         std::wstring espresso_settings_location = PTSettingsHelper::get_module_save_file_location(MODULE_NAME);
 
-        std::wstring executable_args = L"--config " + espresso_settings_location;
-        //executable_args.append(std::to_wstring(powertoys_pid));
+        std::wstring executable_args = L"--config " + espresso_settings_location + L" --ptpid " + std::to_wstring(powertoys_pid);
+        Logger::trace(L"Espresso launching with parameters: " + executable_args);
 
         SHELLEXECUTEINFOW sei{ sizeof(sei) };
         sei.fMask = { SEE_MASK_NOCLOSEPROCESS | SEE_MASK_FLAG_NO_UI };