From 2a11f6b51bdff033da25c5a178a848427262c048 Mon Sep 17 00:00:00 2001
From: Brad Hards <bradh@frogmouth.net>
Date: Thu, 22 Aug 2024 05:43:37 +1000
Subject: [PATCH] heif: make simple identification more robust (#10618)

This adds support for variable length ftyp box, and matches the major brand
and compatible brand checks to libheif.
---
 autotest/gdrivers/heif.py     | 59 +++++++++++++++++++++++++++++
 frmts/heif/heifdrivercore.cpp | 71 ++++++++++++++++++++---------------
 2 files changed, 99 insertions(+), 31 deletions(-)

diff --git a/autotest/gdrivers/heif.py b/autotest/gdrivers/heif.py
index 0c43a93367ad..92fd99640806 100644
--- a/autotest/gdrivers/heif.py
+++ b/autotest/gdrivers/heif.py
@@ -29,6 +29,7 @@
 # DEALINGS IN THE SOFTWARE.
 ###############################################################################
 
+import os
 import shutil
 
 import pytest
@@ -171,3 +172,61 @@ def test_heif_subdatasets(tmp_path):
         gdal.Open("HEIF:1")
     with pytest.raises(Exception):
         gdal.Open("HEIF:1:")
+
+
+def test_heif_identify_no_match():
+
+    drv = gdal.IdentifyDriverEx("data/byte.tif", allowed_drivers=["HEIF"])
+    assert drv is None
+
+
+def test_heif_identify_heic():
+
+    drv = gdal.IdentifyDriverEx("data/heif/subdatasets.heic", allowed_drivers=["HEIF"])
+    assert drv.GetDescription() == "HEIF"
+
+
+@pytest.mark.parametrize(
+    "major_brand,compatible_brands,expect_success",
+    [
+        ("heic", [], True),
+        ("heix", [], True),
+        ("j2ki", [], True),
+        ("j2ki", ["j2ki"], True),
+        ("jpeg", [], True),
+        ("jpg ", [], False),
+        ("miaf", [], True),
+        ("mif1", [], True),
+        ("mif2", [], True),
+        ("mif9", [], False),  # this doesn't exist
+        ("fake", ["miaf"], True),
+        ("j2kj", [], False),
+        ("fake", [], False),
+        ("fake", ["fake", "also"], False),
+        ("fake", ["fake", "avif"], True),
+        ("fake", ["fake", "bvif"], False),
+        ("fake", ["fake", "mif2"], True),
+        ("fake", ["fake", "mif9"], False),
+    ],
+)
+def test_identify_various(major_brand, compatible_brands, expect_success):
+
+    f = gdal.VSIFOpenL("/vsimem/heif_header.bin", "wb")
+    gdal.VSIFSeekL(f, 4, os.SEEK_SET)
+    gdal.VSIFWriteL("ftyp", 1, 4, f)  # box type
+    gdal.VSIFWriteL(major_brand, 1, 4, f)
+    gdal.VSIFWriteL(b"\x00\x00\x00\x00", 1, 4, f)  # minor_version
+    for brand in compatible_brands:
+        gdal.VSIFWriteL(brand, 1, 4, f)
+    length = gdal.VSIFTellL(f)
+    gdal.VSIFSeekL(f, 0, os.SEEK_SET)  # go back and fill in actual box size
+    gdal.VSIFWriteL(length.to_bytes(4, "big"), 1, 4, f)
+    gdal.VSIFCloseL(f)
+
+    drv = gdal.IdentifyDriverEx("/vsimem/heif_header.bin", allowed_drivers=["HEIF"])
+    if expect_success:
+        assert drv.GetDescription() == "HEIF"
+    else:
+        assert drv is None
+
+    gdal.Unlink("/vsimem/heif_header.bin")
diff --git a/frmts/heif/heifdrivercore.cpp b/frmts/heif/heifdrivercore.cpp
index 24e424c3b2f3..e4e293073e97 100644
--- a/frmts/heif/heifdrivercore.cpp
+++ b/frmts/heif/heifdrivercore.cpp
@@ -33,42 +33,51 @@
 /*                    HEIFDriverIdentifySimplified()                    */
 /************************************************************************/
 
-int HEIFDriverIdentifySimplified(GDALOpenInfo *poOpenInfo)
+static const GByte FTYP_4CC[] = {'f', 't', 'y', 'p'};
+static const GByte supportedBrands[][4]{
+    {'a', 'v', 'i', 'f'}, {'h', 'e', 'i', 'c'}, {'h', 'e', 'i', 'x'},
+    {'j', '2', 'k', 'i'}, {'j', 'p', 'e', 'g'}, {'m', 'i', 'a', 'f'},
+    {'m', 'i', 'f', '1'}, {'m', 'i', 'f', '2'}};
 
+int HEIFDriverIdentifySimplified(GDALOpenInfo *poOpenInfo)
 {
     if (STARTS_WITH_CI(poOpenInfo->pszFilename, "HEIF:"))
+    {
         return true;
-
-    if (poOpenInfo->nHeaderBytes < 12 || poOpenInfo->fpL == nullptr)
+    }
+    if (poOpenInfo->nHeaderBytes < 16 || poOpenInfo->fpL == nullptr)
+    {
         return false;
-
-    // Simplistic test...
-    const unsigned char abySig1[] = "\x00"
-                                    "\x00"
-                                    "\x00"
-                                    "\x20"
-                                    "ftypheic";
-    const unsigned char abySig2[] = "\x00"
-                                    "\x00"
-                                    "\x00"
-                                    "\x18"
-                                    "ftypheic";
-    const unsigned char abySig3[] = "\x00"
-                                    "\x00"
-                                    "\x00"
-                                    "\x18"
-                                    "ftypmif1"
-                                    "\x00"
-                                    "\x00"
-                                    "\x00"
-                                    "\x00"
-                                    "mif1heic";
-    return (poOpenInfo->nHeaderBytes >= static_cast<int>(sizeof(abySig1)) &&
-            memcmp(poOpenInfo->pabyHeader, abySig1, sizeof(abySig1)) == 0) ||
-           (poOpenInfo->nHeaderBytes >= static_cast<int>(sizeof(abySig2)) &&
-            memcmp(poOpenInfo->pabyHeader, abySig2, sizeof(abySig2)) == 0) ||
-           (poOpenInfo->nHeaderBytes >= static_cast<int>(sizeof(abySig3)) &&
-            memcmp(poOpenInfo->pabyHeader, abySig3, sizeof(abySig3)) == 0);
+    }
+    if (memcmp(poOpenInfo->pabyHeader + 4, FTYP_4CC, sizeof(FTYP_4CC)) != 0)
+    {
+        return false;
+    }
+    uint32_t lengthBigEndian;
+    memcpy(&lengthBigEndian, poOpenInfo->pabyHeader, sizeof(uint32_t));
+    uint32_t lengthHostEndian = CPL_MSBWORD32(lengthBigEndian);
+    if (lengthHostEndian > static_cast<uint32_t>(poOpenInfo->nHeaderBytes))
+    {
+        lengthHostEndian = static_cast<uint32_t>(poOpenInfo->nHeaderBytes);
+    }
+    for (const GByte *supportedBrand : supportedBrands)
+    {
+        if (memcmp(poOpenInfo->pabyHeader + 8, supportedBrand, 4) == 0)
+        {
+            return true;
+        }
+    }
+    for (uint32_t offset = 16; offset + 4 <= lengthHostEndian; offset += 4)
+    {
+        for (const GByte *supportedBrand : supportedBrands)
+        {
+            if (memcmp(poOpenInfo->pabyHeader + offset, supportedBrand, 4) == 0)
+            {
+                return true;
+            }
+        }
+    }
+    return false;
 }
 
 /************************************************************************/