Awake - VISEGRADRELAY_08152024 (#34316)

* Update with bug fixes for tray icon and support for parent process

* Process information enum

* Update the docs

* Fix spelling

* Make sure that PID is used in PT config flow
This commit is contained in:
Den Delimarsky
2024-08-20 05:24:37 -07:00
committed by GitHub
parent f8269af125
commit 808e6220bc
11 changed files with 154 additions and 51 deletions

View File

@@ -17,6 +17,6 @@ namespace Awake.Core
// Format of the build ID is: CODENAME_MMDDYYYY, where MMDDYYYY
// is representative of the date when the last change was made before
// the pull request is issued.
internal const string BuildId = "DAISY023_04102024";
internal const string BuildId = "VISEGRADRELAY_08152024";
}
}

View File

@@ -5,10 +5,12 @@
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Diagnostics;
using System.Drawing;
using System.Globalization;
using System.IO;
using System.Reactive.Linq;
using System.Runtime.InteropServices;
using System.Text;
using System.Text.Json;
using System.Threading;
@@ -62,7 +64,7 @@ namespace Awake.Core
{
Thread monitorThread = new(() =>
{
Thread.CurrentThread.IsBackground = true;
Thread.CurrentThread.IsBackground = false;
while (true)
{
ExecutionState state = _stateQueue.Take();
@@ -126,10 +128,18 @@ namespace Awake.Core
_stateQueue.Add(ExecutionState.ES_CONTINUOUS);
// Next, make sure that any existing background threads are terminated.
_tokenSource.Cancel();
_tokenSource.Dispose();
if (_tokenSource != null)
{
_tokenSource.Cancel();
_tokenSource.Dispose();
_tokenSource = new CancellationTokenSource();
}
else
{
Logger.LogWarning("The token source was null.");
}
_tokenSource = new CancellationTokenSource();
Logger.LogInfo("Instantiating of new token source and thread token completed.");
}
@@ -137,12 +147,11 @@ namespace Awake.Core
{
PowerToysTelemetry.Log.WriteEvent(new Telemetry.AwakeIndefinitelyKeepAwakeEvent());
CancelExistingThread();
_stateQueue.Add(ComputeAwakeState(keepDisplayOn));
TrayHelper.SetShellIcon(TrayHelper.HiddenWindowHandle, $"{Constants.FullAppName} [{Resources.AWAKE_TRAY_TEXT_INDEFINITE}]", _indefiniteIcon, TrayIconAction.Update);
CancelExistingThread();
_stateQueue.Add(ComputeAwakeState(keepDisplayOn));
if (IsUsingPowerToysConfig)
{
try
@@ -417,5 +426,18 @@ namespace Awake.Core
}
}
}
public static Process? GetParentProcess()
{
return GetParentProcess(Process.GetCurrentProcess().Handle);
}
private static Process? GetParentProcess(IntPtr handle)
{
ProcessBasicInformation pbi = default;
int status = Bridge.NtQueryInformationProcess(handle, 0, ref pbi, Marshal.SizeOf<ProcessBasicInformation>(), out _);
return status != 0 ? null : Process.GetProcessById(pbi.InheritedFromUniqueProcessId.ToInt32());
}
}
}

View File

@@ -0,0 +1,20 @@
// Copyright (c) Microsoft Corporation
// The Microsoft Corporation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.
using System;
using System.Runtime.InteropServices;
namespace Awake.Core.Models
{
[StructLayout(LayoutKind.Sequential)]
internal struct ProcessBasicInformation
{
public IntPtr ExitStatus;
public IntPtr PebAddress;
public IntPtr AffinityMask;
public IntPtr BasePriority;
public IntPtr UniquePID;
public IntPtr InheritedFromUniqueProcessId;
}
}

View File

