[ambz] Adjust OTA2 address during UF2 flashing, improve disconnection

ltchiptool/soc/ambz/flash.py

@@ -1,13 +1,16 @@
 # Copyright (c) Kuba Szczodrzyński 2022-07-29.
 
 from abc import ABC
+from io import BytesIO
 from logging import debug, warning
 from time import sleep
 from typing import IO, Generator, List, Optional, Union
 
 from ltchiptool import SocInterface
+from ltchiptool.soc.amb.system import SystemData
 from ltchiptool.util.flash import FlashConnection
-from ltchiptool.util.intbin import gen2bytes, letoint
+from ltchiptool.util.intbin import gen2bytes
+from ltchiptool.util.logging import verbose
 from ltchiptool.util.streams import ProgressCallback
 from uf2tool import OTAScheme, UploadContext
 
@@ -97,7 +100,7 @@ def flash_connect(self) -> None:
     def flash_disconnect(self) -> None:
         if self.amb:
             try:
-                self.amb.link(disconnect=True)
+                self.amb.disconnect()
             except TimeoutError:
                 pass
             self.amb.close()
@@ -184,26 +187,40 @@ def flash_write_uf2(
     ) -> None:
         # read system data to get active OTA index
         callback.on_message("Checking OTA index...")
-        system = gen2bytes(self.flash_read_raw(0x9000, 256))
-        if len(system) < 256:
+        system_data = gen2bytes(self.flash_read_raw(0x9000, 4096, verify=False))
+        if len(system_data) != 4096:
             raise ValueError(
-                f"Length invalid while reading from 0x9000 - {len(system)}"
+                f"Length invalid while reading from 0x9000 - {len(system_data)}"
             )
+        system = SystemData.unpack(system_data)
+        verbose(f"Realtek System Data: {system}")
 
         # read OTA switch value
-        ota_switch = f"{letoint(system[4:8]):032b}"
+        ota_switch = f"{system.ota2_switch:032b}"
         # count 0-bits
         ota_idx = 1 + (ota_switch.count("0") % 2)
 
-        # validate OTA2 address in system data
-        if ota_idx == 2:
-            ota2_addr = letoint(system[0:4]) & 0xFFFFFF
+        # check OTA2 address
+        try:
+            ota2_addr = system.ota2_address & 0xFFFFFF
             part_addr = ctx.get_offset("ota2", 0)
             if ota2_addr != part_addr:
-                raise ValueError(
-                    f"Invalid OTA2 address on chip - "
-                    f"found {ota2_addr}, expected {part_addr}",
+                # if it differs, correct it
+                system.ota2_address = AMBZ_FLASH_ADDRESS | part_addr
+                # reset OTA switch to use OTA1
+                system.ota2_switch = 0xFFFFFFFF
+                ota_idx = 1
+                # flash new system data
+                system_data = system.pack()
+                callback.on_message("Adjusting OTA2 address...")
+                self.flash_write_raw(
+                    offset=0x9000,
+                    length=len(system_data),
+                    data=BytesIO(system_data),
+                    callback=callback,
                 )
+        except ValueError:
+            warning("OTA2 partition not found in UF2 package")
 
         # collect continuous blocks of data
         parts = ctx.collect_data(

ltchiptool/soc/ambz/util/ambzcode.py

@@ -58,15 +58,26 @@ def print_greeting(delay: float, data: bytes) -> bytes:
 
     @staticmethod
     def download_mode() -> bytes:
+        """Disable booting to SRAM and run download mode again."""
+        # ldr r3, uartimg_boot_sram
+        # ldr r0, [r3]
+        # ldr r1, uartimg_boot_mask
+        # ands r0, r0, r1
+        # str r0, [r3]
         # movs r0, #2
         # ldr r3, UARTIMG_Download
         # blx r3
-        # movs r0, r0
         # UARTIMG_Download: .word 0x900+1
+        # uartimg_boot_sram: .word 0x40000210
+        # uartimg_boot_mask: .word 0xEFFFFFFF
         return (
-            b"\x02\x20\x01\x4b"
-            b"\x98\x47\x00\x00"
+            b"\x04\x4b\x18\x68"
+            b"\x04\x49\x08\x40"
+            b"\x18\x60\x02\x20"
+            b"\x00\x4b\x98\x47"
             b"\x01\x09\x00\x00"  # UARTIMG_Download()
+            b"\x10\x02\x00\x40"  # uartimg_boot_sram
+            b"\xff\xff\xff\xef"  # uartimg_boot_mask
         )
 
     @staticmethod

ltchiptool/soc/ambz/util/ambztool.py

@@ -152,12 +152,22 @@ def xm_putc(self, data, _=1):
     # Basic commands - public low-level API #
     #########################################
 
-    def link(self, disconnect: bool = False) -> None:
+    def disconnect(self) -> None:
+        # try to enter Loud-Handshake mode
+        command = [
+            # - Xmodem -> Handshake
+            # - Loud-Handshake -> Quiet-Handshake
+            AmbZCommand.XMODEM_CAN,
+            # - Handshake -> Xmodem
+            AmbZCommand.XMODEM_HANDSHAKE,
+            # - Xmodem -> Loud-Handshake (resets baud rate)
+            AmbZCommand.XMODEM_CAN,
+        ]
+        self.write(bytes(command))
+
+    def link(self) -> None:
         # clear any data before linking
         self.flush()
-        if disconnect:
-            # break NAK stream to force resetting baudrate
-            self.quiet_handshake()
         handshake = b""
         end = time() + self.link_timeout
         while time() < end:
@@ -169,17 +179,7 @@ def link(self, disconnect: bool = False) -> None:
             handshake = handshake[-4:]
             if len(handshake) == 4 and all(c == NAK[0] for c in handshake):
                 break
-            # try to enter Loud-Handshake mode
-            command = [
-                # - Xmodem -> Handshake
-                # - Loud-Handshake -> Quiet-Handshake
-                AmbZCommand.XMODEM_CAN,
-                # - Handshake -> Xmodem
-                AmbZCommand.XMODEM_HANDSHAKE,
-                # - Xmodem -> Loud-Handshake (resets baud rate)
-                AmbZCommand.XMODEM_CAN,
-            ]
-            self.write(bytes(command))
+            self.disconnect()
             sleep(0.1)
             self.set_baudrate(AMBZ_ROM_BAUDRATE)
         else:
@@ -351,7 +351,6 @@ def memory_write(
 
         self.write(bytes([AmbZCommand.XMODEM_HANDSHAKE]))
         self.expect_ack("Xmodem handshake")
-        self.xm.mode = "xmodem"
         # fake a NAK to make xmodem happy
         self.xm_send_code = NAK
         # fake an ACK after EOT to make xmodem very happy
@@ -499,7 +498,7 @@ def cli(device: str):
         warning(f"Couldn't process flash ID: got {chip_info!r}")
 
     info("Disconnecting...")
-    amb.link(disconnect=True)
+    amb.disconnect()
 
 
 if __name__ == "__main__":