From 15839d8ba126ffcce40229b1f4b0400ba89113f4 Mon Sep 17 00:00:00 2001
From: Jiaqi Gao <jiaqi.gao@intel.com>
Date: Tue, 11 Oct 2022 22:54:04 +0800
Subject: [PATCH] td-shim: add a feature to support optional payload relocation

Signed-off-by: Jiaqi Gao <jiaqi.gao@intel.com>
---
 td-loader/Cargo.toml |  6 +++++-
 td-loader/src/elf.rs | 22 ++++++++++++----------
 td-loader/src/pe.rs  | 26 ++++++++++++++------------
 td-shim/Cargo.toml   |  1 +
 4 files changed, 32 insertions(+), 23 deletions(-)

diff --git a/td-loader/Cargo.toml b/td-loader/Cargo.toml
index 414bfc1a..c9ebb88c 100644
--- 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"
\ No newline at end of file
+env_logger = "0.9.0"
+
+[features]
+default = []
+disable-relocation = []
diff --git a/td-loader/src/elf.rs b/td-loader/src/elf.rs
index 401987f2..18943f54 100644
--- 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
index a6b30588..7e5fe31d 100644
--- 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
index 937e5634..26b2e140 100644
--- 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",