From f8fcc23e2d5abb41bb275628e7e7c78ca0c17517 Mon Sep 17 00:00:00 2001
From: Dave Tucker <dave@dtucker.co.uk>
Date: Sun, 16 Jan 2022 10:40:58 +0000
Subject: [PATCH] aya: Implement subprog relocations

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>
---
 aya/src/obj/mod.rs        | 11 ++++++++++-
 aya/src/obj/relocation.rs | 30 +++++++++++++++++++++++++++++-
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/aya/src/obj/mod.rs b/aya/src/obj/mod.rs
index bdef7d2f..108fc5a8 100644
--- a/aya/src/obj/mod.rs
+++ b/aya/src/obj/mod.rs
@@ -47,6 +47,7 @@ pub struct Object {
     // symbol_offset_by_name caches symbols that could be referenced from a
     // BTF VAR type so the offsets can be fixed up
     pub(crate) symbol_offset_by_name: HashMap<String, u64>,
+    pub(crate) text_section_index: SectionIndex,
 }
 
 #[derive(Debug, Clone, PartialEq)]
@@ -244,6 +245,8 @@ impl Object {
                     .ok_or(BtfError::InvalidSymbolName)?;
                 let sym = Symbol {
                     index: symbol.index().0,
+                    kind: symbol.kind(),
+                    is_local: symbol.is_local(),
                     name: Some(name.clone()),
                     section_index: symbol.section().index(),
                     address: symbol.address(),
@@ -298,6 +301,7 @@ impl Object {
             symbols_by_index: HashMap::new(),
             section_sizes: HashMap::new(),
             symbol_offset_by_name: HashMap::new(),
+            text_section_index: SectionIndex(0),
         }
     }
 
@@ -513,7 +517,10 @@ impl Object {
                 self.maps
                     .insert(section.name.to_string(), parse_map(&section, section.name)?);
             }
-            BpfSectionKind::Text => self.parse_text_section(section)?,
+            BpfSectionKind::Text => {
+                self.text_section_index = section.index;
+                self.parse_text_section(section)?
+            }
             BpfSectionKind::Btf => self.parse_btf(&section)?,
             BpfSectionKind::BtfExt => self.parse_btf_ext(&section)?,
             BpfSectionKind::Maps => {
@@ -1453,6 +1460,8 @@ mod tests {
             1,
             Symbol {
                 index: 1,
+                kind: SymbolKind::Text,
+                is_local: false,
                 section_index: Some(SectionIndex(1)),
                 name: Some("my_config".to_string()),
                 address: 0,
diff --git a/aya/src/obj/relocation.rs b/aya/src/obj/relocation.rs
index e8886bcc..e7320ab1 100644
--- a/aya/src/obj/relocation.rs
+++ b/aya/src/obj/relocation.rs
@@ -1,6 +1,6 @@
 use std::{collections::HashMap, mem};
 
-use object::SectionIndex;
+use object::{SectionIndex, SymbolKind};
 use thiserror::Error;
 
 use crate::{
@@ -52,6 +52,8 @@ pub(crate) struct Relocation {
 #[derive(Debug, Clone)]
 pub(crate) struct Symbol {
     pub(crate) index: usize,
+    pub(crate) kind: SymbolKind,
+    pub(crate) is_local: bool,
     pub(crate) section_index: Option<SectionIndex>,
     pub(crate) name: Option<String>,
     pub(crate) address: u64,
@@ -82,6 +84,7 @@ impl Object {
                     relocations.values(),
                     &maps_by_section,
                     &self.symbols_by_index,
+                    &self.text_section_index,
                 )
                 .map_err(|error| BpfError::RelocationError {
                     function: function.name.clone(),
@@ -114,6 +117,7 @@ fn relocate_maps<'a, I: Iterator<Item = &'a Relocation>>(
     relocations: I,
     maps_by_section: &HashMap<usize, (&str, &Map)>,
     symbol_table: &HashMap<usize, Symbol>,
+    text_section_index: &SectionIndex,
 ) -> Result<(), RelocationError> {
     let section_offset = fun.section_offset;
     let instructions = &mut fun.instructions;
@@ -154,6 +158,19 @@ fn relocate_maps<'a, I: Iterator<Item = &'a Relocation>>(
             None => continue,
         };
 
+        if is_sub_prog(sym, &section_index, text_section_index) {
+            // this isn't a map relocation, it's a subprog relocation
+            // we'll apply this here anyway since iterations are expensive
+            let symbol_offset = sym.address;
+            if symbol_offset % INS_SIZE as u64 != 0 {
+                panic!()
+            }
+            // symbol_offset + insn->imm is the byte offset in the other section
+            instructions[ins_index].imm =
+                (instructions[ins_index].imm + symbol_offset as i32) / INS_SIZE as i32;
+            continue;
+        }
+
         let (name, map) =
             maps_by_section
                 .get(&section_index.0)
@@ -342,3 +359,14 @@ fn is_call(ins: &bpf_insn) -> bool {
         && ins.dst_reg() == 0
         && ins.off == 0
 }
+
+fn is_sub_prog(sym: &Symbol, index: &SectionIndex, text_section_index: &SectionIndex) -> bool {
+    if index != text_section_index {
+        return false;
+    }
+    match sym.kind {
+        SymbolKind::Section if sym.is_local => true,
+        SymbolKind::Text if !sym.is_local => true,
+        _ => false,
+    }
+}