3 files changed, 44 insertions, 34 deletions

diff --git a/Cargo.lock b/Cargo.lock
index b61fd2e98..39f43ba17 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -80,7 +80,7 @@ dependencies = [
  "cfg-if",
  "libc",
  "miniz_oxide",
- "object 0.23.0",
+ "object",
  "rustc-demangle",
 ]
 
@@ -892,6 +892,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0ee1c47aaa256ecabcaea351eae4a9b01ef39ed810004e298d2511ed284b1525"
 
 [[package]]
+name = "memmap"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6585fd95e7bb50d6cc31e20d4cf9afb4e2ba16c5846fc76793f11218da9c475b"
+dependencies = [
+ "libc",
+ "winapi",
+]
+
+[[package]]
 name = "memmap2"
 version = "0.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1010,16 +1020,6 @@ dependencies = [
 
 [[package]]
 name = "object"
-version = "0.22.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8d3b63360ec3cb337817c2dbd47ab4a0f170d285d8e5a2064600f3def1402397"
-dependencies = [
- "flate2",
- "wasmparser",
-]
-
-[[package]]
-name = "object"
 version = "0.23.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a9a7ab5d64814df0fe4a4b5ead45ed6c5f181ee3ff04ba344313a6c80446c5d4"
@@ -1164,7 +1164,8 @@ dependencies = [
  "crossbeam-channel",
  "jod-thread",
  "log",
- "object 0.22.0",
+ "memmap",
+ "object",
  "serde",
  "serde_json",
  "snap",
@@ -1180,7 +1181,7 @@ dependencies = [
  "libloading",
  "mbe",
  "memmap2",
- "object 0.23.0",
+ "object",
  "proc_macro_api",
  "proc_macro_test",
  "serde_derive",
@@ -1899,12 +1900,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "wasmparser"
-version = "0.57.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "32fddd575d477c6e9702484139cf9f23dcd554b06d185ed0f56c857dd3a47aa6"
-
-[[package]]
 name = "winapi"
 version = "0.3.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/crates/proc_macro_api/Cargo.toml b/crates/proc_macro_api/Cargo.toml
index 5b0a05dff..16fd56c7e 100644
--- a/crates/proc_macro_api/Cargo.toml
+++ b/crates/proc_macro_api/Cargo.toml
@@ -21,3 +21,4 @@ base_db = { path = "../base_db", version = "0.0.0" }
 stdx = { path = "../stdx", version = "0.0.0" }
 snap = "1"
 object = { version = "0.23.0", default-features = false, features = ["std", "read_core", "elf", "macho", "pe", "unaligned"] }
+memmap = "0.7.0"
diff --git a/crates/proc_macro_api/src/lib.rs b/crates/proc_macro_api/src/lib.rs
index cec854746..a58e39de0 100644
--- a/crates/proc_macro_api/src/lib.rs
+++ b/crates/proc_macro_api/src/lib.rs
@@ -9,15 +9,22 @@ pub mod msg;
 mod process;
 mod rpc;
 
-use std::{ffi::OsStr, fs::read as fsread, io::{self, Read}, path::{Path, PathBuf}, sync::Arc};
-
 use base_db::{Env, ProcMacro};
+use std::{
+    ffi::OsStr,
+    fs::File,
+    io::{self, Read},
+    path::{Path, PathBuf},
+    sync::Arc,
+};
+
 use tt::{SmolStr, Subtree};
 
 use crate::process::{ProcMacroProcessSrv, ProcMacroProcessThread};
 
 pub use rpc::{ExpansionResult, ExpansionTask, ListMacrosResult, ListMacrosTask, ProcMacroKind};
 
+use memmap::Mmap;
 use object::read::{File as BinaryFile, Object, ObjectSection};
 use snap::read::FrameDecoder as SnapDecoder;
 
@@ -110,13 +117,13 @@ impl ProcMacroClient {
 
     // This is used inside self.read_version() to locate the ".rustc" section
     // from a proc macro crate's binary file.
-    fn read_section<'a>(&self, dylib_binary: &'a [u8], section_name: &str) -> &'a [u8] {
+    fn read_section<'a>(&self, dylib_binary: &'a [u8], section_name: &str) -> io::Result<&'a [u8]> {
         BinaryFile::parse(dylib_binary)
-            .unwrap()
+            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?
             .section_by_name(section_name)
-            .unwrap()
+            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "section read error"))?
             .data()
-            .unwrap()
+            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
     }
 
     // Check the version of rustc that was used to compile a proc macro crate's
@@ -138,10 +145,19 @@ impl ProcMacroClient {
     // * [some more bytes that we don really care but still there] :-)
     // Check this issue for more about the bytes layout:
     // https://github.com/rust-analyzer/rust-analyzer/issues/6174
-    fn read_version(&self, dylib_path: &Path) -> String {
-        let dylib_binary = fsread(dylib_path).unwrap();
+    #[allow(unused)]
+    fn read_version(&self, dylib_path: &Path) -> io::Result<String> {
+        let dylib_file = File::open(dylib_path)?;
+        let dylib_mmaped = unsafe { Mmap::map(&dylib_file) }?;
+
+        let dot_rustc = self.read_section(&dylib_mmaped, ".rustc")?;
 
-        let dot_rustc = self.read_section(&dylib_binary, ".rustc");
+        let header = &dot_rustc[..8];
+        const EXPECTED_HEADER: [u8; 8] = [b'r', b'u', b's', b't', 0, 0, 0, 5];
+        // check if header is valid
+        if !(header == EXPECTED_HEADER) {
+            return Err(io::Error::new(io::ErrorKind::InvalidData, format!(".rustc section should start with header {:?}; header {:?} is actually presented.",EXPECTED_HEADER ,header)));
+        }
 
         let snappy_portion = &dot_rustc[8..];
 
@@ -154,14 +170,12 @@ impl ProcMacroClient {
         // so 13 bytes in total, and we should check the 13th byte
         // to know the length
         let mut bytes_before_version = [0u8; 13];
-        snappy_decoder
-            .read_exact(&mut bytes_before_version)
-            .unwrap();
+        snappy_decoder.read_exact(&mut bytes_before_version)?;
         let length = bytes_before_version[12]; // what? can't use -1 indexing?
 
         let mut version_string_utf8 = vec![0u8; length as usize];
-        snappy_decoder.read_exact(&mut version_string_utf8).unwrap();
-        let version_string = String::from_utf8(version_string_utf8).unwrap();
-        version_string
+        snappy_decoder.read_exact(&mut version_string_utf8)?;
+        let version_string = String::from_utf8(version_string_utf8);
+        version_string.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
     }
 }