Update process handler (#4425)

* Update env expand logic

* Upd process handler args validation

* Remove test ex

* Add copyright + cleanup imports

* Simplify telemetry type

* Simplify PH telemetry tests

* Update debug logs

* Exclude % from allowed symbols

* Fix cmd expand env

* remove useless telemetry

* Add test traits

* Fix env incorrect expanding

* Add more test cases

* Fix exception condition

* Move validation args to public

* Update validation logic

* Add % to test

* Code cleanup

* Hide new logic under knob

* Move constants to class

* move const value back

* Add PublishTelemetry overload

* Update throw logic

* Update invalid args exception

* Update private const letter case

---------

Co-authored-by: Kirill Ivlev <102740624+kirill-ivlev@users.noreply.github.com>

Author: localconst

src/Agent.Sdk/Knob/AgentKnobs.cs

@@ -473,6 +473,13 @@ public class AgentKnobs
             new EnvironmentKnobSource("AZP_75787_ENABLE_COLLECT"),
             new BuiltInDefaultKnobSource("false"));
 
+        public static readonly Knob ProcessHandlerEnableNewLogic = new Knob(
+            nameof(ProcessHandlerEnableNewLogic),
+            "Enables new sanitization logic for process handler",
+            new RuntimeKnobSource("AZP_75787_ENABLE_NEW_PH_LOGIC"),
+            new EnvironmentKnobSource("AZP_75787_ENABLE_NEW_PH_LOGIC"),
+            new BuiltInDefaultKnobSource("false"));
+
         public static readonly Knob DisableDrainQueuesAfterTask = new Knob(
             nameof(DisableDrainQueuesAfterTask),
             "Forces the agent to disable draining queues after each task",

src/Agent.Worker/Handlers/Handler.cs

@@ -308,10 +308,20 @@ protected void RemovePSModulePathFromEnvironment()
             }
         }
 
+        // This overload is to handle specific types some other way.
         protected void PublishTelemetry<T>(
             Dictionary<string, T> telemetryData,
             string feature = "TaskHandler"
         )
