From ea345f485491a6d01f8a59a19d7546beaad5e7de Mon Sep 17 00:00:00 2001 From: qhy040404 Date: Fri, 31 May 2024 23:30:25 +0800 Subject: [PATCH] use client to determine whether to redirect --- .../LifeCycle/HutaoActivationArguments.cs | 5 +- .../InterProcess/PipePacketCommand.cs | 3 + .../LifeCycle/InterProcess/PipePacketType.cs | 2 +- .../InterProcess/PrivateNamedPipeClient.cs | 56 ++++++++++++++----- .../PrivateNamedPipeMessageDispatcher.cs | 7 +++ .../InterProcess/PrivateNamedPipeServer.cs | 41 ++++++-------- 6 files changed, 71 insertions(+), 43 deletions(-) diff --git a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/HutaoActivationArguments.cs b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/HutaoActivationArguments.cs index 052a1a34..41b8dedb 100644 --- a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/HutaoActivationArguments.cs +++ b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/HutaoActivationArguments.cs @@ -8,8 +8,6 @@ namespace Snap.Hutao.Core.LifeCycle; internal sealed class HutaoActivationArguments { - public bool IsElevated { get; set; } - public bool IsRedirectTo { get; set; } public bool IsToastActivated { get; set; } @@ -20,11 +18,10 @@ internal sealed class HutaoActivationArguments public string? LaunchActivatedArguments { get; set; } - public static HutaoActivationArguments FromAppActivationArguments(AppActivationArguments args, bool isRedirected = false, bool isElevated = false) + public static HutaoActivationArguments FromAppActivationArguments(AppActivationArguments args, bool isRedirected = false) { HutaoActivationArguments result = new() { - IsElevated = isElevated, IsRedirectTo = isRedirected, }; diff --git a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PipePacketCommand.cs b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PipePacketCommand.cs index 58f25a1e..643391a4 100644 --- a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PipePacketCommand.cs +++ b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PipePacketCommand.cs @@ -6,6 +6,9 @@ namespace Snap.Hutao.Core.LifeCycle.InterProcess; internal enum PipePacketCommand : byte { None = 0, + Connect = 1, RedirectActivation = 10, + + Exit = 30, } \ No newline at end of file diff --git a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PipePacketType.cs b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PipePacketType.cs index a2cadeca..6d70a112 100644 --- a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PipePacketType.cs +++ b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PipePacketType.cs @@ -8,5 +8,5 @@ internal enum PipePacketType : byte None = 0, Request = 1, Response = 2, - Termination = 3, + SessionTermination = 3, } \ No newline at end of file diff --git a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PrivateNamedPipeClient.cs b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PrivateNamedPipeClient.cs index 25bc8372..34941413 100644 --- a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PrivateNamedPipeClient.cs +++ b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PrivateNamedPipeClient.cs @@ -17,41 +17,67 @@ internal sealed partial class PrivateNamedPipeClient : IDisposable { if (clientStream.TryConnectOnce()) { - bool shouldElevate = false; + Span headerSpan = stackalloc byte[sizeof(PipePacketHeader)]; + bool serverElevated = false; { + // Connect + PipePacketHeader connectPacket = default; + connectPacket.Version = 1; + connectPacket.Type = PipePacketType.Request; + connectPacket.Command = PipePacketCommand.Connect; + + clientStream.Write(new(&connectPacket, sizeof(PipePacketHeader))); + } + + { + // Get previous instance elevated status + clientStream.ReadExactly(headerSpan); + fixed (byte* pHeader = headerSpan) + { + PipePacketHeader* header = (PipePacketHeader*)pHeader; + ReadOnlySpan content = clientStream.GetValidatedContent(header); + serverElevated = JsonSerializer.Deserialize(content); + } + } + + if (!serverElevated && runtimeOptions.IsElevated) + { + // Kill previous instance to use current elevated instance + PipePacketHeader killPacket = default; + killPacket.Version = 1; + killPacket.Type = PipePacketType.SessionTermination; + killPacket.Command = PipePacketCommand.Exit; + + clientStream.Write(new(&killPacket, sizeof(PipePacketHeader))); + clientStream.Flush(); + return false; + } + + { + // Redirect to previous instance PipePacketHeader redirectActivationPacket = default; redirectActivationPacket.Version = 1; redirectActivationPacket.Type = PipePacketType.Request; redirectActivationPacket.Command = PipePacketCommand.RedirectActivation; redirectActivationPacket.ContentType = PipePacketContentType.Json; - HutaoActivationArguments hutaoArgs = HutaoActivationArguments.FromAppActivationArguments(args, isRedirected: true, isElevated: runtimeOptions.IsElevated); + HutaoActivationArguments hutaoArgs = HutaoActivationArguments.FromAppActivationArguments(args, isRedirected: true); byte[] jsonBytes = JsonSerializer.SerializeToUtf8Bytes(hutaoArgs); clientStream.WritePacket(&redirectActivationPacket, jsonBytes); } { - Span headerSpan = stackalloc byte[sizeof(PipePacketHeader)]; - clientStream.ReadExactly(headerSpan); - fixed (byte* pHeader = headerSpan) - { - PipePacketHeader* header = (PipePacketHeader*)pHeader; - ReadOnlySpan content = clientStream.GetValidatedContent(header); - shouldElevate = JsonSerializer.Deserialize(content); - } - } - - { + // Terminate session PipePacketHeader terminationPacket = default; terminationPacket.Version = 1; - terminationPacket.Type = PipePacketType.Termination; + terminationPacket.Type = PipePacketType.SessionTermination; clientStream.Write(new(&terminationPacket, sizeof(PipePacketHeader))); } clientStream.Flush(); - return !shouldElevate; + return true; } return false; diff --git a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PrivateNamedPipeMessageDispatcher.cs b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PrivateNamedPipeMessageDispatcher.cs index 378ca343..9047a5fd 100644 --- a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PrivateNamedPipeMessageDispatcher.cs +++ b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PrivateNamedPipeMessageDispatcher.cs @@ -18,4 +18,11 @@ internal sealed partial class PrivateNamedPipeMessageDispatcher serviceProvider.GetRequiredService().Activate(args); } + + public void Exit() + { + ITaskContext taskContext = serviceProvider.GetRequiredService(); + App app = serviceProvider.GetRequiredService(); + taskContext.BeginInvokeOnMainThread(app.Exit); + } } \ No newline at end of file diff --git a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PrivateNamedPipeServer.cs b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PrivateNamedPipeServer.cs index f7874d37..6e5f20a0 100644 --- a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PrivateNamedPipeServer.cs +++ b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PrivateNamedPipeServer.cs @@ -1,7 +1,6 @@ // Copyright (c) DGP Studio. All rights reserved. // Licensed under the MIT license. -using System.Diagnostics; using System.IO.Pipes; using System.Security.AccessControl; using System.Security.Principal; @@ -13,6 +12,8 @@ internal sealed partial class PrivateNamedPipeServer : IDisposable { private readonly PrivateNamedPipeMessageDispatcher messageDispatcher; private readonly RuntimeOptions runtimeOptions; + private readonly ITaskContext taskContext; + private readonly App app; private readonly CancellationTokenSource serverTokenSource = new(); private readonly SemaphoreSlim serverSemaphore = new(1); @@ -23,6 +24,8 @@ internal sealed partial class PrivateNamedPipeServer : IDisposable { messageDispatcher = serviceProvider.GetRequiredService(); runtimeOptions = serviceProvider.GetRequiredService(); + taskContext = serviceProvider.GetRequiredService(); + app = serviceProvider.GetRequiredService(); PipeSecurity? pipeSecurity = default; @@ -76,7 +79,6 @@ internal sealed partial class PrivateNamedPipeServer : IDisposable private unsafe void RunPacketSession(NamedPipeServerStream serverStream, CancellationToken token) { Span headerSpan = stackalloc byte[sizeof(PipePacketHeader)]; - bool shouldElevate = false; bool sessionTerminated = false; while (serverStream.IsConnected && !sessionTerminated && !token.IsCancellationRequested) { @@ -87,33 +89,26 @@ internal sealed partial class PrivateNamedPipeServer : IDisposable switch ((header->Type, header->Command, header->ContentType)) { - case (PipePacketType.Request, PipePacketCommand.RedirectActivation, PipePacketContentType.Json): - ReadOnlySpan content = serverStream.GetValidatedContent(header); - HutaoActivationArguments? hutaoArgs = JsonSerializer.Deserialize(content); - ArgumentNullException.ThrowIfNull(hutaoArgs); + case (PipePacketType.Request, PipePacketCommand.Connect, _): + PipePacketHeader elevatedPacket = default; + elevatedPacket.Version = 1; + elevatedPacket.Type = PipePacketType.Response; + elevatedPacket.ContentType = PipePacketContentType.Json; - PipePacketHeader responsePacket = default; - responsePacket.Version = 1; - responsePacket.Type = PipePacketType.Response; - responsePacket.ContentType = PipePacketContentType.Json; - - shouldElevate = !runtimeOptions.IsElevated && hutaoArgs.IsElevated; - byte[] jsonBytes = JsonSerializer.SerializeToUtf8Bytes(shouldElevate); - serverStream.WritePacket(&responsePacket, jsonBytes); - serverStream.Flush(); - - if (!shouldElevate) - { - messageDispatcher.RedirectActivation(hutaoArgs); - } + byte[] elevatedBytes = JsonSerializer.SerializeToUtf8Bytes(runtimeOptions.IsElevated); + serverStream.WritePacket(&elevatedPacket, elevatedBytes); break; - case (PipePacketType.Termination, _, _): + case (PipePacketType.Request, PipePacketCommand.RedirectActivation, PipePacketContentType.Json): + ReadOnlySpan content = serverStream.GetValidatedContent(header); + messageDispatcher.RedirectActivation(JsonSerializer.Deserialize(content)); + break; + case (PipePacketType.SessionTermination, _, _): serverStream.Disconnect(); sessionTerminated = true; - if (shouldElevate) + if (header->Command is PipePacketCommand.Exit) { - Process.GetCurrentProcess().Kill(); + messageDispatcher.Exit(); } return;