// 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. #nullable enable using System; using System.Collections.Generic; using System.ComponentModel; using System.Diagnostics; using System.IO; using System.Text; using System.Text.Json; using System.Text.Json.Serialization; using System.Threading; using System.Threading.Tasks; namespace PowerToys.FileLocksmithUI.Services { internal sealed class FileLocksmithQueryService { internal static readonly TimeSpan DefaultTimeout = TimeSpan.FromSeconds(30); private const string WorkerExecutableName = "FileLocksmithCLI.exe"; private const string WorkerArgument = "--worker-json"; private readonly Func _startInfoFactory; private readonly TimeSpan _timeout; private readonly Action? _processStarted; internal FileLocksmithQueryService() : this(CreateWorkerStartInfo, DefaultTimeout, null) { } internal FileLocksmithQueryService( Func startInfoFactory, TimeSpan timeout, Action? processStarted) { ArgumentNullException.ThrowIfNull(startInfoFactory); ArgumentOutOfRangeException.ThrowIfLessThanOrEqual(timeout, TimeSpan.Zero); _startInfoFactory = startInfoFactory; _timeout = timeout; _processStarted = processStarted; } internal async Task FindProcessesAsync( IReadOnlyCollection paths, CancellationToken cancellationToken) { ArgumentNullException.ThrowIfNull(paths); if (paths.Count == 0) { return new FileLocksmithQueryResult(FileLocksmithQueryStatus.Success, Array.Empty()); } using var process = new Process { StartInfo = _startInfoFactory(), }; try { if (!process.Start()) { return FailedToStart(); } } catch (Win32Exception) { return FailedToStart(); } catch (InvalidOperationException) { return FailedToStart(); } _processStarted?.Invoke(process.Id); #pragma warning disable CA2016 // These reads must drain the redirected pipes after timeout cancellation. var outputTask = process.StandardOutput.ReadToEndAsync(CancellationToken.None); var errorTask = process.StandardError.ReadToEndAsync(CancellationToken.None); #pragma warning restore CA2016 var inputClosed = false; try { using var timeoutCancellation = new CancellationTokenSource(_timeout); using var linkedCancellation = CancellationTokenSource.CreateLinkedTokenSource( timeoutCancellation.Token, cancellationToken); try { var request = JsonSerializer.Serialize(new WorkerRequest(paths)); await process.StandardInput.WriteAsync(request.AsMemory(), linkedCancellation.Token); process.StandardInput.Close(); inputClosed = true; await process.WaitForExitAsync(linkedCancellation.Token); } catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested) { await TerminateAsync(process); await Task.WhenAll(outputTask, errorTask); return new FileLocksmithQueryResult( FileLocksmithQueryStatus.TimedOut, Array.Empty()); } catch (IOException) { return Failed(process.HasExited ? process.ExitCode : null); } catch (InvalidOperationException) { return Failed(process.HasExited ? process.ExitCode : null); } if (process.ExitCode != 0) { await Task.WhenAll(outputTask, errorTask); return Failed(process.ExitCode); } var output = await outputTask; await errorTask; try { var response = JsonSerializer.Deserialize(output); if (response?.Processes is null) { return MalformedOutput(); } var processes = new List(response.Processes.Length); foreach (var processInfo in response.Processes) { if (processInfo.Name is null || processInfo.User is null || processInfo.Files is null) { return MalformedOutput(); } processes.Add(new FileLocksmithProcessInfo( processInfo.Name, processInfo.Pid, processInfo.User, processInfo.Files)); } return new FileLocksmithQueryResult( FileLocksmithQueryStatus.Success, processes); } catch (JsonException) { return MalformedOutput(); } } finally { if (!process.HasExited) { await TerminateAsync(process); } if (!inputClosed) { process.StandardInput.BaseStream.Dispose(); } await Task.WhenAll(outputTask, errorTask); } } private static ProcessStartInfo CreateWorkerStartInfo() { var installedPath = Path.Combine(AppContext.BaseDirectory, WorkerExecutableName); var buildOutputPath = Path.GetFullPath(Path.Combine(AppContext.BaseDirectory, "..", WorkerExecutableName)); var workerPath = File.Exists(installedPath) ? installedPath : buildOutputPath; var startInfo = new ProcessStartInfo(workerPath) { CreateNoWindow = true, RedirectStandardError = true, RedirectStandardInput = true, RedirectStandardOutput = true, StandardErrorEncoding = Encoding.UTF8, StandardInputEncoding = Encoding.UTF8, StandardOutputEncoding = Encoding.UTF8, UseShellExecute = false, }; startInfo.ArgumentList.Add(WorkerArgument); return startInfo; } private static async Task TerminateAsync(Process process) { if (process.HasExited) { return; } try { process.Kill(entireProcessTree: true); } catch (InvalidOperationException) when (process.HasExited) { } await process.WaitForExitAsync(CancellationToken.None); } private static FileLocksmithQueryResult FailedToStart() => Failed(null); private static FileLocksmithQueryResult Failed(int? exitCode) => new(FileLocksmithQueryStatus.Failed, Array.Empty(), exitCode); private static FileLocksmithQueryResult MalformedOutput() => new(FileLocksmithQueryStatus.MalformedOutput, Array.Empty()); private sealed record WorkerRequest([property: JsonPropertyName("paths")] IReadOnlyCollection Paths); private sealed class WorkerResponse { [JsonPropertyName("processes")] public WorkerProcessInfo[]? Processes { get; init; } } private sealed class WorkerProcessInfo { [JsonPropertyName("name")] public string? Name { get; init; } [JsonPropertyName("pid")] public uint Pid { get; init; } [JsonPropertyName("user")] public string? User { get; init; } [JsonPropertyName("files")] public string[]? Files { get; init; } } } }