bugfix: bind USB signing to staged PSBT

### shared/auth.py
@@ -393,6 +393,10 @@ async def interact(self):
         from glob import PSRAM
         self.parsed_write_count = PSRAM.txn_write_count
         self.parsed_sha = psram_sha256(self.offset, self.psbt_len)
+        if self.psbt_sha is not None and self.psbt_sha != self.parsed_sha:
+            del self.psbt
+            gc.collect()
+            return await self.failure("PSBT checksum mismatch")
 
         dis.fullscreen("Validating...")
 

### shared/usb.py
@@ -130,10 +130,9 @@ class USBHandler:
     def __init__(self):
         self.dev = pyb.USB_HID()
 
-        # We keep a running hash over whatever has been uploaded
+        # We keep a running hash over whatever has been staged
         # - reset at offset zero, can be read back anytime
-        self.file_checksum = sha256()
-        self.is_fw_upgrade = False
+        self.reset_upload()
 
         # handle simulator
         self.blockable = getattr(self.dev, 'pipe', self.dev)
@@ -195,6 +194,7 @@ async def usb_hid_recv(self):
                 else:
                     # treat zero-length packets as a reset request
                     # do not echo anything back on link.. used to resync connection
+                    self.reset_upload()
                     msg_len = 0
                     continue
 
@@ -267,15 +267,23 @@ async def usb_hid_recv(self):
             except FramingError as exc:
                 reason = exc.args[0]
                 # print("Framing: %s" % reason)
+                self.reset_upload()
                 await self.framing_error(reason)
                 msg_len = 0
 
             except BaseException as exc:
                 # recover from general issues/keep going
                 #print("USB!")
                 #sys.print_exception(exc)
+                self.reset_upload()
                 msg_len = 0
 
+    def reset_upload(self):
+        self.file_checksum = sha256()
+        self.upload_total_size = 0
+        self.upload_next_offset = 0
+        self.is_fw_upgrade = False
+
     def decrypt_inplace(self, msg_len):
         # self.msg is encrypted. decode it in place
         # - seems dangerous to use memview here, but works
@@ -409,10 +417,13 @@ async def handle(self, cmd, args):
             return b'biny' + args
 
         if cmd == 'upld':
-            offset, total_size = unpack_from('<II', args)
-            data = memoryview(args)[4+4:]
-
-            return await self.handle_upload(offset, total_size, data)
+            try:
+                offset, total_size = unpack_from('<II', args)
+                data = memoryview(args)[4+4:]
+                return await self.handle_upload(offset, total_size, data)
+            except:
+                self.reset_upload()
+                raise
 
         if cmd == 'dwld':
             offset, length, fileno = unpack_from('<III', args)
@@ -537,7 +548,11 @@ async def handle(self, cmd, args):
 
             assert 50 < txn_len <= MAX_TXN_LEN, "badlen"
 
-            from auth import sign_transaction
+            from auth import sign_transaction, psram_sha256, TXN_INPUT_OFFSET
+            if txn_len != self.upload_total_size or \
+                    txn_sha != psram_sha256(TXN_INPUT_OFFSET, txn_len):
+                return b'err_Checksum'
+
             sign_transaction(txn_len, (flags & STXN_FLAGS_MASK), txn_sha, input_method="usb")
             return None
 
@@ -823,15 +838,15 @@ async def handle_upload(self, offset, total_size, data):
         # invalidates any previous download lease
         glob.ALLOWED_DOWNLOAD = None
 
-        # maintain a running SHA256 over what's received
+        # maintain a running SHA256 over what's staged
         if offset == 0:
-            self.file_checksum = sha256()
-            self.is_fw_upgrade = False
+            self.reset_upload()
             dis.fullscreen("Receiving...", 0)
         else:
             dis.progress_sofar(offset, total_size)
 
         assert offset % 256 == 0, 'alignment'
