PowerShell · pwshBot · Nov 26, 2024
diff --git a/src/System.Management.Automation/engine/Interop/Windows/WNetGetConnection.cs b/src/System.Management.Automation/engine/Interop/Windows/WNetGetConnection.cs
@@ -6,6 +6,7 @@
 using System;
 using System.Buffers;
 using System.Runtime.InteropServices;
+using System.Management.Automation.Internal;
 
 internal static partial class Interop
 {
@@ -14,7 +15,7 @@ internal static unsafe partial class Windows
         private static bool s_WNetApiNotAvailable;
 
         [LibraryImport("mpr.dll", EntryPoint = "WNetGetConnectionW", StringMarshalling = StringMarshalling.Utf16)]
-        internal static partial int WNetGetConnection(ReadOnlySpan<char> localName, Span<char> remoteName, ref uint remoteNameLength);
+        internal static partial int WNetGetConnection(ReadOnlySpan<char> localName, Span<char> remoteName, ref int remoteNameLength);
 
         internal static int GetUNCForNetworkDrive(char drive, out string? uncPath)
         {
@@ -24,46 +25,43 @@ internal static int GetUNCForNetworkDrive(char drive, out string? uncPath)
                 return ERROR_NOT_SUPPORTED;
             }
 
-            uint bufferSize = MAX_PATH;
-
-#if DEBUG
-            // In Debug mode buffer size is initially set to 3 and if additional buffer is required, the
-            // required buffer size is allocated and the WNetGetConnection API is executed with the newly
-            // allocated buffer size.
-            bufferSize = 3;
-#endif
-
             ReadOnlySpan<char> driveName = stackalloc char[] { drive, ':', '\0' };
-            Span<char> uncBuffer = stackalloc char[(int)bufferSize];
-            int errorCode = ERROR_NO_NETWORK;
-
-            try
+            int bufferSize = MAX_PATH;
+            Span<char> uncBuffer = stackalloc char[MAX_PATH];
+            if (InternalTestHooks.WNetGetConnectionBufferSize > 0 && InternalTestHooks.WNetGetConnectionBufferSize <= MAX_PATH)
             {
-                errorCode = WNetGetConnection(driveName, uncBuffer, ref bufferSize);
-            }
-            catch (System.DllNotFoundException)
-            {
-                s_WNetApiNotAvailable = true;
-                return ERROR_NOT_SUPPORTED;
+                bufferSize = InternalTestHooks.WNetGetConnectionBufferSize;
+                uncBuffer = uncBuffer.Slice(0, bufferSize);
             }
 
-            if (errorCode == ERROR_SUCCESS)
-            {
-                // exclude null terminator
-                uncPath = uncBuffer.Slice(0, (int)bufferSize - 1).ToString();
-            }
-            else if (errorCode == ERROR_MORE_DATA)
+            char[]? rentedArray = null;
+            while (true)
             {
-                char[]? rentedArray = null;
+                int errorCode;
                 try
                 {
-                    uncBuffer = rentedArray = ArrayPool<char>.Shared.Rent((int)bufferSize);
-                    errorCode = WNetGetConnection(driveName, uncBuffer, ref bufferSize);
+                    try
+                    {
+                        errorCode = WNetGetConnection(driveName, uncBuffer, ref bufferSize);
+                    }
+                    catch (DllNotFoundException)
+                    {
+                        s_WNetApiNotAvailable = true;
+                        return ERROR_NOT_SUPPORTED;
+                    }
 
                     if (errorCode == ERROR_SUCCESS)
                     {
-                        // exclude null terminator
-                        uncPath = uncBuffer.Slice(0, (int)bufferSize - 1).ToString();
+                        // Cannot rely on bufferSize as it's only set if
+                        // the first call ended with ERROR_MORE_DATA,
+                        // instead slice at the null terminator.
+                        unsafe
+                        {
+                            fixed (char* uncBufferPtr = uncBuffer)
+                            {
+                                uncPath = new string(uncBufferPtr);
+                            }
+                        }
                     }
                 }
                 finally
@@ -73,9 +71,16 @@ internal static int GetUNCForNetworkDrive(char drive, out string? uncPath)
                         ArrayPool<char>.Shared.Return(rentedArray);
                     }
                 }
