From 36282c20f2a12a580a0b07bc41e3e08264f2b2c3 Mon Sep 17 00:00:00 2001
From: Jose Sua <josesua@microsoft.com>
Date: Wed, 14 Dec 2022 08:42:53 -0800
Subject: [PATCH] Add Instrumentation to AmsiUtil and make the init variable
 readonly (#18727)

* Add Insturmentation to AmsiUtil and make the init variable readonly

* remove unwanted file

* Fix task to hex

* ingnore AsmiInit when not windows

* rename event name

* CR feedback

* the missing dot in the argument documentation because that is very very important.

* CR feedback

Co-authored-by: Jose Sua <josesua@ntdev.microsoft.com>
---
 .../PowerShell.Core.Instrumentation.man       | 52 +++++++++++++++++--
 .../engine/remoting/common/PSETWTracer.cs     |  4 +-
 .../logging/LogProvider.cs                    | 16 ++++++
 .../security/SecuritySupport.cs               | 23 ++++----
 .../utils/tracing/PSEtwLog.cs                 | 10 ++++
 .../utils/tracing/PSEtwLogProvider.cs         | 22 +++++++-
 .../utils/tracing/PSSysLogProvider.cs         | 10 ++++
 7 files changed, 123 insertions(+), 14 deletions(-)

diff --git a/src/PowerShell.Core.Instrumentation/PowerShell.Core.Instrumentation.man b/src/PowerShell.Core.Instrumentation/PowerShell.Core.Instrumentation.man
index ea8fcdc49c..8979ad8ebe 100644
--- a/src/PowerShell.Core.Instrumentation/PowerShell.Core.Instrumentation.man
+++ b/src/PowerShell.Core.Instrumentation/PowerShell.Core.Instrumentation.man
@@ -2208,6 +2208,18 @@
               value="0x6017"
               version="1"
               />
+            <event
+                channel="C_ANALYTIC"
+                keywords="AmsiState"
+                level="win:Verbose"
+                message="$(string.PS_PROVIDER.event.E_A_AmsiState.message)"
+                opcode="Method"
+                symbol="AmsiState"
+                task="Amsi"
+                template="T_AmsiState"
+                value="0x4001"
+                version="1"
+              />
         </events>
         <channels>
           <!--There are three channels defined for PowerShell instrumentation
@@ -2415,6 +2427,12 @@
               symbol="T_ISEOperation"
               value="120"
               />
+            <task
+                message="$(string.PS_PROVIDER.task.T_AmsiState.message)"
+                name="Amsi"
+                symbol="T_Amsi"
+                value="130"
+              />
         </tasks>
         <opcodes>
           <opcode
@@ -2575,6 +2593,12 @@
               name="PSWorkflow"
               symbol="K_PSWORKFLOW"
               />
+            <keyword
+                mask="0x400"
+                message="$(string.PS_PROVIDER.keyword.K_AmsiState.message)"
+                name="AmsiState"
+                symbol="K_AmsiState"
+              />
         </keywords>
         <maps>
           <!-- please keep in sync with SerializationMethod from
@@ -4056,6 +4080,16 @@
                 name="FileName"
                 />
           </template>
+            <template tid="T_AmsiState">
+                <data
+                    inType="win:UnicodeString"
+                    name="Action"
+                />
+                <data
+                    inType="win:UnicodeString"
+                    name="AmsiContext"
+                />
+            </template>
         </templates>
       </provider>
     </events>
@@ -4949,7 +4983,11 @@
             id="PS_PROVIDER.event.E_O_M3PWorkflowExecutionStarted.message"
             value="Workflow execution started. %n %t WorkflowId: %1 %n %t ManagedNodes: %2"
             />
-        <string
+          <string
+              id="PS_PROVIDER.event.E_A_AmsiState.message"
+              value="AmsiUtil state. %n %t state: %1 %n %t Context: %2"
+            />
+          <string
             id="PS_PROVIDER.event.E_O_M3PEndpointRegistered.message"
             value="A new PowerShell endpoint was registered. %n %t EndpointName: %1 %n %t EndpointType: %2 %n %t RegisteredBy: %3"
             />
@@ -5417,7 +5455,11 @@
             id="PS_PROVIDER.keyword.K_PSWORKFLOW.message"
             value="PSWorkflow Hosting And Execution Layer"
             />
-        <string
+          <string
+              id="PS_PROVIDER.keyword.K_AmsiState.message"
+              value="Amsi state"
+            />
+          <string
             id="PS_PROVIDER.keyword.K_SESSION.message"
             value="All session layer"
             />
@@ -5589,7 +5631,11 @@
             id="PS_PROVIDER.task.T_ISEOperation.message"
             value="PowerShell ISE Operation"
             />
-        <string
+          <string
+              id="PS_PROVIDER.task.T_AmsiState.message"
+              value="Amsi State"
+            />
+          <string
             id="PS_PROVIDER.event.E_O_ISEExecuteScript.message"
             value="Windows PowerShell ISE has started to run script file %1."
             />
diff --git a/src/System.Management.Automation/engine/remoting/common/PSETWTracer.cs b/src/System.Management.Automation/engine/remoting/common/PSETWTracer.cs
index 5eb9c5fc33..d59e882883 100644
--- a/src/System.Management.Automation/engine/remoting/common/PSETWTracer.cs
+++ b/src/System.Management.Automation/engine/remoting/common/PSETWTracer.cs
@@ -158,6 +158,7 @@ namespace System.Management.Automation.Internal
         Provider_Lifecycle = 0x1F03,
         Settings = 0x1F04,
         Engine_Trace = 0x1F06,
+        Amsi_Init = 0x4001,
 
         // Experimental Features
         ExperimentalFeature_InvalidName = 0x3001,
@@ -239,7 +240,8 @@ namespace System.Management.Automation.Internal
         ExperimentalFeature = 0x6B,
         ScheduledJob = 0x6E,
         NamedPipe = 0x6F,
-        ISEOperation = 0x78
+        ISEOperation = 0x78,
+        Amsi = 0X82
     }
 
     /// <summary>