+        assert offset == self.upload_next_offset, 'offset'
         assert 1 <= total_size <= MAX_UPLOAD_LEN, 'long'
         assert offset + len(data) <= total_size, 'long'
 
@@ -841,7 +856,7 @@ async def handle_upload(self, offset, total_size, data):
             if offset == 0:
                 assert data[0:5] == b'psbt\xff', 'psbt'
 
-        self.file_checksum.update(data)
+        self.upload_total_size = total_size
 
         for pos in range(offset, offset+len(data), 256):
 
@@ -871,6 +886,12 @@ async def handle_upload(self, offset, total_size, data):
                 hdr = memoryview(here)[-128:]
                 assert hdr == self.is_fw_upgrade        # indicates hacking
 
+                # The legacy USB upgrade protocol includes this duplicate
+                # trailer in its running checksum, even though it is not
+                # written to PSRAM a second time.
+                self.file_checksum.update(here)
+                self.upload_next_offset = offset + len(data)
+
                 # but don't write it, instead offer user a chance to abort
                 from auth import authorize_upgrade
                 authorize_upgrade(self.is_fw_upgrade, pos, psram_offset=0)
@@ -880,6 +901,9 @@ async def handle_upload(self, offset, total_size, data):
 
             # write to PSRAM
             PSRAM.write(pos, here)
+            self.file_checksum.update(here)
+
+        self.upload_next_offset = offset + len(data)
 
         if offset+len(data) >= total_size and not hsm_active:
             # probably done

