Skip to content

Commit

Permalink
Finalize tests and add wsman registration
Browse files Browse the repository at this point in the history
  • Loading branch information
@jborean93
jborean93 committed Feb 1, 2024
1 parent 99f0b5b commit b9bf190
Show file tree
Hide file tree
Showing 19 changed files with 560 additions and 140 deletions.
7 changes: 7 additions & 0 deletions RemoteForge.sln
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@ Project("{2150E333-8FDC-42A3-9474-1A3956D46DE8}") = "tests", "tests", "{97A99A31
EndProject
Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "TestForge", "tests\TestForge\TestForge.csproj", "{3E5F758E-F653-43AD-B642-AF79353610AE}"
EndProject
Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "RemoteForgeTests", "tests\units\RemoteForgeTests.csproj", "{85CB0F60-E896-4E5A-ADAF-6C8AE58E2BCA}"
EndProject
Global
GlobalSection(SolutionConfigurationPlatforms) = preSolution
Debug|Any CPU = Debug|Any CPU
Expand All @@ -26,8 +28,13 @@ Global
{3E5F758E-F653-43AD-B642-AF79353610AE}.Debug|Any CPU.Build.0 = Debug|Any CPU
{3E5F758E-F653-43AD-B642-AF79353610AE}.Release|Any CPU.ActiveCfg = Release|Any CPU
{3E5F758E-F653-43AD-B642-AF79353610AE}.Release|Any CPU.Build.0 = Release|Any CPU
{85CB0F60-E896-4E5A-ADAF-6C8AE58E2BCA}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
{85CB0F60-E896-4E5A-ADAF-6C8AE58E2BCA}.Debug|Any CPU.Build.0 = Debug|Any CPU
{85CB0F60-E896-4E5A-ADAF-6C8AE58E2BCA}.Release|Any CPU.ActiveCfg = Release|Any CPU
{85CB0F60-E896-4E5A-ADAF-6C8AE58E2BCA}.Release|Any CPU.Build.0 = Release|Any CPU
EndGlobalSection
GlobalSection(NestedProjects) = preSolution
{3E5F758E-F653-43AD-B642-AF79353610AE} = {97A99A31-7BF9-4747-A6E4-82497A6240A3}
{85CB0F60-E896-4E5A-ADAF-6C8AE58E2BCA} = {97A99A31-7BF9-4747-A6E4-82497A6240A3}
EndGlobalSection
EndGlobal
31 changes: 0 additions & 31 deletions invoke-command.ps1
View file

This file was deleted.

43 changes: 0 additions & 43 deletions invoke-remote.ps1
View file

This file was deleted.

18 changes: 0 additions & 18 deletions remote-test.ps1
View file

This file was deleted.

43 changes: 31 additions & 12 deletions src/RemoteForge/Commands/InvokeRemoteCommand.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -158,11 +158,29 @@ protected override void BeginProcessing()
parameters["--%"] = usingParameters;
}

Queue<(Runspace?, RunspaceConnectionInfo, string)> connectionQueue = new(ConnectionInfo.Length);
foreach (StringForgeConnectionInfoPSSession conn in ConnectionInfo)
{
RunspaceConnectionInfo? connInfo = conn.GetConnectionInfo(this);
if (connInfo == null)
{
continue;
}
connectionQueue.Enqueue((conn.PSSession?.Runspace, connInfo, conn.ToString()));
}

if (connectionQueue.Count == 0)
{
_outputPipe.CompleteAdding();
return;
}

