Merge remote-tracking branch 'origin/main' into roc-dev-inline-expects

crates/linker/src/pe.rs

@@ -8,8 +8,8 @@ use bincode::{deserialize_from, serialize_into};
 use memmap2::MmapMut;
 use object::{
     pe::{
-        self, ImageFileHeader, ImageImportDescriptor, ImageNtHeaders64, ImageSectionHeader,
-        ImageThunkData64,
+        self, ImageBaseRelocation, ImageFileHeader, ImageImportDescriptor, ImageNtHeaders64,
+        ImageSectionHeader, ImageThunkData64,
     },
     read::pe::ImportTable,
     LittleEndian as LE, Object, RelocationTarget, SectionIndex,
@@ -20,7 +20,8 @@ use roc_collections::{MutMap, VecMap};
 use roc_error_macros::internal_error;
 
 use crate::{
-    generate_dylib::APP_DLL, load_struct_inplace, load_struct_inplace_mut, open_mmap, open_mmap_mut,
+    generate_dylib::APP_DLL, load_struct_inplace, load_struct_inplace_mut,
+    load_structs_inplace_mut, open_mmap, open_mmap_mut,
 };
 
 /// The metadata stores information about/from the host .exe because
@@ -48,7 +49,18 @@ struct PeMetadata {
     dynamic_relocations: DynamicRelocationsPe,
 
     /// File offset for the thunks of our dummy .dll
-    thunks_start_offset: usize,
+    thunks_start_offset_in_file: usize,
+
+    /// Offset for the thunks of our dummy .dll within the .rdata section
+    thunks_start_offset_in_section: usize,
+
+    /// Virtual address of the .rdata section
+    rdata_virtual_address: u32,
+
+    /// The offset into the file of the .reloc section
+    reloc_offset_in_file: usize,
+
+    reloc_section_index: usize,
 
     /// Constants from the host .exe header
     image_base: u64,
@@ -96,7 +108,26 @@ impl PeMetadata {
             .unwrap();
 
         let dynamic_relocations = DynamicRelocationsPe::new(preprocessed_data);
-        let thunks_start_offset = find_thunks_start_offset(preprocessed_data, &dynamic_relocations);
+        let thunks_start_offset_in_file =
+            find_thunks_start_offset(preprocessed_data, &dynamic_relocations);
+
+        let rdata_section = dynhost_obj
+            .sections()
+            .find(|s| s.name() == Ok(".rdata"))
+            .unwrap();
+
+        let thunks_start_offset_in_section =
+            thunks_start_offset_in_file - rdata_section.file_range().unwrap().0 as usize;
+
+        let rdata_virtual_address = rdata_section.address() as u32;
+
+        let (reloc_section_index, reloc_section) = dynhost_obj
+            .sections()
+            .enumerate()
+            .find(|(_, s)| s.name() == Ok(".reloc"))
+            .unwrap();
+
+        let reloc_offset_in_file = reloc_section.file_range().unwrap().0 as usize;
 
         let optional_header = dynhost_obj.nt_headers().optional_header;
         let optional_header_offset = dynhost_obj.dos_header().nt_headers_offset() as usize
@@ -142,7 +173,11 @@ impl PeMetadata {
             imports,
             exports,
             dynamic_relocations,
-            thunks_start_offset,
+            thunks_start_offset_in_file,
+            thunks_start_offset_in_section,
+            rdata_virtual_address,
+            reloc_offset_in_file,
+            reloc_section_index,
         }
     }
 }
@@ -221,6 +256,8 @@ pub(crate) fn surgery_pe(executable_path: &Path, metadata_path: &Path, roc_app_b
     let mut section_header_start = md.dynamic_relocations.section_headers_offset_in_file as usize
         + md.host_section_count * std::mem::size_of::<ImageSectionHeader>();
 
+    relocate_dummy_dll_entries(executable, &md);
+
     let mut code_bytes_added = 0;
     let mut data_bytes_added = 0;
     let mut file_bytes_added = 0;
@@ -249,7 +286,8 @@ pub(crate) fn surgery_pe(executable_path: &Path, metadata_path: &Path, roc_app_b
             }
         }
 
-        let virtual_size = length as u32;
+        // NOTE: sections cannot be zero in size!
+        let virtual_size = u32::max(1, length as u32);
         let size_of_raw_data = next_multiple_of(length, file_alignment) as u32;
 
         match kind {
@@ -289,7 +327,12 @@ pub(crate) fn surgery_pe(executable_path: &Path, metadata_path: &Path, roc_app_b
             let slice = &roc_app_bytes[section.file_range.start..section.file_range.end];
             executable[offset..][..slice.len()].copy_from_slice(slice);
 
-            for (name, app_relocation) in section.relocations.iter() {
+            let it = section
+                .relocations
+                .iter()
+                .flat_map(|(name, rs)| rs.iter().map(move |r| (name, r)));
+
+            for (name, app_relocation) in it {
                 let AppRelocation {
                     offset_in_section,
                     relocation,
@@ -359,7 +402,12 @@ pub(crate) fn surgery_pe(executable_path: &Path, metadata_path: &Path, roc_app_b
         .map(|s| (s.name, s.offset_in_section as u64))
         .collect();
 
-    redirect_dummy_dll_functions(executable, &symbols, &md.imports, md.thunks_start_offset);
+    redirect_dummy_dll_functions(
+        executable,
+        &symbols,
+        &md.imports,
+        md.thunks_start_offset_in_file,
+    );
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -602,9 +650,14 @@ impl Preprocessor {
         result[self.new_headers_size..].copy_from_slice(&data[self.old_headers_size..]);
     }
 
-    fn write_dummy_sections(&self, result: &mut MmapMut, extra_sections: &[[u8; 8]]) {
+    fn write_dummy_sections(&self, result: &mut MmapMut, extra_section_names: &[[u8; 8]]) {
         const W: usize = std::mem::size_of::<ImageSectionHeader>();
 
+        // only correct for the first section, but that is OK because it's overwritten later
+        // anyway. But, this value may be used to check whether a previous section overruns into
+        // the app sections.
+        let pointer_to_raw_data = result.len() - self.additional_length;
+
         let previous_section_header =
             load_struct_inplace::<ImageSectionHeader>(result, self.extra_sections_start - W);
 
@@ -614,8 +667,14 @@ impl Preprocessor {
         let mut next_virtual_address =
             next_multiple_of(previous_section_header_end as usize, self.section_alignment);
 
-        for (i, name) in extra_sections.iter().enumerate() {
-            let header = ImageSectionHeader {
+        let extra_section_headers = load_structs_inplace_mut::<ImageSectionHeader>(
+            result,
+            self.extra_sections_start,
+            extra_section_names.len(),
+        );
+
+        for (header, name) in extra_section_headers.iter_mut().zip(extra_section_names) {
+            *header = ImageSectionHeader {
                 name: *name,
                 // NOTE: the virtual_size CANNOT BE ZERO! the binary is invalid if a section has
                 // zero virtual size. Setting it to 1 works, (because this one byte is not backed
@@ -624,7 +683,7 @@ impl Preprocessor {
                 // NOTE: this must be a valid virtual address, using 0 is invalid!
                 virtual_address: object::U32::new(LE, next_virtual_address as u32),
                 size_of_raw_data: Default::default(),
-                pointer_to_raw_data: Default::default(),
+                pointer_to_raw_data: object::U32::new(LE, pointer_to_raw_data as u32),
                 pointer_to_relocations: Default::default(),
                 pointer_to_linenumbers: Default::default(),
                 number_of_relocations: Default::default(),
@@ -632,10 +691,6 @@ impl Preprocessor {
                 characteristics: Default::default(),
             };
 
-            let header_array: [u8; W] = unsafe { std::mem::transmute(header) };
-
-            result[self.extra_sections_start + i * W..][..W].copy_from_slice(&header_array);
-
             next_virtual_address += self.section_alignment;
         }
     }
@@ -828,7 +883,7 @@ struct Section {
     /// File range of the section (in the app object)
     file_range: Range<usize>,
     kind: SectionKind,
-    relocations: MutMap<String, AppRelocation>,
+    relocations: MutMap<String, Vec<AppRelocation>>,
     app_section_index: SectionIndex,
 }
 
@@ -904,7 +959,7 @@ impl AppSections {
                 _ => continue,
             };
 
-            let mut relocations = MutMap::default();
+            let mut relocations: MutMap<String, Vec<AppRelocation>> = MutMap::default();
 
             for (offset_in_section, relocation) in section.relocations() {
                 match relocation.target() {
@@ -916,14 +971,14 @@ impl AppSections {
                         let address = symbol.as_ref().map(|s| s.address()).unwrap_or_default();
                         let name = symbol.and_then(|s| s.name()).unwrap_or_default();
 
-                        relocations.insert(
-                            name.to_string(),
-                            AppRelocation {
+                        relocations
+                            .entry(name.to_string())
+                            .or_default()
+                            .push(AppRelocation {
                                 offset_in_section,
                                 address,
                                 relocation,
-                            },
-                        );
+                            });
                     }
                     _ => todo!(),
                 }
@@ -1081,6 +1136,175 @@ fn write_section_header(
     data[section_header_start..][..header_array.len()].copy_from_slice(&header_array);
 }
 
+fn write_image_base_relocation(
+    mmap: &mut [u8],
+    reloc_section_start: usize,
+    new_block_va: u32,
+    relocations: &[u16],
+) -> usize {
+    // first, collect the blocks of relocations (each relocation block covers a 4K page)
+    // we will be moving stuff around, and have to work from the back to the front. However, the
+    // relocations are encoded in such a way that we can only decode them from front to back.
+    let mut next_block_start = reloc_section_start as u32;
+    let mut blocks = vec![];
+    let mut block_has_relocation = false;
+
+    loop {
+        let header =
+            load_struct_inplace_mut::<ImageBaseRelocation>(mmap, next_block_start as usize);
+
+        if header.virtual_address.get(LE) == 0 {
+            break;
+        }
+
+        if new_block_va == header.virtual_address.get(LE) {
+            block_has_relocation = true;
+        }
+
+        blocks.push((next_block_start, *header));
+
+        next_block_start += header.size_of_block.get(LE);
+    }
+
+    // this is 8 in practice
+    const HEADER_WIDTH: usize = std::mem::size_of::<ImageBaseRelocation>();
+
+    const ENTRY_WIDTH: usize = std::mem::size_of::<u16>();
+
+    // extra space that we'll use for the new relocations
+    let shift_amount = relocations.len() * ENTRY_WIDTH;
+
+    if block_has_relocation {
+        // now, starting from the back, shift sections that need to be shifted and add the
+        // new relocations to the right block
+        while let Some((block_start, header)) = blocks.pop() {
+            let header_va = header.virtual_address.get(LE);
+            let header_size = header.size_of_block.get(LE);
+
+            let block_start = block_start as usize;
+
+            match header_va.cmp(&new_block_va) {
+                std::cmp::Ordering::Greater => {
+                    // shift this block
+                    mmap.copy_within(
+                        block_start..block_start + header_size as usize,
+                        block_start + shift_amount,
+                    );
+                }
+                std::cmp::Ordering::Equal => {
+                    // extend this block
+                    let header = load_struct_inplace_mut::<ImageBaseRelocation>(mmap, block_start);
+
+                    let new_size = header.size_of_block.get(LE) + shift_amount as u32;
+                    header.size_of_block.set(LE, new_size);
+
+                    let number_of_entries = (new_size as usize - HEADER_WIDTH) / ENTRY_WIDTH;
+                    let entries = load_structs_inplace_mut::<u16>(
+                        mmap,
+                        block_start + HEADER_WIDTH,
+                        number_of_entries,
+                    );
+
+                    entries[number_of_entries - relocations.len()..].copy_from_slice(relocations);
+
+                    // sort by VA. Upper 4 bits store the relocation type
+                    entries.sort_unstable_by_key(|x| x & 0b0000_1111_1111_1111);
+                }
+                std::cmp::Ordering::Less => {
+                    // done
+                    break;
+                }
+            }
+        }
+
+        shift_amount
+    } else {
+        let header =
+            load_struct_inplace_mut::<ImageBaseRelocation>(mmap, next_block_start as usize);
+
+        let size_of_block = HEADER_WIDTH + relocations.len() * ENTRY_WIDTH;
+
+        header.virtual_address.set(LE, new_block_va);
+        header.size_of_block.set(LE, size_of_block as u32);
+
+        let number_of_entries = relocations.len();
+        let entries = load_structs_inplace_mut::<u16>(
+            mmap,
+            next_block_start as usize + HEADER_WIDTH,
+            number_of_entries,
+        );
+
+        entries[number_of_entries - relocations.len()..].copy_from_slice(relocations);
+
+        // sort by VA. Upper 4 bits store the relocation type
+        entries.sort_unstable_by_key(|x| x & 0b0000_1111_1111_1111);
+
+        size_of_block
+    }
+}
+
+/// the roc app functions are called from the host with an indirect call: the code will look
+/// in a table to find the actual address of the app function. This table must be relocated,
+/// because it contains absolute addresses to jump to.
+fn relocate_dummy_dll_entries(executable: &mut [u8], md: &PeMetadata) {
+    let thunks_start_va = (md.rdata_virtual_address - md.image_base as u32)
+        + md.thunks_start_offset_in_section as u32;
+
+    // relocations are defined per 4kb page
+    const BLOCK_SIZE: u32 = 4096;
+
+    let thunks_offset_in_block = thunks_start_va % BLOCK_SIZE;
+    let thunks_relocation_block_va = thunks_start_va - thunks_offset_in_block;
+
+    let relocations: Vec<_> = (0..md.dynamic_relocations.name_by_virtual_address.len())
+        .map(|i| (thunks_offset_in_block as usize + 2 * i) as u16)
+        .collect();
+
+    let added_reloc_bytes = write_image_base_relocation(
+        executable,
+        md.reloc_offset_in_file,
+        thunks_relocation_block_va,
+        &relocations,
+    );
+
+    // the reloc section got bigger, and we need to update the header with the new size
+    let reloc_section_header_start = md.dynamic_relocations.section_headers_offset_in_file as usize
+        + md.reloc_section_index * std::mem::size_of::<ImageSectionHeader>();
+
+    let next_section = load_struct_inplace::<ImageSectionHeader>(
+        executable,
+        reloc_section_header_start + std::mem::size_of::<ImageSectionHeader>(),
+    );
+
+    let next_section_pointer_to_raw_data = next_section.pointer_to_raw_data.get(LE);
+
+    let reloc_section =
+        load_struct_inplace_mut::<ImageSectionHeader>(executable, reloc_section_header_start);
+    let old_section_size = reloc_section.virtual_size.get(LE);
+    let new_virtual_size = old_section_size + added_reloc_bytes as u32;
+    reloc_section.virtual_size.set(LE, new_virtual_size);
+
+    assert!(
+        reloc_section.pointer_to_raw_data.get(LE)
+            + reloc_section.virtual_size.get(LE)
+            + (added_reloc_bytes as u32)
+            < next_section_pointer_to_raw_data,
+        "new .reloc section is too big, and runs into the next section!",
+    );
+
+    //    // in the data directories, update the length of the base relocations
+    //    let dir = load_struct_inplace_mut::<pe::ImageDataDirectory>(
+    //        executable,
+    //        md.dynamic_relocations.data_directories_offset_in_file as usize
+    //            + object::pe::IMAGE_DIRECTORY_ENTRY_BASERELOC
+    //                * std::mem::size_of::<pe::ImageDataDirectory>(),
+    //    );
+    //
+    //    let old_dir_size = dir.size.get(LE);
+    //    debug_assert_eq!(old_section_size, old_dir_size);
+    //    dir.size.set(LE, new_virtual_size);
+}
+
 #[cfg(test)]
 mod test {
     const PE_DYNHOST: &[u8] = include_bytes!("../dynhost_benchmarks_windows.exe") as &[_];
@@ -1560,7 +1784,6 @@ mod test {
 
     #[cfg(windows)]
     #[test]
-    #[ignore = "does not work yet"]
     fn basics_windows() {
         assert_eq!("Hello, 234567 32 1 3!\n", windows_test(test_basics))
     }