Adjust frame IP in SGX relative to image base

Author: mzohreva

src/backtrace/libunwind.rs

@@ -40,7 +40,18 @@ impl Frame {
             Frame::Raw(ctx) => ctx,
             Frame::Cloned { ip, .. } => return ip,
         };
-        unsafe { uw::_Unwind_GetIP(ctx) as *mut c_void }
+        #[allow(unused_mut)]
+        let mut ip = unsafe { uw::_Unwind_GetIP(ctx) as *mut c_void };
+
+        // To reduce TCB size in SGX enclaves, we do not want to implement
+        // symbol resolution functionality. Rather, we can print the offset of
+        // the address here, which could be later mapped to correct function.
+        #[cfg(all(target_env = "sgx", target_vendor = "fortanix"))]
+        {
+            let image_base = super::get_image_base();
+            ip = usize::wrapping_sub(ip as usize, image_base as _) as _;
+        }
+        ip
     }
 
     pub fn sp(&self) -> *mut c_void {

src/backtrace/mod.rs

@@ -125,6 +125,37 @@ impl fmt::Debug for Frame {
     }
 }
 
+#[cfg(all(target_env = "sgx", target_vendor = "fortanix", not(feature = "std")))]
+mod sgx_no_std_image_base {
+    use core::sync::atomic::{AtomicU64, Ordering::SeqCst};
+
+    static IMAGE_BASE: AtomicU64 = AtomicU64::new(0);
+
+    /// Set the image base address. This is only available for Fortanix SGX
+    /// target when the `std` feature is not enabled. This can be used in the
+    /// standard library to set the correct base address.
+    pub fn set_image_base(base_addr: u64) {
+        IMAGE_BASE.store(base_addr, SeqCst);
+    }
+
+    pub(crate) fn get_image_base() -> u64 {
+        IMAGE_BASE.load(SeqCst)
+    }
+}
+
+#[cfg(all(target_env = "sgx", target_vendor = "fortanix", not(feature = "std")))]
+pub use self::sgx_no_std_image_base::set_image_base;
+
+#[cfg(all(target_env = "sgx", target_vendor = "fortanix", not(feature = "std")))]
+#[deny(unused)]
+pub(crate) use self::sgx_no_std_image_base::get_image_base;
+
+#[cfg(all(target_env = "sgx", target_vendor = "fortanix", feature = "std"))]
+#[deny(unused)]
+pub(crate) fn get_image_base() -> u64 {
+    std::os::fortanix_sgx::mem::image_base()
+}
+
 cfg_if::cfg_if! {
     // This needs to come first, to ensure that
     // Miri takes priority over the host platform

src/lib.rs

@@ -134,6 +134,12 @@ cfg_if::cfg_if! {
     }
 }
 
+cfg_if::cfg_if! {
+    if #[cfg(all(target_env = "sgx", target_vendor = "fortanix", not(feature = "std")))] {
+        pub use self::backtrace::set_image_base;
+    }
+}
+
 #[allow(dead_code)]
 struct Bomb {
     enabled: bool,

src/print.rs

@@ -219,7 +219,7 @@ impl BacktraceFrameFmt<'_, '_, '_> {
     #[allow(unused_mut)]
     fn print_raw_generic(
         &mut self,
-        mut frame_ip: *mut c_void,
+        frame_ip: *mut c_void,
         symbol_name: Option<SymbolName<'_>>,
         filename: Option<BytesOrWideString<'_>>,
         lineno: Option<u32>,
@@ -233,15 +233,6 @@ impl BacktraceFrameFmt<'_, '_, '_> {
             }
         }
 
-        // To reduce TCB size in Sgx enclave, we do not want to implement symbol
-        // resolution functionality.  Rather, we can print the offset of the
-        // address here, which could be later mapped to correct function.
-        #[cfg(all(feature = "std", target_env = "sgx", target_vendor = "fortanix"))]
-        {
-            let image_base = std::os::fortanix_sgx::mem::image_base();
-            frame_ip = usize::wrapping_sub(frame_ip as usize, image_base as _) as _;
-        }
-
         // Print the index of the frame as well as the optional instruction
         // pointer of the frame. If we're beyond the first symbol of this frame
         // though we just print appropriate whitespace.

tests/sgx-image-base.rs

@@ -0,0 +1,56 @@
+#![cfg(all(target_env = "sgx", target_vendor = "fortanix"))]
+#![feature(sgx_platform)]
+
+#[cfg(feature = "std")]
+#[test]
+fn sgx_image_base_with_std() {
+    use backtrace::trace;
+
+    let image_base = std::os::fortanix_sgx::mem::image_base();
+
+    let mut frame_ips = Vec::new();
+    trace(|frame| {
+        frame_ips.push(frame.ip());
+        true
+    });
+
+    assert!(frame_ips.len() > 0);
+    for ip in frame_ips {
+        let ip: u64 = ip as _;
+        assert!(ip < image_base);
+    }
+}
+
+#[cfg(not(feature = "std"))]
+#[test]
+fn sgx_image_base_no_std() {
+    use backtrace::trace_unsynchronized;
+
+    fn guess_image_base() -> u64 {
+        let mut top_frame_ip = None;
+        unsafe {
+            trace_unsynchronized(|frame| {
+                top_frame_ip = Some(frame.ip());
+                false
+            });
+        }
+        top_frame_ip.unwrap() as u64 & 0xFFFFFF000000
+    }
+
+    let image_base = guess_image_base();
+    backtrace::set_image_base(image_base);
+
+    let mut frame_ips = Vec::new();
+    unsafe {
+        trace_unsynchronized(|frame| {
+            frame_ips.push(frame.ip());
+            true
+        });
+    }
+
+    assert!(frame_ips.len() > 0);
+    for ip in frame_ips {
+        let ip: u64 = ip as _;
+        assert!(ip < image_base);
+    }
+}