From 835fe7fa7b6fd7552e0e6ed4814641daf696722d Mon Sep 17 00:00:00 2001
From: Addison Crump <addison.crump@cispa.de>
Date: Tue, 1 Aug 2023 23:45:49 +0200
Subject: [PATCH] segment out resolver to a trait

---
 aya/src/maps/stack_trace.rs | 25 ++++++++++++++++++++-----
 aya/src/util.rs             |  3 ++-
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/aya/src/maps/stack_trace.rs b/aya/src/maps/stack_trace.rs
index eff8df47..38d31ccc 100644
--- a/aya/src/maps/stack_trace.rs
+++ b/aya/src/maps/stack_trace.rs
@@ -159,6 +159,24 @@ impl<'a, T: Borrow<MapData>> IntoIterator for &'a StackTraceMap<T> {
     }
 }
 
+/// A resolver for symbols based on an address obtained from a stacktrace.
+pub trait SymbolResolver {
+    /// Resolve a symbol for a given address, if possible.
+    fn resolve_sym(&self, addr: u64) -> Option<String>;
+}
+
+/// The simplest resolver: a direct map from addresses to strings.
+pub type SimpleSymbolResolver = BTreeMap<u64, String>;
+
+impl SymbolResolver for SimpleSymbolResolver {
+    fn resolve_sym(&self, addr: u64) -> Option<String> {
+        self
+            .range(..=addr)
+            .next_back()
+            .map(|(_, s)| s.clone())
+    }
+}
+
 /// A kernel or user space stack trace.
 ///
 /// See the [`StackTraceMap`] documentation for examples.
@@ -174,12 +192,9 @@ impl StackTrace {
     /// You can use [`util::kernel_symbols()`](crate::util::kernel_symbols) to load kernel symbols. For
     /// user-space traces you need to provide the symbols, for example loading
     /// them from debug info.
-    pub fn resolve(&mut self, symbols: &BTreeMap<u64, String>) -> &StackTrace {
+    pub fn resolve<R: SymbolResolver>(&mut self, symbols: &R) -> &StackTrace {
         for frame in self.frames.iter_mut() {
-            frame.symbol_name = symbols
-                .range(..=frame.ip)
-                .next_back()
-                .map(|(_, s)| s.clone())
+            frame.symbol_name = symbols.resolve_sym(frame.ip)
         }
 
         self
diff --git a/aya/src/util.rs b/aya/src/util.rs
index 2e6e01ec..f55c0768 100644
--- a/aya/src/util.rs
+++ b/aya/src/util.rs
@@ -15,6 +15,7 @@ use crate::{
 };
 
 use libc::{if_nametoindex, sysconf, uname, utsname, _SC_PAGESIZE};
+use crate::maps::stack_trace::SimpleSymbolResolver;
 
 /// Represents a kernel version, in major.minor.release version.
 // Adapted from https://docs.rs/procfs/latest/procfs/sys/kernel/struct.Version.html.
@@ -204,7 +205,7 @@ fn parse_cpu_ranges(data: &str) -> Result<Vec<u32>, ()> {
 /// Loads kernel symbols from `/proc/kallsyms`.
 ///
 /// The symbols can be passed to [`StackTrace::resolve`](crate::maps::stack_trace::StackTrace::resolve).
-pub fn kernel_symbols() -> Result<BTreeMap<u64, String>, io::Error> {
+pub fn kernel_symbols() -> Result<SimpleSymbolResolver, io::Error> {
     let mut reader = BufReader::new(File::open("/proc/kallsyms")?);
     parse_kernel_symbols(&mut reader)
 }