Fix 'Set-Location DriveName:' to restore current working directory in the drive (#6774)

Matthew Bobke · daxian-dbw · commit 4e3db1bb93b1 · 2018-05-16T17:13:22.000-07:00
diff --git a/src/System.Management.Automation/engine/SessionStateLocationAPIs.cs b/src/System.Management.Automation/engine/SessionStateLocationAPIs.cs
@@ -5,6 +5,7 @@
 
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
+using System.IO;
 using System.Management.Automation.Provider;
 using Dbg = System.Management.Automation;
 
@@ -262,10 +263,18 @@ internal PathInfo SetLocation(string path, CmdletProviderContext context)
                     // Since the path is an absolute path
                     // we need to change the current working
                     // drive
-
                     PSDriveInfo newWorkingDrive = GetDrive(driveName);
                     CurrentDrive = newWorkingDrive;
 
+                    // If the path is simply a colon-terminated drive,
+                    // not a slash-terminated path to the root of a drive,
+                    // set the path to the current working directory of that drive.
+                    string colonTerminatedVolume = CurrentDrive.Name + ':';
+                    if (CurrentDrive.VolumeSeparatedByColon && (path.Length == colonTerminatedVolume.Length))
+                    {
+                        path = Path.Combine((colonTerminatedVolume + Path.DirectorySeparatorChar), CurrentDrive.CurrentLocation);
+                    }
+
                     // Now that the current working drive is set,
                     // process the rest of the path as a relative path.
                 }
diff --git a/src/System.Management.Automation/namespaces/LocationGlobber.cs b/src/System.Management.Automation/namespaces/LocationGlobber.cs
@@ -2013,6 +2013,9 @@ internal string GetDriveRootRelativePathFromPSPath(
             // Check to see if the path is relative or absolute
             bool isPathForCurrentDrive = false;
 
+            // Check to see if the path is to the root of a drive
+            bool isPathForRootOfDrive = false;
+
             if (IsAbsolutePath(path, out driveName))
             {
                 Dbg.Diagnostics.Assert(
@@ -2080,6 +2083,12 @@ internal string GetDriveRootRelativePathFromPSPath(
                         // this is the default behavior for all windows drives, and all non-filesystem
                         // drives on non-windows
                         path = path.Substring(driveName.Length + 1);
+
+                        if (String.IsNullOrEmpty(path))
+                        {
+                            // path was to the root of a drive such as 'c:'
+                            isPathForRootOfDrive = true;
+                        }
                     }
                 }
             }
@@ -2111,14 +2120,23 @@ internal string GetDriveRootRelativePathFromPSPath(
                 // have access to it.
                 context.Drive = workingDriveForPath;
 
-                string relativePath =
-                    GenerateRelativePath(
-                        workingDriveForPath,
-                        path,
-                        escapeCurrentLocation,
-                        providerInstance,
-                        context);
+                string relativePath = String.Empty;
 
+                if (isPathForRootOfDrive)
+                {
+                    relativePath = context.Drive.Root;
+                }
+                else
+                {
+                    relativePath =
+                        GenerateRelativePath(
+                            workingDriveForPath,
+                            path,
+                            escapeCurrentLocation,
+                            providerInstance,
+                            context);
+                }
+                
                 return relativePath;
             }
             catch (PSNotSupportedException)
diff --git a/test/powershell/Modules/Microsoft.PowerShell.Management/New-Item.Tests.ps1 b/test/powershell/Modules/Microsoft.PowerShell.Management/New-Item.Tests.ps1
@@ -117,6 +117,33 @@ Describe "New-Item" -Tags "CI" {
         $fileInfo.Target | Should -BeNullOrEmpty
         $fileInfo.LinkType | Should -BeExactly "HardLink"
     }
+
+    It "Should create a file at the root of the drive while the current working directory is not the root" {
+        try {
+            New-Item -Name $testfolder -Path "TestDrive:\" -ItemType directory > $null
+            Push-Location -Path "TestDrive:\$testfolder"
+            New-Item -Name $testfile -Path "TestDrive:\" -ItemType file > $null
+            $FullyQualifiedFile | Should -Exist
+        }
+        finally {
+            Pop-Location
+        }
+    }
+
+    It "Should create a folder at the root of the drive while the current working directory is not the root" {
+        $testfolder2 = "newDirectory2"
+        $FullyQualifiedFolder2 = Join-Path -Path $tmpDirectory -ChildPath $testfolder2
+
+        try {
+            New-Item -Name $testfolder -Path "TestDrive:\" -ItemType directory > $null
+            Push-Location -Path "TestDrive:\$testfolder"
+            New-Item -Name $testfolder2 -Path "TestDrive:\" -ItemType directory > $null
+            $FullyQualifiedFolder2 | Should -Exist
+        }
+        finally {
+            Pop-Location
+        }
+    }
 }
 
 # More precisely these tests require SeCreateSymbolicLinkPrivilege.
@@ -186,7 +213,7 @@ Describe "New-Item with links" -Tags @('CI', 'RequireAdminOnWindows') {
     }
 
     It "New-Item -ItemType SymbolicLink should understand directory path ending with slash" {
-        $folderName = [System.IO.Path]::GetRandomFileName()            
+        $folderName = [System.IO.Path]::GetRandomFileName()
         $symbolicLinkPath = New-Item -ItemType SymbolicLink -Path "$tmpDirectory/$folderName/" -Value "/bar/"
         $symbolicLinkPath | Should -Not -BeNullOrEmpty
     }
diff --git a/test/powershell/Modules/Microsoft.PowerShell.Management/Set-Location.Tests.ps1 b/test/powershell/Modules/Microsoft.PowerShell.Management/Set-Location.Tests.ps1
@@ -67,4 +67,30 @@ Describe "Set-Location" -Tags "CI" {
             (Get-Location).Path | Should -BeExactly (Get-PSProvider FileSystem).Home
         }
     }
+
+    It "Should set location to new drive's current working directory when path is the colon-terminated name of a different drive" {
+        try
+        {
+            $oldLocation = Get-Location
+            Set-Location 'TestDrive:\'
+            New-Item -Path 'TestDrive:\' -Name 'Directory1' -ItemType Directory
+            New-PSDrive -Name 'Z' -PSProvider FileSystem -Root 'TestDrive:\Directory1'
+            New-Item -Path 'Z:\' -Name 'Directory2' -ItemType Directory
+
+            Set-Location 'TestDrive:\Directory1'
+            $pathToTest1 = (Get-Location).Path
+            Set-Location 'Z:\Directory2'
+            $pathToTest2 = (Get-Location).Path
+
+            Set-Location 'TestDrive:'
+            (Get-Location).Path | Should -BeExactly $pathToTest1
+            Set-Location 'Z:'
+            (Get-Location).Path | Should -BeExactly $pathToTest2
+        }
+        finally
+        {
+            Set-Location $oldLocation
+            Remove-PSDrive -Name 'Z'
+        }
+    }
 }