diff --git a/src/System.Management.Automation/logging/LogProvider.cs b/src/System.Management.Automation/logging/LogProvider.cs
index a46f181b26..bca5d5fb93 100644
--- a/src/System.Management.Automation/logging/LogProvider.cs
+++ b/src/System.Management.Automation/logging/LogProvider.cs
@@ -102,6 +102,13 @@ namespace System.Management.Automation
         /// <param name="previousValue"></param>
         internal abstract void LogSettingsEvent(LogContext logContext, string variableName, string value, string previousValue);
 
+        /// <summary>
+        /// Provider interface function for logging AmsiUtil State event.
+        /// </summary>
+        /// <param name="state">This the action performed in AmsiUtil class, like init, scan, etc.</param>
+        /// <param name="context">The amsiContext handled - Session pair.</param>
+        internal abstract void LogAmsiUtilStateEvent(string state, string context);
+
         /// <summary>
         /// True if the log provider needs to use logging variables.
         /// </summary>
@@ -370,6 +377,15 @@ namespace System.Management.Automation
         {
         }
 
+        /// <summary>
+        /// Provider interface function for logging provider health event.
+        /// </summary>
+        /// <param name="state">This the action performed in AmsiUtil class, like init, scan, etc.</param>
+        /// <param name="context">The amsiContext handled - Session pair.</param>
+        internal override void LogAmsiUtilStateEvent(string state, string context)
+        {
+        }
+
         #endregion
     }
 }
