snagrecover: Add rockchip support

Add support for booting a rockchip bin file, made with boot_merger.
It can be:
- a file containing the rockchip miniloader, but in this case,
  it's not useful, as rockusb is not implemented
- a file containing uboot tpl and spl. In this case, if u-boot
  is configured for SPL DFU, the snagrecover configuration file
  may specify an u-boot FIT image to load into the first DFU slot.
  This will allow booting u-boot proper and then to fastboot/ums/dfu/...
  to flash.

Signed-off-by: Arnaud Patard <[email protected]>

Author: apatard

src/snagrecover/50-snagboot.rules

@@ -76,3 +76,21 @@ SUBSYSTEM=="usb", ATTRS{idVendor}=="0451", ATTRS{idProduct}=="d022", MODE="0660"
 
 #Xilinx rules
 SUBSYSTEM=="usb", ATTRS{idVendor}=="03fd", ATTRS{idProduct}=="0050", MODE="0660", TAG+="uaccess"
+
+#Rockchip systems
+SUBSYSTEM=="usb", ATTRS{idVendor}=="2207", ATTRS{idProduct}=="2207", MODE="0660", TAG+="uaccess"
+SUBSYSTEM=="usb", ATTRS{idVendor}=="2207", ATTRS{idProduct}=="110a", MODE="0660", TAG+="uaccess"
+SUBSYSTEM=="usb", ATTRS{idVendor}=="2207", ATTRS{idProduct}=="110b", MODE="0660", TAG+="uaccess"
+SUBSYSTEM=="usb", ATTRS{idVendor}=="2207", ATTRS{idProduct}=="300a", MODE="0660", TAG+="uaccess"
+SUBSYSTEM=="usb", ATTRS{idVendor}=="2207", ATTRS{idProduct}=="301a", MODE="0660", TAG+="uaccess"
+SUBSYSTEM=="usb", ATTRS{idVendor}=="2207", ATTRS{idProduct}=="310b", MODE="0660", TAG+="uaccess"
+SUBSYSTEM=="usb", ATTRS{idVendor}=="2207", ATTRS{idProduct}=="310c", MODE="0660", TAG+="uaccess"
+SUBSYSTEM=="usb", ATTRS{idVendor}=="2207", ATTRS{idProduct}=="320a", MODE="0660", TAG+="uaccess"
+SUBSYSTEM=="usb", ATTRS{idVendor}=="2207", ATTRS{idProduct}=="320b", MODE="0660", TAG+="uaccess"
+SUBSYSTEM=="usb", ATTRS{idVendor}=="2207", ATTRS{idProduct}=="320c", MODE="0660", TAG+="uaccess"
+SUBSYSTEM=="usb", ATTRS{idVendor}=="2207", ATTRS{idProduct}=="330a", MODE="0660", TAG+="uaccess"
+SUBSYSTEM=="usb", ATTRS{idVendor}=="2207", ATTRS{idProduct}=="330c", MODE="0660", TAG+="uaccess"
+SUBSYSTEM=="usb", ATTRS{idVendor}=="2207", ATTRS{idProduct}=="330d", MODE="0660", TAG+="uaccess"
+SUBSYSTEM=="usb", ATTRS{idVendor}=="2207", ATTRS{idProduct}=="330e", MODE="0660", TAG+="uaccess"
+SUBSYSTEM=="usb", ATTRS{idVendor}=="2207", ATTRS{idProduct}=="350a", MODE="0660", TAG+="uaccess"
+SUBSYSTEM=="usb", ATTRS{idVendor}=="2207", ATTRS{idProduct}=="350b", MODE="0660", TAG+="uaccess"

src/snagrecover/config.py

@@ -52,6 +52,23 @@
 		"imx8mm": "1fc9:0134",
 		"imx8mq": "1fc9:012b",
 		"imx53" : "15a2:004e",
+	},
+	"rockchip": {
+		"rv1108": "2207:110a",
+		"rv1126": "2207:110b",
+		"rk3066": "2207:300a",
+		"rk3036": "2207:301a",
+		"rk3188": "2207:310b",
+		"rk3128": "2207:310c",
+		"rk3288": "2207:320a",
+		"rk322x": "2207:320b",
+		"rk3328": "2207:320c",
+		"rk3368": "2207:330a",
+		"rk3399": "2207:330c",
+		"px30": "2207:330d",
+		"rk3308": "2207:330e",
+		"rk3568": "2207:350a",
+		"rk3588": "2207:350b"
 	}
 }
 
