Split LOAD segment containing range in three.

This change splits the LOAD segments so the [l, r) range is located in
its separate segment to make removal of it in the later stages
possible.

This change also moves the logic of resizing the segments after adding
new program header to AddPhdr method.

Bug: 998082
Change-Id: Ibec04434caea993ffddb73bcd1771730f2bd1e9a
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1795448
Commit-Queue: Denis Solonkov <solonkovda@google.com>
Reviewed-by: Boris Sazonov <bsazonov@chromium.org>
Reviewed-by: Alex Ilin <alexilin@chromium.org>
Cr-Commit-Position: refs/heads/master@{#697151}

tools/android/elf_compression/compress_section.py

@@ -120,33 +120,24 @@ def _FileRangeToVirtualAddressRange(data, l, r):
   segments is required here.
   """
   elf = elf_headers.ElfHeader(data)
-  for phdr in elf.GetPhdrs():
-    if phdr.p_type == elf_headers.ProgramHeader.Type.PT_LOAD:
-      # Current version of the prototype only supports ranges which are fully
-      # contained inside one LOAD segment. It should cover most of the common
-      # cases.
-      if phdr.p_offset >= r or phdr.p_offset + phdr.p_filesz <= l:
-        # Range doesn't overlap.
-        continue
-      if phdr.p_offset > l or phdr.p_offset + phdr.p_filesz < r:
-        # Range overlap with LOAD segment but isn't fully covered by it.
-        raise RuntimeError('Range is not contained within one LOAD segment')
-      l_virt = phdr.p_vaddr + (l - phdr.p_offset)
-      r_virt = phdr.p_vaddr + (r - phdr.p_offset)
-      return l_virt, r_virt
+  for phdr in elf.GetPhdrsByType(elf_headers.ProgramHeader.Type.PT_LOAD):
+    # Current version of the prototype only supports ranges which are fully
+    # contained inside one LOAD segment. It should cover most of the common
+    # cases.
+    if phdr.p_offset >= r or phdr.FilePositionEnd() <= l:
+      # Range doesn't overlap.
+      continue
+    if phdr.p_offset > l or phdr.FilePositionEnd() < r:
+      # Range overlap with LOAD segment but isn't fully covered by it.
+      raise RuntimeError('Range is not contained within one LOAD segment')
+    l_virt = phdr.p_vaddr + (l - phdr.p_offset)
+    r_virt = phdr.p_vaddr + (r - phdr.p_offset)
+    return l_virt, r_virt
   raise RuntimeError('Specified range is outside of all LOAD segments.')
 
 
 def _CopyRangeIntoCompressedSection(data, l, r):
   """Adds a new section containing compressed version of provided range."""
-  virtual_l, virtual_r = _FileRangeToVirtualAddressRange(data, l, r)
-  # LOAD segments borders are being rounded to the page size so we have to
-  # shrink [l, r) so corresponding virtual addresses are aligned.
-  l += AlignUp(virtual_l) - virtual_l
-  r -= virtual_r - AlignDown(virtual_r)
-  if l >= r:
-    raise RuntimeError('Range collapsed after aligning by page size')
-
   compressed_range = compression.CompressData(data[l:r])
 
   with tempfile.TemporaryDirectory() as tmpdir:
@@ -197,12 +188,16 @@ def _MovePhdrToTheEnd(data):
   if elf_hdr.e_phoff + elf_hdr.e_phnum * elf_hdr.e_phentsize == len(data):
     return
 
-  new_phoff = len(data)
+  old_phoff = elf_hdr.e_phoff
+  new_phoff = elf_hdr.e_phoff = len(data)
+
   unaligned_new_vaddr = _FindNewVaddr(elf_hdr.GetPhdrs())
   new_vaddr = MatchVaddrAlignment(unaligned_new_vaddr, new_phoff)
   # Since we moved the PHDR section to the end of the file, we need to create a
   # new LOAD segment to load it in.
-  new_filesz = (elf_hdr.e_phnum + 1) * elf_hdr.e_phentsize
+  current_filesize = elf_hdr.e_phnum * elf_hdr.e_phentsize
+  # We are using current_filesize while adding new program header due to
+  # AddPhdr handling the increase of size due to addition of new header.
   elf_hdr.AddPhdr(
       elf_headers.ProgramHeader.Create(
           elf_hdr.byte_order,
@@ -211,32 +206,99 @@ def _MovePhdrToTheEnd(data):
           p_offset=new_phoff,
           p_vaddr=new_vaddr,
           p_paddr=new_vaddr,
-          p_filesz=new_filesz,
-          p_memsz=new_filesz,
+          p_filesz=current_filesize,
+          p_memsz=current_filesize,
           p_align=ADDRESS_ALIGN,
       ))
 
   # PHDR segment if it exists should point to the new location.
-  for phdr in elf_hdr.GetPhdrs():
-    if phdr.p_type == elf_headers.ProgramHeader.Type.PT_PHDR:
-      phdr.p_offset = new_phoff
-      phdr.p_vaddr = new_vaddr
-      phdr.p_paddr = new_vaddr
-      phdr.p_filesz = new_filesz
-      phdr.p_memsz = new_filesz
-      phdr.p_align = ADDRESS_ALIGN
+  for phdr in elf_hdr.GetPhdrsByType(elf_headers.ProgramHeader.Type.PT_PHDR):
+    phdr.p_offset = new_phoff
+    phdr.p_vaddr = new_vaddr
+    phdr.p_paddr = new_vaddr
+    phdr.p_align = ADDRESS_ALIGN
 
   # We need to replace the previous phdr placement with zero bytes to fail
   # fast if dynamic linker doesn't like the new program header.
   previous_phdr_size = (elf_hdr.e_phnum - 1) * elf_hdr.e_phentsize
-  data[elf_hdr.e_phoff:elf_hdr.e_phoff +
-       previous_phdr_size] = [0] * previous_phdr_size
+  data[old_phoff:old_phoff + previous_phdr_size] = [0] * previous_phdr_size
 
   # Updating ELF header to point to the new location.
-  elf_hdr.e_phoff = new_phoff
   elf_hdr.PatchData(data)
 
 
+def _SplitLoadSegment(data, l, r):
+  """Find LOAD segment covering [l, r) and splits it into three segments.
+
+  Split is done so one of the LOAD segments contains only [l, r) and nothing
+  else. If the range is located at the start or at the end of the segment less
+  than three segments may be created.
+  """
+  elf_hdr = elf_headers.ElfHeader(data)
+
+  range_phdr = None
+  for phdr in elf_hdr.GetPhdrsByType(elf_headers.ProgramHeader.Type.PT_LOAD):
+    if phdr.p_offset <= l and phdr.FilePositionEnd() >= r:
+      range_phdr = phdr
+      break
+  if range_phdr is None:
+    raise RuntimeError('No LOAD segment covering the range found')
+
+  # The range_phdr will become the LOAD segment containing the [l, r) range
+  # but we need to create the additional two segments.
+  left_segment_size = l - range_phdr.p_offset
+  if left_segment_size > 0:
+    # Creating LOAD segment containing the [phdr.p_offset, l) part.
+    elf_hdr.AddPhdr(
+        elf_headers.ProgramHeader.Create(
+            elf_hdr.byte_order,
+            p_type=range_phdr.p_type,
+            p_flags=range_phdr.p_flags,
+            p_offset=range_phdr.p_offset,
+            p_vaddr=range_phdr.p_vaddr,
+            p_paddr=range_phdr.p_paddr,
+            p_filesz=left_segment_size,
+            p_memsz=left_segment_size,
+            p_align=range_phdr.p_align,
+        ))
+  if range_phdr.p_offset + range_phdr.p_memsz > r:
+    # Creating LOAD segment containing the [r, phdr.p_offset + phdr.p_memsz).
+    right_segment_delta = r - range_phdr.p_offset
+    right_segment_address = range_phdr.p_vaddr + right_segment_delta
+    right_segment_filesize = max(range_phdr.p_filesz - right_segment_delta, 0)
+    right_segment_memsize = range_phdr.p_memsz - right_segment_delta
+    elf_hdr.AddPhdr(
+        elf_headers.ProgramHeader.Create(
+            elf_hdr.byte_order,
+            p_type=range_phdr.p_type,
+            p_flags=range_phdr.p_flags,
+            p_offset=r,
+            p_vaddr=right_segment_address,
+            p_paddr=right_segment_address,
+            p_filesz=right_segment_filesize,
+            p_memsz=right_segment_memsize,
+            p_align=range_phdr.p_align,
+        ))
+  # Modifying the range_phdr
+  central_segment_address = range_phdr.p_vaddr + left_segment_size
+  range_phdr.p_offset = l
+  range_phdr.p_vaddr = central_segment_address
+  range_phdr.p_paddr = central_segment_address
+  range_phdr.p_filesz = r - l
+  range_phdr.p_memsz = r - l
+
+  elf_hdr.PatchData(data)
+
+
+def _ShrinkRangeToAlignVirtualAddress(data, l, r):
+  virtual_l, virtual_r = _FileRangeToVirtualAddressRange(data, l, r)
+  # LOAD segments borders are being rounded to the page size so we have to
+  # shrink [l, r) so corresponding virtual addresses are aligned.
+  l += AlignUp(virtual_l) - virtual_l
+  r -= virtual_r - AlignDown(virtual_r)
+  return l, r
+
+
 def main():
   _SetupLogging()
   args = _ParseArguments()
@@ -245,8 +307,14 @@ def main():
     data = f.read()
     data = bytearray(data)
 
-  _CopyRangeIntoCompressedSection(data, args.left_range, args.right_range)
+  left_range, right_range = _ShrinkRangeToAlignVirtualAddress(
+      data, args.left_range, args.right_range)
+  if left_range >= right_range:
+    raise RuntimeError('Range collapsed after aligning by page size')
+
+  _CopyRangeIntoCompressedSection(data, left_range, right_range)
   _MovePhdrToTheEnd(data)
+  _SplitLoadSegment(data, left_range, right_range)
 
   with open(args.output, 'wb') as f:
     f.write(data)

tools/android/elf_compression/elf_headers.py

@@ -147,6 +147,10 @@ class ProgramHeader(ElfEntry):
     ]
     super(ProgramHeader, self).__init__(byte_order, fields)
 
+  def FilePositionEnd(self):
+    """Returns the end (exclusive) of the segment file range."""
+    return self.p_offset + self.p_filesz
+
 
 class ElfHeader(ElfEntry):
   """This class represents ELFHdr from the ELF standard.
@@ -266,15 +270,57 @@ class ElfHeader(ElfEntry):
     """Returns the list of file's program headers."""
     return self.phdrs
 
+  def GetPhdrsByType(self, phdr_type):
+    """Yields program headers of the given type."""
+    return (phdr for phdr in self.phdrs if phdr.p_type == phdr_type)
+
   def AddPhdr(self, phdr):
     """Adds a new ProgramHeader entry correcting the e_phnum variable.
 
+    This method will increase the size of LOAD segment containing the program
+    headers without correcting the other offsets. It is up to the caller to
+    deal with the results. One way to avoid any problems would be to move
+    program headers to the end of the file.
+
     Args:
       phdr: ProgramHeader. Instance of ProgramHeader to add.
     """
     self.phdrs.append(phdr)
+
+    phdrs_size = self.e_phnum * self.e_phentsize
+    # We need to locate the LOAD segment containing program headers and
+    # increase its size.
+    phdr_found = False
+    for phdr in self.GetPhdrsByType(ProgramHeader.Type.PT_LOAD):
+      if phdr.p_offset > self.e_phoff:
+        continue
+      if phdr.FilePositionEnd() < self.e_phoff + phdrs_size:
+        continue
+      phdr.p_filesz += self.e_phentsize
+      phdr.p_memsz += self.e_phentsize
+      phdr_found = True
+      break
+    if not phdr_found:
+      raise RuntimeError('Failed to increase program headers LOAD segment')
+
+    # If PHDR segment exists it needs to be corrected as well.
+    for phdr in self.GetPhdrsByType(ProgramHeader.Type.PT_PHDR):
+      phdr.p_filesz += self.e_phentsize
+      phdr.p_memsz += self.e_phentsize
     self.e_phnum += 1
 
+  def _OrderPhdrs(self):
+    """Orders program LOAD headers by p_vaddr to comply with standard."""
+
+    def HeaderToKey(phdr):
+      if phdr.p_type == ProgramHeader.Type.PT_LOAD:
+        return (0, phdr.p_vaddr)
+      else:
+        # We want to preserve the order of non LOAD segments.
+        return (1, 0)
+
+    self.phdrs.sort(key=HeaderToKey)
+
   def PatchData(self, data):
     """Patches the given data array to reflect all changes made to the header.
 
@@ -294,6 +340,7 @@ class ElfHeader(ElfEntry):
     elf_bytes = self.ToBytes()
     data[:len(elf_bytes)] = elf_bytes
     current_offset = self.e_phoff
+    self._OrderPhdrs()
     for phdr in self.GetPhdrs():
       phdr_bytes = phdr.ToBytes()
       data[current_offset:current_offset + len(phdr_bytes)] = phdr_bytes