Skip to content

Commit 6db3ee1

Browse files
MJECloudMJECloud
committed
Add cross-platform pwsh task handler
1 parent 8d612d4 commit 6db3ee1

7 files changed

Lines changed: 356 additions & 10 deletions

File tree

‎src/Agent.Worker/Handlers/HandlerFactory.cs‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -76,12 +76,16 @@ public IHandler Create(
7676
}
7777
else if (data is PowerShell3HandlerData)
7878
{
79-
#pragma warning disable CA1416 // PowerShell handlers are Windows only
80-
// PowerShell3.
79+
#pragma warning disable CA1416 // PowerShell3 handler is Windows only
8180
handler = HostContext.CreateService<IPowerShell3Handler>();
8281
(handler as IPowerShell3Handler).Data = data as PowerShell3HandlerData;
8382
#pragma warning restore CA1416
8483
}
84+
else if (data is PwshHandlerData)
85+
{
86+
handler = HostContext.CreateService<IPwshHandler>();
87+
(handler as IPwshHandler).Data = data as PwshHandlerData;
88+
}
8589
else if (data is PowerShellExeHandlerData)
8690
{
8791
#pragma warning disable CA1416 // PowerShell handlers are Windows only
@@ -169,4 +173,4 @@ private List<Guid> getTaskExceptionList()
169173
return exceptionList;
170174
}
171175
}
172-
}
176+
}

