This commit is contained in:
qhy040404
2024-05-31 22:03:57 +08:00
parent 53cda02071
commit 86b04bb5a3
4 changed files with 78 additions and 19 deletions

View File

@@ -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,
};

View File

@@ -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);
}
}

View File

@@ -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<RuntimeOptions>();
}
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<byte> headerSpan = stackalloc byte[sizeof(PipePacketHeader)];
clientStream.ReadExactly(headerSpan);
fixed (byte* pHeader = headerSpan)
{
PipePacketHeader* header = (PipePacketHeader*)pHeader;
ReadOnlySpan<byte> content = clientStream.GetValidatedContent(header);
shouldElevate = JsonSerializer.Deserialize<bool>(content);
}
}
{
@@ -42,7 +56,7 @@ internal sealed class PrivateNamedPipeClient : IDisposable
}
clientStream.Flush();
return true;
return !shouldElevate;
}
return false;

View File

@@ -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<byte> 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<byte> content = GetValidatedContent(serverStream, header);
messageDispatcher.RedirectActivation(JsonSerializer.Deserialize<HutaoActivationArguments>(content));
ReadOnlySpan<byte> content = serverStream.GetValidatedContent(header);
HutaoActivationArguments? hutaoArgs = JsonSerializer.Deserialize<HutaoActivationArguments>(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;
}
}