From 4d38dd77f901a48736c540a9c7499816fa9b4d90 Mon Sep 17 00:00:00 2001
From: Simon Glass <sjg@chromium.org>
Date: Tue, 8 Feb 2022 11:49:55 -0700
Subject: [PATCH] elf: Add a way to read segment information from an ELF file

Add a function which reads the segments and the entry address.

Also fix a comment nit in the tests while we are here.

Signed-off-by: Simon Glass <sjg@chromium.org>
---
 tools/binman/elf.py      | 37 +++++++++++++++++++++++++++++++++++++
 tools/binman/elf_test.py | 31 +++++++++++++++++++++++++++++--
 2 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/tools/binman/elf.py b/tools/binman/elf.py
index bc4966e..5e7d6ae 100644
--- a/tools/binman/elf.py
+++ b/tools/binman/elf.py
@@ -20,6 +20,7 @@ from patman import tout
 ELF_TOOLS = True
 try:
     from elftools.elf.elffile import ELFFile
+    from elftools.elf.elffile import ELFError
     from elftools.elf.sections import SymbolTableSection
 except:  # pragma: no cover
     ELF_TOOLS = False
@@ -369,3 +370,39 @@ def UpdateFile(infile, outfile, start_sym, end_sym, insert):
     newdata += data[syms[end_sym].offset:]
     tools.write_file(outfile, newdata)
     tout.info('Written to offset %#x' % syms[start_sym].offset)
+
+def read_segments(data):
+    """Read segments from an ELF file
+
+    Args:
+        data (bytes): Contents of file
+
+    Returns:
+        tuple:
+            list of segments, each:
+                int: Segment number (0 = first)
+                int: Start address of segment in memory
+                bytes: Contents of segment
+            int: entry address for image
+
+    Raises:
+        ValueError: elftools is not available
+    """
+    if not ELF_TOOLS:
+        raise ValueError('Python elftools package is not available')
+    with io.BytesIO(data) as inf:
+        try:
+            elf = ELFFile(inf)
+        except ELFError as err:
+            raise ValueError(err)
+        entry = elf.header['e_entry']
+        segments = []
+        for i in range(elf.num_segments()):
+            segment = elf.get_segment(i)
+            if segment['p_type'] != 'PT_LOAD' or not segment['p_memsz']:
+                skipped = 1  # To make code-coverage see this line
+                continue
+            start = segment['p_offset']
+            rend = start + segment['p_filesz']
+            segments.append((i, segment['p_paddr'], data[start:rend]))
+    return segments, entry
diff --git a/tools/binman/elf_test.py b/tools/binman/elf_test.py
index 47ebfba..f92352d 100644
--- a/tools/binman/elf_test.py
+++ b/tools/binman/elf_test.py
@@ -56,8 +56,8 @@ class FakeSection:
 def BuildElfTestFiles(target_dir):
     """Build ELF files used for testing in binman
 
-    This compiles and links the test files into the specified directory. It the
-    Makefile and source files in the binman test/ directory.
+    This compiles and links the test files into the specified directory. It uses
+    the Makefile and source files in the binman test/ directory.
 
     Args:
         target_dir: Directory to put the files into
@@ -258,6 +258,33 @@ class TestElf(unittest.TestCase):
         offset = elf.GetSymbolFileOffset(fname, ['missing_sym'])
         self.assertEqual({}, offset)
 
+    def test_read_segments(self):
+        """Test for read_segments()"""
+        if not elf.ELF_TOOLS:
+            self.skipTest('Python elftools not available')
+        fname = self.ElfTestFile('embed_data')
+        segments, entry = elf.read_segments(tools.ReadFile(fname))
+
+    def test_read_segments_fail(self):
+        """Test for read_segments() without elftools"""
+        try:
+            old_val = elf.ELF_TOOLS
+            elf.ELF_TOOLS = False
+            fname = self.ElfTestFile('embed_data')
+            with self.assertRaises(ValueError) as e:
+                elf.read_segments(tools.ReadFile(fname))
+            self.assertIn('Python elftools package is not available',
+                          str(e.exception))
+        finally:
+            elf.ELF_TOOLS = old_val
+
+    def test_read_segments_bad_data(self):
+        """Test for read_segments() with an invalid ELF file"""
+        fname = self.ElfTestFile('embed_data')
+        with self.assertRaises(ValueError) as e:
+            elf.read_segments(tools.GetBytes(100, 100))
+        self.assertIn('Magic number does not match', str(e.exception))
+
 
 if __name__ == '__main__':
     unittest.main()
-- 
2.7.4