‎src/Agent.Worker/Handlers/PwshHandler.cs‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
// Copyright (c) Microsoft Corporation.
2+
// Licensed under the MIT License.
3+
4+
using Agent.Sdk;
5+
using Agent.Sdk.Knob;
6+
using Microsoft.VisualStudio.Services.Agent.Util;
7+
using System;
8+
using System.IO;
9+
using System.Threading.Tasks;
10+
11+
namespace Microsoft.VisualStudio.Services.Agent.Worker.Handlers
12+
{
13+
[ServiceLocator(Default = typeof(PwshHandler))]
14+
public interface IPwshHandler : IHandler
15+
{
16+
PwshHandlerData Data { get; set; }
17+
}
18+
19+
public sealed class PwshHandler : Handler, IPwshHandler
20+
{
21+
public PwshHandlerData Data { get; set; }
22+
23+
public async Task RunAsync()
24+
{
25+
Trace.Entering();
26+
ArgUtil.NotNull(Data, nameof(Data));
27+
ArgUtil.NotNull(ExecutionContext, nameof(ExecutionContext));
28+
ArgUtil.NotNull(Inputs, nameof(Inputs));
29+
ArgUtil.Directory(TaskDirectory, nameof(TaskDirectory));
30+
31+
AddInputsToEnvironment();
32+
AddEndpointsToEnvironment();
33+
AddSecureFilesToEnvironment();
34+
AddVariablesToEnvironment();
35+
AddTaskVariablesToEnvironment();
36+
AddPrependPathToEnvironment();
37+
if (PlatformUtil.RunningOnWindows)
38+
{
39+
RemovePSModulePathFromEnvironment();
40+
}
41+
42+
string scriptFile = ResolveScriptFile();
43+
string scriptDirectory = Path.GetDirectoryName(scriptFile);
44+
string moduleFile = ResolveModuleFile(scriptDirectory);
45+
string pwshArgs = BuildPwshArguments(moduleFile, scriptFile);
46+
string pwsh = ResolvePwshExecutable();
47+
48+
StepHost.OutputDataReceived += OnDataReceived;
49+
StepHost.ErrorDataReceived += OnDataReceived;
50+
51+
var sigintTimeout = TimeSpan.FromMilliseconds(AgentKnobs.ProccessSigintTimeout.GetValue(ExecutionContext).AsInt());
52+
var sigtermTimeout = TimeSpan.FromMilliseconds(AgentKnobs.ProccessSigtermTimeout.GetValue(ExecutionContext).AsInt());
53+
var useGracefulShutdown = AgentKnobs.UseGracefulProcessShutdown.GetValue(ExecutionContext).AsBoolean();
54+
55+
try
56+
{
57+
await StepHost.ExecuteAsync(
58+
workingDirectory: StepHost.ResolvePathForStepHost(scriptDirectory),
59+
fileName: pwsh,
60+
arguments: pwshArgs,
61+
environment: Environment,
62+
requireExitCodeZero: true,
63+
outputEncoding: null,
64+
killProcessOnCancel: false,
65+
inheritConsoleHandler: !ExecutionContext.Variables.Retain_Default_Encoding,
66+
continueAfterCancelProcessTreeKillAttempt: _continueAfterCancelProcessTreeKillAttempt,
67+
sigintTimeout: sigintTimeout,
68+
sigtermTimeout: sigtermTimeout,
69+
useGracefulShutdown: useGracefulShutdown,
70+
cancellationToken: ExecutionContext.CancellationToken);
71+
}
72+
finally
73+
{
74+
StepHost.OutputDataReceived -= OnDataReceived;
75+
StepHost.ErrorDataReceived -= OnDataReceived;
76+
}
77+
}
78+
79+
private void OnDataReceived(object sender, ProcessDataReceivedEventArgs e)
80+
{
81+
if (!CommandManager.TryProcessCommand(ExecutionContext, e.Data))
82+
{
83+
ExecutionContext.Output(e.Data);
84+
}
85+
}
86+
87+
private string ResolveScriptFile()
88+
{
89+
ArgUtil.NotNullOrEmpty(Data.Target, nameof(Data.Target));
90+
string scriptFile = Path.Combine(TaskDirectory, Data.Target);
91+
ArgUtil.File(scriptFile, nameof(scriptFile));
92+
return scriptFile;
93+
}
94+
95+
private string ResolveModuleFile(string scriptDirectory)
96+
{
97+
string moduleFile = Path.Combine(scriptDirectory, "ps_modules", "VstsTaskSdk", "VstsTaskSdk.psd1");
98+
ArgUtil.File(moduleFile, nameof(moduleFile));
99+
return moduleFile;
100+
}
101+
102+
private string BuildPwshArguments(string moduleFile, string scriptFile)
103+
{
104+
if (AgentKnobs.UsePSScriptWrapper.GetValue(ExecutionContext).AsBoolean())
105+
{
106+
return BuildWrapperArguments(moduleFile, scriptFile);
107+
}
108+
109+
return BuildDirectInvocationArguments(moduleFile, scriptFile);
110+
}
111+
112+
private string BuildWrapperArguments(string moduleFile, string scriptFile)
113+
{
114+
return StringUtil.Format(
115+
@"-NoLogo -NoProfile -ExecutionPolicy Unrestricted -Command ""{3}"" -VstsSdkPath {0} -DebugOption {1} -ScriptBlockString ""{2}""",
116+
StepHost.ResolvePathForStepHost(moduleFile).Replace("'", "''"),
117+
ExecutionContext.Variables.System_Debug == true ? "Continue" : "SilentlyContinue",
118+
StepHost.ResolvePathForStepHost(scriptFile).Replace("'", "''''"),
119+
Path.Combine(HostContext.GetDirectory(WellKnownDirectory.Bin), "powershell", "Start-AzpTask.ps1"));
120+
}
121+
122+
private string BuildDirectInvocationArguments(string moduleFile, string scriptFile)
123+
{
124+
return StringUtil.Format(
125+
@"-NoLogo -NoProfile -NonInteractive -ExecutionPolicy Unrestricted -Command "". ([scriptblock]::Create('if ([Console]::InputEncoding -is [Text.UTF8Encoding] -and [Console]::InputEncoding.GetPreamble().Length -ne 0) {{ [Console]::InputEncoding = New-Object Text.UTF8Encoding $false }}')) 2>&1 | ForEach-Object {{ Write-Verbose $_.Exception.Message -Verbose }} ; Import-Module -Name '{0}' -ArgumentList @{{ NonInteractive = $true }} -ErrorAction Stop ; $VerbosePreference = '{1}' ; $DebugPreference = '{1}' ; Invoke-VstsTaskScript -ScriptBlock ([scriptblock]::Create('. ''{2}'''))""",
126+
StepHost.ResolvePathForStepHost(moduleFile).Replace("'", "''"),
127+
ExecutionContext.Variables.System_Debug == true ? "Continue" : "SilentlyContinue",
128+
StepHost.ResolvePathForStepHost(scriptFile).Replace("'", "''''"));
129+
}
130+
131+
private string ResolvePwshExecutable()
132+
{
133+
string pwsh = "pwsh";
134+
if (StepHost is DefaultStepHost)
135+
{
136+
pwsh = HostContext.GetService<IPwshExeUtil>().GetPath();
137+
}
138+
139+
ArgUtil.NotNullOrEmpty(pwsh, nameof(pwsh));
140+
return pwsh;
141+
}
142+
}
143+
}

