From 087ff25c76fc9d82fd81c8273ff8ae82783f99c7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@blender.org>
Date: Thu, 22 Jul 2021 10:28:48 +0200
Subject: [PATCH] Add 'strict pointer mode' to `BlendFile`

Add a 'strict pointer mode' to the `BlendFile` class, which is enabled
by default. This allows users of the `BlendFile` class to decide whether
a bad pointer (i.e. one that points to a non-existing datablock) returns
`None` or raises a `SegmentationFault` exception.
---
 blender_asset_tracer/blendfile/__init__.py | 35 +++++++++++++++++-----
 blender_asset_tracer/pack/__init__.py      |  8 +++++
 tests/test_blendfile_loading.py            |  9 ++++++
 3 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/blender_asset_tracer/blendfile/__init__.py b/blender_asset_tracer/blendfile/__init__.py
index f80ce85..edb66be 100644
--- a/blender_asset_tracer/blendfile/__init__.py
+++ b/blender_asset_tracer/blendfile/__init__.py
@@ -106,6 +106,12 @@ class BlendFile:
 
     log = log.getChild("BlendFile")
 
+    strict_pointer_mode = True
+    """Raise exceptions.SegmentationFault when dereferencing an unknown pointer.
+
+    Set to False to disable this exception, and to return None instead.
+    """
+
     def __init__(self, path: pathlib.Path, mode="rb") -> None:
         """Create a BlendFile instance for the blend file at the path.
 
@@ -401,15 +407,21 @@ class BlendFile:
 
         return abspath
 
-    def dereference_pointer(self, address: int) -> "BlendFileBlock":
-        """Return the pointed-to block, or raise SegmentationFault."""
+    def dereference_pointer(self, address: int) -> typing.Optional["BlendFileBlock"]:
+        """Return the pointed-to block, or raise SegmentationFault.
+
+        When BlendFile.strict_pointer_mode is False, the exception will not be
+        thrown, but None will be returned.
+        """
 
         try:
             return self.block_from_addr[address]
         except KeyError:
-            raise exceptions.SegmentationFault(
-                "address does not exist", address
-            ) from None
+            if self.strict_pointer_mode:
+                raise exceptions.SegmentationFault(
+                    "address does not exist", address
+                ) from None
+            return None
 
     def struct(self, name: bytes) -> dna.Struct:
         index = self.sdna_index_from_id[name]
@@ -755,7 +767,11 @@ class BlendFileBlock:
             address = endian.read_pointer(fileobj, ps)
             if address == 0:
                 continue
-            yield self.bfile.dereference_pointer(address)
+            dereferenced = self.bfile.dereference_pointer(address)
+            if dereferenced is None:
+                # This can happen when strict pointer mode is disabled.
+                continue
+            yield dereferenced
 
     def iter_fixed_array_of_pointers(
         self, path: dna.FieldPath
@@ -786,7 +802,12 @@ class BlendFileBlock:
             if not address:
                 # Fixed-size arrays contain 0-pointers.
                 continue
-            yield self.bfile.dereference_pointer(address)
+
+            dereferenced = self.bfile.dereference_pointer(address)
+            if dereferenced is None:
+                # This can happen when strict pointer mode is disabled.
+                continue
+            yield dereferenced
 
     def __getitem__(self, path: dna.FieldPath):
         return self.get(path)
diff --git a/blender_asset_tracer/pack/__init__.py b/blender_asset_tracer/pack/__init__.py
index 585f8f3..c6d1a8b 100644
--- a/blender_asset_tracer/pack/__init__.py
+++ b/blender_asset_tracer/pack/__init__.py
@@ -512,6 +512,14 @@ class Packer:
 
                 # Find the same block in the newly copied file.
                 block = bfile.dereference_pointer(usage.block.addr_old)
+
+                # Pointers can point to a non-existing data block, in which case
+                # either a SegmentationFault exception is thrown, or None is
+                # returned, based on the strict pointer mode set on the
+                # BlendFile class. Since this block was already meant to be
+                # rewritten, it was found before.
+                assert block is not None
+
                 if usage.path_full_field is None:
                     dir_field = usage.path_dir_field
                     assert dir_field is not None
diff --git a/tests/test_blendfile_loading.py b/tests/test_blendfile_loading.py
index 8e03732..368f29c 100644
--- a/tests/test_blendfile_loading.py
+++ b/tests/test_blendfile_loading.py
@@ -217,6 +217,15 @@ class PointerTest(AbstractBlendFileTest):
         with self.assertRaises(exceptions.SegmentationFault):
             scene.get_pointer(b"ed")
 
+    def test_disabled_strict_pointer_mode(self):
+        scene = self.bf.code_index[b"SC"][0]
+        ed_ptr = scene.get(b"ed")
+        del self.bf.block_from_addr[ed_ptr]
+
+        self.bf.strict_pointer_mode = False
+        dereferenced = scene.get_pointer(b"ed")
+        self.assertIsNone(dereferenced)
+
     def test_abs_offset(self):
         scene = self.bf.code_index[b"SC"][0]
         ed = scene.get_pointer(b"ed")