PowerShell · TravisEz13 · May 1, 2023 · Dec 25, 2019 · Dec 25, 2019 · Dec 27, 2019
diff --git a/src/Microsoft.PowerShell.Commands.Management/commands/management/TestConnectionCommand.cs b/src/Microsoft.PowerShell.Commands.Management/commands/management/TestConnectionCommand.cs
@@ -27,6 +27,7 @@ namespace Microsoft.PowerShell.Commands
     [OutputType(typeof(PingMtuStatus), ParameterSetName = new string[] { MtuSizeDetectParameterSet })]
     [OutputType(typeof(int), ParameterSetName = new string[] { MtuSizeDetectParameterSet })]
     [OutputType(typeof(TraceStatus), ParameterSetName = new string[] { TraceRouteParameterSet })]
+    [OutputType(typeof(TcpPortStatus), ParameterSetName = new string[] { TcpPortParameterSet })]
     public class TestConnectionCommand : PSCmdlet, IDisposable
     {
         #region Parameter Set Names
@@ -134,6 +135,7 @@ public class TestConnectionCommand : PSCmdlet, IDisposable
         /// The default (from Windows) is 4 times.
         /// </summary>
         [Parameter(ParameterSetName = DefaultPingParameterSet)]
+        [Parameter(ParameterSetName = TcpPortParameterSet)]
         [ValidateRange(ValidateRangeKind.Positive)]
         public int Count { get; set; } = 4;
 
@@ -143,6 +145,7 @@ public class TestConnectionCommand : PSCmdlet, IDisposable
         /// </summary>
         [Parameter(ParameterSetName = DefaultPingParameterSet)]
         [Parameter(ParameterSetName = RepeatPingParameterSet)]
+        [Parameter(ParameterSetName = TcpPortParameterSet)]
         [ValidateRange(ValidateRangeKind.Positive)]
         public int Delay { get; set; } = 1;
 
@@ -169,6 +172,7 @@ public class TestConnectionCommand : PSCmdlet, IDisposable
         /// Gets or sets whether to continue pinging until user presses Ctrl-C (or Int.MaxValue threshold reached).
         /// </summary>
         [Parameter(Mandatory = true, ParameterSetName = RepeatPingParameterSet)]
+        [Parameter(ParameterSetName = TcpPortParameterSet)]
         [Alias("Continuous")]
         public SwitchParameter Repeat { get; set; }
 
@@ -180,6 +184,13 @@ public class TestConnectionCommand : PSCmdlet, IDisposable
         [Parameter]
         public SwitchParameter Quiet { get; set; }
 
+        /// <summary>
+        /// Gets or sets whether to enable detailed output mode while running a TCP connection test.
+        /// Without this flag, the TCP test will return a boolean result.
+        /// </summary>
+        [Parameter]
+        public SwitchParameter Detailed;
+
         /// <summary>
         /// Gets or sets the timeout value for an individual ping in seconds.
         /// If a response is not received in this time, no response is assumed.
@@ -227,6 +238,7 @@ public class TestConnectionCommand : PSCmdlet, IDisposable
 
         /// <summary>
         /// BeginProcessing implementation for TestConnectionCommand.
+        /// Sets Count for different types of tests unless specified explicitly.
         /// </summary>
         protected override void BeginProcessing()
         {
@@ -235,6 +247,9 @@ protected override void BeginProcessing()
                 case RepeatPingParameterSet:
                     Count = int.MaxValue;
                     break;
+                case TcpPortParameterSet:
+                    SetCountForTcpTest();
+                    break;
             }
         }
 
@@ -281,6 +296,18 @@ protected override void StopProcessing()
 
         #region ConnectionTest
 
+        private void SetCountForTcpTest()
+        {
+            if (Repeat.IsPresent)
+            {
+                Count = int.MaxValue;
+            }
+            else if (!MyInvocation.BoundParameters.ContainsKey(nameof(Count)))
+            {
+                Count = 1;
+            }
+        }
+
         private void ProcessConnectionByTCPPort(string targetNameOrAddress)
         {
             if (!TryResolveNameOrAddress(targetNameOrAddress, out _, out IPAddress? targetAddress))
@@ -293,42 +320,80 @@ private void ProcessConnectionByTCPPort(string targetNameOrAddress)
                 return;
             }
 
-            TcpClient client = new();
+            int timeoutMilliseconds = TimeoutSeconds * 1000;
+            int delayMilliseconds = Delay * 1000;
 
