Avoid traversing links in remove_virtualenv

crates/uv-fs/src/lib.rs

@@ -855,7 +855,14 @@ pub fn copy_dir_all(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> std::io::Re
 }
 
 /// Perform a safe removal of a virtual environment.
+///
+/// Links at `location` are removed without following them.
 pub fn remove_virtualenv(location: &Path) -> io::Result<()> {
+    let file_type = fs_err::symlink_metadata(location)?.file_type();
+    if file_type.is_symlink() {
+        return remove_symlink(location);
+    }
+
     // On Windows, if the current executable is in the directory, defer self-deletion since Windows
     // won't let you unlink a running executable.
     #[cfg(windows)]
@@ -906,12 +913,28 @@ pub fn remove_virtualenv(location: &Path) -> io::Result<()> {
     Ok(())
 }
 
+/// Prepare an empty virtual environment directory, resolving links when possible.
+///
+/// Returns whether an existing entry was found.
+pub fn clear_virtualenv(location: &Path) -> io::Result<bool> {
+    let location = location
+        .canonicalize()
+        .unwrap_or_else(|_| location.to_path_buf());
+    let cleared = match remove_virtualenv(&location) {
+        Ok(()) => true,
+        Err(err) if err.kind() == io::ErrorKind::NotFound => false,
+        Err(err) => return Err(err),
+    };
+    fs_err::create_dir_all(location)?;
+    Ok(cleared)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
-    fn remove_symlink_removes_directory_link_without_removing_target() -> std::io::Result<()> {
+    fn remove_symlink_removes_directory_link_without_removing_target() -> io::Result<()> {
         let tempdir = tempfile::tempdir()?;
         let target = tempdir.path().join("target");
         fs_err::create_dir(&target)?;
@@ -923,9 +946,39 @@ mod tests {
 
         assert!(matches!(
             fs_err::symlink_metadata(&link),
-            Err(err) if err.kind() == std::io::ErrorKind::NotFound
+            Err(err) if err.kind() == io::ErrorKind::NotFound
         ));
         assert_eq!(fs_err::read_to_string(target.join("file"))?, "content");
         Ok(())
     }
+
+    #[test]
+    fn remove_virtualenv_removes_directory_link_without_removing_target() -> io::Result<()> {
+        let tempdir = tempfile::tempdir()?;
+        let target = tempdir.path().join("target");
+        fs_err::create_dir(&target)?;
+        let marker = target.join("marker");
+        fs_err::write(&marker, "")?;
+        let environment = tempdir.path().join("environment");
+        create_symlink(&target, &environment)?;
+
+        remove_virtualenv(&environment)?;
+
+        assert!(matches!(
+            fs_err::symlink_metadata(environment),
+            Err(err) if err.kind() == io::ErrorKind::NotFound
+        ));
+        assert!(marker.is_file());
+        Ok(())
+    }
+
+    #[test]
+    fn clear_virtualenv_recreates_missing_directory() -> io::Result<()> {
+        let tempdir = tempfile::tempdir()?;
+        let environment = tempdir.path().join("environment");
+
+        assert!(!clear_virtualenv(&environment)?);
+        assert!(environment.is_dir());
+        Ok(())
+    }
 }

crates/uv-virtualenv/src/virtualenv.rs

@@ -150,14 +150,7 @@ pub(crate) fn create(
                         }
                     }
                     debug!("Removing existing {name} ({reason})");
-                    // Before removing the virtual environment, we need to canonicalize the path
-                    // because `Path::metadata` will follow the symlink but we're still operating on
-                    // the unresolved path and will remove the symlink itself.
-                    let location = location
-                        .canonicalize()
-                        .unwrap_or_else(|_| location.to_path_buf());
-                    uv_fs::remove_virtualenv(&location)?;
-                    fs_err::create_dir_all(&location)?;
+                    uv_fs::clear_virtualenv(location)?;
                 }
                 OnExisting::Fail => return err,
                 // If not a virtual environment, fail without prompting.
@@ -166,14 +159,7 @@ pub(crate) fn create(
                     match confirm_clear(location, name)? {
                         Some(true) => {
                             debug!("Removing existing {name} due to confirmation");
-                            // Before removing the virtual environment, we need to canonicalize the
-                            // path because `Path::metadata` will follow the symlink but we're still
-                            // operating on the unresolved path and will remove the symlink itself.
-                            let location = location
-                                .canonicalize()
-                                .unwrap_or_else(|_| location.to_path_buf());
-                            uv_fs::remove_virtualenv(&location)?;
-                            fs_err::create_dir_all(&location)?;
+                            uv_fs::clear_virtualenv(location)?;
                         }
                         Some(false) => return err,
                         // When we don't have a TTY, require `--clear` explicitly.