+        {
+            // JsonConvert.SerializeObject always converts to base object.
+            PublishTelemetry((object)telemetryData, feature);
+        }
+
+        private void PublishTelemetry(
+            object telemetryData,
+            string feature = "TaskHandler"
+        )
         {
             ArgUtil.NotNull(Task, nameof(Task));
 

src/Agent.Worker/Handlers/Helpers/CmdArgsSanitizer.cs

@@ -0,0 +1,96 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System.Collections.Generic;
+using System.Linq;
+using System.Text.RegularExpressions;
+using Microsoft.VisualStudio.Services.Agent.Util;
+
+namespace Agent.Worker.Handlers.Helpers
+{
+    public static class CmdArgsSanitizer
+    {
+        private const string _removedSymbolSign = "_#removed#_";
+        private const string _argsSplitSymbols = "^^";
+        private static readonly Regex _sanitizeRegExp = new("(?<!\\^)([^a-zA-Z0-9\\\\` _''\"\\-=\\/:\\.*,+~?^%])");
+
+        public static (string, CmdArgsSanitizingTelemetry) SanitizeArguments(string inputArgs)
+        {
+            if (inputArgs == null)
+            {
+                return (null, null);
+            }
+
+            var argsChunks = inputArgs.Split(_argsSplitSymbols);
+            var matchesChunks = new List<MatchCollection>();
+
+            for (int i = 0; i < argsChunks.Length; i++)
+            {
+                var matches = _sanitizeRegExp.Matches(argsChunks[i]);
+                if (matches.Count > 0)
+                {
+                    matchesChunks.Add(matches);
+                    argsChunks[i] = _sanitizeRegExp.Replace(argsChunks[i], _removedSymbolSign);
+                }
+            }
+
+            var resultArgs = string.Join(_argsSplitSymbols, argsChunks);
+
+            CmdArgsSanitizingTelemetry telemetry = null;
+
+            if (resultArgs != inputArgs)
+            {
+                var symbolsCount = matchesChunks
+                                    .Select(chunk => chunk.Count)
+                                    .Aggregate(0, (acc, mc) => acc + mc);
+                telemetry = new CmdArgsSanitizingTelemetry
+                    (
+                        RemovedSymbols: CmdArgsSanitizingTelemetry.ToSymbolsDictionary(matchesChunks),
+                        RemovedSymbolsCount: symbolsCount
+                    );
+            }
+
+            return (resultArgs, telemetry);
+        }
+    }
+
+    public record CmdArgsSanitizingTelemetry
+    (
+        Dictionary<string, int> RemovedSymbols,
+        int RemovedSymbolsCount
+    )
+    {
+        public static Dictionary<string, int> ToSymbolsDictionary(List<MatchCollection> matches)
+        {
+            ArgUtil.NotNull(matches, nameof(matches));
+
+            var symbolsDict = new Dictionary<string, int>();
+            foreach (var mc in matches)
+            {
+                foreach (var m in mc.Cast<Match>())
+                {
+                    var symbol = m.Value;
+                    if (symbolsDict.TryGetValue(symbol, out _))
+                    {
+                        symbolsDict[symbol] += 1;
+                    }
+                    else
+                    {
+                        symbolsDict[symbol] = 1;
+                    }
+                }
+            }
+
+            return symbolsDict;
+        }
+
+        public Dictionary<string, object> ToDictionary()
+        {
+            return new()
+            {
+                ["removedSymbols"] = RemovedSymbols,
+                ["removedSymbolsCount"] = RemovedSymbolsCount,
+            };
+        }
+    }
+}

src/Agent.Worker/Handlers/Helpers/ProcessHandlerHelper.cs

@@ -1,66 +1,46 @@
-using System;
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System;
 using System.Collections.Generic;
-using System.Text;
+using Agent.Sdk.Knob;
+using Microsoft.VisualStudio.Services.Agent.Util;
+using Microsoft.VisualStudio.Services.Agent.Worker;
+using Microsoft.VisualStudio.Services.Common;
 
 namespace Agent.Worker.Handlers.Helpers
 {
     public static class ProcessHandlerHelper
     {
-        public static (string, CmdTelemetry) ProcessInputArguments(string inputArgs)
+        private const char _escapingSymbol = '^';
+        private const string _envPrefix = "%";
+        private const string _envPostfix = "%";
+
+        public static (string, CmdTelemetry) ExpandCmdEnv(string inputArgs, Dictionary<string, string> environment)
         {
-            const char quote = '"';
-            const char escapingSymbol = '^';
-            const string envPrefix = "%";
-            const string envPostfix = "%";
+            ArgUtil.NotNull(inputArgs, nameof(inputArgs));
+            ArgUtil.NotNull(environment, nameof(environment));
 
             string result = inputArgs;
             int startIndex = 0;
             var telemetry = new CmdTelemetry();
 
             while (true)
             {
-                int prefixIndex = result.IndexOf(envPrefix, startIndex);
+                int prefixIndex = result.IndexOf(_envPrefix, startIndex);
                 if (prefixIndex < 0)
                 {
                     break;
                 }
 
                 telemetry.FoundPrefixes++;
 
-                if (prefixIndex > 0 && result[prefixIndex - 1] == escapingSymbol)
+                if (prefixIndex > 0 && result[prefixIndex - 1] == _escapingSymbol)
                 {
-                    if (result[prefixIndex - 2] == 0 || result[prefixIndex - 2] != escapingSymbol)
-                    {
-                        startIndex++;
-                        result = result[..(prefixIndex - 1)] + result[prefixIndex..];
-
-                        telemetry.EscapedVariables++;
-
-                        continue;
-                    }
-
-                    telemetry.EscapedEscapingSymbols++;
+                    telemetry.EscapingSymbolBeforeVar++;
                 }
 
-                // We possibly should simplify that part -> if just no close quote, then break
-                int quoteIndex = result.IndexOf(quote, startIndex);
-                if (quoteIndex >= 0 && prefixIndex > quoteIndex)
-                {
-                    int nextQuoteIndex = result.IndexOf(quote, quoteIndex + 1);
-                    if (nextQuoteIndex < 0)
-                    {
-                        telemetry.QuotesNotEnclosed = 1;
-                        break;
-                    }
-
-                    startIndex = nextQuoteIndex + 1;
-
-                    telemetry.QuottedBlocks++;
-
-                    continue;
-                }
-
-                int envStartIndex = prefixIndex + envPrefix.Length;
+                int envStartIndex = prefixIndex + _envPrefix.Length;
                 int envEndIndex = FindEnclosingIndex(result, prefixIndex);
                 if (envEndIndex == 0)
                 {
@@ -69,32 +49,29 @@ public static (string, CmdTelemetry) ProcessInputArguments(string inputArgs)
                 }
 
                 string envName = result[envStartIndex..envEndIndex];
-
-                telemetry.BracedVariables++;
-
-                if (envName.StartsWith(escapingSymbol))
+                if (envName.StartsWith(_escapingSymbol))
                 {
-                    var sanitizedEnvName = envPrefix + envName[1..] + envPostfix;
-
-                    result = result[..prefixIndex] + sanitizedEnvName + result[(envEndIndex + envPostfix.Length)..];
-                    startIndex = prefixIndex + sanitizedEnvName.Length;
-
                     telemetry.VariablesStartsFromES++;
-
-                    continue;
                 }
 
                 var head = result[..prefixIndex];
-                if (envName.Contains(escapingSymbol))
+                if (envName.Contains(_escapingSymbol, StringComparison.Ordinal))
                 {
-                    head += envName.Split(escapingSymbol)[1];
-                    envName = envName.Split(escapingSymbol)[0];
-
                     telemetry.VariablesWithESInside++;
                 }
 
-                var envValue = System.Environment.GetEnvironmentVariable(envName) ?? "";
-                var tail = result[(envEndIndex + envPostfix.Length)..];
+                // Since Windows have case-insensetive environment, and Process handler is windows-specific, we should allign this behavior.
+                var windowsEnvironment = new Dictionary<string, string>(environment, StringComparer.OrdinalIgnoreCase);
+
+                // In case we don't have such variable, we just leave it as is
+                if (!windowsEnvironment.TryGetValue(envName, out string envValue) || string.IsNullOrEmpty(envValue))
+                {
+                    telemetry.NotExistingEnv++;
+                    startIndex = prefixIndex + 1;
+                    continue;
+                }
+
+                var tail = result[(envEndIndex + _envPostfix.Length)..];
 
                 result = head + envValue + tail;
                 startIndex = prefixIndex + envValue.Length;
@@ -119,49 +96,88 @@ private static int FindEnclosingIndex(string input, int targetIndex)
 
             return 0;
         }
+
+        public static (bool, Dictionary<string, object>) ValidateInputArguments(
+            string inputArgs,
+            Dictionary<string, string> environment,
+            IExecutionContext context)
+        {
+            var enableValidation = AgentKnobs.ProcessHandlerSecureArguments.GetValue(context).AsBoolean();
+            context.Debug($"Enable args validation: '{enableValidation}'");
+            var enableAudit = AgentKnobs.ProcessHandlerSecureArgumentsAudit.GetValue(context).AsBoolean();
+            context.Debug($"Enable args validation audit: '{enableAudit}'");
+            var enableTelemetry = AgentKnobs.ProcessHandlerTelemetry.GetValue(context).AsBoolean();
+            context.Debug($"Enable telemetry: '{enableTelemetry}'");
+
+            if (enableValidation || enableAudit || enableTelemetry)
+            {
+                context.Debug("Starting args env expansion");
+                var (expandedArgs, envExpandTelemetry) = ExpandCmdEnv(inputArgs, environment);
+                context.Debug($"Expanded args={expandedArgs}");
+
+                context.Debug("Starting args sanitization");
+                var (sanitizedArgs, sanitizeTelemetry) = CmdArgsSanitizer.SanitizeArguments(expandedArgs);
+
+                Dictionary<string, object> telemetry = null;
+                if (sanitizedArgs != inputArgs)
+                {
+                    if (enableTelemetry)
+                    {
+                        telemetry = envExpandTelemetry.ToDictionary();
+                        if (sanitizeTelemetry != null)
+                        {
+                            telemetry.AddRange(sanitizeTelemetry.ToDictionary());
+                        }
+                    }
+                    if (sanitizedArgs != expandedArgs)
+                    {
+                        if (enableAudit && !enableValidation)
+                        {
+                            context.Warning(StringUtil.Loc("ProcessHandlerInvalidScriptArgs"));
+                        }
+                        if (enableValidation)
+                        {
+                            return (false, telemetry);
+                        }
+
+                        return (true, telemetry);
+                    }
+                }
+
+                return (true, null);
+            }
+            else
+            {
+                context.Debug("Args sanitization skipped.");
+                return (true, null);
+            }
+        }
     }
 
     public class CmdTelemetry
     {
         public int FoundPrefixes { get; set; } = 0;
-        public int QuottedBlocks { get; set; } = 0;
         public int VariablesExpanded { get; set; } = 0;
-        public int EscapedVariables { get; set; } = 0;
-        public int EscapedEscapingSymbols { get; set; } = 0;
+        public int EscapingSymbolBeforeVar { get; set; } = 0;
         public int VariablesStartsFromES { get; set; } = 0;
-        public int BraceSyntaxEntries { get; set; } = 0;
-        public int BracedVariables { get; set; } = 0;
         public int VariablesWithESInside { get; set; } = 0;
         public int QuotesNotEnclosed { get; set; } = 0;
         public int NotClosedEnvSyntaxPosition { get; set; } = 0;
+        public int NotExistingEnv { get; set; } = 0;
 
-        public Dictionary<string, int> ToDictionary()
+        public Dictionary<string, object> ToDictionary()
         {
-            return new Dictionary<string, int>
+            return new Dictionary<string, object>
             {
                 ["foundPrefixes"] = FoundPrefixes,
-                ["quottedBlocks"] = QuottedBlocks,
                 ["variablesExpanded"] = VariablesExpanded,
-                ["escapedVariables"] = EscapedVariables,
-                ["escapedEscapingSymbols"] = EscapedEscapingSymbols,
+                ["escapedVariables"] = EscapingSymbolBeforeVar,
                 ["variablesStartsFromES"] = VariablesStartsFromES,
-                ["braceSyntaxEntries"] = BraceSyntaxEntries,
-                ["bracedVariables"] = BracedVariables,
                 ["bariablesWithESInside"] = VariablesWithESInside,
                 ["quotesNotEnclosed"] = QuotesNotEnclosed,
-                ["notClosedBraceSyntaxPosition"] = NotClosedEnvSyntaxPosition
+                ["notClosedBraceSyntaxPosition"] = NotClosedEnvSyntaxPosition,
+                ["notExistingEnv"] = NotExistingEnv
             };
         }
-
-        public Dictionary<string, string> ToStringsDictionary()
-        {
-            var dict = ToDictionary();
-            var result = new Dictionary<string, string>();
-            foreach (var key in dict.Keys)
-            {
-                result.Add(key, dict[key].ToString());
-            }
-            return result;
-        }
     };
 }

src/Agent.Worker/Handlers/NodeHandler.cs

@@ -156,7 +156,6 @@ public async Task RunAsync()
             {
                 Trace.Info("AZP_AGENT_IGNORE_VSTSTASKLIB enabled, ignoring fix");
             }
-            
 
             StepHost.OutputDataReceived += OnDataReceived;
             StepHost.ErrorDataReceived += OnDataReceived;