@@ -80,6 +97,8 @@ def init_config(args: list):
 		if args.usb_path is None:
 			if soc_family == "imx":
 				usb_ids = default_usb_ids["imx"][soc_model]
+			elif soc_family == "rockchip":
+				usb_ids = default_usb_ids["rockchip"][soc_model]
 			else:
 				usb_ids = default_usb_ids[soc_family]
 

src/snagrecover/firmware/firmware.py

@@ -146,6 +146,9 @@ def run_firmware(port, fw_name: str, subfw_name: str = ""):
 	elif soc_family == "zynqmp":
 		from snagrecover.firmware.zynqmp_fw import zynqmp_run
 		zynqmp_run(port, fw_name, fw_blob, subfw_name)
+	elif soc_family == "rockchip":
+		from snagrecover.firmware.rk_fw import rockchip_run
+		rockchip_run(port, fw_name, fw_blob)
 	else:
 		raise Exception(f"Unsupported SoC family {soc_family}")
 	logger.info(f"Done installing firmware {fw_name}")

src/snagrecover/firmware/rk_fw.py

@@ -12,28 +12,29 @@
 import logging
 import struct
 import time
-from cryptography.hazmat.primitives.ciphers import Cipher
-from cryptography.hazmat.decrepit.ciphers.algorithms import ARC4
 from crccheck import crc
 from dataclasses import dataclass
 
 logger = logging.getLogger("snagrecover")
 from snagrecover.protocols import rockchip
+from snagrecover.protocols import dfu
 from snagrecover.utils import BinFileHeader
 from snagrecover.config import recovery_config
 
+
 # List generated with a grep on rkbin repository
-NEWIDB_LIST = [ "rk3506", "rk3506b", "rk3528", "rk3562", "rk3566", "rk3568", "rk3576", "rk3583", "rk3588", "rv1103b", "rv1106" ]
+NEWIDB_LIST = ["rk3506", "rk3506b", "rk3528", "rk3562", "rk3566", "rk3568", "rk3576", "rk3583", "rk3588", "rv1103b", "rv1106"]
 
 BOOTTAG = b"BOOT"
 LDRTAG = b"LDR "
-TAG_LIST = [ BOOTTAG, LDRTAG ]
+TAG_LIST = [BOOTTAG, LDRTAG]
 BOOTENTRYSIZE = 57
 BOOTHEADERENTRYSIZE = 6
 BOOTHEADERSIZE = 102
 BOOTHEADERTIMESIZE = 7
 RC4_KEY = bytearray([124, 78, 3, 4, 85, 5, 9, 7, 45, 44, 123, 56, 23, 13, 23, 17])
 
+
 @dataclass
 class BootEntry(BinFileHeader):
 	size: int
@@ -76,6 +77,7 @@ class BootReleaseTime(BinFileHeader):
 	def __str__(self):
 		return f"{self.year}/{self.month}/{self.day} {self.hour}:{self.minute}:{self.second}"
 
+
 class LoaderFileError(Exception):
 	def __init__(self, message):
 		self.message = message
@@ -84,6 +86,7 @@ def __init__(self, message):
 	def __str__(self):
 		return f"File format error: {self.message}"
 
+
 @dataclass
 class BootHeader(BinFileHeader):
 	tag: bytes
@@ -105,7 +108,7 @@ class BootHeader(BinFileHeader):
 
 	def __post_init__(self):
 		if self.tag not in TAG_LIST:
-			raise LoaderFileError(f"Invalid tag {self.header.tag}")
+			raise LoaderFileError(f"Invalid tag {self.tag}")
 		# not sure how to exactly parse version/merge_version
 		self.maj_ver = self.version >> 8
 		self.min_ver = self.version & 0xff
