From dfe995534652e23b3c14bd5dfd28445233d9fc6b Mon Sep 17 00:00:00 2001
From: "Joel Sallow (/u/ta11ow)" <32407840+vexx32@users.noreply.github.com>
Date: Sat, 28 Mar 2020 06:50:32 -0400
Subject: [PATCH] Don't write DNS resolution errors on Test-Connection -Quiet
 (#12204)

---
 .../management/TestConnectionCommand.cs       | 73 +++++++++++++++----
 .../Test-Connection.Tests.ps1                 | 13 ++--
 2 files changed, 66 insertions(+), 20 deletions(-)

diff --git a/src/Microsoft.PowerShell.Commands.Management/commands/management/TestConnectionCommand.cs b/src/Microsoft.PowerShell.Commands.Management/commands/management/TestConnectionCommand.cs
index a4044e3d5c..a6ff78d8c0 100644
--- a/src/Microsoft.PowerShell.Commands.Management/commands/management/TestConnectionCommand.cs
+++ b/src/Microsoft.PowerShell.Commands.Management/commands/management/TestConnectionCommand.cs
@@ -58,6 +58,8 @@ namespace Microsoft.PowerShell.Commands
 
         private static byte[]? s_DefaultSendBuffer;
 
+        private readonly CancellationTokenSource _dnsLookupCancel = new CancellationTokenSource();
+
         private bool _disposed;
 
         private Ping? _sender;
@@ -275,6 +277,7 @@ namespace Microsoft.PowerShell.Commands
         protected override void StopProcessing()
         {
             _sender?.SendAsyncCancel();
+            _dnsLookupCancel.Cancel();
         }
 
         #region ConnectionTest
@@ -283,6 +286,11 @@ namespace Microsoft.PowerShell.Commands
         {
             if (!TryResolveNameOrAddress(targetNameOrAddress, out _, out IPAddress? targetAddress))
             {
+                if (Quiet.IsPresent)
+                {
+                    WriteObject(false);
+                }
+
                 return;
             }
 
@@ -334,6 +342,11 @@ namespace Microsoft.PowerShell.Commands
 
             if (!TryResolveNameOrAddress(targetNameOrAddress, out string resolvedTargetName, out IPAddress? targetAddress))
             {
+                if (!Quiet.IsPresent)
+                {
+                    WriteObject(false);
+                }
+
                 return;
             }
 
@@ -471,6 +484,11 @@ namespace Microsoft.PowerShell.Commands
             PingReply? reply, replyResult = null;
             if (!TryResolveNameOrAddress(targetNameOrAddress, out string resolvedTargetName, out IPAddress? targetAddress))
             {
+                if (Quiet.IsPresent)
+                {
+                    WriteObject(-1);
+                }
+
                 return;
             }
 
@@ -574,6 +592,11 @@ namespace Microsoft.PowerShell.Commands
         {
             if (!TryResolveNameOrAddress(targetNameOrAddress, out string resolvedTargetName, out IPAddress? targetAddress))
             {
+                if (Quiet.IsPresent)
+                {
+                    WriteObject(false);
+                }
+
                 return;
             }
 
@@ -667,7 +690,7 @@ namespace Microsoft.PowerShell.Commands
 
                 if (ResolveDestination)
                 {
-                    hostEntry = Dns.GetHostEntry(targetNameOrAddress);
+                    hostEntry = GetCancellableHostEntry(targetNameOrAddress);
                     resolvedTargetName = hostEntry.HostName;
                 }
                 else
@@ -679,27 +702,35 @@ namespace Microsoft.PowerShell.Commands
             {
                 try
                 {
-                    hostEntry = Dns.GetHostEntry(targetNameOrAddress);
+                    hostEntry = GetCancellableHostEntry(targetNameOrAddress);
 
                     if (ResolveDestination)
                     {
                         resolvedTargetName = hostEntry.HostName;
-                        hostEntry = Dns.GetHostEntry(hostEntry.HostName);
+                        hostEntry = GetCancellableHostEntry(hostEntry.HostName);
                     }
                 }
+                catch (PipelineStoppedException)
+                {
+                    throw;
+                }
                 catch (Exception ex)
                 {
-                    string message = StringUtil.Format(
-                        TestConnectionResources.NoPingResult,
-                        resolvedTargetName,
-                        TestConnectionResources.CannotResolveTargetName);
-                    Exception pingException = new PingException(message, ex);
-                    ErrorRecord errorRecord = new ErrorRecord(
-                        pingException,
-                        TestConnectionExceptionId,
-                        ErrorCategory.ResourceUnavailable,
-                        resolvedTargetName);
-                    WriteError(errorRecord);
+                    if (!Quiet.IsPresent)
+                    {
+                        string message = StringUtil.Format(
+                            TestConnectionResources.NoPingResult,
+                            resolvedTargetName,
+                            TestConnectionResources.CannotResolveTargetName);
+                        Exception pingException = new PingException(message, ex);
+                        ErrorRecord errorRecord = new ErrorRecord(
+                            pingException,
+                            TestConnectionExceptionId,
+                            ErrorCategory.ResourceUnavailable,
+                            resolvedTargetName);
+                        WriteError(errorRecord);
+                    }
+
                     return false;
                 }
 
@@ -732,6 +763,20 @@ namespace Microsoft.PowerShell.Commands
             return true;
         }
 
+        private IPHostEntry GetCancellableHostEntry(string targetNameOrAddress)
+        {
+            var task = Dns.GetHostEntryAsync(targetNameOrAddress);
+            var waitHandles = new[] { ((IAsyncResult)task).AsyncWaitHandle, _dnsLookupCancel.Token.WaitHandle };
+
+            // WaitAny() returns the index of the first signal it gets; 1 is our cancellation token.
+            if (WaitHandle.WaitAny(waitHandles) == 1)
+            {
+                throw new PipelineStoppedException();
+            }
+
+            return task.GetAwaiter().GetResult();
+        }
+
         private IPAddress? GetHostAddress(IPHostEntry hostEntry)
         {
             AddressFamily addressFamily = IPv6 ? AddressFamily.InterNetworkV6 : AddressFamily.InterNetwork;
diff --git a/test/powershell/Modules/Microsoft.PowerShell.Management/Test-Connection.Tests.ps1 b/test/powershell/Modules/Microsoft.PowerShell.Management/Test-Connection.Tests.ps1
index 5e78380b49..31e14a0ce7 100644
--- a/test/powershell/Modules/Microsoft.PowerShell.Management/Test-Connection.Tests.ps1
+++ b/test/powershell/Modules/Microsoft.PowerShell.Management/Test-Connection.Tests.ps1
@@ -59,18 +59,19 @@ Describe "Test-Connection" -tags "CI" {
             $result2 | Should -BeFalse
         }
 
-        It "Ping fake host" {
+        It 'returns false without errors for an unresolvable address when using -Quiet' {
+            Test-Connection -Quiet -ErrorAction Stop -Count 1 -TargetName "fakeHost" | Should -BeFalse
+        }
 
-            { $result = Test-Connection "fakeHost" -Count 1 -Quiet -ErrorAction Stop } |
+        It "Ping fake host" {
+            { Test-Connection "fakeHost" -Count 1 -ErrorAction Stop } |
                 Should -Throw -ErrorId "TestConnectionException,Microsoft.PowerShell.Commands.TestConnectionCommand"
             # Error code = 11001 - Host not found.
             if ((Get-PlatformInfo).Platform -match "raspbian") {
                 $code = 11
-            }
-            elseif (!$IsWindows) {
+            } elseif (!$IsWindows) {
                 $code = -131073
-            }
-            else {
+            } else {
                 $code = 11001
             }
             $error[0].Exception.InnerException.ErrorCode | Should -Be $code