refactor(core/prodtest): unify binary update function

diff --git a/core/embed/projects/prodtest/cmd/common.c b/core/embed/projects/prodtest/cmd/common.c
index 82968a6b..6f0ff457 100644
--- a/core/embed/projects/prodtest/cmd/common.c
+++ b/core/embed/projects/prodtest/cmd/common.c
@@ -1,3 +1,25 @@
+/*
+ * This file is part of the Trezor project, https://trezor.io/
+ *
+ * Copyright (c) SatoshiLabs
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#include <trezor_model.h>
+#include <trezor_rtl.h>
+
 #include "common.h"
 
 #include "buffer.h"
@@ -432,3 +454,98 @@ bool check_cert_chain(cli_t* cli, const uint8_t* chain, size_t chain_size,
 
   return true;
 }
+
+#ifdef USE_NRF
+#define BINARY_MAXSIZE \
+  (0x50000 > BOOTLOADER_MAXSIZE ? 0x50000 : BOOTLOADER_MAXSIZE)
+#else
+#define BINARY_MAXSIZE BOOTLOADER_MAXSIZE
+#endif
+
+__attribute__((section(".buf"))) static uint8_t binary_buffer[BINARY_MAXSIZE];
+static size_t binary_len = 0;
+static bool binary_update_in_progress = false;
+
+void binary_update(cli_t* cli, bool (*finalize)(uint8_t* data, size_t len)) {
+  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);
+      goto cleanup;
+    }
+
+    // Reset our state
+    binary_len = 0;
+    binary_update_in_progress = true;
+    cli_trace(cli, "Begin");
+    cli_ok(cli, "");
+
+  } else if (0 == strcmp(phase, "chunk")) {
+    if (cli_arg_count(cli) < 2) {
+      cli_error_arg_count(cli);
+      goto cleanup;
+    }
+
+    if (!binary_update_in_progress) {
+      cli_error(cli, CLI_ERROR, "Update not started. Use 'begin' first.");
+      goto cleanup;
+    }
+
+    // Receive next piece of the image
+    size_t chunk_len = 0;
+
+    if (!cli_arg_hex(cli, "hex-data", &binary_buffer[binary_len],
+                     sizeof(binary_buffer) - binary_len, &chunk_len)) {
+      cli_error_arg(cli, "Expecting hex data for chunk.");
+      goto cleanup;
+    }
+
+    binary_len += chunk_len;
+
+    cli_ok(cli, "%u %u", (unsigned)chunk_len, (unsigned)binary_len);
+
+  } else if (0 == strcmp(phase, "end")) {
+    if (cli_arg_count(cli) != 1) {
+      cli_error_arg_count(cli);
+      goto cleanup;
+    }
+
+    if (binary_len == 0) {
+      cli_error(cli, CLI_ERROR, "No data received");
+      goto cleanup;
+    }
+
+    if (!finalize(binary_buffer, binary_len)) {
+      binary_len = 0;
+      cli_error(cli, CLI_ERROR, "Error while finalizing the update");
+      goto cleanup;
+    }
+
+    cli_trace(cli, "Update successful (%u bytes)", (unsigned)binary_len);
+    cli_ok(cli, "");
+
+    // Reset state so next begin must come before chunks
+    binary_len = 0;
+    binary_update_in_progress = false;
+
+  } else {
+    cli_error(cli, CLI_ERROR, "Unknown phase '%s' (begin|chunk|end)", phase);
+    goto cleanup;
+  }
+
+  return;
+
+cleanup:
+  binary_update_in_progress = false;
+  binary_len = 0;
+}
diff --git a/core/embed/projects/prodtest/cmd/common.h b/core/embed/projects/prodtest/cmd/common.h
index 9efb916c..0b1f8e0a 100644
--- a/core/embed/projects/prodtest/cmd/common.h
+++ b/core/embed/projects/prodtest/cmd/common.h
@@ -17,8 +17,7 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-#ifndef __COMMON_H
-#define __COMMON_H
+#pragma once
 
 #include <rtl/cli.h>
 
@@ -28,4 +27,4 @@ bool check_cert_chain(cli_t* cli, const uint8_t* chain, size_t chain_size,
                       const uint8_t* sig, size_t sig_size,
                       const uint8_t challenge[CHALLENGE_SIZE]);
 
-#endif  // __COMMON_H
+void binary_update(cli_t* cli, bool (*finalize)(uint8_t* data, size_t len));
diff --git a/core/embed/projects/prodtest/cmd/prodtest_bootloader.c b/core/embed/projects/prodtest/cmd/prodtest_bootloader.c
index 7dc575a4..37d96eab 100644
--- a/core/embed/projects/prodtest/cmd/prodtest_bootloader.c
+++ b/core/embed/projects/prodtest/cmd/prodtest_bootloader.c
@@ -26,6 +26,8 @@
 #include <util/boot_image.h>
 #include <util/image.h>
 
+#include "common.h"
+
 #ifdef USE_BOOT_UCB
 #include <util/boot_header.h>
 #endif
@@ -79,9 +81,6 @@ static void prodtest_bootloader_version(cli_t *cli) {
 }
 
 #ifndef TREZOR_MODEL_T2T1
-__attribute__((
-    section(".buf"))) static uint8_t bootloader_buffer[BOOTLOADER_MAXSIZE];
-static size_t bootloader_len = 0;
 
 #if USE_BOOT_UCB
 // Writes boot header and bootloader code to the BOOTUPDATE_AREA
@@ -110,104 +109,38 @@ static bool write_to_bootupdate_area(const uint8_t *data, size_t size) {
 }
 #endif  // USE_BOOT_UCB
 
-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_trace(cli, "Begin");
-    cli_ok(cli, "");
-
-  } 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, "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, "%u %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, "No data received");
-      return;
-    }
-
+static bool prodtest_bootloader_update_finalize(uint8_t *data, size_t len) {
 #if USE_BOOT_UCB
-    mpu_mode_t mpu_mode = mpu_reconfig(MPU_MODE_BOOTUPDATE);
+  mpu_mode_t mpu_mode = mpu_reconfig(MPU_MODE_BOOTUPDATE);
 
-    if (!write_to_bootupdate_area(bootloader_buffer, bootloader_len)) {
-      mpu_restore(mpu_mode);
-      cli_error(cli, CLI_ERROR, "Failed to flash bootloader");
-      return;
-    }
+  if (!write_to_bootupdate_area(data, len)) {
+    mpu_restore(mpu_mode);
+    return false;
+  }
 
-    boot_image_t bootloader_image = {
-        .image_ptr = (const void *)BOOTUPDATE_START,
-        .image_size = bootloader_len,
-    };
+  boot_image_t bootloader_image = {
+      .image_ptr = (const void *)BOOTUPDATE_START,
+      .image_size = len,
+  };
 
-    boot_image_replace(&bootloader_image);
+  boot_image_replace(&bootloader_image);
 
-    mpu_restore(mpu_mode);
+  mpu_restore(mpu_mode);
 
 #else
-    boot_image_t bootloader_image = {
-        .image_ptr = bootloader_buffer,
-        .image_size = bootloader_len,
-    };
+  boot_image_t bootloader_image = {
+      .image_ptr = data,
+      .image_size = len,
+  };
 
-    boot_image_replace(&bootloader_image);
+  boot_image_replace(&bootloader_image);
 #endif
 
-    // Reset state so next begin must come before chunks
-    bootloader_len = 0;
-
-    cli_trace(cli, "Update successful (%u bytes)", (unsigned)bootloader_len);
-    cli_ok(cli, "");
+  return true;
+}
 
-  } else {
-    cli_error(cli, CLI_ERROR, "Unknown phase '%s' (begin|chunk|end)", phase);
-  }
+static void prodtest_bootloader_update(cli_t *cli) {
+  binary_update(cli, prodtest_bootloader_update_finalize);
 }
 #endif
 
diff --git a/core/embed/projects/prodtest/cmd/prodtest_nrf.c b/core/embed/projects/prodtest/cmd/prodtest_nrf.c
index 012e1d18..f2999f62 100644
--- a/core/embed/projects/prodtest/cmd/prodtest_nrf.c
+++ b/core/embed/projects/prodtest/cmd/prodtest_nrf.c
@@ -26,6 +26,7 @@
 #include <rtl/cli.h>
 #include <util/flash_otp.h>
 
+#include "common.h"
 #include "prodtest_optiga.h"
 
 static void prodtest_nrf_communication(cli_t* cli) {
@@ -73,105 +74,12 @@ static void prodtest_nrf_version(cli_t* cli) {
          info.version_patch, info.version_tweak);
 }
 
-#define NRF_UPDATE_MAXSIZE 0x50000
-__attribute__((section(".buf"))) static uint8_t nrf_buffer[NRF_UPDATE_MAXSIZE];
+static bool prodtest_nrf_update_finalize(uint8_t* data, size_t len) {
+  return nrf_update(data, len);
+}
 
 static void prodtest_nrf_update(cli_t* cli) {
-  static size_t nrf_len = 0;
-  static bool nrf_update_in_progress = false;
-
-  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).");
-    return;
-  }
-
-  if (0 == strcmp(phase, "begin")) {
-    if (cli_arg_count(cli) != 1) {
-      cli_error_arg_count(cli);
-      goto cleanup;
-    }
-
-    // Reset our state
-    nrf_len = 0;
-    nrf_update_in_progress = true;
-    cli_trace(cli, "begin");
-    cli_ok(cli, "");
-
-  } else if (0 == strcmp(phase, "chunk")) {
-    if (cli_arg_count(cli) < 2) {
-      cli_error_arg_count(cli);
-      goto cleanup;
-    }
-
-    if (!nrf_update_in_progress) {
-      cli_error(cli, CLI_ERROR, "Update not started. Use 'begin' first.");
-      goto cleanup;
-    }
-
-    // Receive next piece of the image
-    size_t chunk_len = 0;
-    uint8_t chunk_buf[512];  // tune this if you like
-
-    if (!cli_arg_hex(cli, "hex-data", chunk_buf, sizeof(chunk_buf),
-                     &chunk_len)) {
-      cli_error_arg(cli, "Expecting hex-data for chunk.");
-      goto cleanup;
-    }
-
-    if (nrf_len + chunk_len > NRF_UPDATE_MAXSIZE) {
-      cli_error(cli, CLI_ERROR, "Buffer overflow (have %u, need %u)",
-                (unsigned)nrf_len, (unsigned)chunk_len);
-      goto cleanup;
-    }
-
-    memcpy(&nrf_buffer[nrf_len], chunk_buf, chunk_len);
-    nrf_len += chunk_len;
-
-    cli_ok(cli, "%u %u", (unsigned)chunk_len, (unsigned)nrf_len);
-
-  } else if (0 == strcmp(phase, "end")) {
-    if (cli_arg_count(cli) != 1) {
-      cli_error_arg_count(cli);
-      goto cleanup;
-    }
-
-    if (!nrf_update_in_progress) {
-      cli_error(cli, CLI_ERROR, "Update not started. Use 'begin' first.");
-      goto cleanup;
-    }
-
-    if (nrf_len == 0) {
-      cli_error(cli, CLI_ERROR, "No data received");
-      goto cleanup;
-    }
-
-    // Hand off to your firmware update routine
-    if (!nrf_update(nrf_buffer, nrf_len)) {
-      cli_error(cli, CLI_ERROR, "Update failed");
-      goto cleanup;
-    }
-
-    // Clear state so next begin is required
-    nrf_len = 0;
-    nrf_update_in_progress = false;
-
-    cli_trace(cli, "Update successful");
-    cli_ok(cli, "");
-
-  } else {
-    cli_error(cli, CLI_ERROR, "Unknown phase '%s' (begin|chunk|end)", phase);
-  }
-
-  return;
-
-cleanup:
-  nrf_update_in_progress = false;
+  binary_update(cli, prodtest_nrf_update_finalize);
 }
 
 static void prodtest_nrf_pair(cli_t* cli) {