@@ -128,6 +131,7 @@ def __post_init__(self):
 	def __str__(self):
 		return f"{self.tag}, {self.size} ,{self.maj_ver}.{self.min_ver}, 0x{self.merge_version:0x}, {self.releasetime}, {self.chip}, {self.entry471}, {self.entry472}, {self.loader}, sign: {self.sign}, enc: {self.rc4}"
 
+
 class RkCrc32(crc.Crc32Base):
 	"""CRC-32/ROCKCHIP
 	"""
@@ -139,6 +143,7 @@ class RkCrc32(crc.Crc32Base):
 	_reflect_output = False
 	_xor_output = 0
 
+
 class LoaderFile():
 	def __init__(self, blob):
 		self.blob = blob
@@ -170,7 +175,7 @@ def __init__(self, blob):
 		(self.crc32,) = struct.unpack("<I", crc32)
 		assert self.crc32 == calc_crc32
 
-	def entry_data(self, name, idx = 0):
+	def entry_data(self, name, idx=0):
 		entry = None
 		if name == "471":
 			entry = self.entry471
@@ -190,6 +195,38 @@ def entry_data(self, name, idx = 0):
 	def __str__(self):
 		return f"{self.header} crc: {self.crc32:02x}"
 
+
+# Manual implementation of RC4, from Wikipedia's page
+# Not very secure, so don't use it elsewhere.
+class rc4():
+	def __init__(self):
+		self.S = list(range(256))
+		self.i = 0
+		self.j = 0
+
+	def ksa(self, key):
+		keylength = len(key)
+		self.S = list(range(256))
+		j = 0
+		for i in range(256):
+			j = (j + self.S[i] + key[i % keylength]) % 256
+			self.S[i], self.S[j] = self.S[j], self.S[i]
+
+	def prga(self):
+		self.i = (self.i + 1) % 256
+		self.j = (self.j + self.S[self.i]) % 256
+		self.S[self.i], self.S[self.j] = self.S[self.j], self.S[self.i]
+		K = self.S[(self.S[self.i] + self.S[self.j]) % 256]
+		return K
+
+	def encrypt(self, buf):
+		obuf = bytearray(len(buf))
+
+		for offset in list(range(len(buf))):
+			obuf[offset] = buf[offset] ^ self.prga()
+		return obuf
+
+
 def rc4_encrypt(fw_blob):
 
 	soc_model = recovery_config["soc_model"]
@@ -201,28 +238,39 @@ def rc4_encrypt(fw_blob):
 	padded_len = (blob_len+4095)//4096 * 4096
 	fw_blob = bytearray(fw_blob)
 	fw_blob += bytearray([0]*(padded_len - blob_len))
-	a = ARC4(RC4_KEY)
-	c = Cipher(a, mode=None)
-	d = c.encryptor()
+	encoder = rc4()
+	encoder.ksa(RC4_KEY)
 	obuf = bytearray()
 	for i in range(padded_len):
-		obuf += d.update(fw_blob[i*512:(i+1)*512])
+		obuf += encoder.encrypt(fw_blob[i*512:(i+1)*512])
 	return obuf
 
+
 def rockchip_run(dev, fw_name, fw_blob):
-	rom = rockchip.RochipBootRom(dev)
 
 	if fw_name == 'code471':
 		logger.info("Downloading code471...")
+		rom = rockchip.RochipBootRom(dev)
 		blob = rc4_encrypt(fw_blob)
 		rom.write_blob(blob, 0x471)
 	elif fw_name == 'code472':
 		logger.info("Downloading code472...")
+		rom = rockchip.RochipBootRom(dev)
 		blob = rc4_encrypt(fw_blob)
 		rom.write_blob(blob, 0x472)
+	elif fw_name == "u-boot-fit":
+		id = dfu.search_partid(dev, "u-boot.itb")
+		if id is None:
+			logger.error("Missing u-boot.itb DFU partition")
+		dfu_cmd = dfu.DFU(dev, stm32=False)
+		dfu_cmd.get_status()
+		dfu_cmd.download_and_run(fw_blob, id, 0, len(fw_blob))
+		dfu_cmd.get_status()
+		dfu_cmd.detach(id)
 	else:
 		fw = LoaderFile(fw_blob)
 		logger.info(f"{fw}")