@@ -56,13 +56,13 @@ namespace Awake.Core.Native
[return: MarshalAs(UnmanagedType.Bool)]
internal static extern bool DestroyMenu(IntPtr hMenu);
[DllImport("user32.dll")]
[DllImport("user32.dll", SetLastError = true)]
internal static extern bool DestroyWindow(IntPtr hWnd);
[DllImport("user32.dll")]
[DllImport("user32.dll", SetLastError = true)]
internal static extern void PostQuitMessage(int nExitCode);
[DllImport("shell32.dll")]
[DllImport("shell32.dll", SetLastError = true)]
internal static extern bool Shell_NotifyIcon(int dwMessage, ref NotifyIconData pnid);
[DllImport("user32.dll", SetLastError = true)]
@@ -83,14 +83,14 @@ namespace Awake.Core.Native
[DllImport("user32.dll", SetLastError = true)]
internal static extern bool ShowWindow(IntPtr hWnd, int nCmdShow);
[DllImport("user32.dll")]
[DllImport("user32.dll", SetLastError = true)]
[return: MarshalAs(UnmanagedType.Bool)]
internal static extern bool GetCursorPos(out Point lpPoint);
[DllImport("user32.dll")]
[DllImport("user32.dll", SetLastError = true)]
internal static extern bool ScreenToClient(IntPtr hWnd, ref Point lpPoint);
[DllImport("user32.dll")]
[DllImport("user32.dll", SetLastError = true)]
internal static extern bool GetMessage(out Msg lpMsg, IntPtr hWnd, uint wMsgFilterMin, uint wMsgFilterMax);
[DllImport("user32.dll", SetLastError = true)]
@@ -100,7 +100,10 @@ namespace Awake.Core.Native
[return: MarshalAs(UnmanagedType.Bool)]
internal static extern bool SetMenuInfo(IntPtr hMenu, ref MenuInfo lpcmi);
[DllImport("user32.dll")]
[DllImport("user32.dll", SetLastError = true)]
internal static extern bool SetForegroundWindow(IntPtr hWnd);
[DllImport("ntdll.dll")]
internal static extern int NtQueryInformationProcess(IntPtr processHandle, int processInformationClass, ref ProcessBasicInformation processInformation, int processInformationLength, out int returnLength);
}
}

View File

