From 3f691ba5ebae6f89c4984ca076e50b57ad73e04d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marvin=20Ha=CC=88user?=
 <8659494+mhaeuser@users.noreply.github.com>
Date: Sat, 15 Apr 2023 17:03:24 +0200
Subject: [PATCH] ImageTool: Determine PE section characteristics based on
 permissions

---
 BaseTools/ImageTool/ElfScan.c   |  2 --
 BaseTools/ImageTool/ImageTool.h |  7 ----
 BaseTools/ImageTool/PeEmit.c    | 57 ++++++++++++---------------------
 BaseTools/ImageTool/PeScan.c    | 20 ++----------
 4 files changed, 23 insertions(+), 63 deletions(-)

diff --git a/BaseTools/ImageTool/ElfScan.c b/BaseTools/ImageTool/ElfScan.c
index 622a2d8966..43a0333eee 100644
--- a/BaseTools/ImageTool/ElfScan.c
+++ b/BaseTools/ImageTool/ElfScan.c
@@ -631,7 +631,6 @@ CreateIntermediate (
       Segments[SIndex].Read         = true;
       Segments[SIndex].Write        = false;
       Segments[SIndex].Execute      = true;
-      Segments[SIndex].Type         = ToolImageSectionTypeCode;
       ++SIndex;
       continue;
     } else if (IsDataShdr (Shdr)) {
@@ -665,7 +664,6 @@ CreateIntermediate (
       Segments[SIndex].Read         = true;
       Segments[SIndex].Write        = true;
       Segments[SIndex].Execute      = false;
-      Segments[SIndex].Type         = ToolImageSectionTypeInitialisedData;
       ++SIndex;
       continue;
     } else if (IsHiiRsrcShdr (Shdr)) {
diff --git a/BaseTools/ImageTool/ImageTool.h b/BaseTools/ImageTool/ImageTool.h
index dc261294d6..db7e5c6a3a 100644
--- a/BaseTools/ImageTool/ImageTool.h
+++ b/BaseTools/ImageTool/ImageTool.h
@@ -82,12 +82,6 @@ typedef struct {
   bool     IsXip;
 } image_tool_header_info_t;
 
-typedef enum {
-  ToolImageSectionTypeCode,
-  ToolImageSectionTypeInitialisedData,
-  ToolImageSectionTypeUninitialisedData
-} image_tool_type_t;
-
 typedef struct {
   char     *Name;
   uint8_t  *Data;
@@ -97,7 +91,6 @@ typedef struct {
   bool     Read;
   bool     Write;
   bool     Execute;
-  image_tool_type_t Type;
 } image_tool_segment_t;
 
 typedef struct {
diff --git a/BaseTools/ImageTool/PeEmit.c b/BaseTools/ImageTool/PeEmit.c
index f0cfeefadd..15bf8b8110 100644
--- a/BaseTools/ImageTool/PeEmit.c
+++ b/BaseTools/ImageTool/PeEmit.c
@@ -275,20 +275,7 @@ ToolImageEmitPeSectionHeaders (
   assert (SectionHeadersSize <= *BufferSize);
 
   for (Index = 0; Index < Image->SegmentInfo.NumSegments; ++Index) {
-    switch (Image->SegmentInfo.Segments[Index].Type) {
-      case ToolImageSectionTypeCode:
-        Sections[Index].Characteristics = EFI_IMAGE_SCN_CNT_CODE;
-        break;
-      case ToolImageSectionTypeInitialisedData:
-        Sections[Index].Characteristics = EFI_IMAGE_SCN_CNT_INITIALIZED_DATA;
-        break;
-      case ToolImageSectionTypeUninitialisedData:
-        Sections[Index].Characteristics = EFI_IMAGE_SCN_CNT_UNINITIALIZED_DATA;
-        break;
-      default:
-        assert (false);
-        break;
-    }
+    assert (Sections[Index].Characteristics == 0);
 
     if (Image->SegmentInfo.Segments[Index].Read) {
       Sections[Index].Characteristics |= EFI_IMAGE_SCN_MEM_READ;
@@ -299,7 +286,9 @@ ToolImageEmitPeSectionHeaders (
     }
 
     if (Image->SegmentInfo.Segments[Index].Execute) {
-      Sections[Index].Characteristics |= EFI_IMAGE_SCN_MEM_EXECUTE;
+      Sections[Index].Characteristics |= EFI_IMAGE_SCN_MEM_EXECUTE | EFI_IMAGE_SCN_CNT_CODE;
+    } else {
+      Sections[Index].Characteristics |= EFI_IMAGE_SCN_CNT_INITIALIZED_DATA;
     }
 
     Sections[Index].PointerToRawData = SectionOffset;
@@ -547,15 +536,18 @@ ToolImageEmitPeSections (
 
     Context->PeHdr->SizeOfImage += Segment->ImageSize;
 
-    if (FirstCode && (Segment->Type == ToolImageSectionTypeCode)) {
-      Context->PeHdr->BaseOfCode = (UINT32)Segment->ImageAddress;
-      FirstCode = false;
+    if (Segment->Execute) {
+      if (FirstCode) {
+        Context->PeHdr->BaseOfCode = (UINT32)Segment->ImageAddress;
+        FirstCode = false;
+      }
     }
-
 #if defined(EFI_TARGET32)
-    if (FirstData && (Segment->Type == ToolImageSectionTypeInitialisedData)) {
-      Context->PeHdr->BaseOfData = (UINT32)Segment->ImageAddress;
-      FirstData = false;
+    else {
+      if (FirstData) {
+        Context->PeHdr->BaseOfData = (UINT32)Segment->ImageAddress;
+        FirstData = false;
+      }
     }
 #endif
 
@@ -578,23 +570,14 @@ ToolImageEmitPeSections (
     *Buffer      += SectionPadding;
     SectionsSize += SectionPadding;
 
-    switch (Segment->Type) {
-      case ToolImageSectionTypeCode:
-        Context->PeHdr->BaseOfCode = MIN(Context->PeHdr->BaseOfCode, (UINT32)Segment->ImageAddress);
-        Context->PeHdr->SizeOfCode += Segment->ImageSize;
-        break;
-      case ToolImageSectionTypeInitialisedData:
+    if (Segment->Execute) {
+      Context->PeHdr->BaseOfCode = MIN(Context->PeHdr->BaseOfCode, (UINT32)Segment->ImageAddress);
+      Context->PeHdr->SizeOfCode += Segment->ImageSize;
+    } else {
 #if defined(EFI_TARGET32)
-        Context->PeHdr->BaseOfData = MIN(Context->PeHdr->BaseOfData, (UINT32)Segment->ImageAddress);
+      Context->PeHdr->BaseOfData = MIN(Context->PeHdr->BaseOfData, (UINT32)Segment->ImageAddress);
 #endif
-        Context->PeHdr->SizeOfInitializedData += Segment->ImageSize;
-        break;
-      case ToolImageSectionTypeUninitialisedData:
-        Context->PeHdr->SizeOfUninitializedData += Segment->ImageSize;
-        break;
-      default:
-        assert (false);
-        break;
+      Context->PeHdr->SizeOfInitializedData += Segment->ImageSize;
     }
   }
 
diff --git a/BaseTools/ImageTool/PeScan.c b/BaseTools/ImageTool/PeScan.c
index 381c9bc6b8..c8569f17e3 100644
--- a/BaseTools/ImageTool/PeScan.c
+++ b/BaseTools/ImageTool/PeScan.c
@@ -240,23 +240,9 @@ ScanPeGetSegmentInfo (
       ImageSegment->ImageAddress = Section->VirtualAddress;
       ImageSegment->ImageSize    = Size;
 
-      if ((Section->Characteristics & EFI_IMAGE_SCN_MEM_READ) != 0) {
-        ImageSegment->Read = true;
-      }
-
-      if ((Section->Characteristics & EFI_IMAGE_SCN_MEM_WRITE) != 0) {
-        ImageSegment->Write = true;
-      }
-
-      if ((Section->Characteristics & EFI_IMAGE_SCN_MEM_EXECUTE) != 0) {
-        ImageSegment->Execute = true;
-      }
-
-      if (ImageSegment->Execute) {
-        ImageSegment->Type = ToolImageSectionTypeCode;
-      } else {
-        ImageSegment->Type = ToolImageSectionTypeInitialisedData;
-      }
+      ImageSegment->Read    = (Section->Characteristics & EFI_IMAGE_SCN_MEM_READ) != 0;
+      ImageSegment->Write   = (Section->Characteristics & EFI_IMAGE_SCN_MEM_WRITE) != 0;
+      ImageSegment->Execute = (Section->Characteristics & EFI_IMAGE_SCN_MEM_EXECUTE) != 0;
 
       ++SegmentInfo->NumSegments;
       ++ImageSegment;