+		rom = rockchip.RochipBootRom(dev)
 		for i in range(fw.header.entry471.count):
 			logger.info(f"Downloading entry 471 {i}...")
 			(data, delay) = fw.entry_data("471", i)

src/snagrecover/recoveries/rockchip.py

@@ -0,0 +1,51 @@
+import usb
+import logging
+logger = logging.getLogger("snagrecover")
+from snagrecover.firmware.firmware import run_firmware
+from snagrecover.utils import get_usb
+from snagrecover.utils import cli_error
+from snagrecover.config import recovery_config
+import time
+
+
+def main():
+	usb_addr = recovery_config["usb_path"]
+	dev = get_usb(usb_addr)
+
+	# Blob made with boot_merger
+	if "xpl" in recovery_config["firmware"]:
+		try:
+			run_firmware(dev, "xpl")
+			usb.util.dispose_resources(dev)
+		except Exception as e:
+			cli_error(f"Failed to run firmware: {e}")
+
+	# u-boot binaries.
+	elif "code471" in recovery_config["firmware"] and "code472" in recovery_config["firmware"]:
+		try:
+			run_firmware(dev, "code471")
+			usb.util.dispose_resources(dev)
+		except Exception as e:
+			cli_error(f"Failed to run code471 firmware: {e}")
+		if "delay" in recovery_config["firmware"]["code471"]:
+			delay = recovery_config["firmware"]["code471"]["delay"]
+			logger.info(f"Sleeping {delay}ms")
+			time.sleep(delay / 1000)
+		try:
+			run_firmware(dev, "code472")
+			usb.util.dispose_resources(dev)
+		except Exception as e:
+			cli_error(f"Failed to run code472 firmware: {e}")
+		if "delay" in recovery_config["firmware"]["code472"]:
+			delay = recovery_config["firmware"]["code472"]["delay"]
+			logger.info(f"Sleeping {delay}ms")
+			time.sleep(delay / 1000)
+	else:
+		cli_error("Missing xpl or code471/code472 binary configuration")
+
+	if "u-boot-fit" in recovery_config["firmware"]:
+		try:
+			dev = get_usb(usb_addr)
+			run_firmware(dev, "u-boot-fit")
+		except Exception as e:
+			cli_error(f"Failed to load u-boot.itb: {e}")

src/snagrecover/supported_socs.yaml

@@ -57,6 +57,10 @@ tested:
         family: stm32mp
   zynqmp:
         family: zynqmp
+  rk3399:
+        family: rockchip
+  rk3588:
+        family: rockchip
 untested:
   a10:
         family: sunxi
@@ -144,4 +148,29 @@ untested:
         family: sunxi
   v853:
         family: sunxi
-
+  rv1108:
+        family: rockchip
+  rv1126:
+        family: rockchip
+  rk3066:
+        family: rockchip
+  rk3036:
+        family: rockchip
+  rk3188:
+        family: rockchip
+  rk3128:
+        family: rockchip
+  rk3288:
+        family: rockchip
+  rk322x:
+        family: rockchip
+  rk3328:
+        family: rockchip
+  rk3368:
+        family: rockchip
+  px30:
+        family: rockchip
+  rk3308:
+        family: rockchip
+  rk3568:
+        family: rockchip

src/snagrecover/templates/rockchip-merger.yaml

@@ -0,0 +1,6 @@
+xpl:
+  # u-boot TPL+SPL with SPL configured for DFU boot
+  path: rk3399_uboot.bin
+u-boot-fit:
+  # proper u-boot
+  path: u-boot.itb

src/snagrecover/templates/rockchip-tpl-spl.yaml

@@ -0,0 +1,5 @@
+code471:
+  path: u-boot-rockchip-usb471.bin
+  delay: 1
+code472:
+  path: u-boot-rockchip-usb472.bin

src/snagrecover/utils.py

@@ -191,6 +191,9 @@ def get_recovery(soc_family: str):
 	elif soc_family == "zynqmp":
 		from snagrecover.recoveries.zynqmp import main as zynqmp_recovery
 		return zynqmp_recovery
+	elif soc_family == "rockchip":
+		from snagrecover.recoveries.rockchip import main as rockchip_recovery
+		return rockchip_recovery
 	else:
 		cli_error(f"unsupported board family {soc_family}")