bugfix: use write counter for firmware approval

### releases/Next-ChangeLog.md
@@ -11,6 +11,8 @@ This lists the new changes that have not yet been published in a normal release.
 - Bugfix: Reject malformed PSBTs containing P2SH-P2WSH inputs with a missing or
   incorrect redeem script, preventing transactions with an unknown fee from
   proceeding to approval.
+- Bugfix: Abort a pending firmware upgrade if its staged image is overwritten before
+  approval. Thanks to Huzaifa Jawaid.
 - Bugfix: Restore the ability to view the device-generated seed before adding user
   entropy, which was available in the previous dice-roll workflow but was inadvertently
   removed in 5.6.1/1.5.1Q. The new **View TRNG Words** menu item displays the full

### shared/auth.py
@@ -1568,6 +1568,7 @@ def __init__(self, hdr, length, hdr_check=False, psram_offset=None):
         self.length = length
         self.hdr_check = hdr_check
         self.psram_offset = psram_offset
+        self.upgrade_write_count = glob.PSRAM.upgrade_write_count
 
     async def interact(self):
         from version import decode_firmware_header
@@ -1585,10 +1586,6 @@ async def interact(self):
                 self.pop_menu()
                 return
 
-        # bind this request to the exact bytes staged in PSRAM right now;
-        # a second upload may land while the consent screen is up
-        self.staged_sha = psram_sha256(self.psram_offset, self.length)
-
         # Get informed consent to upgrade.
         date, version, _ = decode_firmware_header(self.hdr)
 
@@ -1605,8 +1602,7 @@ async def interact(self):
             ch = await ux_show_story(msg)
 
             if ch == 'y':
-                # re-verify the staged bytes are unchanged since approval
-                assert psram_sha256(self.psram_offset, self.length) == self.staged_sha
+                assert glob.PSRAM.upgrade_write_count == self.upgrade_write_count
 
                 # Accepted:
                 # - write final file header, so bootloader will see it

### shared/psram.py
@@ -15,6 +15,10 @@ class PSRAMWrapper:
     # mutation of the transaction being signed
     txn_write_count = 0
 
+    # kept separate from transaction signing: firmware approval captures this
+    # value and aborts if its staged image is overwritten
+    upgrade_write_count = 0
+
     def __init__(self):
         self._wr = uctypes.bytearray_at(self.base, self.length)
 
@@ -31,6 +35,7 @@ def write_at(self, offset, ln):
         if offset < 2 * version.MAX_TXN_LEN:
             # covers both TXN staging regions (input and output)
             self.txn_write_count += 1
+            self.upgrade_write_count += 1
 
         return memoryview(self._wr)[offset:offset+ln]
 

### shared/vdisk.py
@@ -109,6 +109,9 @@ def import_file(self, filename, sz):
         glob.ALLOWED_DOWNLOAD = None
         glob.PSRAM.txn_write_count += 1
 
+        # native copy bypasses PSRAMWrapper.write_at()
+        glob.PSRAM.upgrade_write_count += 1
+
         # I could not resist doing this in C... since we already have the
         # data in memory, why mess around with file concepts?
         actual = VBLKDEV.copy_file(0, filename.split('/')[-1])

### testing/test_upgrades.py
@@ -26,11 +26,12 @@ def doit(data, pkt_len=2048):
 
 @pytest.fixture
 def make_firmware(src_root_dir):
-    def doit(hw_compat, fname=f'{src_root_dir}/stm32/firmware-signed.bin', outname='tmp-firmware.bin'):
+    def doit(hw_compat, fname=f'{src_root_dir}/stm32/firmware-signed.bin',
+             outname='tmp-firmware.bin', version='3.0.99'):
         # os.system(f'signit sign 3.0.99 --keydir ../stm32/keys -r {fname} -o {outname} --hw-compat=0x{hw_compat:02x}')
         p = subprocess.run(
             [
-                'signit', 'sign', '3.0.99',
+                'signit', 'sign', version,
                  '--keydir', f'{src_root_dir}/stm32/keys',
                  '-r', f'{fname}',
                  '-o', f'{outname}',
@@ -136,34 +137,41 @@ def test_hacky_upgrade(mode, cap_story, transport, dev, sim_exec, make_firmware,
     #         assert a == data[pos:pos+128], repr(pos)
 
 
-def test_upgrade_staged_image_tamper(dev, make_firmware, upload_file, cap_story,
-                                     need_keypress, sim_exec, is_q1, is_mark5):
+def test_upgrade_staged_image_tamper(make_firmware, upload_file, cap_story,
+                                     press_select, sim_exec, sim_eval, is_q1, is_mark5):
     # a second upload may land while the upgrade approval is on screen
     # (check_busy allow-lists FirmwareUpgradeRequest); the staged bytes
     # must be re-verified before flashing, not just the header snapshot
     hw = "q1" if is_q1 else (5 if is_mark5 else 4)
-    data_a = make_firmware(hw)
+    data_a = make_firmware(hw, version='3.0.98')
     hdr_a = data_a[FW_HEADER_OFFSET:FW_HEADER_OFFSET+FW_HEADER_SIZE]
 
+    sim_exec("import glob; from pincodes import pa; "
+             "glob._fw_upgrade = pa.firmware_upgrade; "
+             "glob._fw_upgrade_called = False; "
+             "pa.firmware_upgrade = lambda *a: setattr(glob, '_fw_upgrade_called', True)")
+
     # upload image A with trailer -> fires authorize_upgrade
     upload_file(data_a + hdr_a)
-    _, story = cap_story()
-    assert "Install this new firmware?" in story
 
     # upload image B as a raw image (no trailer) -> no re-auth, but
     # overwrites the staging area via PSRAM.write
-    data_b = make_firmware(hw, outname='tmp-firmware-b.bin')
+    data_b = make_firmware(hw, outname='tmp-firmware-b.bin', version='3.0.99')
     assert len(data_b) == len(data_a)
+    assert data_b != data_a
     upload_file(data_b)
 
-    # approve what was displayed (image A); the pre-flash assert fires,
-    # caught by interact()'s except -> self.failed, cleanup, pop_menu
-    need_keypress('y')
-    time.sleep(1)
-    # must not have upgraded: request done and cleaned up, no reboot
-    rv = sim_exec("from auth import UserAuthorizedAction; "
-                  "print(UserAuthorizedAction.active_request is None)")
-    assert 'True' in rv
+    _, story = cap_story()
+    assert "Install this new firmware?" in story
+    assert "3.0.98" in story
+
+    try:
+        press_select()
+        time.sleep(1)
+        assert sim_eval("glob._fw_upgrade_called") == 'False'
+    finally:
+        sim_exec("from pincodes import pa; import glob; "
+                 "pa.firmware_upgrade = glob._fw_upgrade")
 
 
 # EOF