refactor(core): prepare fw for differently sized fw chunks

[no changelog]

core/embed/lib/image.c

@@ -224,12 +224,13 @@ secbool check_image_contents(const image_header *const hdr, uint32_t firstskip,
     return secfalse;
   }
 
-  const void *data =
+  int remaining = hdr->codelen;
-      flash_area_get_address(area, firstskip, IMAGE_CHUNK_SIZE - firstskip);
+
+  const void *data = flash_area_get_address(
+      area, firstskip, MIN(remaining, IMAGE_CHUNK_SIZE - firstskip));
   if (!data) {
     return secfalse;
   }
-  int remaining = hdr->codelen;
   if (sectrue !=
       check_single_hash(hdr->hashes, data,
                         MIN(remaining, IMAGE_CHUNK_SIZE - firstskip))) {
@@ -242,7 +243,7 @@ secbool check_image_contents(const image_header *const hdr, uint32_t firstskip,
 
   while (remaining > 0) {
     data = flash_area_get_address(area, chunk * IMAGE_CHUNK_SIZE,
-                                  IMAGE_CHUNK_SIZE);
+                                  MIN(remaining, IMAGE_CHUNK_SIZE));
     if (!data) {
       return secfalse;
     }

core/embed/lib/image.h

@@ -22,11 +22,11 @@
 
 #include <stdint.h>
 #include "flash.h"
+#include "model.h"
 #include "secbool.h"
 
 #define IMAGE_HEADER_SIZE 0x400  // size of the bootloader or firmware header
 #define IMAGE_SIG_SIZE 65
-#define IMAGE_CHUNK_SIZE (128 * 1024)
 #define IMAGE_INIT_CHUNK_SIZE (16 * 1024)
 
 #define BOOTLOADER_IMAGE_MAGIC 0x425A5254  // TRZB

core/embed/models/model_D001.h

@@ -22,6 +22,7 @@
 #define BOOTLOADER_START 0x08020000
 #define FIRMWARE_START 0x08040000
 
+#define IMAGE_CHUNK_SIZE (128 * 1024)
 #define BOOTLOADER_IMAGE_MAXSIZE (128 * 1024 * 1)  // 128 KB
 #define FIRMWARE_IMAGE_MAXSIZE (128 * 1024 * 13)   // 1664 KB
 #define NORCOW_SECTOR_SIZE (64 * 1024)

core/embed/models/model_T1B1.h

@@ -10,6 +10,7 @@
 #define BOOTLOADER_START 0x08000000
 #define FIRMWARE_START 0x08010000
 
+#define IMAGE_CHUNK_SIZE (64 * 1024)
 #define BOOTLOADER_IMAGE_MAXSIZE (32 * 1024 * 1)  // 32 KB
 #define FIRMWARE_IMAGE_MAXSIZE (64 * 1024 * 15)   // 960 KB
 #define NORCOW_SECTOR_SIZE (16 * 1024)

core/embed/models/model_T2B1.h

@@ -22,6 +22,7 @@
 #define BOOTLOADER_START 0x08020000
 #define FIRMWARE_START 0x08040000
 
+#define IMAGE_CHUNK_SIZE (128 * 1024)
 #define BOOTLOADER_IMAGE_MAXSIZE (128 * 1024 * 1)  // 128 KB
 #define FIRMWARE_IMAGE_MAXSIZE (128 * 1024 * 13)   // 1664 KB
 #define NORCOW_SECTOR_SIZE (64 * 1024)

core/embed/models/model_T2T1.h

@@ -22,6 +22,7 @@
 #define BOOTLOADER_START 0x08020000
 #define FIRMWARE_START 0x08040000
 
+#define IMAGE_CHUNK_SIZE (128 * 1024)
 #define BOOTLOADER_IMAGE_MAXSIZE (128 * 1024 * 1)  // 128 KB
 #define FIRMWARE_IMAGE_MAXSIZE (128 * 1024 * 13)   // 1664 KB
 #define NORCOW_SECTOR_SIZE (64 * 1024)

python/src/trezorlib/firmware/consts.py

@@ -18,10 +18,6 @@ from . import models
 
 V1_SIGNATURE_SLOTS = 3
 
-ONEV2_CHUNK_SIZE = 1024 * 64
-V2_CHUNK_SIZE = 1024 * 128
-
-
 # === KEYS KEPT FOR COMPATIBILITY ===
 # use `trezorlib.firmware.models` directly
 

python/src/trezorlib/firmware/core.py

@@ -14,7 +14,6 @@
 # You should have received a copy of the License along with this library.
 # If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
 
-import hashlib
 import typing as t
 from copy import copy
 from enum import Enum
@@ -117,11 +116,8 @@ class FirmwareImage(Struct):
         c.Terminated,
     )
 
-    HASH_PARAMS = util.FirmwareHashParameters(
+    def get_hash_params(self) -> "util.FirmwareHashParameters":
-        hash_function=hashlib.blake2s,
+        return Model.from_hw_model(self.header.hw_model).hash_params()
-        chunk_size=consts.V2_CHUNK_SIZE,
-        padding_byte=None,
-    )
 
     def code_hashes(self) -> t.List[bytes]:
         """Calculate hashes of chunks of `code`.
@@ -129,26 +125,23 @@ class FirmwareImage(Struct):
         Assume that the first `code_offset` bytes of `code` are taken up by the header.
         """
         hashes = []
+
+        hash_params = self.get_hash_params()
+
         # End offset for each chunk. Normally this would be (i+1)*chunk_size for i-th chunk,
         # but the first chunk is shorter by code_offset, so all end offsets are shifted.
-        ends = [
+        ends = [(i + 1) * hash_params.chunk_size - self._code_offset for i in range(16)]
-            (i + 1) * self.HASH_PARAMS.chunk_size - self._code_offset for i in range(16)
-        ]
         start = 0
         for end in ends:
             chunk = self.code[start:end]
             # padding for last non-empty chunk
-            if (
+            if hash_params.padding_byte is not None and start < len(self.code) < end:
-                self.HASH_PARAMS.padding_byte is not None
+                chunk += hash_params.padding_byte[0:1] * (end - start - len(chunk))
-                and start < len(self.code)
-                and end > len(self.code)
-            ):
-                chunk += self.HASH_PARAMS.padding_byte[0:1] * (end - start - len(chunk))
 
             if not chunk:
                 hashes.append(b"\0" * 32)
             else:
-                hashes.append(self.HASH_PARAMS.hash_function(chunk).digest())
+                hashes.append(hash_params.hash_function(chunk).digest())
 
             start = end
 
@@ -159,13 +152,16 @@ class FirmwareImage(Struct):
             raise util.FirmwareIntegrityError("Invalid firmware data.")
 
     def digest(self) -> bytes:
+
+        hash_params = self.get_hash_params()
+
         header = copy(self.header)
         header.hashes = self.code_hashes()
         header.signature = b"\x00" * 64
         header.sigmask = 0
         header.v1_key_indexes = [0] * consts.V1_SIGNATURE_SLOTS
         header.v1_signatures = [b"\x00" * 64] * consts.V1_SIGNATURE_SLOTS
-        return self.HASH_PARAMS.hash_function(header.build()).digest()
+        return hash_params.hash_function(header.build()).digest()
 
 
 class VendorFirmware(Struct):

python/src/trezorlib/firmware/legacy.py

@@ -24,6 +24,7 @@ from construct_classes import Struct, subcon
 
 from . import consts, models, util
 from .core import FirmwareImage
+from .models import Model
 
 __all__ = [
     "LegacyFirmware",
@@ -98,14 +99,11 @@ def check_sig_signmessage(
 class LegacyV2Firmware(FirmwareImage):
     """Firmware image in the format used by Trezor One 1.8.0 and newer."""
 
-    HASH_PARAMS = util.FirmwareHashParameters(
-        hash_function=hashlib.sha256,
-        chunk_size=consts.ONEV2_CHUNK_SIZE,
-        padding_byte=b"\xff",
-    )
-
     V3_FIRST_VERSION = (1, 12, 0)
 
+    def get_hash_params(self) -> "util.FirmwareHashParameters":
+        return Model.ONE.hash_params()
+
     def verify_v2(self, dev_keys: bool) -> None:
         if not dev_keys:
             public_keys = models.LEGACY_V1V2.firmware_keys

python/src/trezorlib/firmware/models.py

@@ -14,10 +14,13 @@
 # You should have received a copy of the License along with this library.
 # If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
 
+import hashlib
 import typing as t
 from dataclasses import dataclass
 from enum import Enum
 
+from .util import FirmwareHashParameters
+
 if t.TYPE_CHECKING:
     from typing_extensions import Self
 
@@ -26,12 +29,13 @@ class Model(Enum):
     T1B1 = b"T1B1"
     T2T1 = b"T2T1"
     T2B1 = b"T2B1"
-    DISC1 = b"D001"
+    D001 = b"D001"
 
     # legacy aliases
     ONE = T1B1
     T = T2T1
     R = T2B1
+    DISC1 = D001
 
     @classmethod
     def from_hw_model(cls, hw_model: t.Union["Self", bytes]) -> "Self":
@@ -48,6 +52,9 @@ class Model(Enum):
             model_map = MODEL_MAP
         return model_map[self]
 
+    def hash_params(self) -> "FirmwareHashParameters":
+        return MODEL_HASH_PARAMS_MAP[self]
+
 
 @dataclass
 class ModelKeys:
@@ -206,18 +213,38 @@ T2B1 = ModelKeys(
     firmware_sigs_needed=-1,
 )
 
+
+LEGACY_HASH_PARAMS = FirmwareHashParameters(
+    hash_function=hashlib.sha256,
+    chunk_size=1024 * 64,
+    padding_byte=b"\xff",
+)
+
+T2T1_HASH_PARAMS = FirmwareHashParameters(
+    hash_function=hashlib.blake2s,
+    chunk_size=1024 * 128,
+    padding_byte=None,
+)
+
 MODEL_MAP = {
     Model.T1B1: LEGACY_V3,
     Model.T2T1: T2T1,
     Model.T2B1: T2B1,
-    Model.DISC1: TREZOR_CORE_DEV,
+    Model.D001: TREZOR_CORE_DEV,
 }
 
 MODEL_MAP_DEV = {
     Model.T1B1: LEGACY_V3_DEV,
     Model.T2T1: TREZOR_CORE_DEV,
     Model.T2B1: TREZOR_CORE_DEV,
-    Model.DISC1: TREZOR_CORE_DEV,
+    Model.D001: TREZOR_CORE_DEV,
+}
+
+MODEL_HASH_PARAMS_MAP = {
+    Model.T1B1: LEGACY_HASH_PARAMS,
+    Model.T2T1: T2T1_HASH_PARAMS,
+    Model.T2B1: T2T1_HASH_PARAMS,
+    Model.D001: T2T1_HASH_PARAMS,
 }
 
 # aliases
@@ -234,3 +261,5 @@ TREZOR_R_DEV = TREZOR_CORE_DEV
 
 DISC1 = TREZOR_CORE_DEV
 DISC1_DEV = TREZOR_CORE_DEV
+D001 = TREZOR_CORE_DEV
+D001_DEV = TREZOR_CORE_DEV