PR feedback

Signed-off-by: Alan Jowett <[email protected]>
vbpf · Jan 8, 2025 · 768692b · 768692b
1 parent 14b2b3f
commit 768692b
Showing 1 changed file with 19 additions and 12 deletions.
diff --git a/src/asm_files.cpp b/src/asm_files.cpp
@@ -162,15 +162,19 @@ get_program_name_and_size(const ELFIO::section& sec, const ELFIO::Elf_Xword star
     return {program_name, size};
 }
 
+void verify_load_instruction(const ebpf_inst& instruction, const std::string& symbol_name, ELFIO::Elf64_Addr offset) {
+    if ((instruction.opcode & INST_CLS_MASK) != INST_CLS_LD) {
+        throw UnmarshalError("Illegal operation on symbol " + symbol_name + " at location " +
+                             std::to_string(offset / sizeof(ebpf_inst)));
+    }
+}
+
 void relocate_map(ebpf_inst& reloc_inst, const std::string& symbol_name,
                   const std::variant<size_t, std::map<std::string, size_t>>& map_record_size_or_map_offsets,
                   const program_info& info, const ELFIO::Elf64_Addr offset, const ELFIO::Elf_Word index,
                   const ELFIO::const_symbol_section_accessor& symbols) {
     // Only permit loading the address of the map.
-    if ((reloc_inst.opcode & INST_CLS_MASK) != INST_CLS_LD) {
-        throw UnmarshalError("Illegal operation on symbol " + symbol_name + " at location " +
-                             std::to_string(offset / sizeof(ebpf_inst)));
-    }
+    verify_load_instruction(reloc_inst, symbol_name, offset);
     reloc_inst.src = INST_LD_MODE_MAP_FD;
 
     // Relocation value is an offset into the "maps" or ".maps" section.
@@ -187,6 +191,8 @@ void relocate_map(ebpf_inst& reloc_inst, const std::string& symbol_name,
         const auto it = map_descriptors_offsets.find(symbol_name);
         if (it != map_descriptors_offsets.end()) {
             reloc_value = it->second;
+        } else {
+            throw UnmarshalError("Map descriptor not found for symbol " + symbol_name);
         }
     }
     if (reloc_value >= info.map_descriptors.size()) {
@@ -201,10 +207,7 @@ void relocate_global_variable(ebpf_inst& reloc_inst, ebpf_inst& next_reloc_inst,
                               const std::variant<size_t, std::map<std::string, size_t>>& map_record_size_or_map_offsets,
                               const ELFIO::Elf64_Addr offset) {
     // Only permit loading the address of the global variable.
-    if ((reloc_inst.opcode & INST_CLS_MASK) != INST_CLS_LD) {
-        throw UnmarshalError("Illegal operation on symbol " + symbol_name + " at location " +
-                             std::to_string(offset / sizeof(ebpf_inst)));
-    }
+    verify_load_instruction(reloc_inst, symbol_name, offset);
 
     // Copy the immediate value to the next instruction.
     next_reloc_inst.imm = reloc_inst.imm;
@@ -215,6 +218,8 @@ void relocate_global_variable(ebpf_inst& reloc_inst, ebpf_inst& next_reloc_inst,
     const auto it = map_descriptors_offsets.find(symbol_name);
     if (it != map_descriptors_offsets.end()) {
         reloc_value = it->second;
+    } else {
+        throw UnmarshalError("Map descriptor not found for symbol " + symbol_name);
     }
 
     if (reloc_value >= info.map_descriptors.size()) {
@@ -470,13 +475,15 @@ vector<raw_program> read_elf(std::istream& input_stream, const std::string& path
                         continue;
                     }
                     offset -= program_offset;
-                    if (offset / sizeof(ebpf_inst) >= prog.prog.size()) {
+
+                    // Relocation is always on a double width instruction, so check that the offset is within the
+                    // program and that offset + 1 is within the program.
+                    if (((offset / sizeof(ebpf_inst)) + 1) >= prog.prog.size()) {
                         throw UnmarshalError("Invalid relocation data");
                     }
+
                     ebpf_inst& reloc_inst = prog.prog[offset / sizeof(ebpf_inst)];
-                    ebpf_inst& next_reloc_inst = offset / sizeof(ebpf_inst) + 1 < prog.prog.size()
-                                                     ? prog.prog[offset / sizeof(ebpf_inst) + 1]
-                                                     : reloc_inst;
+                    ebpf_inst& next_reloc_inst = prog.prog[offset / sizeof(ebpf_inst) + 1];
 
                     auto [symbol_name, symbol_section_index] = get_symbol_name_and_section_index(symbols, index);