@@ -43,11 +43,6 @@ namespace Awake.Core
HiddenWindowHandle = IntPtr.Zero;
}
public static void InitializeTray(string text, Icon icon)
{
CreateHiddenWindow(icon, text);
}
private static void ShowContextMenu(IntPtr hWnd)
{
if (TrayMenu != IntPtr.Zero)
@@ -88,7 +83,7 @@ namespace Awake.Core
}
}
private static void CreateHiddenWindow(Icon icon, string text)
public static void InitializeTray(Icon icon, string text)
{
IntPtr hWnd = IntPtr.Zero;
@@ -143,7 +138,13 @@ namespace Awake.Core
Bridge.UpdateWindow(hWnd);
SetShellIcon(hWnd, text, icon);
});
}).Wait();
Task.Run(() =>
{
RunOnMainThread(() =>
{
RunMessageLoop();
});
});
@@ -151,6 +152,8 @@ namespace Awake.Core
internal static void SetShellIcon(IntPtr hWnd, string text, Icon? icon, TrayIconAction action = TrayIconAction.Add)
{
Logger.LogInfo($"Setting the shell icon.\nText: {text}\nAction: {action}");
int message = Native.Constants.NIM_ADD;
switch (action)
@@ -168,6 +171,8 @@ namespace Awake.Core
if (action == TrayIconAction.Add || action == TrayIconAction.Update)
{
Logger.LogInfo($"Adding or updating tray icon. HIcon handle is {icon?.Handle}\nHWnd: {hWnd}");
_notifyIconData = new NotifyIconData
{
CbSize = Marshal.SizeOf(typeof(NotifyIconData)),

View File

@@ -44,6 +44,7 @@ namespace Awake
internal static readonly string[] AliasesTimeOption = ["--time-limit", "-t"];
internal static readonly string[] AliasesPidOption = ["--pid", "-p"];
internal static readonly string[] AliasesExpireAtOption = ["--expire-at", "-e"];
internal static readonly string[] AliasesParentPidOption = ["--use-parent-pid", "-u"];
private static readonly Icon _defaultAwakeIcon = new(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "Assets/Awake/awake.ico"));
@@ -116,6 +117,12 @@ namespace Awake
IsRequired = false,
};
var parentPidOption = new Option<bool>(AliasesParentPidOption, () => true, Resources.AWAKE_CMD_PARENT_PID_OPTION)
{
Arity = ArgumentArity.ZeroOrOne,
IsRequired = false,
};
RootCommand? rootCommand =
[
configOption,
@@ -123,10 +130,11 @@ namespace Awake
timeOption,
pidOption,
expireAtOption,
parentPidOption,
];
rootCommand.Description = Core.Constants.AppName;
rootCommand.SetHandler(HandleCommandLineArguments, configOption, displayOption, timeOption, pidOption, expireAtOption);
rootCommand.SetHandler(HandleCommandLineArguments, configOption, displayOption, timeOption, pidOption, expireAtOption, parentPidOption);
return rootCommand.InvokeAsync(args).Result;
}
@@ -153,9 +161,9 @@ namespace Awake
Manager.CompleteExit(exitCode);
}
private static void HandleCommandLineArguments(bool usePtConfig, bool displayOn, uint timeLimit, int pid, string expireAt)
private static void HandleCommandLineArguments(bool usePtConfig, bool displayOn, uint timeLimit, int pid, string expireAt, bool useParentPid)
{
if (pid == 0)
if (pid == 0 && !useParentPid)
{
Logger.LogInfo("No PID specified. Allocating console...");
Manager.AllocateConsole();
@@ -175,11 +183,12 @@ namespace Awake
Logger.LogInfo($"The value for --time-limit is: {timeLimit}");
Logger.LogInfo($"The value for --pid is: {pid}");
Logger.LogInfo($"The value for --expire-at is: {expireAt}");
Logger.LogInfo($"The value for --use-parent-pid is: {useParentPid}");
// Start the monitor thread that will be used to track the current state.
Manager.StartMonitor();
TrayHelper.InitializeTray(Core.Constants.FullAppName, _defaultAwakeIcon);
TrayHelper.InitializeTray(_defaultAwakeIcon, Core.Constants.FullAppName);
var eventHandle = new EventWaitHandle(false, EventResetMode.ManualReset, PowerToys.Interop.Constants.AwakeExitEvent());
new Thread(() =>
@@ -209,12 +218,47 @@ namespace Awake
}
ScaffoldConfiguration(settingsPath);
if (pid != 0)
{
Logger.LogInfo($"Bound to target process while also using PowerToys settings: {pid}");
RunnerHelper.WaitForPowerToysRunner(pid, () =>
{
Logger.LogInfo($"Triggered PID-based exit handler for PID {pid}.");
Exit(Resources.AWAKE_EXIT_BINDING_HOOK_MESSAGE, 0);
});
}
}
catch (Exception ex)
{
Logger.LogError($"There was a problem with the configuration file. Make sure it exists.\n{ex.Message}");
}
}
else if (pid != 0 || useParentPid)
{
// Second, we snap to process-based execution. Because this is something that
// is snapped to a running entity, we only want to enable the ability to set
// indefinite keep-awake with the display settings that the user wants to set.
int targetPid = pid != 0 ? pid : useParentPid ? Manager.GetParentProcess()?.Id ?? 0 : 0;
if (targetPid != 0)
{
Logger.LogInfo($"Bound to target process: {targetPid}");
Manager.SetIndefiniteKeepAwake(displayOn);
RunnerHelper.WaitForPowerToysRunner(targetPid, () =>
{
Logger.LogInfo($"Triggered PID-based exit handler for PID {targetPid}.");
Exit(Resources.AWAKE_EXIT_BINDING_HOOK_MESSAGE, 0);
});
}
else
{
Logger.LogError("Not binding to any process.");
}
}
else
{
// Date-based binding takes precedence over timed configuration, so we want to
@@ -247,15 +291,6 @@ namespace Awake
}
}
}
if (pid != 0)
{
RunnerHelper.WaitForPowerToysRunner(pid, () =>
{
Logger.LogInfo($"Triggered PID-based exit handler for PID {pid}.");
Exit(Resources.AWAKE_EXIT_BINDING_HOOK_MESSAGE, 0);
});
}
}
private static void ScaffoldConfiguration(string settingsPath)

View File

@@ -114,6 +114,15 @@ namespace Awake.Properties {
}
}
/// <summary>
/// Looks up a localized string similar to Uses the parent process as the bound target - once the process terminates, Awake stops..
/// </summary>
internal static string AWAKE_CMD_PARENT_PID_OPTION {
get {
return ResourceManager.GetString("AWAKE_CMD_PARENT_PID_OPTION", resourceCulture);
}
}
/// <summary>
/// Looks up a localized string similar to Exit.
/// </summary>

View File

@@ -205,4 +205,7 @@
<value>s</value>
<comment>Used to display number of seconds in the system tray tooltip.</comment>
</data>
<data name="AWAKE_CMD_PARENT_PID_OPTION" xml:space="preserve">
<value>Uses the parent process as the bound target - once the process terminates, Awake stops.</value>
</data>
</root>