improve filesearch::get_or_default_sysroot r=ozkanonur

Signed-off-by: Onur Özkan <work@onurozkan.dev>

compiler/rustc_session/src/config.rs

@@ -2431,7 +2431,7 @@ pub fn build_session_options(matches: &getopts::Matches) -> Options {
     let sysroot = match &sysroot_opt {
         Some(s) => s,
         None => {
-            tmp_buf = crate::filesearch::get_or_default_sysroot();
+            tmp_buf = crate::filesearch::get_or_default_sysroot().expect("Failed finding sysroot");
             &tmp_buf
         }
     };

compiler/rustc_session/src/filesearch.rs

@@ -1,5 +1,6 @@
 //! A module for searching for libraries
 
+use smallvec::{smallvec, SmallVec};
 use std::env;
 use std::fs;
 use std::iter::FromIterator;
@@ -62,9 +63,99 @@ pub fn make_target_lib_path(sysroot: &Path, target_triple: &str) -> PathBuf {
     PathBuf::from_iter([sysroot, Path::new(&rustlib_path), Path::new("lib")])
 }
 
+#[cfg(unix)]
+fn current_dll_path() -> Result<PathBuf, String> {
+    use std::ffi::{CStr, OsStr};
+    use std::os::unix::prelude::*;
+
+    unsafe {
+        let addr = current_dll_path as usize as *mut _;
+        let mut info = std::mem::zeroed();
+        if libc::dladdr(addr, &mut info) == 0 {
+            return Err("dladdr failed".into());
+        }
+        if info.dli_fname.is_null() {
+            return Err("dladdr returned null pointer".into());
+        }
+        let bytes = CStr::from_ptr(info.dli_fname).to_bytes();
+        let os = OsStr::from_bytes(bytes);
+        Ok(PathBuf::from(os))
+    }
+}
+
+#[cfg(windows)]
+fn current_dll_path() -> Result<PathBuf, String> {
+    use std::ffi::OsString;
+    use std::io;
+    use std::os::windows::prelude::*;
+    use std::ptr;
+
+    use winapi::um::libloaderapi::{
+        GetModuleFileNameW, GetModuleHandleExW, GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS,
+    };
+
+    unsafe {
+        let mut module = ptr::null_mut();
+        let r = GetModuleHandleExW(
+            GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS,
+            current_dll_path as usize as *mut _,
+            &mut module,
+        );
+        if r == 0 {
+            return Err(format!("GetModuleHandleExW failed: {}", io::Error::last_os_error()));
+        }
+        let mut space = Vec::with_capacity(1024);
+        let r = GetModuleFileNameW(module, space.as_mut_ptr(), space.capacity() as u32);
+        if r == 0 {
+            return Err(format!("GetModuleFileNameW failed: {}", io::Error::last_os_error()));
+        }
+        let r = r as usize;
+        if r >= space.capacity() {
+            return Err(format!("our buffer was too small? {}", io::Error::last_os_error()));
+        }
+        space.set_len(r);
+        let os = OsString::from_wide(&space);
+        Ok(PathBuf::from(os))
+    }
+}
+
+pub fn sysroot_candidates() -> SmallVec<[PathBuf; 2]> {
+    let target = crate::config::host_triple();
+    let mut sysroot_candidates: SmallVec<[PathBuf; 2]> =
+        smallvec![get_or_default_sysroot().expect("Failed finding sysroot")];
+    let path = current_dll_path().and_then(|s| Ok(s.canonicalize().map_err(|e| e.to_string())?));
+    if let Ok(dll) = path {
+        // use `parent` twice to chop off the file name and then also the
+        // directory containing the dll which should be either `lib` or `bin`.
+        if let Some(path) = dll.parent().and_then(|p| p.parent()) {
+            // The original `path` pointed at the `rustc_driver` crate's dll.
+            // Now that dll should only be in one of two locations. The first is
+            // in the compiler's libdir, for example `$sysroot/lib/*.dll`. The
+            // other is the target's libdir, for example
+            // `$sysroot/lib/rustlib/$target/lib/*.dll`.
+            //
+            // We don't know which, so let's assume that if our `path` above
+            // ends in `$target` we *could* be in the target libdir, and always
+            // assume that we may be in the main libdir.
+            sysroot_candidates.push(path.to_owned());
+
+            if path.ends_with(target) {
+                sysroot_candidates.extend(
+                    path.parent() // chop off `$target`
+                        .and_then(|p| p.parent()) // chop off `rustlib`
+                        .and_then(|p| p.parent()) // chop off `lib`
+                        .map(|s| s.to_owned()),
+                );
+            }
+        }
+    }
+
+    return sysroot_candidates;
+}
+
 /// This function checks if sysroot is found using env::args().next(), and if it
-/// is not found, uses env::current_exe() to imply sysroot.
-pub fn get_or_default_sysroot() -> PathBuf {
+/// is not found, finds sysroot from current rustc_driver dll.
+pub fn get_or_default_sysroot() -> Result<PathBuf, String> {
     // Follow symlinks.  If the resolved path is relative, make it absolute.
     fn canonicalize(path: PathBuf) -> PathBuf {
         let path = fs::canonicalize(&path).unwrap_or(path);
@@ -74,17 +165,32 @@ pub fn get_or_default_sysroot() -> PathBuf {
         fix_windows_verbatim_for_gcc(&path)
     }
 
-    // Use env::current_exe() to get the path of the executable following
-    // symlinks/canonicalizing components.
-    fn from_current_exe() -> PathBuf {
-        match env::current_exe() {
-            Ok(exe) => {
-                let mut p = canonicalize(exe);
-                p.pop();
-                p.pop();
-                p
-            }
-            Err(e) => panic!("failed to get current_exe: {e}"),
+    fn default_from_rustc_driver_dll() -> Result<PathBuf, String> {
+        let dll = current_dll_path().and_then(|s| Ok(canonicalize(s)))?;
+
+        // `dll` will be in one of the following two:
+        // - compiler's libdir: $sysroot/lib/*.dll
+        // - target's libdir: $sysroot/lib/rustlib/$target/lib/*.dll
+        //
+        // use `parent` twice to chop off the file name and then also the
+        // directory containing the dll
+        let dir = dll.parent().and_then(|p| p.parent()).ok_or(format!(
+            "Could not move 2 levels upper using `parent()` on {}",
+            dll.display()
+        ))?;
+
+        // if `dir` points target's dir, move up to the sysroot
+        if dir.ends_with(crate::config::host_triple()) {
+            dir.parent() // chop off `$target`
+                .and_then(|p| p.parent()) // chop off `rustlib`
+                .and_then(|p| p.parent()) // chop off `lib`
+                .map(|s| s.to_owned())
+                .ok_or(format!(
+                    "Could not move 3 levels upper using `parent()` on {}",
+                    dir.display()
+                ))
+        } else {
+            Ok(dir.to_owned())
         }
     }
 
@@ -118,7 +224,5 @@ pub fn get_or_default_sysroot() -> PathBuf {
         }
     }
 
-    // Check if sysroot is found using env::args().next(), and if is not found,
-    // use env::current_exe() to imply sysroot.
-    from_env_args_next().unwrap_or_else(from_current_exe)
+    Ok(from_env_args_next().unwrap_or(default_from_rustc_driver_dll()?))
 }

compiler/rustc_session/src/session.rs

@@ -1280,7 +1280,7 @@ pub fn build_session(
 
     let sysroot = match &sopts.maybe_sysroot {
         Some(sysroot) => sysroot.clone(),
-        None => filesearch::get_or_default_sysroot(),
+        None => filesearch::get_or_default_sysroot().expect("Failed finding sysroot"),
     };
 
     let target_cfg = config::build_target_config(&sopts, target_override, &sysroot);