From 2d5d7201d591c5fc08dd4821195072390654ef0f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Vejpustek?= <ondrej.vejpustek@satoshilabs.com>
Date: Wed, 2 Oct 2024 18:06:32 +0200
Subject: [PATCH] feat(core): allow benchmark only in debug

[no changelog]
---
 core/SConscript.firmware            | 4 ++++
 core/SConscript.unix                | 4 ++++
 core/src/apps/benchmark/__init__.py | 4 ++++
 core/src/apps/workflow_handlers.py  | 4 ++--
 4 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/core/SConscript.firmware b/core/SConscript.firmware
index 7a8185ccfb..4253dc087c 100644
--- a/core/SConscript.firmware
+++ b/core/SConscript.firmware
@@ -19,6 +19,10 @@ NEW_RENDERING = ARGUMENTS.get('NEW_RENDERING', '1') == '1'
 MODEL_IDENTIFIER = models.get_model_identifier(TREZOR_MODEL)
 BENCHMARK = ARGUMENTS.get('BENCHMARK', '0') == '1'
 
+if BENCHMARK and PYOPT != '0':
+  print("BENCHMARK=1 works only with PYOPT=0.")
+  exit(1)
+
 FEATURE_FLAGS = {
     "RDI": True,
     "SECP256K1_ZKP": True,  # required for trezor.crypto.curve.bip340 (BIP340/Taproot)
diff --git a/core/SConscript.unix b/core/SConscript.unix
index a5d7ffee6b..bc791f69d2 100644
--- a/core/SConscript.unix
+++ b/core/SConscript.unix
@@ -17,6 +17,10 @@ FROZEN = ARGUMENTS.get('TREZOR_EMULATOR_FROZEN', 0)
 RASPI = os.getenv('TREZOR_EMULATOR_RASPI') == '1'
 
 
+if BENCHMARK and PYOPT != '0':
+  print("BENCHMARK=1 works only with PYOPT=0.")
+  exit(1)
+
 FEATURES_WANTED = ["input", "sd_card", "dma2d", "optiga", "sbu"]
 if NEW_RENDERING:
     FEATURES_WANTED.append("new_rendering")
diff --git a/core/src/apps/benchmark/__init__.py b/core/src/apps/benchmark/__init__.py
index e69de29bb2..2ff83459bb 100644
--- a/core/src/apps/benchmark/__init__.py
+++ b/core/src/apps/benchmark/__init__.py
@@ -0,0 +1,4 @@
+if not __debug__:
+    from trezor import utils
+
+    utils.halt("Disabled in production mode")
diff --git a/core/src/apps/workflow_handlers.py b/core/src/apps/workflow_handlers.py
index e581aeacbe..973d807134 100644
--- a/core/src/apps/workflow_handlers.py
+++ b/core/src/apps/workflow_handlers.py
@@ -207,9 +207,9 @@ def _find_message_handler_module(msg_type: int) -> str:
             return "apps.solana.sign_tx"
 
         # benchmark
-        if msg_type == MessageType.BenchmarkListNames:
+        if __debug__ and msg_type == MessageType.BenchmarkListNames:
             return "apps.benchmark.list_names"
-        if msg_type == MessageType.BenchmarkRun:
+        if __debug__ and msg_type == MessageType.BenchmarkRun:
             return "apps.benchmark.run"
 
     raise ValueError