diff --git a/src/System.Management.Automation/security/SecuritySupport.cs b/src/System.Management.Automation/security/SecuritySupport.cs
index f396ca2043..16d6ca3afe 100644
--- a/src/System.Management.Automation/security/SecuritySupport.cs
+++ b/src/System.Management.Automation/security/SecuritySupport.cs
@@ -854,6 +854,7 @@ namespace Microsoft.PowerShell.Commands
 
 namespace System.Management.Automation
 {
+    using System.Management.Automation.Tracing;
     using System.Security.Cryptography.Pkcs;
 
     /// <summary>
@@ -1334,6 +1335,14 @@ namespace System.Management.Automation
 
     internal static class AmsiUtils
     {
+        static AmsiUtils()
+        {
+#if !UNIX
+            s_amsiInitFailed = !CheckAmsiInit();
+            PSEtwLog.LogAmsiUtilStateEvent($"init-{s_amsiInitFailed}", $"{s_amsiContext}-{s_amsiSession}");
+#endif
+        }
+
         internal static int Init()
         {
             Diagnostics.Assert(s_amsiContext == IntPtr.Zero, "Init should be called just once");
@@ -1355,11 +1364,6 @@ namespace System.Management.Automation
                 AppDomain.CurrentDomain.ProcessExit += CurrentDomain_ProcessExit;
 
                 var hr = AmsiNativeMethods.AmsiInitialize(appName, ref s_amsiContext);
-                if (!Utils.Succeeded(hr))
-                {
-                    s_amsiInitFailed = true;
-                }
-
                 return hr;
             }
         }
@@ -1403,6 +1407,7 @@ namespace System.Management.Automation
             // If we had a previous initialization failure, just return the neutral result.
             if (s_amsiInitFailed)
             {
+                PSEtwLog.LogAmsiUtilStateEvent("ScanContent-InitFail", $"{s_amsiContext}-{s_amsiSession}");
                 return AmsiNativeMethods.AMSI_RESULT.AMSI_RESULT_NOT_DETECTED;
             }
 
@@ -1410,6 +1415,7 @@ namespace System.Management.Automation
             {
                 if (s_amsiInitFailed)
                 {
+                    PSEtwLog.LogAmsiUtilStateEvent("ScanContent-InitFail", $"{s_amsiContext}-{s_amsiSession}");
                     return AmsiNativeMethods.AMSI_RESULT.AMSI_RESULT_NOT_DETECTED;
                 }
 
@@ -1449,6 +1455,7 @@ namespace System.Management.Automation
                     if (!Utils.Succeeded(hr))
                     {
                         // If we got a failure, just return the neutral result ("AMSI_RESULT_NOT_DETECTED")
+                        PSEtwLog.LogAmsiUtilStateEvent($"AmsiScanBuffer-{hr}", $"{s_amsiContext}-{s_amsiSession}");
                         return AmsiNativeMethods.AMSI_RESULT.AMSI_RESULT_NOT_DETECTED;
                     }
 
@@ -1456,7 +1463,7 @@ namespace System.Management.Automation
                 }
                 catch (DllNotFoundException)
                 {
-                    s_amsiInitFailed = true;
+                    PSEtwLog.LogAmsiUtilStateEvent("DllNotFoundException", $"{s_amsiContext}-{s_amsiSession}");
                     return AmsiNativeMethods.AMSI_RESULT.AMSI_RESULT_NOT_DETECTED;
                 }
             }
@@ -1557,7 +1564,6 @@ namespace System.Management.Automation
 
                 if (!Utils.Succeeded(hr))
                 {
-                    s_amsiInitFailed = true;
                     return false;
                 }
             }
@@ -1571,7 +1577,6 @@ namespace System.Management.Automation
 
                 if (!Utils.Succeeded(hr))
                 {
-                    s_amsiInitFailed = true;
                     return false;
                 }
             }
@@ -1593,7 +1598,7 @@ namespace System.Management.Automation
         [SuppressMessage("Microsoft.Reliability", "CA2006:UseSafeHandleToEncapsulateNativeResources")]
         private static IntPtr s_amsiSession = IntPtr.Zero;
 
-        private static bool s_amsiInitFailed = false;
+        private static readonly bool s_amsiInitFailed = false;
         private static bool s_amsiNotifyFailed = false;
         private static readonly object s_amsiLockObject = new object();
 
diff --git a/src/System.Management.Automation/utils/tracing/PSEtwLog.cs b/src/System.Management.Automation/utils/tracing/PSEtwLog.cs
index ca8471749e..8212cf55cd 100644
--- a/src/System.Management.Automation/utils/tracing/PSEtwLog.cs
+++ b/src/System.Management.Automation/utils/tracing/PSEtwLog.cs
@@ -118,6 +118,16 @@ namespace System.Management.Automation.Tracing
             provider.LogProviderLifecycleEvent(logContext, providerName, newState);
         }
 
