Improve checkpoint superblock selection (#3)

Author: Schamper

dissect/apfs/apfs.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import logging
+import os
 from typing import TYPE_CHECKING, BinaryIO
 
 from dissect.apfs.c_apfs import c_apfs
@@ -11,6 +13,9 @@
     from dissect.apfs.objects.fs import FS
     from dissect.apfs.objects.keybag import ContainerKeybag
 
+log = logging.getLogger(__name__)
+log.setLevel(os.getenv("DISSECT_LOG_APFS", "CRITICAL"))
+
 
 class APFS:
     """Container class for APFS operations.
@@ -24,8 +29,29 @@ def __init__(self, fh: BinaryIO):
         self.fh.seek(0)
 
         self.sb = NxSuperblock.from_block(self, 0, self.fh.read(c_apfs.NX_DEFAULT_BLOCK_SIZE))
-        self.sbs = [self.sb] + [obj for obj in self.sb.checkpoint_objects if isinstance(obj, NxSuperblock)]
-        self.sb = sorted(self.sbs, key=lambda obj: obj.xid)[-1]
+        self.sb.check()
+
+        self.sbs = sorted(
+            [obj for obj in self.sb.checkpoint_objects if isinstance(obj, NxSuperblock)],
+            key=lambda obj: obj.xid,
+            reverse=True,
+        )
+
+        # TODO: Do more accurate checkpoint traversal
+        for sb in self.sbs:
+            try:
+                sb.check()
+                sb.compare(self.sb)
+            except Exception as e:
+                log.debug("Skipping superblock xid=%d: %s", sb.xid, e)
+                continue
+
+            if not sb.omap.is_valid():
+                log.debug("Skipping superblock xid=%d: invalid OMAP", sb.xid)
+                continue
+
+            self.sb = sb
+            break
 
     @property
     def block_size(self) -> int:

dissect/apfs/objects/nx_superblock.py

@@ -29,15 +29,41 @@ class NxSuperblock(Object):
     __struct__ = c_apfs.nx_superblock
     object: c_apfs.nx_superblock
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def check(self) -> None:
+        """Check the validity of the superblock."""
+        if not self.is_valid():
+            raise Error("Invalid nx_superblock checksum")
+
+        if self.type != c_apfs.OBJECT_TYPE.NX_SUPERBLOCK:
+            raise Error("Invalid nx_superblock type")
+
+        if not self.is_ephemeral:
+            raise Error("Invalid nx_superblock storage type")
 
         if self.object.nx_magic.to_bytes(4, "big") != c_apfs.NX_MAGIC:
             raise Error(
                 "Invalid nx_superblock magic "
                 f"(expected {c_apfs.NX_MAGIC!r}, got {self.object.nx_magic.to_bytes(4, 'big')!r})"
             )
 
+    def compare(self, other: NxSuperblock) -> None:
+        """Compare this superblock to another superblock."""
+        if self.header.o_xid < other.header.o_xid:
+            raise Error("Lower xid than other superblock")
+
+        for attr in (
+            "nx_uuid",
+            "nx_fusion_uuid",
+            "nx_block_size",
+            "nx_block_count",
+            "nx_xp_desc_blocks",
+            "nx_xp_data_blocks",
+            "nx_xp_desc_base",
+            "nx_xp_data_base",
+        ):
+            if getattr(self.object, attr) != getattr(other.object, attr):
+                raise Error(f"Mismatch on {attr}")
+
     @cached_property
     def block_size(self) -> int:
         """The block size of the container."""
@@ -66,14 +92,20 @@ def uuid(self) -> UUID:
     @cached_property
     def checkpoint_objects(self) -> list[CheckpointMap | NxSuperblock]:
         """All checkpoint objects in the container."""
-        return list(_read_checkpoint_objects(self.container, self.object.nx_xp_desc_base, self.object.nx_xp_desc_len))
+        # TODO: Rework this a bit to be more accurate
+        return list(
+            _read_checkpoint_objects(self.container, self.object.nx_xp_desc_base, self.object.nx_xp_desc_blocks)
+        )
 
     @cached_property
     def ephemeral_objects(self) -> dict[int, Object]:
         """All ephemeral objects in the container."""
+        # TODO: I don't think this is correct
         return {
             obj.oid: obj
-            for obj in _read_checkpoint_objects(self.container, self.object.nx_xp_data_base, self.object.nx_xp_data_len)
+            for obj in _read_checkpoint_objects(
+                self.container, self.object.nx_xp_data_base, self.object.nx_xp_data_blocks
+            )
         }
 
     @cached_property

dissect/apfs/objects/omap.py

@@ -20,6 +20,15 @@ def __init__(self, *args, **kwargs):
 
         self.lookup = lru_cache(128)(self.lookup)
 
+    def is_valid(self) -> bool:
+        return (
+            super().is_valid()
+            and self.type == c_apfs.OBJECT_TYPE.OMAP
+            and self.subtype == 0
+            and self.is_physical
+            and self.oid == self.address
+        )
+
     @cached_property
     def btree(self) -> BTree:
         """The B-tree of the object map."""

tests/_data/corrupt.bin.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:672c556cb370c4ae547e69b822d1896a2327fcb392eeba19f3d73414382cc80f
+size 595750

tests/test_apfs.py

@@ -6,7 +6,7 @@
 
 import pytest
 
-from dissect.apfs.apfs import APFS
+from dissect.apfs.apfs import APFS, log
 from dissect.apfs.c_apfs import c_apfs
 from tests.conftest import absolute_path
 
@@ -357,3 +357,15 @@ def test_snapshots() -> None:
         for i, snapshot in enumerate(volume.snapshots):
             assert snapshot.name == f"Snapshot {i}"
             assert snapshot.open().get("file").open().read() == f"Snapshot {i}\n".encode()
+
+
+def test_corrupt_checkpoints(caplog: pytest.LogCaptureFixture) -> None:
+    """Test APFS volumes with corrupt checkpoints."""
+    with gzip.open(absolute_path("_data/corrupt.bin.gz"), "rb") as fh, caplog.at_level("DEBUG", log.name):
+        container = APFS(fh)
+
+        assert container.sb.xid == 302
+        assert len(container.volumes) == 1
+
+    assert caplog.messages[0] == "Skipping superblock xid=304: invalid OMAP"
+    assert caplog.messages[1] == "Skipping superblock xid=303: Invalid nx_superblock checksum"