Skip to content

Commit

Permalink
Add using implementation
Browse files Browse the repository at this point in the history
  • Loading branch information
jborean93 committed Jan 25, 2024
1 parent 39d8b85 commit ac1af7c
Show file tree
Hide file tree
Showing 2 changed files with 82 additions and 4 deletions.
17 changes: 16 additions & 1 deletion remote-test.ps1
Original file line number Diff line number Diff line change
@@ -1,3 +1,18 @@
Import-Module $PSScriptRoot/output/RemoteForge
Import-Module -Name "$PSScriptRoot/tests/TestForge/bin/Release/net7.0/TestForge.dll"
Invoke-Remote $args { 'foo' }
# Invoke-Remote $args { 'foo' }

$a = @(1, 2, 3)
$b = @{foo = 'test'}
$c = 'foo'
$aIdx = 0
# $s = New-RemoteForgeSession PipeTest:?logmessages=true
$s = New-RemoteForgeSession PipeTest:
Invoke-Remote PipeTest: {
# Invoke-Command -Session $s -ScriptBlock {
$aIdx = 'o'
$using:b["Foo"]
$using:c
}

$s | Remove-PSSession
69 changes: 66 additions & 3 deletions src/RemoteForge/Commands/InvokeRemoteCommand.cs
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,15 @@
using System.Collections.Specialized;
using System.IO;
using System.Management.Automation;
using System.Management.Automation.Language;
using System.Management.Automation.Remoting;
using System.Management.Automation.Runspaces;
using System.Threading;
using System.Threading.Tasks;
using System.Reflection;
using Microsoft.PowerShell.Commands;
using System.Text;
using System.Linq;

namespace RemoteForge.Commands;

Expand Down Expand Up @@ -53,7 +56,7 @@ private enum PipelineType
Position = 1,
ParameterSetName = "ScriptBlockParam"
)]
public ScriptBlock? ScriptBlock { get; set; }
public string? ScriptBlock { get; set; }

[Parameter(
Mandatory = true,
Expand Down Expand Up @@ -109,7 +112,7 @@ protected override void BeginProcessing()
string commandToRun;
if (ScriptBlock != null)
{
commandToRun = ScriptBlock.ToString();
commandToRun = ScriptBlock;
}
else
{
Expand Down Expand Up @@ -140,7 +143,6 @@ protected override void BeginProcessing()
commandToRun = File.ReadAllText(resolvedPath);
}

// TODO: Check if `$using:...` is used in commandToRun
OrderedDictionary? parameters = null;
if (ParamSplat?.Count > 0)
{
Expand All @@ -160,6 +162,14 @@ protected override void BeginProcessing()
parameters.Add(kvp.Key, value);
}
}

Hashtable usingParameters = GetUsingParameters(commandToRun);
if (usingParameters.Count > 0)
{
parameters ??= new();
parameters["--%"] = usingParameters;
}

_worker = Task.Run(async () =>
{
try
Expand All @@ -176,6 +186,59 @@ await RunWorker(
});
}

private Hashtable GetUsingParameters(string script)
{
ScriptBlock? indexAccesor = null;

Hashtable usingParams = new();

// ParseInput will not fail on errors, to allow providing anything as
// the script to run we don't care about that, FindAll will only work
// if the script is valid PowerShell.
ScriptBlockAst sbkAst = Parser.ParseInput(script, out var _1, out var _2);

foreach (var usingStatement in sbkAst.FindAll((a) => a is UsingExpressionAst, true))
{
UsingExpressionAst usingAst = (UsingExpressionAst)usingStatement;
VariableExpressionAst backingVariableAst = UsingExpressionAst.ExtractUsingVariable(usingAst);
string varPath = backingVariableAst.VariablePath.UserPath;

// The non-index $using variable ensures the key is lowercased.
// The index variant is not in case the index itself is a string
// as that could be case sensitive.
string varText = usingAst.ToString();
if (usingAst.SubExpression is VariableExpressionAst)
{
varText = varText.ToLowerInvariant();
}
string key = Convert.ToBase64String(Encoding.Unicode.GetBytes(varText));

if (!usingParams.ContainsKey(key))
{
object? value = SessionState.PSVariable.GetValue(varPath);

if (usingAst.SubExpression is IndexExpressionAst usingIndex)
{
// ConstantExpressionAst, StringConstantExpressionAst, and
// VariableExpressionAst with IsConstantVariable() all
// should return a value here. Those 3 scenarios are the
// only index scenarios that $using supports.
object index = usingIndex.Index.SafeGetValue();

// PowerShell doesn't have a nice API to dynamically
// invoke the Item[] lookup so to avoid complex reflection
// code we just fallback to doing it in pwsh itself.
indexAccesor ??= System.Management.Automation.ScriptBlock.Create("$args[0][$args[1]]");
value = indexAccesor.Invoke(value, index).FirstOrDefault();
}

usingParams.Add(key, value);
}
}

return usingParams;
}

protected override void ProcessRecord()
{
foreach (PSObject? input in InputObject)
Expand Down

0 comments on commit ac1af7c

Please sign in to comment.