From 67c9cc03597e6f00bb6917a493cd9bb405a05b4d Mon Sep 17 00:00:00 2001
From: Alessandro Decina <alessandro.d@gmail.com>
Date: Sat, 20 Mar 2021 23:20:51 +0000
Subject: [PATCH] aya: add util::kernel_symbols()

kernel_symbols() can be used to load /proc/kallsyms in a BTreeMap.
Useful for looking up symbols from stack addresses.
---
 aya/src/util.rs | 43 ++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 42 insertions(+), 1 deletion(-)

diff --git a/aya/src/util.rs b/aya/src/util.rs
index ef692886..c5bc773c 100644
--- a/aya/src/util.rs
+++ b/aya/src/util.rs
@@ -1,5 +1,12 @@
 //! Utility functions.
-use std::{fs, io, str::FromStr};
+use std::{
+    collections::{BTreeMap, HashMap},
+    fs::{self, File},
+    io::{self, BufReader},
+    str::FromStr,
+};
+
+use io::BufRead;
 
 const ONLINE_CPUS: &str = "/sys/devices/system/cpu/online";
 pub(crate) const POSSIBLE_CPUS: &str = "/sys/devices/system/cpu/possible";
@@ -50,6 +57,26 @@ fn parse_cpu_ranges(data: &str) -> Result<Vec<u32>, ()> {
     Ok(cpus)
 }
 
+pub fn kernel_symbols() -> Result<BTreeMap<u64, String>, io::Error> {
+    let mut reader = BufReader::new(File::open("/proc/kallsyms")?);
+    parse_kernel_symbols(&mut reader)
+}
+
+fn parse_kernel_symbols(reader: &mut dyn BufRead) -> Result<BTreeMap<u64, String>, io::Error> {
+    let mut syms = BTreeMap::new();
+
+    for line in reader.lines() {
+        let line = line?;
+        let parts = line.splitn(4, ' ').collect::<Vec<_>>();
+        let addr = u64::from_str_radix(parts[0], 16)
+            .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, line.clone()))?;
+        let name = parts[2].to_owned();
+        syms.insert(addr, name);
+    }
+
+    Ok(syms)
+}
+
 #[cfg(test)]
 mod tests {
     use std::iter::FromIterator;
@@ -68,4 +95,18 @@ mod tests {
         assert!(parse_cpu_ranges("0-1,2-").is_err());
         assert!(parse_cpu_ranges("foo").is_err());
     }
+
+    #[test]
+    fn test_parse_kernel_symbols() {
+        let data = "0000000000002000 A irq_stack_backing_store\n\
+                          0000000000006000 A cpu_tss_rw [foo bar]\n"
+            .as_bytes();
+        let syms = parse_kernel_symbols(&mut BufReader::new(data)).unwrap();
+        assert_eq!(syms.keys().collect::<Vec<_>>(), vec![&0x2000, &0x6000]);
+        assert_eq!(
+            syms.get(&0x2000u64).unwrap().as_str(),
+            "irq_stack_backing_store"
+        );
+        assert_eq!(syms.get(&0x6000u64).unwrap().as_str(), "cpu_tss_rw");
+    }
 }