From 1a22792ee75642a3998634cc16abda7d813b45bb Mon Sep 17 00:00:00 2001
From: ajwerner <awerner32@gmail.com>
Date: Wed, 28 Dec 2022 13:20:33 -0500
Subject: [PATCH] Resolve symbol address for PIE executables

See https://github.com/foniod/redbpf/pull/308 for a similar change.
---
 aya/src/programs/uprobe.rs | 34 ++++++++++++++++++++++++++++++----
 1 file changed, 30 insertions(+), 4 deletions(-)

diff --git a/aya/src/programs/uprobe.rs b/aya/src/programs/uprobe.rs
index 751dc77a..d201ff5c 100644
--- a/aya/src/programs/uprobe.rs
+++ b/aya/src/programs/uprobe.rs
@@ -1,6 +1,6 @@
 //! User space probes.
 use libc::pid_t;
-use object::{Object, ObjectSymbol};
+use object::{Object, ObjectSection, ObjectSymbol};
 use std::{
     error::Error,
     ffi::CStr,
@@ -319,15 +319,41 @@ enum ResolveSymbolError {
 
     #[error("unknown symbol `{0}`")]
     Unknown(String),
+
+    #[error("symbol `{0}` does not appear in section")]
+    NotInSection(String),
+
+    #[error("symbol `{0}` in section `{1:?}` which has no offset")]
+    SectionFileRangeNone(String, Result<String, object::Error>),
 }
 
 fn resolve_symbol(path: &str, symbol: &str) -> Result<u64, ResolveSymbolError> {
     let data = fs::read(path)?;
     let obj = object::read::File::parse(&*data)?;
 
-    obj.dynamic_symbols()
+    let sym = obj
+        .dynamic_symbols()
         .chain(obj.symbols())
         .find(|sym| sym.name().map(|name| name == symbol).unwrap_or(false))
-        .map(|s| s.address())
-        .ok_or_else(|| ResolveSymbolError::Unknown(symbol.to_string()))
+        .ok_or_else(|| ResolveSymbolError::Unknown(symbol.to_string()))?;
+
+    let needs_addr_translation = matches!(
+        obj.kind(),
+        object::ObjectKind::Dynamic | object::ObjectKind::Executable
+    );
+    if !needs_addr_translation {
+        Ok(sym.address())
+    } else {
+        let index = sym
+            .section_index()
+            .ok_or_else(|| ResolveSymbolError::NotInSection(symbol.to_string()))?;
+        let section = obj.section_by_index(index)?;
+        let (offset, _size) = section.file_range().ok_or_else(|| {
+            ResolveSymbolError::SectionFileRangeNone(
+                symbol.to_string(),
+                section.name().map(str::to_owned),
+            )
+        })?;
+        Ok(sym.address() - section.address() + offset)
+    }
 }