+        /// <summary>
+        /// Provider interface function for logging AmsiUtil State event.
+        /// </summary>
+        /// <param name="state">This the action performed in AmsiUtil class, like init, scan, etc.</param>
+        /// <param name="context">The amsiContext handled - Session pair.</param>
+        internal static void LogAmsiUtilStateEvent(string state, string context)
+        {
+            provider.LogAmsiUtilStateEvent(state, context);
+        }
+
         /// <summary>
         /// Provider interface function for logging settings event.
         /// </summary>
diff --git a/src/System.Management.Automation/utils/tracing/PSEtwLogProvider.cs b/src/System.Management.Automation/utils/tracing/PSEtwLogProvider.cs
index a13dac02e1..09b1cd9b0f 100755
--- a/src/System.Management.Automation/utils/tracing/PSEtwLogProvider.cs
+++ b/src/System.Management.Automation/utils/tracing/PSEtwLogProvider.cs
@@ -124,7 +124,17 @@ namespace System.Management.Automation.Tracing
                     }
                     else
                     {
-                        payload.AppendLine(StringUtil.Format(EtwLoggingStrings.CommandStateChange, logContext.CommandName, newState.ToString()));
+                        if (newState == CommandState.Stopped ||
+                            newState == CommandState.Terminated)
+                        {
+                            // When state is stopped or termianted only log the CommandName
+                            payload.AppendLine(StringUtil.Format(EtwLoggingStrings.CommandStateChange, logContext, newState.ToString()));
+                        }
+                        else
+                        {
+                            // When state is Start log the CommandLine which has arguments for completeness. 
+                            payload.AppendLine(StringUtil.Format(EtwLoggingStrings.CommandStateChange, logContext.CommandLine, newState.ToString()));
+                        }
                     }
                 }
 
@@ -182,6 +192,16 @@ namespace System.Management.Automation.Tracing
             WriteEvent(PSEventId.Provider_Health, PSChannel.Operational, PSOpcode.Exception, PSTask.ExecutePipeline, logContext, payload.ToString());
         }
 
+        /// <summary>
+        /// Provider interface function for logging provider health event.
+        /// </summary>
+        /// <param name="state">This the action performed in AmsiUtil class, like init, scan, etc.</param>
+        /// <param name="context">The amsiContext handled - Session pair.</param>
+        internal override void LogAmsiUtilStateEvent(string state, string context)
+        {
+            WriteEvent(PSEventId.Amsi_Init, PSChannel.Analytic, PSOpcode.Method, PSLevel.Informational, PSTask.Amsi, (PSKeyword)0x0, state, context);
+        }
+
         /// <summary>
         /// Provider interface function for logging provider lifecycle event.
         /// </summary>
diff --git a/src/System.Management.Automation/utils/tracing/PSSysLogProvider.cs b/src/System.Management.Automation/utils/tracing/PSSysLogProvider.cs
index 5a55d2f4cc..629cdd6447 100755
--- a/src/System.Management.Automation/utils/tracing/PSSysLogProvider.cs
+++ b/src/System.Management.Automation/utils/tracing/PSSysLogProvider.cs
@@ -92,6 +92,16 @@ namespace System.Management.Automation.Tracing
             WriteEvent(PSEventId.Engine_Health, PSChannel.Operational, PSOpcode.Exception, PSTask.ExecutePipeline, logContext, payload.ToString());
         }
 
+        /// <summary>
+        /// Provider interface function for logging provider health event.
+        /// </summary>
+        /// <param name="state">This the action performed in AmsiUtil class, like init, scan, etc</param>
+        /// <param name="context">The amsiContext handled - Session pair</param>
+        internal override void LogAmsiUtilStateEvent(string state, string context)
+        {
+            WriteEvent(PSEventId.Amsi_Init, PSChannel.Analytic, PSOpcode.Method, PSLevel.Informational, PSTask.Amsi, (PSKeyword)0x0, state, context);
+        }
+
         /// <summary>
         /// Provider interface function for logging engine lifecycle event.
         /// </summary>