diff --git a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/HutaoActivationArguments.cs b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/HutaoActivationArguments.cs index 41b8dedb..052a1a34 100644 --- a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/HutaoActivationArguments.cs +++ b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/HutaoActivationArguments.cs @@ -8,6 +8,8 @@ namespace Snap.Hutao.Core.LifeCycle; internal sealed class HutaoActivationArguments { + public bool IsElevated { get; set; } + public bool IsRedirectTo { get; set; } public bool IsToastActivated { get; set; } @@ -18,10 +20,11 @@ internal sealed class HutaoActivationArguments public string? LaunchActivatedArguments { get; set; } - public static HutaoActivationArguments FromAppActivationArguments(AppActivationArguments args, bool isRedirected = false) + public static HutaoActivationArguments FromAppActivationArguments(AppActivationArguments args, bool isRedirected = false, bool isElevated = false) { HutaoActivationArguments result = new() { + IsElevated = isElevated, IsRedirectTo = isRedirected, }; diff --git a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PipeStreamExtension.cs b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PipeStreamExtension.cs new file mode 100644 index 00000000..1ac4c516 --- /dev/null +++ b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PipeStreamExtension.cs @@ -0,0 +1,28 @@ +// Copyright (c) DGP Studio. All rights reserved. +// Licensed under the MIT license. + +using Snap.Hutao.Core.ExceptionService; +using System.IO.Hashing; +using System.IO.Pipes; + +namespace Snap.Hutao.Core.LifeCycle.InterProcess; + +internal static class PipeStreamExtension +{ + public static unsafe byte[] GetValidatedContent(this PipeStream stream, PipePacketHeader* header) + { + byte[] content = new byte[header->ContentLength]; + stream.ReadAtLeast(content, header->ContentLength, false); + HutaoException.ThrowIf(XxHash64.HashToUInt64(content) != header->Checksum, "PipePacket Content Hash incorrect"); + return content; + } + + public static unsafe void WritePacket(this PipeStream stream, PipePacketHeader* header, byte[] content) + { + header->ContentLength = content.Length; + header->Checksum = XxHash64.HashToUInt64(content); + + stream.Write(new(header, sizeof(PipePacketHeader))); + stream.Write(content); + } +} 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 fdecdfef..1fc1926d 100644 --- a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PrivateNamedPipeClient.cs +++ b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PrivateNamedPipeClient.cs @@ -10,12 +10,19 @@ namespace Snap.Hutao.Core.LifeCycle.InterProcess; [Injection(InjectAs.Singleton)] internal sealed class PrivateNamedPipeClient : IDisposable { + private readonly RuntimeOptions runtimeOptions; private readonly NamedPipeClientStream clientStream = new(".", "Snap.Hutao.PrivateNamedPipe", PipeDirection.InOut, PipeOptions.Asynchronous | PipeOptions.WriteThrough); + public PrivateNamedPipeClient(IServiceProvider serviceProvider) + { + runtimeOptions = serviceProvider.GetRequiredService(); + } + public unsafe bool TryRedirectActivationTo(AppActivationArguments args) { if (clientStream.TryConnectOnce()) { + bool shouldElevate = false; { PipePacketHeader redirectActivationPacket = default; redirectActivationPacket.Version = 1; @@ -23,14 +30,21 @@ internal sealed class PrivateNamedPipeClient : IDisposable redirectActivationPacket.Command = PipePacketCommand.RedirectActivation; redirectActivationPacket.ContentType = PipePacketContentType.Json; - HutaoActivationArguments hutaoArgs = HutaoActivationArguments.FromAppActivationArguments(args, isRedirected: true); + HutaoActivationArguments hutaoArgs = HutaoActivationArguments.FromAppActivationArguments(args, isRedirected: true, isElevated: runtimeOptions.IsElevated); byte[] jsonBytes = JsonSerializer.SerializeToUtf8Bytes(hutaoArgs); - redirectActivationPacket.ContentLength = jsonBytes.Length; - redirectActivationPacket.Checksum = XxHash64.HashToUInt64(jsonBytes); + clientStream.WritePacket(&redirectActivationPacket, jsonBytes); + } - clientStream.Write(new(&redirectActivationPacket, sizeof(PipePacketHeader))); - clientStream.Write(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); + } } { @@ -42,7 +56,7 @@ internal sealed class PrivateNamedPipeClient : IDisposable } clientStream.Flush(); - return true; + return !shouldElevate; } return false; 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 1dc0ba16..f7874d37 100644 --- a/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PrivateNamedPipeServer.cs +++ b/src/Snap.Hutao/Snap.Hutao/Core/LifeCycle/InterProcess/PrivateNamedPipeServer.cs @@ -1,8 +1,7 @@ // Copyright (c) DGP Studio. All rights reserved. // Licensed under the MIT license. -using Snap.Hutao.Core.ExceptionService; -using System.IO.Hashing; +using System.Diagnostics; using System.IO.Pipes; using System.Security.AccessControl; using System.Security.Principal; @@ -74,17 +73,10 @@ internal sealed partial class PrivateNamedPipeServer : IDisposable } } - private static unsafe byte[] GetValidatedContent(NamedPipeServerStream serverStream, PipePacketHeader* header) - { - byte[] content = new byte[header->ContentLength]; - serverStream.ReadAtLeast(content, header->ContentLength, false); - HutaoException.ThrowIf(XxHash64.HashToUInt64(content) != header->Checksum, "PipePacket Content Hash incorrect"); - return content; - } - 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) { @@ -96,12 +88,34 @@ internal sealed partial class PrivateNamedPipeServer : IDisposable switch ((header->Type, header->Command, header->ContentType)) { case (PipePacketType.Request, PipePacketCommand.RedirectActivation, PipePacketContentType.Json): - ReadOnlySpan content = GetValidatedContent(serverStream, header); - messageDispatcher.RedirectActivation(JsonSerializer.Deserialize(content)); + ReadOnlySpan content = serverStream.GetValidatedContent(header); + HutaoActivationArguments? hutaoArgs = JsonSerializer.Deserialize(content); + ArgumentNullException.ThrowIfNull(hutaoArgs); + + 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); + } + break; case (PipePacketType.Termination, _, _): serverStream.Disconnect(); sessionTerminated = true; + if (shouldElevate) + { + Process.GetCurrentProcess().Kill(); + } + return; } }