-            }
 
-            return errorCode;
+                if (errorCode == ERROR_MORE_DATA)
+                {
+                    uncBuffer = rentedArray = ArrayPool<char>.Shared.Rent(bufferSize);
+                }
+                else
+                {
+                    return errorCode;
+                }
+            }
         }
     }
 }
diff --git a/src/System.Management.Automation/engine/Utils.cs b/src/System.Management.Automation/engine/Utils.cs
@@ -1602,6 +1602,9 @@ public static class InternalTestHooks
         internal static bool OneDriveTestRecurseOn;
         internal static string OneDriveTestSymlinkName = "link-Beta";
 
+        // Test out smaller connection buffer size when calling WNetGetConnection.
+        internal static int WNetGetConnectionBufferSize = -1;
+
         /// <summary>This member is used for internal test purposes.</summary>
         public static void SetTestHook(string property, object value)
         {

diff --git a/test/powershell/Modules/Microsoft.PowerShell.Management/Get-PSDrive.Tests.ps1 b/test/powershell/Modules/Microsoft.PowerShell.Management/Get-PSDrive.Tests.ps1
@@ -66,7 +66,7 @@ Describe "Temp: drive" -Tag Feature {
     }
 }
 
-Describe "Get-PSDrive for network path" -Tags "Feature","RequireAdminOnWindows" {
+Describe "Get-PSDrive for network path" -Tags "CI","RequireAdminOnWindows" {
 
     It 'Check P/Invoke GetDosDevice/QueryDosDevice' -Skip:(-not $IsWindows) {
         $UsedDrives  = Get-PSDrive | Select-Object -ExpandProperty Name
@@ -78,4 +78,39 @@ Describe "Get-PSDrive for network path" -Tags "Feature","RequireAdminOnWindows"
         $drive.DisplayRoot | Should -BeExactly '\\localhost\c$\Windows'
         subst "$($PSDriveName):" /D
     }
+
+    It 'Check P/Invoke for WNetGetConnection with small buffer: <SmallBuffer>' -Skip:(-not $IsWindows) -TestCases @(
+        @{ SmallBuffer = $false }
+        @{ SmallBuffer = $true }
+    ) {
+        param ($SmallBuffer)
+
+        $UsedDrives  = Get-PSDrive | Select-Object -ExpandProperty Name
+        $PSDriveName = 'D'..'Z' | Where-Object -FilterScript {$_ -notin $UsedDrives} | Get-Random
+
+        $drive = New-PSDrive -Name $PSDriveName -PSProvider FileSystem -Root \\localhost\c$\Windows -Persist
+        try {
+            if ($SmallBuffer) {
+                [System.Management.Automation.Internal.InternalTestHooks]::SetTestHook('WNetGetConnectionBufferSize', 4)
+            }
+
+            # The result is cached in the current instance, use a new
+            # PowerShell instance to test out WNetGetConnection code.
+            $ps = [PowerShell]::Create()
+            $actual = $ps.AddCommand('Get-PSDrive').AddParameter('Name', $PSDriveName).Invoke()
+            if ($ps.HadErrors) {
+                throw $ps.Streams.Error[0]
+            }
+
+            $actual.Name | Should -BeExactly $PSDriveName
+            $actual.DisplayRoot | Should -BeExactly '\\localhost\c$\Windows'
+        }
+        finally {
+            $drive | Remove-PSDrive
+
+            if ($SmallBuffer) {
+                [System.Management.Automation.Internal.InternalTestHooks]::SetTestHook('WNetGetConnectionBufferSize', -1)
+            }
+        }
+    }
 }