Move program headers to the end of the file.

This change implements step 2 of the script: moving all of its program
headers to the end of the file and zeroing the previous position.

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

tools/android/elf_compression/compress_section.py

@@ -43,6 +43,48 @@ OBJCOPY_PATH = pathlib.Path(__file__).resolve().parents[3].joinpath(
     'third_party/llvm-build/Release+Asserts/bin/llvm-objcopy')
 
 
+def AlignUp(addr, page_size=ADDRESS_ALIGN):
+  """Rounds up given address to be aligned to page_size.
+
+  Args:
+    addr: int. Virtual address to be aligned.
+    page_size: int. Page size to be used for the alignment.
+  """
+  if addr % page_size != 0:
+    addr += page_size - (addr % page_size)
+  return addr
+
+
+def AlignDown(addr, page_size=ADDRESS_ALIGN):
+  """Round down given address to be aligned to page_size.
+
+  Args:
+    addr: int. Virtual address to be aligned.
+    page_size: int. Page size to be used for the alignment.
+  """
+  return addr - addr % page_size
+
+
+def MatchVaddrAlignment(vaddr, offset, align=ADDRESS_ALIGN):
+  """Align vaddr to comply with ELF standard binary alignment.
+
+  Increases vaddr until the following is true:
+    vaddr % align == offset % align
+
+  Args:
+    vaddr: virtual address to be aligned.
+    offset: file offset to be aligned.
+    align: alignment value.
+
+  Returns:
+    Aligned virtual address, bigger or equal than the vaddr.
+  """
+  delta = offset % align - vaddr % align
+  if delta < 0:
+    delta += align
+  return vaddr + delta
+
+
 def _SetupLogging():
   logging.basicConfig(
       format='%(asctime)s %(filename)s:%(lineno)s %(levelname)s] %(message)s',
@@ -79,7 +121,7 @@ def _FileRangeToVirtualAddressRange(data, l, r):
   """
   elf = elf_headers.ElfHeader(data)
   for phdr in elf.GetPhdrs():
-    if phdr.p_type == elf_headers.ProgramHeader.Type.PT_LOAD.value:
+    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.
@@ -100,9 +142,8 @@ def _CopyRangeIntoCompressedSection(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.
-  if virtual_l % ADDRESS_ALIGN != 0:
-    l += ADDRESS_ALIGN - (virtual_l % ADDRESS_ALIGN)
-  r -= virtual_r % ADDRESS_ALIGN
+  l += AlignUp(virtual_l) - virtual_l
+  r -= virtual_r - AlignDown(virtual_r)
   if l >= r:
     raise RuntimeError('Range collapsed after aligning by page size')
 
@@ -134,6 +175,68 @@ def _CopyRangeIntoCompressedSection(data, l, r):
       data[:] = bytearray(f.read())
 
 
+def _FindNewVaddr(phdrs):
+  """Returns the virt address that is safe to use for insertion of new data."""
+  max_vaddr = 0
+  # Strictly speaking it should be sufficient to look through only LOAD
+  # segments, but better be safe than sorry.
+  for phdr in phdrs:
+    max_vaddr = max(max_vaddr, phdr.p_vaddr + phdr.p_memsz)
+  # When the mapping occurs end address is increased to be a multiple
+  # of page size. To ensure compatibility with anything relying on this
+  # behaviour we take this increase into account.
+  max_vaddr = AlignUp(max_vaddr)
+  return max_vaddr
+
+
+def _MovePhdrToTheEnd(data):
+  """Moves Phdrs to the end of the file and adjusts all references to it."""
+  elf_hdr = elf_headers.ElfHeader(data)
+
+  # If program headers are already in the end of the file, nothing to do.
+  if elf_hdr.e_phoff + elf_hdr.e_phnum * elf_hdr.e_phentsize == len(data):
+    return
+
+  new_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
+  elf_hdr.AddPhdr(
+      elf_headers.ProgramHeader.Create(
+          elf_hdr.byte_order,
+          p_type=elf_headers.ProgramHeader.Type.PT_LOAD,
+          p_flags=elf_headers.ProgramHeader.Flags.PF_R,
+          p_offset=new_phoff,
+          p_vaddr=new_vaddr,
+          p_paddr=new_vaddr,
+          p_filesz=new_filesz,
+          p_memsz=new_filesz,
+          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
+
+  # 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
+
+  # Updating ELF header to point to the new location.
+  elf_hdr.e_phoff = new_phoff
+  elf_hdr.PatchData(data)
+
+
 def main():
   _SetupLogging()
   args = _ParseArguments()
@@ -143,6 +246,7 @@ def main():
     data = bytearray(data)
 
   _CopyRangeIntoCompressedSection(data, args.left_range, args.right_range)
+  _MovePhdrToTheEnd(data)
 
   with open(args.output, 'wb') as f:
     f.write(data)

tools/android/elf_compression/elf_headers.py

@@ -66,6 +66,28 @@ class ElfEntry(object):
         continue
       setattr(self, field_name, kwargs[field_name])
 
+  def ToBytes(self):
+    """Returns byte representation of ELF entry."""
+    bytearr = bytearray()
+    for field_name, field_size in self._fields:
+      field_bytes = getattr(self, field_name).to_bytes(
+          field_size, byteorder=self.byte_order)
+      bytearr.extend(field_bytes)
+    return bytearr
+
+  @classmethod
+  def Create(cls, byte_order, **kwargs):
+    """Static wrapper around ApplyKwargs method.
+
+    Args:
+      byte_order: str. Either 'little' for little endian or 'big' for big
+        endian.
+      **kwargs: will be passed directly to the ApplyKwargs method.
+    """
+    obj = cls(byte_order)
+    obj.ApplyKwargs(**kwargs)
+    return obj
+
   @classmethod
   def FromBytes(cls, byte_order, data, offset):
     """Static wrapper around ParseBytes method.
@@ -93,6 +115,11 @@ class ProgramHeader(ElfEntry):
     PT_SHLIB = 5
     PT_PHDR = 6
 
+  class Flags(enum.IntFlag):
+    PF_X = 1
+    PF_W = 2
+    PF_R = 4
+
   def __init__(self, byte_order):
     """ProgramHeader constructor.
 
@@ -238,3 +265,36 @@ class ElfHeader(ElfEntry):
   def GetPhdrs(self):
     """Returns the list of file's program headers."""
     return self.phdrs
+
+  def AddPhdr(self, phdr):
+    """Adds a new ProgramHeader entry correcting the e_phnum variable.
+
+    Args:
+      phdr: ProgramHeader. Instance of ProgramHeader to add.
+    """
+    self.phdrs.append(phdr)
+    self.e_phnum += 1
+
+  def PatchData(self, data):
+    """Patches the given data array to reflect all changes made to the header.
+
+    This method doesn't completely rewrite the data, instead it patches
+    inplace. Not only the ElfHeader is patched but all of its ProgramHeader
+    as well.
+
+    The important limitation is that this method doesn't take changes of sizes
+    and offsets into account. As example, if new ProgramHeader is added, this
+    method will override whatever data is located under its placement so the
+    user has to move the headers to the end beforehand or the user mustn't
+    change header's size.
+
+    Args:
+      data: bytearray. The data array to be patched.
+    """
+    elf_bytes = self.ToBytes()
+    data[:len(elf_bytes)] = elf_bytes
+    current_offset = self.e_phoff
+    for phdr in self.GetPhdrs():
+      phdr_bytes = phdr.ToBytes()
+      data[current_offset:current_offset + len(phdr_bytes)] = phdr_bytes
+      current_offset += self.e_phentsize

tools/android/elf_compression/test/compression_script_test.py

@@ -2,17 +2,18 @@
 # Copyright 2019 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
-"""Basic test for the compress_section script.
-
-Runs the script on a sample library and verifies that it still works.
-"""
+"""Tests for the compress_section script."""
 
 import os
 import pathlib
 import subprocess
+import sys
 import tempfile
 import unittest
 
+sys.path.append(str(pathlib.Path(__file__).resolve().parents[1]))
+import compress_section
+
 LIBRARY_CC_NAME = 'libtest.cc'
 OPENER_CC_NAME = 'library_opener.cc'
 
@@ -118,6 +119,28 @@ class CompressionScriptTest(unittest.TestCase):
     opener_output = self._RunOpener(opener_path, patched_library_path)
     self.assertEqual(opener_output, '1046506\n')
 
+  def testAlignUp(self):
+    """Tests for AlignUp method of the script."""
+    self.assertEqual(compress_section.AlignUp(1024, 1024), 1024)
+    self.assertEqual(compress_section.AlignUp(1023, 1024), 1024)
+    self.assertEqual(compress_section.AlignUp(1025, 1024), 2048)
+    self.assertEqual(compress_section.AlignUp(5555, 4096), 8192)
+
+  def testAlignDown(self):
+    """Tests for AlignDown method of the script."""
+    self.assertEqual(compress_section.AlignDown(1024, 1024), 1024)
+    self.assertEqual(compress_section.AlignDown(1023, 1024), 0)
+    self.assertEqual(compress_section.AlignDown(1025, 1024), 1024)
+    self.assertEqual(compress_section.AlignDown(5555, 4096), 4096)
+
+  def testMatchVaddrAlignment(self):
+    """Tests for MatchVaddrAlignment method of the script."""
+    self.assertEqual(compress_section.MatchVaddrAlignment(100, 100, 1024), 100)
+    self.assertEqual(compress_section.MatchVaddrAlignment(99, 100, 1024), 100)
+    self.assertEqual(compress_section.MatchVaddrAlignment(101, 100, 1024), 1124)
+    self.assertEqual(
+        compress_section.MatchVaddrAlignment(1024, 2049, 1024), 1025)
+
 
 if __name__ == '__main__':
   unittest.main()

tools/android/elf_compression/test/elf_headers_test.py

@@ -69,6 +69,52 @@ class ElfHeaderTest(unittest.TestCase):
     self.assertEqual(load_phdr.p_vaddr, 0x1000)
     self.assertEqual(load_phdr.p_paddr, 0x1000)
 
+  def testElfHeaderNoopPatching(self):
+    """Patching the ELF without any changes."""
+    with open(self.library_path, 'rb') as f:
+      data = bytearray(f.read())
+    data_copy = data[:]
+    elf = elf_headers.ElfHeader(data)
+    elf.PatchData(data)
+    self.assertEqual(data, data_copy)
+
+  def testElfHeaderPatchingAndParsing(self):
+    """Patching the ELF and validating that it worked."""
+    with open(self.library_path, 'rb') as f:
+      data = bytearray(f.read())
+    elf = elf_headers.ElfHeader(data)
+    # Changing some values.
+    elf.e_ehsize = 42
+    elf.GetPhdrs()[0].p_align = 1
+    elf.GetPhdrs()[0].p_filesz = 10
+    elf.PatchData(data)
+
+    updated_elf = elf_headers.ElfHeader(data)
+    # Validating all of the ELF header fields.
+    self.assertEqual(updated_elf.e_type, elf.e_type)
+    self.assertEqual(updated_elf.e_entry, elf.e_entry)
+    self.assertEqual(updated_elf.e_phoff, elf.e_phoff)
+    self.assertEqual(updated_elf.e_shoff, elf.e_shoff)
+    self.assertEqual(updated_elf.e_flags, elf.e_flags)
+    self.assertEqual(updated_elf.e_ehsize, 42)
+    self.assertEqual(updated_elf.e_phentsize, elf.e_phentsize)
+    self.assertEqual(updated_elf.e_phnum, elf.e_phnum)
+    self.assertEqual(updated_elf.e_shentsize, elf.e_shentsize)
+    self.assertEqual(updated_elf.e_shnum, elf.e_shnum)
+    self.assertEqual(updated_elf.e_shstrndx, elf.e_shstrndx)
+
+    # Validating all of the fields of the first segment.
+    load_phdr = elf.GetPhdrs()[0]
+    updated_load_phdr = updated_elf.GetPhdrs()[0]
+
+    self.assertEqual(updated_load_phdr.p_offset, load_phdr.p_offset)
+    self.assertEqual(updated_load_phdr.p_vaddr, load_phdr.p_vaddr)
+    self.assertEqual(updated_load_phdr.p_paddr, load_phdr.p_paddr)
+    self.assertEqual(updated_load_phdr.p_filesz, 10)
+    self.assertEqual(updated_load_phdr.p_memsz, load_phdr.p_memsz)
+    self.assertEqual(updated_load_phdr.p_flags, load_phdr.p_flags)
+    self.assertEqual(updated_load_phdr.p_align, 1)
+
 
 if __name__ == '__main__':
   unittest.main()