MegaBites-AI/Windows-powershell

1// Copyright (c) Microsoft Corporation.2// Licensed under the MIT License.3 4#nullable enable5 6using System;7using System.Buffers;8using System.ComponentModel;9using System.Management.Automation;10using System.Runtime.InteropServices;11 12internal static partial class Interop13{14    internal static partial class Windows15    {16        [LibraryImport(PinvokeDllNames.QueryDosDeviceDllName, EntryPoint = "QueryDosDeviceW", StringMarshalling = StringMarshalling.Utf16, SetLastError = true)]17        internal static partial int QueryDosDevice(Span<char> lpDeviceName, Span<char> lpTargetPath, uint ucchMax);18 19        internal static string GetDosDeviceForNetworkPath(char deviceName)20        {21            // By default buffer size is set to 300 which would generally be sufficient in most of the cases.22            const int StartLength =23#if DEBUG24                // In debug, validate ArrayPool growth.25                1;26#else27                300;28#endif29 30            Span<char> buffer = stackalloc char[StartLength + 1];31            Span<char> fullDeviceName = [deviceName, ':', '\0'];32            char[]? rentedArray = null;33 34            try35            {36                while (true)37                {38                    uint length = (uint)buffer.Length;39                    int retValue = QueryDosDevice(fullDeviceName, buffer, length);40                    if (retValue > 0)41                    {42                        if (buffer.StartsWith("\\??\\"))43                        {44                            // QueryDosDevice always return array of NULL-terminating strings with additional final NULL45                            // so the buffer has always two NULL-s on end.46                            //47                            // "\\??\\UNC\\localhost\\c$\\tmp\0\0" -> "UNC\\localhost\\c$\\tmp\0\0"48                            Span<char> res = buffer.Slice(4);49                            if (res.StartsWith("UNC"))50                            {51                                // -> "C\\localhost\\c$\\tmp\0\0" -> "\\\\localhost\\c$\\tmp"52                                //53                                // We need to take only first null-terminated string as QueryDosDevice() docs say.54                                int i = 3;55                                for (; i < res.Length; i++)56                                {57                                    if (res[i] == '\0')58                                    {59                                        break;60                                    }61                                }62 63                                Diagnostics.Assert(i < res.Length, "Broken QueryDosDevice() buffer.");64 65                                res = res.Slice(2, i);66                                res[0] = '\\';67 68                                // If we want always to have terminating slash -> "\\\\localhost\\c$\\tmp\\"69                                // res = res.Slice(2, retValue - 3);70                                // res[0] = '\\';71                                // res[^1] = '\\';72                            }73                            // else if (res[^3] == ':')74                            // {75                            //     Diagnostics.Assert(false, "Really it is a dead code since GetDosDevice() is called only if PSDrive.DriveType == DriveType.Network");76 77                            //     // The substed path is the root path of a drive. For example: subst Y: C:\78                            //     // -> "C:\0\0" -> "C:\"79                            //     res = res.Slice(0, retValue - 1);80                            //     res[^1] = '\\';81                            // }82                            else83                            {84                                throw new Exception("GetDosDeviceForNetworkPath() can be called only if PSDrive.DriveType == DriveType.Network.");85                            }86 87                            return res.ToString();88                        }89                        else90                        {91                            Diagnostics.Assert(false, "Really it is a dead code since GetDosDevice() is called only if PSDrive.DriveType == DriveType.Network");92 93                            // The drive name is not a substed path, then we return the root path of the drive94                            // "C:\0" -> "C:\\"95                            fullDeviceName[^1] = '\\';96                            return fullDeviceName.ToString();97                        }98                    }99 100                    const int ERROR_INSUFFICIENT_BUFFER = 122;101                    int errorCode = Marshal.GetLastPInvokeError();102                    if (errorCode != ERROR_INSUFFICIENT_BUFFER)103                    {104                        throw new Win32Exception((int)errorCode);105                    }106 107                    char[]? toReturn = rentedArray;108                    buffer = rentedArray = ArrayPool<char>.Shared.Rent(buffer.Length * 2);109                    if (toReturn is not null)110                    {111                        ArrayPool<char>.Shared.Return(toReturn);112                    }113                }114            }115            finally116            {117                if (rentedArray is not null)118                {119                    ArrayPool<char>.Shared.Return(rentedArray);120                }121            }122        }123    }124}125