td-shim: add a feature to support optional payload relocation

Signed-off-by: Jiaqi Gao <[email protected]>
confidential-containers · Oct 17, 2022 · 15839d8 · 15839d8
1 parent ddf3a07
commit 15839d8
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 23 deletions.
diff --git a/td-loader/Cargo.toml b/td-loader/Cargo.toml
@@ -14,4 +14,8 @@ log = "0.4.13"
 scroll = { version = "0.10", default-features=false, features = ["derive"] }
 
 [dev-dependencies]
-env_logger = "0.9.0"
+env_logger = "0.9.0"
+
+[features]
+default = []
+disable-relocation = []
diff --git a/td-loader/src/elf.rs b/td-loader/src/elf.rs
@@ -67,16 +67,18 @@ pub fn relocate_elf_with_per_program_header(
         }
     }
 
-    // relocate to base
-    for reloc in elf.relocations()? {
-        if reloc.r_type() == R_X86_64_RELATIVE {
-            let r_addend = reloc.r_addend;
-            loaded_buffer
-                .pwrite::<u64>(
-                    new_image_base.checked_add(r_addend as usize)? as u64,
-                    reloc.r_offset as usize,
-                )
-                .ok()?;
+    if !cfg!(feature = "disable-relocation") {
+        // relocate to base
+        for reloc in elf.relocations()? {
+            if reloc.r_type() == R_X86_64_RELATIVE {
+                let r_addend = reloc.r_addend;
+                loaded_buffer
+                    .pwrite::<u64>(
+                        new_image_base.checked_add(r_addend as usize)? as u64,
+                        reloc.r_offset as usize,
+                    )
+                    .ok()?;
+            }
         }
     }
 

diff --git a/td-loader/src/pe.rs b/td-loader/src/pe.rs
@@ -172,8 +172,6 @@ pub fn relocate_with_per_section(
     let coff_optional_offset = coff_standard_end;
     let coff_optional_end = coff_header_end.checked_add(coff_optional_size)?;
     image_buffer.len().checked_sub(coff_optional_end)?;
-    let coff_optional_region = &image_buffer[coff_optional_offset..coff_optional_end];
-    let image_base = coff_optional_region.pread::<u64>(0).ok()?;
 
     // Validate section header region
     // There's no "Data Directories", so "Section Table" follows COFF Optional Fields.
@@ -210,16 +208,20 @@ pub fn relocate_with_per_section(
         }
     }
 
-    let sections = Sections::parse(sections_buffer, num_sections as usize)?;
-    for section in sections {
-        if &section.name == b".reloc\0\0" && image_base != new_image_base as u64 {
-            reloc_to_base(
-                loaded_buffer,
-                image_buffer,
-                &section,
-                image_base as usize,
-                new_image_base as usize,
-            )?;
+    if !cfg!(feature = "disable-relocation") {
+        let coff_optional_region = &image_buffer[coff_optional_offset..coff_optional_end];
+        let image_base = coff_optional_region.pread::<u64>(0).ok()?;
+        let sections = Sections::parse(sections_buffer, num_sections as usize)?;
+        for section in sections {
+            if &section.name == b".reloc\0\0" && image_base != new_image_base as u64 {
+                reloc_to_base(
+                    loaded_buffer,
+                    image_buffer,
+                    &section,
+                    image_base as usize,
+                    new_image_base as usize,
+                )?;
+            }
         }
     }
 

diff --git a/td-shim/Cargo.toml b/td-shim/Cargo.toml
@@ -51,6 +51,7 @@ boot-kernel = ["td-layout/boot-kernel"]
 secure-boot = ["der", "ring"]
 tdx = ["tdx-tdcall", "td-exception/tdx", "td-logger/tdx", "x86"]
 lazy-accept = ["tdx"]
+disable-relocation = ["td-loader/disable-relocation"]
 main = [
     "td-loader",
     "linked_list_allocator",