From a7b21df34f7b1875fdefd646a48d36d8b1fd5542 Mon Sep 17 00:00:00 2001
From: Martin Liska <martin.liska@hey.com>
Date: Wed, 4 Sep 2024 15:51:25 +0200
Subject: [PATCH] Factor out functionality to write_section_raw

---
 wild_lib/src/elf_writer.rs | 65 ++++++++++++++++++--------------------
 1 file changed, 30 insertions(+), 35 deletions(-)

diff --git a/wild_lib/src/elf_writer.rs b/wild_lib/src/elf_writer.rs
index 267d3c50..5a9e1dae 100644
--- a/wild_lib/src/elf_writer.rs
+++ b/wild_lib/src/elf_writer.rs
@@ -52,6 +52,7 @@ use crate::threading::prelude::*;
 use ahash::AHashMap;
 use anyhow::anyhow;
 use anyhow::bail;
+use anyhow::ensure;
 use anyhow::Context;
 use linker_utils::elf::rel_type_to_string;
 use linker_utils::elf::shf;
@@ -1142,14 +1143,15 @@ impl<'data> ObjectLayout<'data> {
                 SectionSlot::Loaded(sec) => {
                     self.write_section(layout, sec, buffers, table_writer)?
                 }
-                SectionSlot::LoadedDebugInfo(part_id, section_id, section_size) => self
-                    .write_debug_section(
+                SectionSlot::LoadedDebugInfo(part_id, section_id, section_size) => {
+                    self.write_section_raw(
                         layout,
                         part_id,
                         *section_id,
                         *section_size as usize,
                         buffers,
-                    )?,
+                    )?;
+                }
                 SectionSlot::EhFrameData(section_index) => {
                     self.write_eh_frame_data(*section_index, layout, table_writer)?;
                 }
@@ -1203,21 +1205,15 @@ impl<'data> ObjectLayout<'data> {
             .output_sections
             .has_data_in_file(sec.output_section_id())
         {
-            let section_buffer = buffers.get_mut(sec.output_part_id());
-            let allocation_size = sec.capacity() as usize;
-            if section_buffer.len() < allocation_size {
-                bail!(
-                    "Insufficient space allocated to section `{}`. Tried to take {} bytes, but only {} remain",
-                    self.object.section_display_name(sec.index),
-                    allocation_size, section_buffer.len()
-                );
-            }
-            let out = slice_take_prefix_mut(section_buffer, allocation_size);
-            // Cut off any padding so that our output buffer is the size of our input buffer.
             let object_section = self.object.section(sec.index)?;
-            let section_size = self.object.section_size(object_section)?;
-            let out = &mut out[..section_size as usize];
-            self.object.copy_section_data(object_section, out)?;
+            let section_size = self.object.section_size(object_section)? as usize;
+            let out = self.write_section_raw(
+                layout,
+                &sec.output_part_id(),
+                sec.index,
+                section_size,
+                buffers,
+            )?;
             self.apply_relocations(out, sec, layout, table_writer)
                 .with_context(|| {
                     format!(
@@ -1238,35 +1234,34 @@ impl<'data> ObjectLayout<'data> {
         Ok(())
     }
 
-    fn write_debug_section(
+    fn write_section_raw(
         &self,
         layout: &Layout,
         part_id: &PartId,
         section_index: SectionIndex,
         section_size: usize,
-        buffers: &mut OutputSectionPartMap<&mut [u8]>,
-    ) -> Result {
-        if layout
+        buffers: &'data mut OutputSectionPartMap<&mut [u8]>,
+    ) -> Result<&'data mut [u8]> {
+        ensure!(layout
             .output_sections
-            .has_data_in_file(part_id.output_section_id())
-        {
-            let section_buffer = buffers.get_mut(*part_id);
-            let allocation_size = section_size;
-            if section_buffer.len() < allocation_size {
-                bail!(
+            .has_data_in_file(part_id.output_section_id()));
+
+        let section_buffer = buffers.get_mut(*part_id);
+        let allocation_size = section_size;
+        if section_buffer.len() < allocation_size {
+            bail!(
                     "Insufficient space allocated to section `{}`. Tried to take {} bytes, but only {} remain",
                     self.object.section_display_name(section_index),
                     allocation_size, section_buffer.len()
                 );
-            }
-            let out = slice_take_prefix_mut(section_buffer, allocation_size);
-            // Cut off any padding so that our output buffer is the size of our input buffer.
-            let object_section = self.object.section(section_index)?;
-            let section_size = self.object.section_size(object_section)?;
-            let out = &mut out[..section_size as usize];
-            self.object.copy_section_data(object_section, out)?;
         }
-        Ok(())
+        let out = slice_take_prefix_mut(section_buffer, allocation_size);
+        // Cut off any padding so that our output buffer is the size of our input buffer.
+        let object_section = self.object.section(section_index)?;
+        let section_size = self.object.section_size(object_section)?;
+        let out: &'data mut [u8] = &mut out[..section_size as usize];
+        self.object.copy_section_data(object_section, out)?;
+        Ok(out)
     }
 
     /// Writes debug symbols.