### testing/test_hsm.py
@@ -1803,4 +1803,16 @@ def test_hsm_sign_download_lease(dev, quick_start_hsm, fake_txn, load_hsm_users,
         dev.send_recv(CCProtocolPacker.download(0, 256, 0))
     assert 'not allowed' in str(e.value)
 
+def test_hsm_rejects_psbt_sha_mismatch(dev, quick_start_hsm, fake_txn, sim_exec):
+    quick_start_hsm(DICT(warnings_ok=True, rules=[{}]))
+
+    psbt = fake_txn(1, 2, segwit_in=True)
+    txn_len, _ = dev.upload_file(psbt)
+    sim_exec("from auth import sign_transaction; "
+             "sign_transaction(%d, psbt_sha=bytes(32), input_method='usb')" % txn_len)
+
+    with pytest.raises(CCProtoError) as e:
+        wait_til_signed(dev)
+    assert 'PSBT checksum mismatch' in str(e.value)
+
 # EOF

### testing/test_usb.py
@@ -10,6 +10,7 @@
 from constants import simulator_fixed_tprv
 from ckcc_protocol.protocol import MAX_MSG_LEN, CCProtocolPacker, CCProtoError
 from ckcc_protocol.constants import MSG_SIGNING_MAX_LENGTH
+from sigheader import FW_HEADER_OFFSET, FW_HEADER_SIZE, FW_HEADER_MAGIC
 
 @pytest.mark.skip
 def test_usb_fuzz(dev):
@@ -130,9 +131,6 @@ def test_upload_short(dev, data_len):
 
     assert chk == hashlib.sha256(data).digest(), 'bad hash'
 
-    # clear screen / test a degerate case
-    dev.send_recv(CCProtocolPacker.upload(256, 256, b''))
-
 @pytest.mark.parametrize('pkt_len', [256, 1024, 2048])
 def test_upload_long(dev, pkt_len, count=5, data=None):
     # upload a larger "file"
@@ -145,9 +143,6 @@ def test_upload_long(dev, pkt_len, count=5, data=None):
         chk = dev.send_recv(CCProtocolPacker.sha256())
         assert chk == hashlib.sha256(data[0:pos+pkt_len]).digest(), 'bad hash'
 
-    # clear screen / test a degerate case
-    dev.send_recv(CCProtocolPacker.upload(256, 256, b''))
-
 @pytest.mark.parametrize('data_len', [0x3f01, 0x3f02, 0x3f03])
 def test_upload_psbt_at_firmware_probe_boundary(dev, data_len):
     dev.upload_file(b'psbt\xff' + bytes(data_len - 5))
@@ -165,6 +160,43 @@ def test_upload_fails(dev):
         # bad position
         v = dev.send_recv(CCProtocolPacker.upload(1000, 3, data))
 
+def test_upload_rejects_sparse_offset(dev):
+    dev.send_recv(CCProtocolPacker.upload(0, 768, bytes(256)))
+
+    with pytest.raises(CCProtoError) as e:
+        dev.send_recv(CCProtocolPacker.upload(512, 768, bytes(256)))
+    assert 'offset' in str(e.value)
+    assert dev.send_recv(CCProtocolPacker.sha256()) == hashlib.sha256(b'').digest()
+
+def test_upload_error_resets_checksum(dev):
+    # This block looks like a firmware header, but fails validation before
+    # it can be written to PSRAM. It must not remain in the upload checksum.
+    total_size = FW_HEADER_OFFSET + FW_HEADER_SIZE
+    for pos in range(0, FW_HEADER_OFFSET & ~255, 1024):
+        here = bytes(min(1024, (FW_HEADER_OFFSET & ~255) - pos))
+        dev.send_recv(CCProtocolPacker.upload(pos, total_size, here))
+
+    bad_hdr = bytearray(256)
+    struct.pack_into('<I', bad_hdr, FW_HEADER_OFFSET & 255, FW_HEADER_MAGIC)
+    with pytest.raises(CCProtoError):
+        dev.send_recv(CCProtocolPacker.upload(FW_HEADER_OFFSET & ~255, total_size, bad_hdr))
+
+    assert dev.send_recv(CCProtocolPacker.sha256()) == hashlib.sha256(b'').digest()
+
+def test_stxn_binds_uploaded_size_and_psram(dev, fake_txn, sim_exec):
+    psbt = fake_txn(1, 2, segwit_in=True)
+    txn_len, txn_sha = dev.upload_file(psbt)
+
+    with pytest.raises(CCProtoError) as e:
+        dev.send_recv(CCProtocolPacker.sign_transaction(txn_len - 1, txn_sha))
+    assert 'Checksum' in str(e.value)
+
+    txn_len, txn_sha = dev.upload_file(psbt)
+    sim_exec("from glob import PSRAM; PSRAM.write(0, b'x')")
+    with pytest.raises(CCProtoError) as e:
+        dev.send_recv(CCProtocolPacker.sign_transaction(txn_len, txn_sha))
+    assert 'Checksum' in str(e.value)
+
 def test_encryption(dev):
     "Setup session key and test link encryption works"
 
@@ -280,8 +312,9 @@ def test_download_lease(dev, fake_txn, start_sign, end_sign):
     in_psbt = fake_txn(1, 2, segwit_in=True)
     start_sign(in_psbt, finalize=False)
     end_sign(accept=True, finalize=False)
-    rv = dev.send_recv(CCProtocolPacker.upload(256, 1024, bytes(256)))
-    assert rv == 256
+    with pytest.raises(CCProtoError) as e:
+        dev.send_recv(CCProtocolPacker.upload(256, 1024, bytes(256)))
+    assert 'offset' in str(e.value)
     with pytest.raises(CCProtoError) as e:
         dev.send_recv(CCProtocolPacker.download(0, 256, 1))
     assert 'not allowed' in str(e.value)
@@ -373,10 +406,12 @@ def test_upld_zero_total_size(dev):
     assert 'long' in str(e.value)
 
 def test_upld_short_args(dev):
+    dev.upload_file(b'previous upload')
     msg = b'upld' + struct.pack('<I', 0)
     with pytest.raises(CCProtoError) as e:
         dev.send_recv(msg, encrypt=False)
     assert 'buffer too small' in str(e.value)
+    assert dev.send_recv(CCProtocolPacker.sha256()) == hashlib.sha256(b'').digest()
 
 def test_ncry_short_args(dev):
     msg = b'ncry' + struct.pack('<I', 1)