-            try
+            for (var i = 1; i <= Count; i++)
             {
-                Task connectionTask = client.ConnectAsync(targetAddress, TcpPort);
-                string targetString = targetAddress.ToString();
+                long latency = 0;
+                SocketError status = SocketError.SocketError;
+
+                Stopwatch stopwatch = new Stopwatch();
 
-                for (var i = 1; i <= TimeoutSeconds; i++)
+                using var client = new TcpClient();
+
+                try
                 {
-                    Task timeoutTask = Task.Delay(millisecondsDelay: 1000);
-                    Task.WhenAny(connectionTask, timeoutTask).Result.Wait();
+                    stopwatch.Start();
 
-                    if (timeoutTask.Status == TaskStatus.Faulted || timeoutTask.Status == TaskStatus.Canceled)
+                    if (client.ConnectAsync(targetAddress, TcpPort).Wait(timeoutMilliseconds, _dnsLookupCancel.Token))
                     {
-                        // Waiting is interrupted by Ctrl-C.
-                        WriteObject(false);
-                        return;
+                        latency = stopwatch.ElapsedMilliseconds;
+                        status = SocketError.Success;
                     }
-
-                    if (connectionTask.Status == TaskStatus.RanToCompletion)
+                    else
                     {
-                        WriteObject(true);
-                        return;
+                        status = SocketError.TimedOut;
                     }
                 }
-            }
-            catch
-            {
-                // Silently ignore connection errors.
-            }
-            finally
-            {
-                client.Close();
-            }
+                catch (AggregateException ae)
+                {
+                    ae.Handle((ex) =>
+                    {
+                        if (ex is TaskCanceledException)
+                        {
+                            throw new PipelineStoppedException();
+                        }
+                        if (ex is SocketException socketException)
+                        {
+                            status = socketException.SocketErrorCode;
+                            return true;
+                        }
+                        else
+                        {
+                            return false;
+                        }
+                    });
+                }
+                finally
+                {
+                    stopwatch.Reset();
+                }
+
+                if (!Detailed.IsPresent)
+                {
+                    WriteObject(status == SocketError.Success);
+                    return;
+                }
+                else
+                {
+                    WriteObject(new TcpPortStatus(
+                        i,
+                        Source,
+                        targetNameOrAddress,
+                        targetAddress,
+                        TcpPort,
+                        latency,
+                        status == SocketError.Success,
+                        status
+                    ));
+                }
 
-            WriteObject(false);
+                if (i < Count)
+                {
+                    Task.Delay(delayMilliseconds).Wait(_dnsLookupCancel.Token);
+                }
+            }
         }
 
         #endregion ConnectionTest
@@ -875,6 +940,75 @@ private PingReply SendCancellablePing(
             }
         }
 
+        /// <summary>
+        /// The class contains information about the TCP connection test.
+        /// </summary>
+        public class TcpPortStatus
+        {
+            /// <summary>
+            /// Initializes a new instance of the <see cref="TcpPortStatus"/> class.
+            /// </summary>
+            /// <param name="id">The number of this test.</param>
+            /// <param name="source">The source machine name or IP of the test.</param>
+            /// <param name="target">The target machine name or IP of the test.</param>
+            /// <param name="targetAddress">The resolved IP from the target.</param>
+            /// <param name="port">The port used for the connection.</param>
+            /// <param name="latency">The latency of the test.</param>
+            /// <param name="connected">If the test connection succeeded.</param>
+            /// <param name="status">Status of the underlying socket.</param>
+            internal TcpPortStatus(int id, string source, string target, IPAddress targetAddress, int port, long latency, bool connected, SocketError status)
+            {
+                Id = id;
+                Source = source;
+                Target = target;
+                TargetAddress = targetAddress;
+                Port = port;
+                Latency = latency;
+                Connected = connected;
+                Status = status;
+            }
+
+            /// <summary>
+            /// Gets and sets the count of the test.
+            /// </summary>
+            public int Id { get; set; }
+
+            /// <summary>
+            /// Gets the source from which the test was sent.
+            /// </summary>
+            public string Source { get; }
+
+            /// <summary>
+            /// Gets the target name.
+            /// </summary>
+            public string Target { get; }
+
+            /// <summary>
+            /// Gets the resolved address for the target.
+            /// </summary>
+            public IPAddress TargetAddress { get; }
+
+            /// <summary>
+            /// Gets the port used for the test.
+            /// </summary>
+            public int Port { get; }
+
+            /// <summary>
+            /// Gets or sets the latancy of the connection.
+            /// </summary>
+            public long Latency { get; set; }
+
+            /// <summary>
+            /// Gets or sets the result of the test.
+            /// </summary>
+            public bool Connected { get; set; }
+
+            /// <summary>
+            /// Gets or sets the state of the socket after the test.
+            /// </summary>
+            public SocketError Status { get; set; }
+        }
+
         /// <summary>
         /// The class contains information about the source, the destination and ping results.
         /// </summary>

diff --git a/...m.Management.Automation/FormatAndOutput/DefaultFormatters/PowerShellCore_format_ps1xml.cs b/...m.Management.Automation/FormatAndOutput/DefaultFormatters/PowerShellCore_format_ps1xml.cs
@@ -248,6 +248,10 @@ internal static IEnumerable<ExtendedTypeDefinition> GetFormatData()
                 "Microsoft.PowerShell.MarkdownRender.PSMarkdownOptionInfo",
                 ViewsOf_Microsoft_PowerShell_MarkdownRender_MarkdownOptionInfo());
 
