Handle vmlinux and kernel modules, including KASLR

https://blogs.oracle.com/linux/post/whats-inside-a-linux-kernel-core-dump
was invaluable to understand where to read the kernel ASLR offset from.

Test Plan
=========

Checked a couple of symbols on 3 different machines (Ubuntu, RHEL,
Fedora) and the KASLR offset was calculated correctly. Everything worked
on x86 and arm64 too.

src/process.rs

@@ -22,6 +22,7 @@ pub enum ExecutableMappingType {
     Anonymous,
     /// Special mapping to optimise certain system calls.
     Vdso,
+    Kernel,
 }
 
 #[derive(Clone)]
@@ -46,6 +47,7 @@ pub struct ExecutableMapping {
     pub kind: ExecutableMappingType,
     pub start_addr: u64,
     pub end_addr: u64,
+    // kaslr info etc etc
     pub offset: u64,
     pub load_address: u64,
     pub main_exec: bool,
@@ -166,6 +168,10 @@ impl ObjectFileInfo {
         virtual_address: u64,
         mapping: &ExecutableMapping,
     ) -> Option<u64> {
+        if mapping.kind == ExecutableMappingType::Kernel {
+            return Some(virtual_address - mapping.offset);
+        }
+
         let offset = virtual_address - mapping.start_addr + mapping.offset;
 
         for segment in &self.elf_load_segments {

src/profile/aggregated.rs

@@ -6,6 +6,7 @@ use lightswitch_object::ExecutableId;
 use tracing::error;
 
 use crate::bpf::profiler_bindings::native_stack_t;
+use crate::kernel::KERNEL_PID;
 use crate::process::ObjectFileInfo;
 use crate::process::Pid;
 use crate::process::ProcessInfo;
@@ -30,10 +31,6 @@ impl RawAggregatedSample {
         procs: &HashMap<Pid, ProcessInfo>,
         objs: &HashMap<ExecutableId, ObjectFileInfo>,
     ) -> Result<AggregatedSample, anyhow::Error> {
-        let Some(info) = procs.get(&self.pid) else {
-            return Err(anyhow!("process not found"));
-        };
-
         let mut processed_sample = AggregatedSample {
             pid: self.pid,
             tid: self.tid,
@@ -43,6 +40,10 @@ impl RawAggregatedSample {
         };
 
         if let Some(native_stack) = self.ustack {
+            let Some(info) = procs.get(&self.pid) else {
+                return Err(anyhow!("process not found"));
+            };
+
             for (i, virtual_address) in native_stack.addresses.into_iter().enumerate() {
                 if native_stack.len <= i.try_into().unwrap() {
                     break;
@@ -53,11 +54,7 @@ impl RawAggregatedSample {
                 };
 
                 let file_offset = match objs.get(&mapping.executable_id) {
-                    Some(obj) => {
-                        // We need the normalized address to symbolize on-disk object files
-                        // and might need the absolute addresses for JIT
-                        obj.normalized_address(virtual_address, &mapping)
-                    }
+                    Some(obj) => obj.normalized_address(virtual_address, &mapping),
                     None => {
                         error!("executable with id {} not found", mapping.executable_id);
                         None
@@ -72,16 +69,33 @@ impl RawAggregatedSample {
             }
         }
 
-        // The kernel stacks are not normalized yet.
         if let Some(kernel_stack) = self.kstack {
+            let Some(info) = procs.get(&KERNEL_PID) else {
+                return Err(anyhow!("kernel process not found"));
+            };
+
             for (i, virtual_address) in kernel_stack.addresses.into_iter().enumerate() {
                 if kernel_stack.len <= i.try_into().unwrap() {
                     break;
                 }
 
+                let Some(mapping) = info.mappings.for_address(virtual_address) else {
+                    continue;
+                };
+
+                let file_offset = match objs.get(&mapping.executable_id) {
+                    Some(obj) => obj.normalized_address(virtual_address, &mapping),
+                    None => {
+                        error!("executable with id {} not found", mapping.executable_id);
+                        None
+                    }
+                };
+
+                // todo: revisit this as the file offset calculation won't work
+                // for kaslr
                 processed_sample.kstack.push(Frame {
                     virtual_address,
-                    file_offset: None,
+                    file_offset,
                     symbolization_result: None,
                 });
             }

src/profile/convert.rs

@@ -4,14 +4,14 @@ use lightswitch_metadata::types::{MetadataLabelValue, TaskKey};
 
 use lightswitch_proto::profile::pprof::Label;
 use lightswitch_proto::profile::{pprof, LabelStringOrNumber, PprofBuilder};
-
 use std::collections::HashMap;
 use std::fmt::Write;
 use std::path::PathBuf;
 use std::time::Duration;
 use std::time::SystemTime;
 use tracing::{debug, error, span, Level};
 
+use crate::kernel::KERNEL_PID;
 use crate::ksym::Ksym;
 use crate::ksym::KsymIter;
 use crate::process::ObjectFileInfo;
@@ -55,32 +55,62 @@ pub fn to_pprof(
         let mut location_ids = Vec::new();
 
         for kframe in kstack {
-            // TODO: Add real values, read kernel build ID, etc.
-            let mapping_id: u64 = pprof.add_mapping(
-                0x1000000,
-                0xFFFFFFFF,
-                0xFFFFFFFF,
-                0x0,
-                "[kernel]", // This is a special marker.
-                "fake_kernel_build_id",
-            );
-
-            let mut lines = vec![];
-
-            match kframe.symbolization_result {
-                Some(Ok((name, _))) => {
-                    let (line, _) = pprof.add_line(&name);
-                    lines.push(line);
+            let virtual_address = kframe.virtual_address;
+
+            let Some(info) = procs.get(&KERNEL_PID) else {
+                continue;
+            };
+
+            let Some(mapping) = info.mappings.for_address(virtual_address) else {
+                // todo: maybe append an error frame for debugging?
+                continue;
+            };
+
+            match objs.get(&mapping.executable_id) {
+                Some(obj) => {
+                    let normalized_addr = kframe
+                        .file_offset
+                        .or_else(|| obj.normalized_address(virtual_address, &mapping));
+
+                    if normalized_addr.is_none() {
+                        debug!("normalized address is none");
+                        continue;
+                    }
+
+                    let mapping_id = pprof.add_mapping(
+                        mapping.executable_id,
+                        mapping.start_addr,
+                        mapping.end_addr,
+                        0x0,
+                        obj.path.to_str().expect("will always be valid"), // should this be named name?,
+                        &mapping
+                            .build_id
+                            .expect("this should never happen")
+                            .to_string(),
+                    );
+
+                    let mut lines = Vec::new();
+
+                    match kframe.symbolization_result {
+                        Some(Ok((name, _))) => {
+                            let (line, _) = pprof.add_line(&name);
+                            lines.push(line);
+                        }
+                        Some(Err(e)) => {
+                            let (line, _) = pprof.add_line(&e.to_string());
+                            lines.push(line);
+                        }
+                        None => {}
+                    }
+
+                    let location =
+                        pprof.add_location(normalized_addr.unwrap_or(0), mapping_id, lines);
+                    location_ids.push(location);
                 }
-                Some(Err(e)) => {
-                    let (line, _) = pprof.add_line(&e.to_string());
-                    lines.push(line);
+                None => {
+                    error!("executable with id {} not found", mapping.executable_id);
                 }
-                None => {}
             }
-
-            let location = pprof.add_location(kframe.virtual_address, mapping_id, lines);
-            location_ids.push(location);
         }
 
         for uframe in ustack {
@@ -146,7 +176,7 @@ pub fn to_pprof(
                     location_ids.push(location);
                 }
                 None => {
-                    debug!("build id not found");
+                    debug!("executable with id {} not found", mapping.executable_id);
                 }
             }
         }

src/profiler.rs

@@ -5,6 +5,7 @@ use std::collections::hash_map::OccupiedEntry;
 use std::collections::HashMap;
 use std::env::temp_dir;
 use std::fs;
+use std::fs::File;
 use std::io::ErrorKind;
 use std::mem::size_of;
 use std::mem::ManuallyDrop;
@@ -36,6 +37,8 @@ use crate::bpf::tracers_skel::{TracersSkel, TracersSkelBuilder};
 use crate::collector::*;
 use crate::debug_info::DebugInfoBackendNull;
 use crate::debug_info::DebugInfoManager;
+use crate::kernel::get_all_kernel_modules;
+use crate::kernel::KERNEL_PID;
 use crate::perf_events::setup_perf_event;
 use crate::process::{
     ExecutableMapping, ExecutableMappingType, ExecutableMappings, ObjectFileInfo, Pid, ProcessInfo,
@@ -197,7 +200,7 @@ fn fetch_vdso_info(
     cache_dir: &Path,
 ) -> Result<(PathBuf, ObjectFile)> {
     // Read raw memory
-    let file = fs::File::open(format!("/proc/{}/mem", pid))?;
+    let file = File::open(format!("/proc/{}/mem", pid))?;
     let size = end_addr - start_addr;
     let mut buf: Vec<u8> = vec![0; size as usize];
     file.read_exact_at(&mut buf, start_addr + offset)?;
@@ -435,6 +438,76 @@ impl Profiler {
         self.profile_send.send(profile).expect("handle send");
     }
 
+    pub fn add_kernel_modules(&mut self) {
+        let kaslr_offset = match lightswitch_object::kernel::kaslr_offset() {
+            Ok(kaslr_offset) => {
+                debug!("kaslr offset: 0x{:x}", kaslr_offset);
+                kaslr_offset
+            }
+            Err(e) => {
+                error!(
+                    "fetching the kaslr offset failed with {:?}, assuming it is 0",
+                    e
+                );
+                0
+            }
+        };
+
+        match get_all_kernel_modules() {
+            Ok(kernel_code_ranges) => {
+                self.procs.write().insert(
+                    KERNEL_PID,
+                    ProcessInfo {
+                        status: ProcessStatus::Running,
+                        mappings: ExecutableMappings(
+                            kernel_code_ranges
+                                .iter()
+                                .map(|e| {
+                                    debug!(
+                                        "kernel module {} 0x{:x} - 0x{:x}",
+                                        e.name, e.start, e.end
+                                    );
+                                    ExecutableMapping {
+                                        executable_id: e.build_id.id().expect("should never fail"),
+                                        build_id: Some(e.build_id.clone()),
+                                        kind: ExecutableMappingType::Kernel,
+                                        start_addr: e.start,
+                                        end_addr: e.end,
+                                        offset: kaslr_offset,
+                                        load_address: 0,
+                                        main_exec: false,
+                                        soft_delete: false,
+                                    }
+                                })
+                                .collect(),
+                        ),
+                    },
+                );
+
+                for kernel_code_range in kernel_code_ranges {
+                    self.object_files.write().insert(
+                        kernel_code_range
+                            .build_id
+                            .id()
+                            .expect("should never happen"),
+                        ObjectFileInfo {
+                            file: File::open("/").expect("should never fail"), // TODO: placeholder, but this will be removed soon
+                            path: PathBuf::from(kernel_code_range.name),
+                            elf_load_segments: vec![],
+                            is_dyn: false,
+                            references: 1,
+                            native_unwind_info_size: None,
+                            is_vdso: false,
+                        },
+                    );
+                }
+            }
+            Err(e) => {
+                error!("Fetching kernel code ranges failed with: {:?}", e);
+            }
+        }
+    }
+
     pub fn run(mut self, collector: ThreadSafeCollector) -> Duration {
         // In this case, we only want to calculate maximum sampling buffer sizes based on the
         // number of "online" CPUs, not "possible" CPUs, which they sometimes differ.
@@ -446,6 +519,7 @@ impl Profiler {
 
         self.setup_perf_events();
         self.set_bpf_map_info();
+        self.add_kernel_modules();
 
         self.tracers.attach().expect("attach tracers");
 
@@ -1569,7 +1643,7 @@ impl Profiler {
 
                     // We want to open the file as quickly as possible to minimise the chances of races
                     // if the file is deleted.
-                    let file = match fs::File::open(&abs_path) {
+                    let file = match File::open(&abs_path) {
                         Ok(f) => f,
                         Err(e) => {
                             debug!("failed to open file {} due to {:?}", abs_path, e);
@@ -1710,7 +1784,7 @@ impl Profiler {
                             debug!("vDSO object file id failed");
                             continue;
                         };
-                        let Ok(file) = std::fs::File::open(&vdso_path) else {
+                        let Ok(file) = File::open(&vdso_path) else {
                             debug!("vDSO object file open failed");
                             continue;
                         };