Fix copied proc-macros not being cleaned up on exit

2 files changed, 49 insertions, 59 deletions

diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib.rs
index 49a249f2cb6..828d49e6a21 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib.rs
@@ -3,17 +3,12 @@
 mod version;
 
 use proc_macro::bridge;
-use std::{
-    fmt,
-    fs::{self, File},
-    io,
-    time::SystemTime,
-};
+use std::{fmt, fs, io, time::SystemTime};
 
 use libloading::Library;
 use memmap2::Mmap;
 use object::Object;
-use paths::{AbsPath, Utf8Path, Utf8PathBuf};
+use paths::{Utf8Path, Utf8PathBuf};
 use proc_macro_api::ProcMacroKind;
 
 use crate::ProcMacroSrvSpan;
@@ -28,14 +23,9 @@ fn is_derive_registrar_symbol(symbol: &str) -> bool {
     symbol.contains(NEW_REGISTRAR_SYMBOL)
 }
 
-fn find_registrar_symbol(file: &Utf8Path) -> io::Result<Option<String>> {
-    let file = File::open(file)?;
-    let buffer = unsafe { Mmap::map(&file)? };
-
-    Ok(object::File::parse(&*buffer)
-        .map_err(invalid_data_err)?
-        .exports()
-        .map_err(invalid_data_err)?
+fn find_registrar_symbol(buffer: &[u8]) -> object::Result<Option<String>> {
+    Ok(object::File::parse(buffer)?
+        .exports()?
         .into_iter()
         .map(|export| export.name())
         .filter_map(|sym| String::from_utf8(sym.into()).ok())
@@ -118,17 +108,17 @@ struct ProcMacroLibraryLibloading {
 }
 
 impl ProcMacroLibraryLibloading {
-    fn open(file: &Utf8Path) -> Result<Self, LoadProcMacroDylibError> {
-        let symbol_name = find_registrar_symbol(file)?.ok_or_else(|| {
-            invalid_data_err(format!("Cannot find registrar symbol in file {file}"))
-        })?;
+    fn open(path: &Utf8Path) -> Result<Self, LoadProcMacroDylibError> {
+        let buffer = unsafe { Mmap::map(&fs::File::open(path)?)? };
+        let symbol_name =
+            find_registrar_symbol(&buffer).map_err(invalid_data_err)?.ok_or_else(|| {
+                invalid_data_err(format!("Cannot find registrar symbol in file {path}"))
+            })?;
 
-        let abs_file: &AbsPath = file
-            .try_into()
-            .map_err(|_| invalid_data_err(format!("expected an absolute path, got {file}")))?;
-        let version_info = version::read_dylib_info(abs_file)?;
+        let version_info = version::read_dylib_info(&buffer)?;
+        drop(buffer);
 
-        let lib = load_library(file).map_err(invalid_data_err)?;
+        let lib = load_library(path).map_err(invalid_data_err)?;
         let proc_macros = crate::proc_macros::ProcMacros::from_lib(
             &lib,
             symbol_name,
@@ -138,20 +128,22 @@ impl ProcMacroLibraryLibloading {
     }
 }
 
-pub(crate) struct Expander {
-    inner: ProcMacroLibraryLibloading,
-    path: Utf8PathBuf,
-    modified_time: SystemTime,
-}
-
-impl Drop for Expander {
+struct RemoveFileOnDrop(Utf8PathBuf);
+impl Drop for RemoveFileOnDrop {
     fn drop(&mut self) {
         #[cfg(windows)]
-        std::fs::remove_file(&self.path).ok();
-        _ = self.path;
+        std::fs::remove_file(&self.0).unwrap();
+        _ = self.0;
     }
 }
 
+// Drop order matters as we can't remove the dylib before the library is unloaded
+pub(crate) struct Expander {
+    inner: ProcMacroLibraryLibloading,
+    _remove_on_drop: RemoveFileOnDrop,
+    modified_time: SystemTime,
+}
+
 impl Expander {
     pub(crate) fn new(lib: &Utf8Path) -> Result<Expander, LoadProcMacroDylibError> {
         // Some libraries for dynamic loading require canonicalized path even when it is
@@ -160,10 +152,9 @@ impl Expander {
         let modified_time = fs::metadata(&lib).and_then(|it| it.modified())?;
 
         let path = ensure_file_with_lock_free_access(&lib)?;
-
         let library = ProcMacroLibraryLibloading::open(path.as_ref())?;
 
-        Ok(Expander { inner: library, path, modified_time })
+        Ok(Expander { inner: library, _remove_on_drop: RemoveFileOnDrop(path), modified_time })
     }
 
     pub(crate) fn expand<S: ProcMacroSrvSpan>(
@@ -205,20 +196,23 @@ fn ensure_file_with_lock_free_access(path: &Utf8Path) -> io::Result<Utf8PathBuf>
     }
 
     let mut to = Utf8PathBuf::from_path_buf(std::env::temp_dir()).unwrap();
+    to.push("rust-analyzer-proc-macros");
+    _ = fs::create_dir(&to);
 
     let file_name = path.file_name().ok_or_else(|| {
         io::Error::new(io::ErrorKind::InvalidInput, format!("File path is invalid: {path}"))
     })?;
 
-    // Generate a unique number by abusing `HashMap`'s hasher.
-    // Maybe this will also "inspire" a libs team member to finally put `rand` in libstd.
-    let t = RandomState::new().build_hasher().finish();
-
-    let mut unique_name = t.to_string();
-    unique_name.push_str(file_name);
-
-    to.push(unique_name);
-    std::fs::copy(path, &to)?;
+    to.push({
+        // Generate a unique number by abusing `HashMap`'s hasher.
+        // Maybe this will also "inspire" a libs team member to finally put `rand` in libstd.
+        let t = RandomState::new().build_hasher().finish();
+        let mut unique_name = t.to_string();
+        unique_name.push_str(file_name);
+        unique_name.push('-');
+        unique_name
+    });
+    fs::copy(path, &to)?;
     Ok(to)
 }
 
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib/version.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib/version.rs
index 7f0e95c50de..c1804e4fef7 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib/version.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib/version.rs
@@ -1,13 +1,8 @@
 //! Reading proc-macro rustc version information from binary data
 
-use std::{
-    fs::File,
-    io::{self, Read},
-};
+use std::io::{self, Read};
 
-use memmap2::Mmap;
 use object::read::{File as BinaryFile, Object, ObjectSection};
-use paths::AbsPath;
 
 #[derive(Debug)]
 #[allow(dead_code)]
@@ -21,14 +16,14 @@ pub struct RustCInfo {
 }
 
 /// Read rustc dylib information
-pub fn read_dylib_info(dylib_path: &AbsPath) -> io::Result<RustCInfo> {
+pub fn read_dylib_info(buffer: &[u8]) -> io::Result<RustCInfo> {
     macro_rules! err {
         ($e:literal) => {
             io::Error::new(io::ErrorKind::InvalidData, $e)
         };
     }
 
-    let ver_str = read_version(dylib_path)?;
+    let ver_str = read_version(buffer)?;
     let mut items = ver_str.split_whitespace();
     let tag = items.next().ok_or_else(|| err!("version format error"))?;
     if tag != "rustc" {
@@ -106,11 +101,8 @@ fn read_section<'a>(dylib_binary: &'a [u8], section_name: &str) -> io::Result<&'
 ///
 /// Check this issue for more about the bytes layout:
 /// <https://github.com/rust-lang/rust-analyzer/issues/6174>
-pub fn read_version(dylib_path: &AbsPath) -> io::Result<String> {
-    let dylib_file = File::open(dylib_path)?;
-    let dylib_mmapped = unsafe { Mmap::map(&dylib_file) }?;
-
-    let dot_rustc = read_section(&dylib_mmapped, ".rustc")?;
+pub fn read_version(buffer: &[u8]) -> io::Result<String> {
+    let dot_rustc = read_section(buffer, ".rustc")?;
 
     // check if magic is valid
     if &dot_rustc[0..4] != b"rust" {
@@ -159,8 +151,12 @@ pub fn read_version(dylib_path: &AbsPath) -> io::Result<String> {
 
 #[test]
 fn test_version_check() {
-    let path = paths::AbsPathBuf::assert(crate::proc_macro_test_dylib_path());
-    let info = read_dylib_info(&path).unwrap();
+    let info = read_dylib_info(&unsafe {
+        memmap2::Mmap::map(&std::fs::File::open(crate::proc_macro_test_dylib_path()).unwrap())
+            .unwrap()
+    })
+    .unwrap();
+
     assert_eq!(
         info.version_string,
         crate::RUSTC_VERSION_STRING,