Parsing LLVM stackmaps.

lib/llvm-backend/Cargo.toml

@@ -14,6 +14,7 @@ goblin = "0.0.20"
 libc = "0.2.49"
 nix = "0.14.0"
 capstone = { version = "0.5.0", optional = true }
+byteorder = "1"
 
 [build-dependencies]
 cc = "1.0"

lib/llvm-backend/cpp/object_loader.cpp

@@ -9,6 +9,22 @@ struct MemoryManager : llvm::RuntimeDyld::MemoryManager {
 public:
     MemoryManager(callbacks_t callbacks) : callbacks(callbacks) {}
 
+    uint8_t * get_stack_map_ptr() const {
+        return stack_map_ptr;
+    }
+
+    size_t get_stack_map_size() const {
+        return stack_map_size;
+    }
+
+    uint8_t * get_code_ptr() const {
+        return (uint8_t *) code_start_ptr;
+    }
+
+    size_t get_code_size() const {
+        return code_size;
+    }
+
     virtual ~MemoryManager() override {
         deregisterEHFrames();
         // Deallocate all of the allocated memory.
@@ -24,11 +40,17 @@ public:
     virtual uint8_t* allocateDataSection(uintptr_t size, unsigned alignment, unsigned section_id, llvm::StringRef section_name, bool read_only) override {
         // Allocate from the read-only section or the read-write section, depending on if this allocation
         // should be read-only or not.
+        uint8_t *ret;
         if (read_only) {
-            return allocate_bump(read_section, read_bump_ptr, size, alignment);
+            ret = allocate_bump(read_section, read_bump_ptr, size, alignment);
         } else {
-            return allocate_bump(readwrite_section, readwrite_bump_ptr, size, alignment);
+            ret = allocate_bump(readwrite_section, readwrite_bump_ptr, size, alignment);
         }
+        if(section_name.equals(llvm::StringRef("__llvm_stackmaps")) || section_name.equals(llvm::StringRef(".llvm_stackmaps"))) {
+            stack_map_ptr = ret;
+            stack_map_size = size;
+        }
+        return ret;
     }
 
     virtual void reserveAllocationSpace(
@@ -53,6 +75,8 @@ public:
         assert(code_result == RESULT_OK);
         code_section = Section { code_ptr_out, code_size_out };
         code_bump_ptr = (uintptr_t)code_ptr_out;
+        code_start_ptr = (uintptr_t)code_ptr_out;
+        this->code_size = code_size;
 
         uint8_t *read_ptr_out = nullptr;
         size_t read_size_out = 0;
@@ -127,12 +151,17 @@ private:
     }
 
     Section code_section, read_section, readwrite_section;
+    uintptr_t code_start_ptr;
+    size_t code_size;
     uintptr_t code_bump_ptr, read_bump_ptr, readwrite_bump_ptr;
     uint8_t* eh_frame_ptr;
     size_t eh_frame_size;
     bool eh_frames_registered = false;
 
     callbacks_t callbacks;
+
+    uint8_t *stack_map_ptr = nullptr;
+    size_t stack_map_size = 0;
 };
 
 struct SymbolLookup : llvm::JITSymbolResolver {
@@ -201,4 +230,29 @@ WasmModule::WasmModule(
 void* WasmModule::get_func(llvm::StringRef name) const {
     auto symbol = runtime_dyld->getSymbol(name);
     return (void*)symbol.getAddress();
-}
+}
+
+uint8_t * WasmModule::get_stack_map_ptr() const {
+    llvm::RuntimeDyld::MemoryManager& mm = *memory_manager;
+    MemoryManager *local_mm = dynamic_cast<MemoryManager *>(&mm);
+    return local_mm->get_stack_map_ptr();
+}
+
+size_t WasmModule::get_stack_map_size() const {
+    llvm::RuntimeDyld::MemoryManager& mm = *memory_manager;
+    MemoryManager *local_mm = dynamic_cast<MemoryManager *>(&mm);
+    return local_mm->get_stack_map_size();
+}
+
+
+uint8_t * WasmModule::get_code_ptr() const {
+    llvm::RuntimeDyld::MemoryManager& mm = *memory_manager;
+    MemoryManager *local_mm = dynamic_cast<MemoryManager *>(&mm);
+    return local_mm->get_code_ptr();
+}
+
+size_t WasmModule::get_code_size() const {
+    llvm::RuntimeDyld::MemoryManager& mm = *memory_manager;
+    MemoryManager *local_mm = dynamic_cast<MemoryManager *>(&mm);
+    return local_mm->get_code_size();
+}

lib/llvm-backend/cpp/object_loader.hh

@@ -156,6 +156,10 @@ struct WasmModule
         callbacks_t callbacks);
 
     void *get_func(llvm::StringRef name) const;
+    uint8_t *get_stack_map_ptr() const;
+    size_t get_stack_map_size() const;
+    uint8_t *get_code_ptr() const;
+    size_t get_code_size() const;
 
     bool _init_failed = false;
   private:
@@ -233,4 +237,20 @@ extern "C"
     {
         return module->get_func(llvm::StringRef(name));
     }
-}
+
+    const uint8_t *llvm_backend_get_stack_map_ptr(const WasmModule *module) {
+        return module->get_stack_map_ptr();
+    }
+
+    size_t llvm_backend_get_stack_map_size(const WasmModule *module) {
+        return module->get_stack_map_size();
+    }
+
+    const uint8_t *llvm_backend_get_code_ptr(const WasmModule *module) {
+        return module->get_code_ptr();
+    }
+
+    size_t llvm_backend_get_code_size(const WasmModule *module) {
+        return module->get_code_size();
+    }
+}

