[llvm] [llvm-objcopy] Fix unaligned `p_offset` after copy. (PR #79889)

[Dmitriy Chestnykh via llvm-commits]

https://github.com/chestnykh updated https://github.com/llvm/llvm-project/pull/79889

>From c95797b5c1df4391f66846a1dc74cc7b6e319d84 Mon Sep 17 00:00:00 2001
From: Dmitry Chestnykh <dm.chestnykh at gmail.com>
Date: Mon, 29 Jan 2024 22:55:33 +0300
Subject: [PATCH] [llvm-objcopy] Fix unaligned `p_offset` after copy.

- When copying the ELF whose layout was described in the test
the llvm-objcopy breaks the last PT_LOAD segment alignment
of p_offset by assignment the `p_filesz` value of prev segment to it.
This patch fixes this case. #79887
---
 llvm/lib/ObjCopy/ELF/ELFObject.cpp            |   8 +-
 .../llvm-objcopy/ELF/aligned-offset.test      | 124 ++++++++++++++++++
 2 files changed, 128 insertions(+), 4 deletions(-)
 create mode 100644 llvm/test/tools/llvm-objcopy/ELF/aligned-offset.test

diff --git a/llvm/lib/ObjCopy/ELF/ELFObject.cpp b/llvm/lib/ObjCopy/ELF/ELFObject.cpp
index c8b66d6fcb5ebf..f4f1dc3e117b8b 100644
--- a/llvm/lib/ObjCopy/ELF/ELFObject.cpp
+++ b/llvm/lib/ObjCopy/ELF/ELFObject.cpp
@@ -2280,17 +2280,17 @@ static uint64_t layoutSegments(std::vector<Segment *> &Segments,
   // segments. So we can simply layout segments one after the other accounting
   // for alignment.
   for (Segment *Seg : Segments) {
+    uint64_t SegAlign = std::max<uint64_t>(Seg->Align, 1);
     // We assume that segments have been ordered by OriginalOffset and Index
     // such that a parent segment will always come before a child segment in
     // OrderedSegments. This means that the Offset of the ParentSegment should
     // already be set and we can set our offset relative to it.
     if (Seg->ParentSegment != nullptr) {
       Segment *Parent = Seg->ParentSegment;
-      Seg->Offset =
-          Parent->Offset + Seg->OriginalOffset - Parent->OriginalOffset;
+      Seg->Offset = alignTo(Parent->Offset, SegAlign) + Seg->OriginalOffset -
+                    Parent->OriginalOffset;
     } else {
-      Seg->Offset =
-          alignTo(Offset, std::max<uint64_t>(Seg->Align, 1), Seg->VAddr);
+      Seg->Offset = alignTo(Offset, SegAlign, Seg->VAddr);
     }
     Offset = std::max(Offset, Seg->Offset + Seg->FileSize);
   }
diff --git a/llvm/test/tools/llvm-objcopy/ELF/aligned-offset.test b/llvm/test/tools/llvm-objcopy/ELF/aligned-offset.test
new file mode 100644
index 00000000000000..4c66433428f081
--- /dev/null
+++ b/llvm/test/tools/llvm-objcopy/ELF/aligned-offset.test
@@ -0,0 +1,124 @@
+# This test tests that PT_LOAD segment has properly aligned p_offset
+# in the presented layout (The presense of PT_PHDR and PT_INTERP is important).
+
+# RUN: yaml2obj %s -o %t
+# RUN: llvm-objcopy %t %t2
+# RUN: llvm-readobj --program-headers %t2 | FileCheck %s
+
+!ELF
+FileHeader:
+  Class:           ELFCLASS64
+  Data:            ELFDATA2LSB
+  Type:            ET_DYN
+  Machine:         EM_X86_64
+Sections:
+  - Name:            .text
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    AddressAlign:    0x1000
+    Size:            4096
+  - Name:            .text2
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    AddressAlign:    0x1000
+    Size:            0x1938c9
+  - Name:            .text3
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    AddressAlign:    0x1000
+    Size:            4096
+  - Name:            .text4
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    AddressAlign:    0x1000
+    Size:            4096
+  - Name:            .text5
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    AddressAlign:    0x1000
+    Size:            1658880
+ProgramHeaders:
+  - Type:     PT_PHDR
+    Flags:    [ PF_R ]
+    Offset:   0x40
+    VAddr:    0x40
+    PAddr:    0x40
+    FileSize: 0x268
+    MemSize:  0x268
+    Align:    0x8
+  - Type:     PT_INTERP
+    Flags:    [ PF_R ]
+    Offset:   0x195000
+    VAddr:    0x195000
+    PAddr:    0x195000
+    FileSize: 0xf
+    MemSize:  0xf
+    Align:    0x1
+  - Type:     PT_LOAD
+    Flags:    [ PF_X, PF_R ]
+    FirstSec: .text
+    LastSec:  .text2
+    FileSize: 0x1948c9
+    MemSize:  0x1948c9
+    Offset:   0
+  - Type:     PT_LOAD
+    Flags:    [ PF_R ]
+    FirstSec: .text3
+    LastSec:  .text4
+    FileSize: 0xdf960
+    MemSize:  0xdf960
+    Offset:   0x195000
+    VAddr:    0x195000
+    PAddr:    0x195000
+
+#CHECK: ProgramHeaders [
+#CHECK-NEXT:   ProgramHeader {
+#CHECK-NEXT:     Type: PT_PHDR (0x6)
+#CHECK-NEXT:     Offset: 0x40
+#CHECK-NEXT:     VirtualAddress: 0x40
+#CHECK-NEXT:     PhysicalAddress: 0x40
+#CHECK-NEXT:     FileSize: 616
+#CHECK-NEXT:     MemSize: 616
+#CHECK-NEXT:     Flags [ (0x4)
+#CHECK-NEXT:       PF_R (0x4)
+#CHECK-NEXT:     ]
+#CHECK-NEXT:     Alignment: 8
+#CHECK-NEXT:   }
+#CHECK-NEXT:   ProgramHeader {
+#CHECK-NEXT:     Type: PT_INTERP (0x3)
+#CHECK-NEXT:     Offset: 0x1948C9
+#CHECK-NEXT:     VirtualAddress: 0x195000
+#CHECK-NEXT:     PhysicalAddress: 0x195000
+#CHECK-NEXT:     FileSize: 15
+#CHECK-NEXT:     MemSize: 15
+#CHECK-NEXT:     Flags [ (0x4)
+#CHECK-NEXT:       PF_R (0x4)
+#CHECK-NEXT:     ]
+#CHECK-NEXT:     Alignment: 1
+#CHECK-NEXT:   }
+#CHECK-NEXT:   ProgramHeader {
+#CHECK-NEXT:     Type: PT_LOAD (0x1)
+#CHECK-NEXT:     Offset: 0x0
+#CHECK-NEXT:     VirtualAddress: 0x0
+#CHECK-NEXT:     PhysicalAddress: 0x0
+#CHECK-NEXT:     FileSize: 1657033
+#CHECK-NEXT:     MemSize: 1657033
+#CHECK-NEXT:     Flags [ (0x5)
+#CHECK-NEXT:       PF_R (0x4)
+#CHECK-NEXT:       PF_X (0x1)
+#CHECK-NEXT:     ]
+#CHECK-NEXT:     Alignment: 4096
+#CHECK-NEXT:   }
+#CHECK-NEXT:   ProgramHeader {
+#CHECK-NEXT:     Type: PT_LOAD (0x1)
+#CHECK-NEXT:     Offset: 0x195000
+#CHECK-NEXT:     VirtualAddress: 0x195000
+#CHECK-NEXT:     PhysicalAddress: 0x195000
+#CHECK-NEXT:     FileSize: 915808
+#CHECK-NEXT:     MemSize: 915808
+#CHECK-NEXT:     Flags [ (0x4)
+#CHECK-NEXT:       PF_R (0x4)
+#CHECK-NEXT:     ]
+#CHECK-NEXT:     Alignment: 4096
+#CHECK-NEXT:   }
+#CHECK-NEXT: ]