Skip to content

Commit

Permalink
Update process handler (#4425)
Browse files Browse the repository at this point in the history 
* 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 <[email protected]>
  • Loading branch information
@KonstantinTyukalov @kirill-ivlev
KonstantinTyukalov and kirill-ivlev authored Oct 16, 2023
1 parent e5dca16 commit 34e2c30
Show file tree
Hide file tree
Showing 10 changed files with 548 additions and 177 deletions.
7 changes: 7 additions & 0 deletions src/Agent.Sdk/Knob/AgentKnobs.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
10 changes: 10 additions & 0 deletions src/Agent.Worker/Handlers/Handler.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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));

Expand Down
96 changes: 96 additions & 0 deletions src/Agent.Worker/Handlers/Helpers/CmdArgsSanitizer.cs
View file
Original file line number Diff line number Diff line change
@@ -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,
};
}
}
}
176 changes: 96 additions & 80 deletions src/Agent.Worker/Handlers/Helpers/ProcessHandlerHelper.cs
View file
Original file line number Diff line number Diff line change
@@ -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)
{
Expand All @@ -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;
Expand All @@ -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;
}
};
}
1 change: 0 additions & 1 deletion src/Agent.Worker/Handlers/NodeHandler.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -156,7 +156,6 @@ public async Task RunAsync()
{
Trace.Info("AZP_AGENT_IGNORE_VSTSTASKLIB enabled, ignoring fix");
}


StepHost.OutputDataReceived += OnDataReceived;
StepHost.ErrorDataReceived += OnDataReceived;
Expand Down
Loading

0 comments on commit 34e2c30

Please sign in to comment.