From 22beba4a47da2c84f01d80d431b194c6865df734 Mon Sep 17 00:00:00 2001 From: Jordan Borean Date: Tue, 26 Nov 2024 16:59:34 +1000 Subject: [PATCH] Fix up buffer management getting network roots (#24600) Fixes the logic used when getting the PSDrive DisplayRoot for network PSDrive paths to exclude null terminators. --- .../Interop/Windows/WNetGetConnection.cs | 71 ++++++++++--------- .../engine/Utils.cs | 3 + .../Get-PSDrive.Tests.ps1 | 37 +++++++++- 3 files changed, 77 insertions(+), 34 deletions(-) diff --git a/src/System.Management.Automation/engine/Interop/Windows/WNetGetConnection.cs b/src/System.Management.Automation/engine/Interop/Windows/WNetGetConnection.cs index fcdaeb65d3..88ec4386a1 100644 --- 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 partial class Interop private static bool s_WNetApiNotAvailable; [LibraryImport("mpr.dll", EntryPoint = "WNetGetConnectionW", StringMarshalling = StringMarshalling.Utf16)] - internal static partial int WNetGetConnection(ReadOnlySpan localName, Span remoteName, ref uint remoteNameLength); + internal static partial int WNetGetConnection(ReadOnlySpan localName, Span remoteName, ref int remoteNameLength); internal static int GetUNCForNetworkDrive(char drive, out string? uncPath) { @@ -24,46 +25,43 @@ internal static partial class Interop 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 driveName = stackalloc char[] { drive, ':', '\0' }; - Span uncBuffer = stackalloc char[(int)bufferSize]; - int errorCode = ERROR_NO_NETWORK; - - try + int bufferSize = MAX_PATH; + Span 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) + char[]? rentedArray = null; + while (true) { - // exclude null terminator - uncPath = uncBuffer.Slice(0, (int)bufferSize - 1).ToString(); - } - else if (errorCode == ERROR_MORE_DATA) - { - char[]? rentedArray = null; + int errorCode; try { - uncBuffer = rentedArray = ArrayPool.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 partial class Interop ArrayPool.Shared.Return(rentedArray); } } - } - return errorCode; + if (errorCode == ERROR_MORE_DATA) + { + uncBuffer = rentedArray = ArrayPool.Shared.Rent(bufferSize); + } + else + { + return errorCode; + } + } } } } diff --git a/src/System.Management.Automation/engine/Utils.cs b/src/System.Management.Automation/engine/Utils.cs index c01bbbc598..2a2091af31 100644 --- a/src/System.Management.Automation/engine/Utils.cs +++ b/src/System.Management.Automation/engine/Utils.cs @@ -1609,6 +1609,9 @@ namespace System.Management.Automation.Internal internal static bool OneDriveTestRecurseOn; internal static string OneDriveTestSymlinkName = "link-Beta"; + // Test out smaller connection buffer size when calling WNetGetConnection. + internal static int WNetGetConnectionBufferSize = -1; + /// This member is used for internal test purposes. 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 index f2684773a1..4ad46de7cb 100644 --- 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: ' -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) + } + } + } }