feat(core/prodtest): add command for bootloader update, along with script for sending the data

Author: TychoVrahe

core/SConscript.prodtest

@@ -66,6 +66,7 @@ SOURCE_MOD_CRYPTO += [
     'vendor/trezor-crypto/bignum.c',
     'vendor/trezor-crypto/blake256.c',
     'vendor/trezor-crypto/blake2b.c',
+    'vendor/trezor-crypto/blake2s.c',
     'vendor/trezor-crypto/buffer.c',
     'vendor/trezor-crypto/chacha_drbg.c',
     'vendor/trezor-crypto/chacha20poly1305/chacha_merged.c',
@@ -123,6 +124,7 @@ SOURCE_MOD += [
     'embed/gfx/gfx_draw.c',
     'embed/gfx/terminal.c',
     'embed/io/display/display_utils.c',
+    'embed/util/bl_check/bl_check.c',
     'embed/util/image/image.c',
     'embed/util/rsod/rsod.c',
     'embed/util/scm_revision/scm_revision.c',
@@ -205,6 +207,7 @@ ALLPATHS = [
         'embed/models',
         'embed/gfx/inc',
         'embed/sys/bsp/inc',
+        'embed/util/bl_check/inc',
         'embed/util/image/inc',
         'embed/util/rsod/inc',
         'embed/util/scm_revision/inc',

core/SConscript.prodtest_emu

@@ -58,6 +58,7 @@ SOURCE_MOD_CRYPTO += [
     'vendor/trezor-crypto/bignum.c',
     'vendor/trezor-crypto/blake256.c',
     'vendor/trezor-crypto/blake2b.c',
+    'vendor/trezor-crypto/blake2s.c',
     'vendor/trezor-crypto/buffer.c',
     'vendor/trezor-crypto/chacha_drbg.c',
     'vendor/trezor-crypto/chacha20poly1305/chacha_merged.c',
@@ -115,6 +116,7 @@ SOURCE_MOD += [
     'embed/gfx/gfx_draw.c',
     'embed/gfx/terminal.c',
     'embed/io/display/display_utils.c',
+    'embed/util/bl_check/bl_check.c',
     'embed/util/image/image.c',
     'embed/util/rsod/rsod.c',
     'embed/util/scm_revision/scm_revision.c',
@@ -194,6 +196,7 @@ ALLPATHS = ['embed/rust',
            'embed/models',
            'embed/projects/unix',
            'embed/gfx/inc',
+           'embed/util/bl_check/inc',
            'embed/util/image/inc',
            'embed/util/rsod/inc',
            'embed/util/scm_revision/inc',

core/embed/projects/prodtest/.changelog.d/5227.added

@@ -0,0 +1 @@
+Added command for updating bootloader.

core/embed/projects/prodtest/README.md

@@ -115,7 +115,7 @@ OK
 
 
 ### reboot-to-bootloader
-This command initiates device reboot to bootlaoder.
+This command initiates device reboot to bootloader.
 
 Example:
 ```
@@ -134,14 +134,18 @@ OK 0.2.6
 ```
 
 ### bootloader-version
-Retrives the version of the bootlaoder. The command returns `OK` followed by the version in the format `<major>.<minor>.<patch>`.
+Retrieves the version of the bootloader. The command returns `OK` followed by the version in the format `<major>.<minor>.<patch>`.
 
 Example:
 ```
 bootloader-version
 OK 2.1.7
 ```
 
+### bootloader-update
+Updates the bootloader to the supplied binary file.
+
+
 ### ble-adv-start
 Starts BLE advertising. Accepts one parameter, advertising name. The command returns `OK` if the operation is successful.
 

core/embed/projects/prodtest/cmd/prodtest_bootloader.c

@@ -18,9 +18,12 @@
  */
 
 #include <trezor_rtl.h>
+#include <trezor_model.h>
 
 #include <rtl/cli.h>
 #include <sys/mpu.h>
+#include <util/bl_check.h>
+#include <util/board_capabilities.h>
 #include <util/image.h>
 
 static void prodtest_bootloader_version(cli_t *cli) {
@@ -51,6 +54,88 @@ static void prodtest_bootloader_version(cli_t *cli) {
   cli_ok(cli, "%d.%d.%d", v & 0xFF, (v >> 8) & 0xFF, (v >> 16) & 0xFF);
 }
 
+#ifndef TREZOR_MODEL_T2T1
+__attribute__((
+    section(".buf"))) static uint8_t bootloader_buffer[BOOTLOADER_MAXSIZE];
+static size_t bootloader_len = 0;
+
+static void prodtest_bootloader_update(cli_t *cli) {
+  if (cli_arg_count(cli) < 1) {
+    cli_error_arg_count(cli);
+    return;
+  }
+
+  const char *phase = cli_arg(cli, "phase");
+
+  if (phase == NULL) {
+    cli_error_arg(cli, "Expecting phase (begin|chunk|end).");
+  }
+
+  if (0 == strcmp(phase, "begin")) {
+    if (cli_arg_count(cli) != 1) {
+      cli_error_arg_count(cli);
+      return;
+    }
+
+    // Reset our state
+    bootloader_len = 0;
+    cli_ok(cli, "BOOT: begin");
+
+  } else if (0 == strcmp(phase, "chunk")) {
+    if (cli_arg_count(cli) < 2) {
+      cli_error_arg_count(cli);
+      return;
+    }
+
+    // Receive next piece of the image
+    size_t chunk_len = 0;
+    // Temporary buffer for this chunk; tweak max if you like
+    uint8_t chunk_buf[1024];
+
+    if (!cli_arg_hex(cli, "hex-data", chunk_buf, sizeof(chunk_buf),
+                     &chunk_len)) {
+      cli_error_arg(cli, "Expecting hex data for chunk.");
+      return;
+    }
+
+    if (bootloader_len + chunk_len > BOOTLOADER_MAXSIZE) {
+      cli_error(cli, CLI_ERROR, "BOOT: buffer overflow (have %u, %u more)",
+                (unsigned)bootloader_len, (unsigned)chunk_len);
+      return;
+    }
+
+    memcpy(&bootloader_buffer[bootloader_len], chunk_buf, chunk_len);
+    bootloader_len += chunk_len;
+
+    cli_ok(cli, "BOOT: got %u bytes (total %u)", (unsigned)chunk_len,
+           (unsigned)bootloader_len);
+
+  } else if (0 == strcmp(phase, "end")) {
+    if (cli_arg_count(cli) != 1) {
+      cli_error_arg_count(cli);
+      return;
+    }
+
+    if (bootloader_len == 0) {
+      cli_error(cli, CLI_ERROR, "BOOT: no data received");
+      return;
+    }
+
+    parse_boardloader_capabilities();
+    bl_check_replace(bootloader_buffer, bootloader_len);
+
+    cli_ok(cli, "BOOT: update successful (%u bytes)", (unsigned)bootloader_len);
+
+    // Optionally reset state so next begin must come before chunks
+    bootloader_len = 0;
+
+  } else {
+    cli_error(cli, CLI_ERROR, "BOOT: unknown phase '%s' (begin|chunk|end)",
+              phase);
+  }
+}
+#endif
+
 // clang-format off
 
 PRODTEST_CLI_CMD(
@@ -59,3 +144,13 @@ PRODTEST_CLI_CMD(
   .info = "Retrieve the bootloader version",
   .args = ""
 );
+
+#ifndef TREZOR_MODEL_T2T1
+PRODTEST_CLI_CMD(
+  .name = "bootloader-update",
+  .func = prodtest_bootloader_update,
+  .info = "Update bootloader",
+  .args = "<phase> <hex-data>"
+);
+#endif
+

core/embed/sys/linker/stm32u58/prodtest.ld

@@ -5,6 +5,7 @@ MEMORY {
 
   MAIN_RAM (wal) : ORIGIN = MAIN_RAM_START, LENGTH = MAIN_RAM_SIZE
   AUX1_RAM (wal) : ORIGIN = AUX1_RAM_START, LENGTH = AUX1_RAM_SIZE
+  AUX2_RAM (wal) : ORIGIN = AUX2_RAM_START, LENGTH = AUX2_RAM_SIZE
   BOOT_ARGS (wal) : ORIGIN = BOOTARGS_START, LENGTH = BOOTARGS_SIZE
   FB1_RAM (wal) : ORIGIN = FB1_RAM_START, LENGTH = FB1_RAM_SIZE
   FB2_RAM (wal) : ORIGIN = FB2_RAM_START, LENGTH = FB2_RAM_SIZE
@@ -77,7 +78,7 @@ SECTIONS {
     . = ALIGN(4);
     *(.no_dma_buffers*);
     . = ALIGN(4);
-  } >AUX1_RAM
+  } >AUX2_RAM
 
   .stack : ALIGN(8) {
     . = 16K; /* Overflow causes UsageFault */

core/tools/bld_update.py

@@ -0,0 +1,86 @@
+#!/usr/bin/env python3
+import time
+from pathlib import Path
+import sys
+import zlib
+
+import click
+import serial
+
+def _compress(data: bytes) -> bytes:
+    """
+    Compress data with zlib at max compression, raw deflate.
+    """
+    compressor = zlib.compressobj(level=9, wbits=-10)
+    return compressor.compress(data) + compressor.flush()
+
+def send_cmd(ser, cmd, expect_ok=True):
+    """Send a line, read response, and abort on CLI_ERROR."""
+    ser.write((cmd + "\r\n").encode())
+    # Give the device a moment to process
+    time.sleep(0.05)
+    resp = ser.readline().decode(errors='ignore').strip()
+    click.echo(f"> {cmd}")
+    click.echo(f"< {resp}")
+    if expect_ok and resp.startswith("CLI_ERROR"):
+        click.echo("Error from device, aborting.", err=True)
+        sys.exit(1)
+    return resp
+
+def upload_bootloader(port, bin_path, chunk_size):
+    # Read binary file
+    data = Path(bin_path).read_bytes()
+    orig_size = len(data)
+    click.echo(f"Read {orig_size} bytes from {bin_path!r}")
+
+    # Compress the data
+    comp_data = _compress(data)
+    comp_size = len(comp_data)
+    click.echo(f"Compressed to {comp_size} bytes ({comp_size*100//orig_size}% of original)")
+
+    # Open USB-VCP port
+    ser = serial.Serial(port, timeout=2)
+    time.sleep(0.1)
+    ser.reset_input_buffer()
+    ser.reset_output_buffer()
+
+    # 1) Begin transfer
+    send_cmd(ser, "bootloader-update begin")
+
+    # 2) Stream compressed chunks
+    offset = 0
+    while offset < comp_size:
+        chunk = comp_data[offset:offset+chunk_size]
+        hexstr = chunk.hex()
+        send_cmd(ser, f"bootloader-update chunk {hexstr}")
+        offset += len(chunk)
+        pct = offset * 100 // comp_size
+        click.echo(f"  Uploaded {offset}/{comp_size} bytes ({pct}%)")
+
+    # 3) Finish transfer
+    send_cmd(ser, "bootloader-update end")
+    click.echo("Bootloader upload complete.")
+
+@click.command(context_settings={"help_option_names": ["-h", "--help"]})
+@click.argument("port", metavar="<serial-port>")
+@click.argument("binary", metavar="<bootloader-binary>", type=click.Path(exists=True, dir_okay=False))
+@click.option("--chunk-size", "-c", default=512, show_default=True,
+              help="Max bytes per chunk (in compressed form)")
+def main(port, binary, chunk_size):
+    """
+    Upload a (compressed) bootloader image via USB-VCP CLI.
+
+    <serial-port> e.g. /dev/ttyUSB0 or COM3
+    <bootloader-binary> path to the .bin file
+    """
+    try:
+        upload_bootloader(port, binary, chunk_size)
+    except serial.SerialException as e:
+        click.echo(f"Serial error: {e}", err=True)
+        sys.exit(1)
+    except KeyboardInterrupt:
+        click.echo("Interrupted by user.", err=True)
+        sys.exit(1)
+
+if __name__ == "__main__":
+    main()