From 2d26003df79839d7f6b5e30eaa49e191dc9e6c87 Mon Sep 17 00:00:00 2001
From: Simon Glass <sjg@chromium.org>
Date: Mon, 8 Jul 2019 14:25:45 -0600
Subject: [PATCH] binman: Support locating an image header

Add support for locating an image header in an image.

Signed-off-by: Simon Glass <sjg@chromium.org>
---
 tools/binman/etype/image_header.py | 23 +++++++++++++++++++++++
 tools/binman/ftest.py              | 27 +++++++++++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/tools/binman/etype/image_header.py b/tools/binman/etype/image_header.py
index b1c4f8a07e..8f9c5aa5d9 100644
--- a/tools/binman/etype/image_header.py
+++ b/tools/binman/etype/image_header.py
@@ -15,6 +15,29 @@ from entry import Entry
 import fdt_util
 
 IMAGE_HEADER_MAGIC = b'BinM'
+IMAGE_HEADER_LEN   = 8
+
+def LocateHeaderOffset(data):
+    """Search an image for an image header
+
+    Args:
+        data: Data to search
+
+    Returns:
+        Offset of image header in the image, or None if not found
+    """
+    hdr_pos = data.find(IMAGE_HEADER_MAGIC)
+    if hdr_pos != -1:
+        size = len(data)
+        hdr = data[hdr_pos:hdr_pos + IMAGE_HEADER_LEN]
+        if len(hdr) == IMAGE_HEADER_LEN:
+            offset = struct.unpack('<I', hdr[4:])[0]
+            if hdr_pos == len(data) - IMAGE_HEADER_LEN:
+                pos = size + offset - (1 << 32)
+            else:
+                pos = offset
+            return pos
+    return None
 
 class Entry_image_header(Entry):
     """An entry which contains a pointer to the FDT map
diff --git a/tools/binman/ftest.py b/tools/binman/ftest.py
index d800ba1e9d..ce66e3a2f2 100644
--- a/tools/binman/ftest.py
+++ b/tools/binman/ftest.py
@@ -25,6 +25,7 @@ import control
 import elf
 import fdt
 from etype import fdtmap
+from etype import image_header
 import fdt_util
 import fmap_util
 import test_util
@@ -2282,6 +2283,32 @@ class TestFunctional(unittest.TestCase):
         data = self._DoReadFile('005_simple.dts')
         self.assertEqual(None, fdtmap.LocateFdtmap(data))
 
+    def testFindImageHeader(self):
+        """Test locating a image header"""
+        self._CheckLz4()
+        data = self._DoReadFileDtb('128_decode_image.dts', use_real_dtb=True,
+                                   update_dtb=True)[0]
+        image = control.images['image']
+        entries = image.GetEntries()
+        entry = entries['fdtmap']
+        # The header should point to the FDT map
+        self.assertEqual(entry.image_pos, image_header.LocateHeaderOffset(data))
+
+    def testFindImageHeaderStart(self):
+        """Test locating a image header located at the start of an image"""
+        data = self._DoReadFileDtb('117_fdtmap_hdr_start.dts',
+                                   use_real_dtb=True, update_dtb=True)[0]
+        image = control.images['image']
+        entries = image.GetEntries()
+        entry = entries['fdtmap']
+        # The header should point to the FDT map
+        self.assertEqual(entry.image_pos, image_header.LocateHeaderOffset(data))
+
+    def testFindImageHeaderMissing(self):
+        """Test failing to locate an image header"""
+        data = self._DoReadFile('005_simple.dts')
+        self.assertEqual(None, image_header.LocateHeaderOffset(data))
+
 
 if __name__ == "__main__":
     unittest.main()
-- 
2.25.1