‎src/Agent.Worker/TaskManager.cs‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -616,6 +616,7 @@ public sealed class ExecutionData
616616
private Node16HandlerData _node16;
617617
private Node20_1HandlerData _node20_1;
618618
private Node24HandlerData _node24;
619+
private PwshHandlerData _pwsh;
619620
private PowerShellHandlerData _powerShell;
620621
private PowerShell3HandlerData _powerShell3;
621622
private PowerShellExeHandlerData _powerShellExe;
@@ -713,6 +714,21 @@ public Node24HandlerData Node24
713714
}
714715
}
715716

717+
[JsonProperty(NullValueHandling = NullValueHandling.Ignore)]
718+
public PwshHandlerData Pwsh
719+
{
720+
get
721+
{
722+
return _pwsh;
723+
}
724+
725+
set
726+
{
727+
_pwsh = value;
728+
Add(value);
729+
}
730+
}
731+
716732
[JsonProperty(NullValueHandling = NullValueHandling.Ignore)]
717733
public PowerShellHandlerData PowerShell
718734
{
@@ -916,6 +932,11 @@ public sealed class PowerShell3HandlerData : HandlerData
916932
public override int Priority => 106;
917933
}
918934

935+
public sealed class PwshHandlerData : HandlerData
936+
{
937+
public override int Priority => 106;
938+
}
939+
919940
public sealed class PowerShellHandlerData : HandlerData
920941
{
921942
public string ArgumentFormat

‎src/Agent.Worker/TaskRunner.cs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -202,7 +202,7 @@ private async Task RunAsyncInternal()
202202
runtimeVariables = new Variables(HostContext, variableCopy, out expansionWarnings);
203203
expansionWarnings?.ForEach(x => ExecutionContext.Warning(x));
204204
}
205-
else if (handlerData is BaseNodeHandlerData || handlerData is PowerShell3HandlerData)
205+
else if (handlerData is BaseNodeHandlerData || handlerData is PowerShell3HandlerData || handlerData is PwshHandlerData)
206206
{
207207
// Only the node, node10, and powershell3 handlers support running inside container.
208208
// Make sure required container is already created.
@@ -580,7 +580,7 @@ public HandlerData GetHandlerData(IExecutionContext ExecutionContext, ExecutionD
580580
targetOS = stepTarget.ExecutionOS;
581581
if (stepTarget is ContainerInfo)
582582
{
583-
if ((currentExecution.All.Any(x => x is PowerShell3HandlerData)) &&
583+
if ((currentExecution.All.Any(x => x is PowerShell3HandlerData || x is PwshHandlerData)) &&
584584
(currentExecution.All.Any(x => x is BaseNodeHandlerData)))
585585
{
586586
Trace.Info($"Since we are targeting a container, we will prefer a node handler if one is available");
@@ -590,7 +590,7 @@ public HandlerData GetHandlerData(IExecutionContext ExecutionContext, ExecutionD
590590
}
591591
Trace.Info($"Get handler data for target platform {targetOS.ToString()}");
592592
return currentExecution.All
593-
.OrderBy(x => !(x.PreferredOnPlatform(targetOS) && (preferPowershellHandler || !(x is PowerShell3HandlerData)))) // Sort true to false.
593+
.OrderBy(x => !(x.PreferredOnPlatform(targetOS) && (preferPowershellHandler || !(x is PowerShell3HandlerData || x is PwshHandlerData)))) // Sort true to false.
594594
.ThenBy(x => x.Priority)
595595
.FirstOrDefault();
596596
}
@@ -742,9 +742,9 @@ private bool IsCorrelationIdRequired(IHandler handler, Definition task)
742742
Trace.Info($"Node SDK version: {nodeSdkVer}. Correlation ID is required: {isIdRequired}.");
743743
}
744744
}
745-
else if (handler is IPowerShell3Handler)
745+
else if (handler is IPowerShell3Handler || handler is IPwshHandler)
746746
{
747-
Trace.Info("Current handler is PowerShell3. Trying to determing the SDK version.");
747+
Trace.Info("Current handler is PowerShell. Trying to determing the SDK version.");
748748
var psSdkVer = task.GetPowerShellSDKVersion();
749749
if (psSdkVer == null)
750750
{
@@ -766,4 +766,4 @@ private bool IsCorrelationIdRequired(IHandler handler, Definition task)
766766
return isIdRequired;
767767
}
768768
}
769-
}
769+
}