_worker = Task.Run(async () =>
{
try
{
await RunWorker(
connectionQueue,
commandToRun,
arguments: arguments,
parameters: parameters);
Expand Down Expand Up @@ -298,12 +316,12 @@ private void WriteResult(PipelineType pipelineType, object? data)
}

private async Task RunWorker(
Queue<(Runspace?, RunspaceConnectionInfo, string)> connections,
string script,
PSObject?[] arguments,
IDictionary? parameters)
{
Queue<StringForgeConnectionInfoPSSession> connections = new(ConnectionInfo);
List<Task> tasks = new(Math.Min(ThrottleLimit, ConnectionInfo.Length));
List<Task> tasks = new(Math.Min(ThrottleLimit, connections.Count));
do
{
if (connections.Count == 0 || tasks.Count >= tasks.Capacity)
Expand All @@ -313,22 +331,22 @@ private async Task RunWorker(
await doneTask;
}

if (connections.TryDequeue(out StringForgeConnectionInfoPSSession? info))
if (connections.TryDequeue(out (Runspace?, RunspaceConnectionInfo, string) info))
{
Task t = Task.Run(async () =>
{
try
{
await RunScript(info, script, arguments, parameters);
await RunScript(info.Item1, info.Item2, info.Item3, script, arguments, parameters);
}
catch (Exception e)
{
string msg = $"Failed to run script on '{info}': {e.Message}";
string msg = $"Failed to run script on '{info.Item3}': {e.Message}";
ErrorRecord errorRecord = new(
e,
"ExecuteException",
ErrorCategory.NotSpecified,
info.ToString())
info.Item3)
{
ErrorDetails = new(msg),
};
Expand All @@ -342,18 +360,19 @@ private async Task RunWorker(
}

private async Task RunScript(
StringForgeConnectionInfoPSSession info,
Runspace? runspace,
RunspaceConnectionInfo connInfo,
string connId,
string script,
PSObject?[] arguments,
IDictionary? parameters)
{
bool disposeRunspace = false;
Runspace? runspace = info.PSSession?.Runspace;
if (runspace == null)
{
disposeRunspace = true;
runspace = await RunspaceHelper.CreateRunspaceAsync(
info.ConnectionInfo,
connInfo,
_cancelToken.Token,
host: Host,
typeTable: null,
Expand All @@ -373,7 +392,7 @@ private async Task RunScript(
PSObject? obj = outputPipe[e.Index];
if (obj != null)
{
obj.Properties.Add(new PSNoteProperty("PSComputerName", info.ToString()));
obj.Properties.Add(new PSNoteProperty("PSComputerName", connId));
obj.Properties.Add(new PSNoteProperty("RunspaceId", runspace.InstanceId));
obj.Properties.Add(new PSNoteProperty("PSShowComputerName", true));
}
Expand All @@ -390,9 +409,9 @@ private async Task RunScript(
_errorRecordSetTargetObject ??= typeof(ErrorRecord).GetMethod(
"SetTargetObject",
BindingFlags.Instance | BindingFlags.NonPublic);
_errorRecordSetTargetObject?.Invoke(errorRecord, new[] { info.ToString() });
_errorRecordSetTargetObject?.Invoke(errorRecord, new[] { connId });
OriginInfo oi = new(info.ToString(), runspace.InstanceId);
OriginInfo oi = new(connId, runspace.InstanceId);
RemotingErrorRecord remoteError = new(errorRecord, oi);
_outputPipe.Add((PipelineType.Error, remoteError));
Expand Down
100 changes: 88 additions & 12 deletions src/RemoteForge/Commands/RemoteForgeSession.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
using System.Collections.Generic;
using System.Management.Automation;
using System.Management.Automation.Runspaces;
using System.Net;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

Expand Down Expand Up @@ -32,10 +34,16 @@ protected override void ProcessRecord()
}
else
{
RunspaceConnectionInfo? connInfo = connection.GetConnectionInfo(this);
if (connInfo == null)
{
continue;
}

creationTasks.Add((connection.ToString(), Task.Run(async () =>
{
Runspace rs = await RunspaceHelper.CreateRunspaceAsync(
connection.ConnectionInfo,
connInfo,
_cancelTokenSource.Token,
host: Host);
Expand Down Expand Up @@ -87,45 +95,113 @@ public void Dispose()
public sealed class StringForgeConnectionInfoPSSession
{
private readonly string _originalString;
internal RunspaceConnectionInfo ConnectionInfo { get; }
internal readonly Func<RunspaceConnectionInfo> _connInfoFactory;
internal PSSession? PSSession { get; }

public StringForgeConnectionInfoPSSession(string info)
{
_originalString = info;
ConnectionInfo = RemoteForgeRegistration.CreateForgeConnectionInfo(info);
_connInfoFactory = () => RemoteForgeRegistration.CreateForgeConnectionInfo(info);
}

public StringForgeConnectionInfoPSSession(IRemoteForge forge)
{
_originalString = forge.GetTransportString();
ConnectionInfo = new RemoteForgeConnectionInfo(forge);
_connInfoFactory = () => new RemoteForgeConnectionInfo(forge);
}

public StringForgeConnectionInfoPSSession(RunspaceConnectionInfo info)
{
_originalString = GetConnectionInfoString(info);
ConnectionInfo = info;
_connInfoFactory = () => info;
}

public StringForgeConnectionInfoPSSession(PSSession session)
{
ConnectionInfo = session.Runspace.OriginalConnectionInfo;
_originalString = GetConnectionInfoString(ConnectionInfo);
_originalString = GetConnectionInfoString(session.Runspace.OriginalConnectionInfo);
_connInfoFactory = () => session.Runspace.OriginalConnectionInfo;
PSSession = session;
}

// We use a Func so the parameter binding works and we can provide a
// better error at runtime.
internal RunspaceConnectionInfo? GetConnectionInfo(PSCmdlet cmdlet)
{
try
{
return _connInfoFactory();
}
catch (Exception e)
{
ErrorRecord err = new(
e,
"InvalidForgeConnection",
ErrorCategory.InvalidArgument,
_originalString);
cmdlet.WriteError(err);
return null;
}
}

private static string GetConnectionInfoString(RunspaceConnectionInfo info) => info switch
{
RemoteForgeConnectionInfo f => f.ConnectionUri,
SSHConnectionInfo => $"ssh:{info.ComputerName}",
WSManConnectionInfo => $"wsman:{info.ComputerName}",
NamedPipeConnectionInfo p => $"pipe:{p.CustomPipeName}",
ContainerConnectionInfo => $"container:{info.ComputerName}",
VMConnectionInfo => $"vm:{info.ComputerName}",
SSHConnectionInfo s => GetSSHConnectionInfoString(s),
WSManConnectionInfo w => GetWSManConnectionInfoString(w),
_ => $"{info.GetType().Name}:{info.ComputerName}",
};

private static string GetSSHConnectionInfoString(SSHConnectionInfo connInfo)
{
StringBuilder connString = new("ssh:");
if (!string.IsNullOrWhiteSpace(connInfo.UserName))
{
connString.Append(connInfo.UserName).Append('@');
}

if (
IPAddress.TryParse(connInfo.ComputerName, out IPAddress? addr) &&
addr.AddressFamily == System.Net.Sockets.AddressFamily.InterNetworkV6)
{
connString.Append('[').Append(connInfo.ComputerName).Append(']');
}
else
{
connString.Append(connInfo.ComputerName);
}

if (connInfo.Port != 22)
{
connString.Append(':').Append(connInfo.Port);
}

return connString.ToString();
}

private static string GetWSManConnectionInfoString(WSManConnectionInfo connInfo)
{
StringBuilder connString = new();

if (connInfo.AppName == "wsman")
{
if (connInfo.Scheme == "http" && connInfo.Port == 5985)
{
connString.Append(connInfo.ComputerName);
}
else if (connInfo.Scheme == "https" && connInfo.Port == 5986)
{
connString.Append("https://").Append(connInfo.ComputerName);
}
}

if (connString.Length == 0)
{
connString.Append(connInfo.ConnectionUri);
}

return $"wsman:{connString}";
}

public override string ToString()
=> _originalString;
}
Loading

0 comments on commit b9bf190

Please sign in to comment.