+            yield return new ExtendedTypeDefinition(
+                "Microsoft.PowerShell.Commands.TestConnectionCommand+TcpPortStatus",
+                ViewsOf_Microsoft_PowerShell_Commands_TestConnectionCommand_TcpPortStatus());
+
             yield return new ExtendedTypeDefinition(
                 "Microsoft.PowerShell.Commands.TestConnectionCommand+PingStatus",
                 ViewsOf_Microsoft_PowerShell_Commands_TestConnectionCommand_PingStatus());
@@ -1924,6 +1928,31 @@ private static IEnumerable<FormatViewDefinition> ViewsOf_Microsoft_PowerShell_Ma
                 .EndList());
         }
 
+        private static IEnumerable<FormatViewDefinition> ViewsOf_Microsoft_PowerShell_Commands_TestConnectionCommand_TcpPortStatus()
+        {
+            yield return new FormatViewDefinition(
+                "Microsoft.PowerShell.Commands.TestConnectionCommand+TcpPortStatus",
+                TableControl.Create()
+                    .AddHeader(Alignment.Right, label: "Id", width: 4)
+                    .AddHeader(Alignment.Left, label: "Source", width: 16)
+                    .AddHeader(Alignment.Left, label: "Address", width: 25)
+                    .AddHeader(Alignment.Right, label: "Port", width: 7)
+                    .AddHeader(Alignment.Right, label: "Latency(ms)", width: 7)
+                    .AddHeader(Alignment.Left, label: "Connected", width: 10)
+                    .AddHeader(Alignment.Left, label: "Status", width: 24)
+                    .StartRowDefinition()
+                        .AddPropertyColumn("Id")
+                        .AddPropertyColumn("Source")
+                        .AddPropertyColumn("TargetAddress")
+                        .AddPropertyColumn("Port")
+                        .AddPropertyColumn("Latency")
+                        .AddPropertyColumn("Connected")
+                        .AddPropertyColumn("Status")
+                    .EndRowDefinition()
+                    .GroupByProperty("Target")
+                .EndTable());
+        }
+
         private static IEnumerable<FormatViewDefinition> ViewsOf_Microsoft_PowerShell_Commands_TestConnectionCommand_PingStatus()
         {
             yield return new FormatViewDefinition(

diff --git a/test/powershell/Modules/Microsoft.PowerShell.Management/Test-Connection.Tests.ps1 b/test/powershell/Modules/Microsoft.PowerShell.Management/Test-Connection.Tests.ps1
@@ -320,13 +320,49 @@ Describe "Connection" -Tag "CI", "RequireAdminOnWindows" {
         $UnreachableAddress = "10.11.12.13"
     }
 
-    It "Test connection to local host port 80" {
+    It "Test connection to local host on working port" {
         Test-Connection '127.0.0.1' -TcpPort $WebListener.HttpPort | Should -BeTrue
     }
 
     It "Test connection to unreachable host port 80" {
         Test-Connection $UnreachableAddress -TcpPort 80 -TimeOut 1 | Should -BeFalse
     }
+
+    It "Test detailed connection to local host on working port" {
+        $result = Test-Connection '127.0.0.1' -TcpPort $WebListener.HttpPort -Detailed
+
+        $result.Count | Should -Be 1
+        $result[0].Id | Should -BeExactly 1
+        $result[0].TargetAddress | Should -BeExactly '127.0.0.1'
+        $result[0].Port | Should -Be $WebListener.HttpPort
+        $result[0].Latency | Should -BeGreaterOrEqual 0
+        $result[0].Connected | Should -BeTrue
+        $result[0].Status | Should -BeExactly 'Success'
+    }
+
+    It "Test detailed connection to local host on working port with modified count" {
+        $result = Test-Connection '127.0.0.1' -TcpPort $WebListener.HttpPort -Detailed -Count 2
+
+        $result.Count | Should -Be 2
+        $result[0].Id | Should -BeExactly 1
+        $result[0].TargetAddress | Should -BeExactly '127.0.0.1'
+        $result[0].Port | Should -Be $WebListener.HttpPort
+        $result[0].Latency | Should -BeGreaterOrEqual 0
+        $result[0].Connected | Should -BeTrue
+        $result[0].Status | Should -BeExactly 'Success'
+    }
+
+    It "Test detailed connection to unreachable host port 80" {
+        $result = Test-Connection $UnreachableAddress -TcpPort 80 -Detailed -TimeOut 1
+
+        $result.Count | Should -Be 1
+        $result[0].Id | Should -BeExactly 1
+        $result[0].TargetAddress | Should -BeExactly $UnreachableAddress
+        $result[0].Port | Should -Be 80
+        $result[0].Latency | Should -BeExactly 0
+        $result[0].Connected | Should -BeFalse
+        $result[0].Status | Should -Not -BeExactly 'Success'
+    }
 }
 
 Describe "Test-Connection should run in the default synchronization context (threadpool)" -Tag "CI" {