lib/llvm-backend/src/backend.rs

@@ -1,3 +1,4 @@
+use super::stackmap::{self, StackmapRegistry};
 use crate::intrinsics::Intrinsics;
 use inkwell::{
     memory_buffer::MemoryBuffer,
@@ -25,6 +26,7 @@ use wasmer_runtime_core::{
     },
     cache::Error as CacheError,
     module::ModuleInfo,
+    state::ModuleStateMap,
     structures::TypedIndex,
     typed_func::{Wasm, WasmTrapInfo},
     types::{LocalFuncIndex, SigIndex},
@@ -76,6 +78,10 @@ extern "C" {
     ) -> LLVMResult;
     fn module_delete(module: *mut LLVMModule);
     fn get_func_symbol(module: *mut LLVMModule, name: *const c_char) -> *const vm::Func;
+    fn llvm_backend_get_stack_map_ptr(module: *const LLVMModule) -> *const u8;
+    fn llvm_backend_get_stack_map_size(module: *const LLVMModule) -> usize;
+    fn llvm_backend_get_code_ptr(module: *const LLVMModule) -> *const u8;
+    fn llvm_backend_get_code_size(module: *const LLVMModule) -> usize;
 
     fn throw_trap(ty: i32);
 
@@ -231,10 +237,15 @@ pub struct LLVMBackend {
     module: *mut LLVMModule,
     #[allow(dead_code)]
     buffer: Arc<Buffer>,
+    msm: Option<ModuleStateMap>,
 }
 
 impl LLVMBackend {
-    pub fn new(module: Module, _intrinsics: Intrinsics) -> (Self, LLVMCache) {
+    pub fn new(
+        module: Module,
+        _intrinsics: Intrinsics,
+        _stackmaps: &StackmapRegistry,
+    ) -> (Self, LLVMCache) {
         Target::initialize_x86(&InitializationConfig {
             asm_parser: true,
             asm_printer: true,
@@ -285,10 +296,24 @@ impl LLVMBackend {
 
         let buffer = Arc::new(Buffer::LlvmMemory(memory_buffer));
 
+        let raw_stackmap = unsafe {
+            ::std::slice::from_raw_parts(
+                llvm_backend_get_stack_map_ptr(module),
+                llvm_backend_get_stack_map_size(module),
+            )
+        };
+        if raw_stackmap.len() > 0 {
+            let map = stackmap::StackMap::parse(raw_stackmap);
+            eprintln!("{:?}", map);
+        } else {
+            eprintln!("WARNING: No stack map");
+        }
+
         (
             Self {
                 module,
                 buffer: Arc::clone(&buffer),
+                msm: None,
             },
             LLVMCache { buffer },
         )
@@ -318,6 +343,7 @@ impl LLVMBackend {
             Self {
                 module,
                 buffer: Arc::clone(&buffer),
+                msm: None,
             },
             LLVMCache { buffer },
         ))
@@ -372,6 +398,19 @@ impl RunnableModule for LLVMBackend {
         Some(unsafe { Wasm::from_raw_parts(trampoline, invoke_trampoline, None) })
     }
 
+    fn get_code(&self) -> Option<&[u8]> {
+        Some(unsafe {
+            ::std::slice::from_raw_parts(
+                llvm_backend_get_code_ptr(self.module),
+                llvm_backend_get_code_size(self.module),
+            )
+        })
+    }
+
+    fn get_module_state_map(&self) -> Option<ModuleStateMap> {
+        self.msm.clone()
+    }
+
     unsafe fn do_early_trap(&self, data: Box<dyn Any>) -> ! {
         throw_any(Box::leak(data))
     }

lib/llvm-backend/src/code.rs

@@ -27,6 +27,7 @@ use wasmparser::{BinaryReaderError, MemoryImmediate, Operator, Type as WpType};
 use crate::backend::LLVMBackend;
 use crate::intrinsics::{CtxType, GlobalCache, Intrinsics, MemoryCache};
 use crate::read_info::{blocktype_to_type, type_to_type};
+use crate::stackmap::{StackmapEntry, StackmapEntryKind, StackmapRegistry};
 use crate::state::{ControlFrame, IfElseState, State};
 use crate::trampolines::generate_trampolines;
 
@@ -342,26 +343,6 @@ pub struct CodegenError {
     pub message: String,
 }
 
-#[derive(Default, Debug, Clone)]
-pub struct StackmapRegistry {
-    entries: Vec<StackmapEntry>,
-}
-
-#[derive(Debug, Clone)]
-pub struct StackmapEntry {
-    kind: StackmapEntryKind,
-    local_function_id: usize,
-    local_count: usize,
-    stack_count: usize,
-}
-
-#[derive(Debug, Clone, Copy)]
-pub enum StackmapEntryKind {
-    Loop,
-    Call,
-    Trappable,
-}
-
 pub struct LLVMModuleCodeGenerator {
     context: Option<Context>,
     builder: Option<Builder>,
@@ -2653,7 +2634,10 @@ impl ModuleCodeGenerator<LLVMFunctionCodeGenerator, LLVMBackend, CodegenError>
 
         // self.module.print_to_stderr();
 
-        let (backend, cache_gen) = LLVMBackend::new(self.module, self.intrinsics.take().unwrap());
+        let stackmaps = self.stackmaps.borrow();
+
+        let (backend, cache_gen) =
+            LLVMBackend::new(self.module, self.intrinsics.take().unwrap(), &*stackmaps);
         Ok((backend, Box::new(cache_gen)))
     }
 

lib/llvm-backend/src/lib.rs

@@ -6,6 +6,7 @@ mod code;
 mod intrinsics;
 mod platform;
 mod read_info;
+mod stackmap;
 mod state;
 mod trampolines;
 

lib/llvm-backend/src/stackmap.rs

@@ -0,0 +1,208 @@
+// https://llvm.org/docs/StackMaps.html#stackmap-section
+
+use byteorder::{LittleEndian, ReadBytesExt};
+use std::io::{self, Cursor};
+
+#[derive(Default, Debug, Clone)]
+pub struct StackmapRegistry {
+    pub entries: Vec<StackmapEntry>,
+}
+
+#[derive(Debug, Clone)]
+pub struct StackmapEntry {
+    pub kind: StackmapEntryKind,
+    pub local_function_id: usize,
+    pub local_count: usize,
+    pub stack_count: usize,
+}
+
+#[derive(Debug, Clone, Copy)]
+pub enum StackmapEntryKind {
+    Loop,
+    Call,
+    Trappable,
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct StackMap {
+    pub version: u8,
+    pub stk_size_records: Vec<StkSizeRecord>,
+    pub constants: Vec<Constant>,
+    pub stk_map_records: Vec<StkMapRecord>,
+}
+
+#[derive(Copy, Clone, Debug, Default)]
+pub struct StkSizeRecord {
+    pub function_address: u64,
+    pub stack_size: u64,
+    pub record_count: u64,
+}
+
+#[derive(Copy, Clone, Debug, Default)]
+pub struct Constant {
+    pub large_constant: u64,
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct StkMapRecord {
+    pub patchpoint_id: u64,
+    pub instruction_offset: u32,
+    pub locations: Vec<Location>,
+    pub live_outs: Vec<LiveOut>,
+}
+
+#[derive(Copy, Clone, Debug)]
+pub struct Location {
+    pub ty: LocationType,
+    pub location_size: u16,
+    pub dwarf_regnum: u16,
+    pub offset_or_small_constant: i32,
+}
+
+#[derive(Copy, Clone, Debug, Default)]
+pub struct LiveOut {
+    pub dwarf_regnum: u16,
+    pub size_in_bytes: u8,
+}
+
+#[derive(Copy, Clone, Debug)]
+pub enum LocationType {
+    Register,
+    Direct,
+    Indirect,
+    Constant,
+    ConstantIndex,
+}
+
+impl StackMap {
+    pub fn parse(raw: &[u8]) -> io::Result<StackMap> {
+        let mut reader = Cursor::new(raw);
+        let mut map = StackMap::default();
+
+        let version = reader.read_u8()?;
+        if version != 3 {
+            return Err(io::Error::new(io::ErrorKind::Other, "version is not 3"));
+        }
+        map.version = version;
+        if reader.read_u8()? != 0 {
+            return Err(io::Error::new(
+                io::ErrorKind::Other,
+                "reserved field is not zero (1)",
+            ));
+        }
+        if reader.read_u16::<LittleEndian>()? != 0 {
+            return Err(io::Error::new(
+                io::ErrorKind::Other,
+                "reserved field is not zero (2)",
+            ));
+        }
+        let num_functions = reader.read_u32::<LittleEndian>()?;
+        let num_constants = reader.read_u32::<LittleEndian>()?;
+        let num_records = reader.read_u32::<LittleEndian>()?;
+        for _ in 0..num_functions {
+            let mut record = StkSizeRecord::default();
+            record.function_address = reader.read_u64::<LittleEndian>()?;
+            record.stack_size = reader.read_u64::<LittleEndian>()?;
+            record.record_count = reader.read_u64::<LittleEndian>()?;
+            map.stk_size_records.push(record);
+        }
+        for _ in 0..num_constants {
+            map.constants.push(Constant {
+                large_constant: reader.read_u64::<LittleEndian>()?,
+            });
+        }
+        for _ in 0..num_records {
+            let mut record = StkMapRecord::default();
+
+            record.patchpoint_id = reader.read_u64::<LittleEndian>()?;
+            record.instruction_offset = reader.read_u32::<LittleEndian>()?;
+            if reader.read_u16::<LittleEndian>()? != 0 {
+                return Err(io::Error::new(
+                    io::ErrorKind::Other,
+                    "reserved field is not zero (3)",
+                ));
+            }
+            let num_locations = reader.read_u16::<LittleEndian>()?;
+            for _ in 0..num_locations {
+                let ty = reader.read_u8()?;
+
+                let mut location = Location {
+                    ty: match ty {
+                        1 => LocationType::Register,
+                        2 => LocationType::Direct,
+                        3 => LocationType::Indirect,
+                        4 => LocationType::Constant,
+                        5 => LocationType::ConstantIndex,
+                        _ => {
+                            return Err(io::Error::new(
+                                io::ErrorKind::Other,
+                                "unknown location type",
+                            ))
+                        }
+                    },
+                    location_size: 0,
+                    dwarf_regnum: 0,
+                    offset_or_small_constant: 0,
+                };
+
+                if reader.read_u8()? != 0 {
+                    return Err(io::Error::new(
+                        io::ErrorKind::Other,
+                        "reserved field is not zero (4)",
+                    ));
+                }
+                location.location_size = reader.read_u16::<LittleEndian>()?;
+                location.dwarf_regnum = reader.read_u16::<LittleEndian>()?;
+                if reader.read_u16::<LittleEndian>()? != 0 {
+                    return Err(io::Error::new(
+                        io::ErrorKind::Other,
+                        "reserved field is not zero (5)",
+                    ));
+                }
+                location.offset_or_small_constant = reader.read_i32::<LittleEndian>()?;
+
+                record.locations.push(location);
+            }
+            if reader.position() % 8 != 0 {
+                if reader.read_u32::<LittleEndian>()? != 0 {
+                    return Err(io::Error::new(
+                        io::ErrorKind::Other,
+                        "reserved field is not zero (6)",
+                    ));
+                }
+            }
+            if reader.read_u16::<LittleEndian>()? != 0 {
+                return Err(io::Error::new(
+                    io::ErrorKind::Other,
+                    "reserved field is not zero (7)",
+                ));
+            }
+            let num_live_outs = reader.read_u16::<LittleEndian>()?;
+            for _ in 0..num_live_outs {
+                let mut liveout = LiveOut::default();
+
+                liveout.dwarf_regnum = reader.read_u16::<LittleEndian>()?;
+                if reader.read_u8()? != 0 {
+                    return Err(io::Error::new(
+                        io::ErrorKind::Other,
+                        "reserved field is not zero (8)",
+                    ));
+                }
+                liveout.size_in_bytes = reader.read_u8()?;
+
+                record.live_outs.push(liveout);
+            }
+            if reader.position() % 8 != 0 {
+                if reader.read_u32::<LittleEndian>()? != 0 {
+                    return Err(io::Error::new(
+                        io::ErrorKind::Other,
+                        "reserved field is not zero (9)",
+                    ));
+                }
+            }
+
+            map.stk_map_records.push(record);
+        }
+        Ok(map)
+    }
+}