winget-cli

PowerShellCmdlet.cs (23583B)

---

 // -----------------------------------------------------------------------------
 // <copyright file="PowerShellCmdlet.cs" company="Microsoft Corporation">
 //     Copyright (c) Microsoft Corporation. Licensed under the MIT License.
 // </copyright>
 // -----------------------------------------------------------------------------
 
 namespace Microsoft.WinGet.Common.Command
 {
     using System;
     using System.Collections.Concurrent;
     using System.Collections.Generic;
     using System.Management.Automation;
     using System.Runtime.ExceptionServices;
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.WinGet.Resources;
     using Microsoft.WinGet.SharedLib.Exceptions;
     using Microsoft.WinGet.SharedLib.PolicySettings;
 
     /// <summary>
     /// This must be the base class for every cmdlet for winget PowerShell modules.
     /// It supports:
     ///  - Async operations.
     ///  - Execute on an MTA. If the thread is already running on an MTA it will executed it, otherwise
     ///    it will create a new MTA thread.
     /// Wait must be used to synchronously wait con the task.
     /// </summary>
     public abstract class PowerShellCmdlet
     {
         private const string Debug = "Debug";
         private static readonly string[] WriteInformationTags = new string[] { "PSHOST" };
 
         private readonly PSCmdlet psCmdlet;
         private readonly Thread pwshThread;
 
         private readonly CancellationTokenSource source = new ();
         private readonly SemaphoreSlim semaphore = new (1, 1);
         private readonly ManualResetEventSlim pwshThreadActionReady = new (false);
         private readonly ManualResetEventSlim pwshThreadActionCompleted = new (false);
 
         private BlockingCollection<QueuedStream> queuedStreams = new ();
         private int progressActivityId = 0;
         private ConcurrentDictionary<int, ProgressRecordType> progressRecords = new ();
         private Action? pwshThreadAction = null;
         private ExceptionDispatchInfo? pwshThreadEdi = null;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="PowerShellCmdlet"/> class.
         /// </summary>
         /// <param name="psCmdlet">PSCmdlet.</param>
         /// <param name="policies">Policies.</param>
         public PowerShellCmdlet(PSCmdlet psCmdlet, HashSet<Policy> policies)
         {
             // Passing Debug will make all the message actions to be Inquire. For async operations
             // and the current queue message implementation this doesn't make sense.
             // PowerShell will inquire for any message giving the impression that the task is
             // paused, but the async operation is still running.
             if (psCmdlet.MyInvocation.BoundParameters.ContainsKey(Debug))
             {
                 throw new NotSupportedException(Resources.DebugNotSupported);
             }
 
             this.ValidatePolicies(policies);
 
             this.psCmdlet = psCmdlet;
             this.pwshThread = Thread.CurrentThread;
         }
 
         /// <summary>
         /// Request cancellation for this command.
         /// </summary>
         public void Cancel()
         {
             this.source.Cancel();
         }
 
         /// <summary>
         /// Execute the delegate in a MTA thread.
         /// Caller must wait on task.
         /// </summary>
         /// <param name="func">Function to execute.</param>
         /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
         internal Task RunOnMTA(Func<Task> func)
         {
             // .NET 4.8 doesn't support TaskCompletionSource.
 #if POWERSHELL_WINDOWS
             throw new NotImplementedException();
 #else
             // This must be called in the main thread.
             if (this.pwshThread != Thread.CurrentThread)
             {
                 throw new InvalidOperationException();
             }
 
             if (Thread.CurrentThread.GetApartmentState() == ApartmentState.MTA)
             {
                 this.Write(StreamType.Verbose, "Already running on MTA");
                 try
                 {
                     Task result = func();
                     result.ContinueWith((task) => this.Complete(), TaskContinuationOptions.ExecuteSynchronously);
                     return result;
                 }
                 catch
                 {
                     this.Complete();
                     throw;
                 }
             }
 
             this.Write(StreamType.Verbose, "Creating MTA thread");
             var tcs = new TaskCompletionSource();
             var thread = new Thread(() =>
             {
                 try
                 {
                     func().GetAwaiter().GetResult();
                     tcs.SetResult();
                 }
                 catch (Exception e)
                 {
                     tcs.SetException(e);
                 }
                 finally
                 {
                     this.Complete();
                 }
             });
 
             thread.SetApartmentState(ApartmentState.MTA);
             thread.Start();
             return tcs.Task;
 #endif
         }
 
         /// <summary>
         /// Execute the delegate in a MTA thread.
         /// Caller must wait on task.
         /// </summary>
         /// <param name="func">Function to execute.</param>
         /// <typeparam name="TResult">Return type of function.</typeparam>
         /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
         internal Task<TResult> RunOnMTA<TResult>(Func<Task<TResult>> func)
         {
             // This must be called in the main thread.
             if (this.pwshThread != Thread.CurrentThread)
             {
                 throw new InvalidOperationException();
             }
 
             if (Thread.CurrentThread.GetApartmentState() == ApartmentState.MTA)
             {
                 this.Write(StreamType.Verbose, "Already running on MTA");
                 try
                 {
                     Task<TResult> result = func();
                     result.ContinueWith((task) => this.Complete(), TaskContinuationOptions.ExecuteSynchronously);
                     return result;
                 }
                 catch
                 {
                     this.Complete();
                     throw;
                 }
             }
 
             this.Write(StreamType.Verbose, "Creating MTA thread");
             var tcs = new TaskCompletionSource<TResult>();
             var thread = new Thread(() =>
             {
                 try
                 {
                     var result = func().GetAwaiter().GetResult();
                     tcs.SetResult(result);
                 }
                 catch (Exception e)
                 {
                     tcs.SetException(e);
                 }
                 finally
                 {
                     this.Complete();
                 }
             });
 
             thread.SetApartmentState(ApartmentState.MTA);
             thread.Start();
             return tcs.Task;
         }
 
         /// <summary>
         /// Execute the delegate in a MTA thread.
         /// Synchronous call.
         /// </summary>
         /// <param name="func">Function to execute.</param>
         /// <typeparam name="TResult">Return type of function.</typeparam>
         /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
         internal TResult RunOnMTA<TResult>(Func<TResult> func)
         {
             // This must be called in the main thread.
             if (this.pwshThread != Thread.CurrentThread)
             {
                 throw new InvalidOperationException();
             }
 
             if (Thread.CurrentThread.GetApartmentState() == ApartmentState.MTA)
             {
                 this.Write(StreamType.Verbose, "Already running on MTA");
                 try
                 {
                     return func();
                 }
                 finally
                 {
                     this.Complete();
                 }
             }
 
             this.Write(StreamType.Verbose, "Creating MTA thread");
             var tcs = new TaskCompletionSource<TResult>();
             var thread = new Thread(() =>
             {
                 try
                 {
                     var result = func();
                     tcs.SetResult(result);
                 }
                 catch (Exception e)
                 {
                     tcs.SetException(e);
                 }
                 finally
                 {
                     this.Complete();
                 }
             });
 
             thread.SetApartmentState(ApartmentState.MTA);
             thread.Start();
             this.Wait(tcs.Task);
             return tcs.Task.Result;
         }
 
         /// <summary>
         /// Executes an action in the main thread.
         /// Blocks until call is executed.
         /// </summary>
         /// <param name="action">Action to perform.</param>
         internal void ExecuteInPowerShellThread(Action action)
         {
             if (this.pwshThread == Thread.CurrentThread)
             {
                 action();
                 return;
             }
 
             this.WaitForOurTurn();
 
             this.pwshThreadAction = action;
             this.pwshThreadActionReady.Set();
             this.WaitMainThreadActionCompletion();
         }
 
         /// <summary>
         /// Waits for the task to be completed. This MUST be called from the main thread.
         /// </summary>
         /// <param name="runningTask">Task to wait for.</param>
         /// <param name="writeCmdlet">The cmdlet that can write to PowerShell.</param>
         internal void Wait(Task runningTask, PowerShellCmdlet? writeCmdlet = null)
         {
             writeCmdlet ??= this;
 
             // This must be called in the main thread.
             if (this.pwshThread != Thread.CurrentThread)
             {
                 throw new InvalidOperationException();
             }
 
             do
             {
                 if (this.pwshThreadActionReady.IsSet)
                 {
                     // Someone needs the main thread.
                     this.pwshThreadActionReady.Reset();
 
                     if (this.pwshThreadAction != null)
                     {
                         try
                         {
                             this.pwshThreadEdi = null;
                             this.pwshThreadAction();
                         }
                         catch (Exception e)
                         {
                             // Make sure we don't throw in the PowerShell thread, this way
                             // we'll get a more meaningful stack by Get-Error.
                             this.pwshThreadEdi = ExceptionDispatchInfo.Capture(e);
                         }
 
                         this.pwshThreadAction = null;
                     }
 
                     // Done.
                     this.pwshThreadActionCompleted.Set();
                 }
 
                 // Take from the blocking collection.
                 if (!this.queuedStreams.IsCompleted && this.queuedStreams.Count > 0)
                 {
                     try
                     {
                         var queuedOutput = this.queuedStreams.Take();
                         if (queuedOutput != null)
                         {
                             this.CmdletWrite(queuedOutput.Type, queuedOutput.Data, writeCmdlet);
                         }
                     }
                     catch (InvalidOperationException)
                     {
                         // An InvalidOperationException means that Take() was called on a completed collection.
                     }
                 }
             }
             while (!(runningTask.IsCompleted && this.queuedStreams.IsCompleted));
 
             if (runningTask.IsFaulted)
             {
                 // If IsFaulted is true, the task's Status is equal to Faulted,
                 // and its Exception property will be non-null.
                 AggregateException? ae = runningTask.Exception! as AggregateException;
                 if (ae != null && ae.InnerExceptions.Count == 1)
                 {
                     ExceptionDispatchInfo.Capture(ae.InnerExceptions[0]).Throw();
                 }
 
                 throw runningTask.Exception!;
             }
         }
 
         /// <summary>
         /// Writes into the corresponding stream if running on the main thread.
         /// Otherwise queue the message.
         /// Is the caller responsibility to use the correct types.
         /// </summary>
         /// <param name="type">Stream type.</param>
         /// <param name="data">Data.</param>
         internal void Write(StreamType type, object data)
         {
             if (type == StreamType.Progress)
             {
                 ProgressRecord progressRecord = (ProgressRecord)data;
                 if (progressRecord.RecordType == ProgressRecordType.Completed)
                 {
                     throw new NotSupportedException("Use CompleteProgress");
                 }
 
                 // Keep track of all progress activity.
                 _ = this.progressRecords.TryAdd(progressRecord.ActivityId, progressRecord.RecordType);
             }
 
             if (this.pwshThread == Thread.CurrentThread)
             {
                 this.CmdletWrite(type, data, this);
                 return;
             }
 
             this.queuedStreams.Add(new QueuedStream(type, data));
         }
 
         /// <summary>
         /// Helper to compute percentage and write progress for processing activities.
         /// </summary>
         /// <param name="activityId">Activity id.</param>
         /// <param name="activity">The activity in progress.</param>
         /// <param name="status">The status of the activity.</param>
         /// <param name="completed">Number of completed actions.</param>
         /// <param name="total">The expected total.</param>
         internal void WriteProgressWithPercentage(int activityId, string activity, string status, int completed, int total)
         {
             double percentComplete = (double)completed / total;
             var record = new ProgressRecord(activityId, activity, status)
             {
                 RecordType = ProgressRecordType.Processing,
                 PercentComplete = (int)(100.0 * percentComplete),
             };
             this.Write(StreamType.Progress, record);
         }
 
         /// <summary>
         /// Helper to complete progress records.
         /// </summary>
         /// <param name="activityId">Activity id.</param>
         /// <param name="activity">The activity in progress.</param>
         /// <param name="status">The status of the activity.</param>
         /// <param name="force">Force write complete progress.</param>
         internal void CompleteProgress(int activityId, string activity, string status, bool force = false)
         {
             var record = new ProgressRecord(activityId, activity, status)
             {
                 RecordType = ProgressRecordType.Completed,
                 PercentComplete = 100,
             };
 
             if (!this.progressRecords.TryAdd(activityId, record.RecordType))
             {
                 _ = this.progressRecords.TryUpdate(activityId, record.RecordType, ProgressRecordType.Processing);
             }
 
             if (this.pwshThread == Thread.CurrentThread)
             {
                 this.CmdletWrite(StreamType.Progress, record, this);
             }
             else
             {
                 // You should only use force if you know the cmdlet that is completing this progress is a sync cmdlet that
                 // is running in an async context. A sync cmdlet is anything that doesn't start with Start-*
                 if (force)
                 {
                     this.ExecuteInPowerShellThread(() => this.CmdletWrite(StreamType.Progress, record, this));
                 }
                 else
                 {
                     this.queuedStreams.Add(new QueuedStream(StreamType.Progress, record));
                 }
             }
         }
 
         /// <summary>
         /// Writes to PowerShell streams.
         /// This method must be called in the original thread.
         /// WARNING: You must only call this when the task is completed.
         /// </summary>
         /// <param name="writeCmdlet">The cmdlet that can write to PowerShell.</param>
         internal void ConsumeAndWriteStreams(PowerShellCmdlet writeCmdlet)
         {
             // This must be called in the main thread.
             if (this.pwshThread != Thread.CurrentThread)
             {
                 throw new InvalidOperationException();
             }
 
             // Take from the blocking collection until is completed.
             try
             {
                 while (true)
                 {
                     var queuedOutput = this.queuedStreams.Take();
                     if (queuedOutput != null)
                     {
                         this.CmdletWrite(queuedOutput.Type, queuedOutput.Data, writeCmdlet);
                     }
                 }
             }
             catch (InvalidOperationException)
             {
                 // We are done.
                 // An InvalidOperationException means that Take() was called on a completed collection.
             }
         }
 
         /// <summary>
         /// Gets a new progress activity id.
         /// </summary>
         /// <returns>The new progress record id.</returns>
         internal int GetNewProgressActivityId()
         {
             return Interlocked.Increment(ref this.progressActivityId);
         }
 
         /// <summary>
         /// Gets the cancellation token.
         /// </summary>
         /// <returns>CancellationToken.</returns>
         internal CancellationToken GetCancellationToken()
         {
             return this.source.Token;
         }
 
         /// <summary>
         /// Gets the current file system location from the cmdlet.
         /// </summary>
         /// <returns>Path.</returns>
         internal string GetCurrentFileSystemLocation()
         {
             return this.psCmdlet.SessionState.Path.CurrentFileSystemLocation.Path;
         }
 
         /// <summary>
         /// Sets a variable.
         /// </summary>
         /// <param name="variableName">Variable name.</param>
         /// <param name="value">Value.</param>
         internal void SetVariable(string variableName, object value)
         {
             this.psCmdlet.SessionState.PSVariable.Set(variableName, value);
         }
 
         /// <summary>
         /// Prompts the user if it should continue processing if possible.
         /// </summary>
         /// <param name="target">Message.</param>
         /// <returns>If the operation should continue.</returns>
         internal bool ShouldProcess(string target)
         {
             // If not on the main thread just continue.
             if (this.pwshThread != Thread.CurrentThread)
             {
                 return true;
             }
 
             return this.psCmdlet.ShouldProcess(target);
         }
 
         private void Complete()
         {
             this.queuedStreams.CompleteAdding();
         }
 
         private void CmdletWrite(StreamType streamType, object data, PowerShellCmdlet writeCmdlet)
         {
             switch (streamType)
             {
                 case StreamType.Debug:
                     throw new NotSupportedException();
                 case StreamType.Verbose:
                     writeCmdlet.psCmdlet.WriteVerbose((string)data);
                     break;
                 case StreamType.Warning:
                     writeCmdlet.psCmdlet.WriteWarning((string)data);
                     break;
                 case StreamType.Error:
                     writeCmdlet.psCmdlet.WriteError((ErrorRecord)data);
                     break;
                 case StreamType.Progress:
                     // If the activity is already completed don't write progress.
                     var progressRecord = (ProgressRecord)data;
                     if (this.progressRecords[progressRecord.ActivityId] == progressRecord.RecordType)
                     {
                         writeCmdlet.psCmdlet.WriteProgress(progressRecord);
                     }
 
                     break;
                 case StreamType.Object:
                     writeCmdlet.psCmdlet.WriteObject(data);
                     break;
                 case StreamType.Information:
                     writeCmdlet.psCmdlet.WriteInformation(data, WriteInformationTags);
                     break;
             }
         }
 
         private void ValidatePolicies(HashSet<Policy> policies)
         {
             GroupPolicy groupPolicy = GroupPolicy.GetInstance();
 
             if (policies.Contains(Policy.WinGet))
             {
                 if (!groupPolicy.IsEnabled(Policy.WinGet))
                 {
                     throw new GroupPolicyException(Policy.WinGet, GroupPolicyFailureType.BlockedByPolicy);
                 }
 
                 policies.Remove(Policy.WinGet);
             }
 
             if (policies.Contains(Policy.Configuration))
             {
                 if (!groupPolicy.IsEnabled(Policy.Configuration))
                 {
                     throw new GroupPolicyException(Policy.Configuration, GroupPolicyFailureType.BlockedByPolicy);
                 }
 
                 policies.Remove(Policy.Configuration);
             }
 
             if (policies.Contains(Policy.WinGetCommandLineInterfaces))
             {
                 if (!groupPolicy.IsEnabled(Policy.WinGetCommandLineInterfaces))
                 {
                     throw new GroupPolicyException(Policy.WinGetCommandLineInterfaces, GroupPolicyFailureType.BlockedByPolicy);
                 }
 
                 policies.Remove(Policy.WinGetCommandLineInterfaces);
             }
 
             if (policies.Count > 0)
             {
                 throw new NotSupportedException($"Invalid policies {string.Join(",", policies)}");
             }
         }
 
         private void WaitForOurTurn()
         {
             this.semaphore.Wait(this.GetCancellationToken());
             this.pwshThreadActionCompleted.Reset();
         }
 
         private void WaitMainThreadActionCompletion()
         {
             WaitHandle.WaitAny(new[]
             {
                 this.GetCancellationToken().WaitHandle,
                 this.pwshThreadActionCompleted.WaitHandle,
             });
 
             try
             {
                 if (this.pwshThreadEdi != null)
                 {
                     this.pwshThreadEdi.Throw();
                 }
             }
             finally
             {
                 this.semaphore.Release();
             }
         }
 
         private class QueuedStream
         {
             public QueuedStream(StreamType type, object data)
             {
                 this.Type = type;
                 this.Data = data;
             }
 
             public StreamType Type { get; }
 
             public object Data { get; }
         }
     }
 }