‎src/Microsoft.VisualStudio.Services.Agent/Util/PwshExeUtil.cs‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
// Copyright (c) Microsoft Corporation.
2+
// Licensed under the MIT License.
3+
4+
using Agent.Sdk;
5+
using Agent.Sdk.Util;
6+
using System;
7+
using System.Diagnostics;
8+
using System.Text.RegularExpressions;
9+
10+
namespace Microsoft.VisualStudio.Services.Agent.Util
11+
{
12+
[ServiceLocator(Default = typeof(PwshExeUtil))]
13+
public interface IPwshExeUtil : IAgentService
14+
{
15+
string GetPath();
16+
}
17+
18+
public sealed class PwshExeUtil : AgentService, IPwshExeUtil
19+
{
20+
private static readonly Version MinimumVersion = new Version(7, 0);
21+
22+
public string GetPath()
23+
{
24+
Trace.Entering();
25+
26+
string commandName = PlatformUtil.RunningOnWindows ? "pwsh.exe" : "pwsh";
27+
string pwshPath = WhichUtil.Which(commandName, trace: Trace);
28+
if (string.IsNullOrEmpty(pwshPath))
29+
{
30+
throw new InvalidOperationException(StringUtil.Loc("FileNotFound", commandName));
31+
}
32+
33+
Version version = GetVersion(pwshPath);
34+
if (version < MinimumVersion)
35+
{
36+
throw new InvalidOperationException($"A compatible version of pwsh was not found. Minimum required version is {MinimumVersion}.");
37+
}
38+
39+
return pwshPath;
40+
}
41+
42+
private Version GetVersion(string pwshPath)
43+
{
44+
ArgUtil.NotNullOrEmpty(pwshPath, nameof(pwshPath));
45+
46+
string output = string.Empty;
47+
using (var process = new Process())
48+
{
49+
process.StartInfo = new ProcessStartInfo
50+
{
51+
FileName = pwshPath,
52+
Arguments = "-NoLogo -NoProfile -NonInteractive -Command \"$PSVersionTable.PSVersion.ToString()\"",
53+
RedirectStandardOutput = true,
54+
RedirectStandardError = true,
55+
UseShellExecute = false,
56+
CreateNoWindow = true,
57+
};
58+
59+
if (!process.Start())
60+
{
61+
throw new InvalidOperationException($"Unable to start '{pwshPath}'.");
62+
}
63+
64+
output = process.StandardOutput.ReadToEnd();
65+
string error = process.StandardError.ReadToEnd();
66+
process.WaitForExit();
67+
68+
if (process.ExitCode != 0)
69+
{
70+
throw new InvalidOperationException($"Unable to determine pwsh version from '{pwshPath}'. {error}".Trim());
71+
}
72+
}
73+
74+
string versionString = output?.Trim();
75+
if (Version.TryParse(versionString, out Version version))
76+
{
77+
return version;
78+
}
79+
80+
Match match = Regex.Match(versionString ?? string.Empty, @"\d+\.\d+(\.\d+)?(\.\d+)?", RegexOptions.CultureInvariant);
81+
if (match.Success && Version.TryParse(match.Value, out version))
82+
{
83+
return version;
84+
}
85+
86+
throw new InvalidOperationException($"Unable to parse pwsh version '{versionString}'.");
87+
}
88+
}
89+
}

0 commit comments

Comments
 (0)