From d48222cb5dd6d2c0e45a966f1b4fce7a00262f53 Mon Sep 17 00:00:00 2001
From: Romain Gantois <romain.gantois@bootlin.com>
Date: Wed, 18 Dec 2024 16:40:35 +0100
Subject: [PATCH] snagrecover: add Xilinx ZynqMP support

Add support for Xilinx ZynqMP platforms. The USB recovery procedure should
be identical for all ZynqMP SoCs so only one SoC model is registered.

Signed-off-by: Romain Gantois <romain.gantois@bootlin.com>
---
 README.md                                     |   2 +-
 docs/fw_binaries.md                           |  27 +++
 src/snagrecover/50-snagboot.rules             |   3 +
 src/snagrecover/config.py                     |   1 +
 src/snagrecover/firmware/firmware.py          |   3 +
 src/snagrecover/firmware/zynqmp_fw.py         | 175 ++++++++++++++++++
 src/snagrecover/protocols/dfu.py              |  10 +
 src/snagrecover/recoveries/zynqmp.py          |  33 ++++
 src/snagrecover/supported_socs.yaml           |   2 +
 src/snagrecover/templates/zynqmp-generic.yaml |   2 +
 src/snagrecover/utils.py                      |   3 +
 11 files changed, 260 insertions(+), 1 deletion(-)
 create mode 100644 src/snagrecover/firmware/zynqmp_fw.py
 create mode 100644 src/snagrecover/recoveries/zynqmp.py
 create mode 100644 src/snagrecover/templates/zynqmp-generic.yaml

diff --git a/README.md b/README.md
index 783c518..00ec0aa 100644
--- a/README.md
+++ b/README.md
@@ -16,7 +16,7 @@ ISP, UUU, and sunxi-fel. Snagboot is made of two separate parts:
 </p>
 
 The currently supported SoC families are ST STM32MP1/2, Microchip SAMA5, NXP
-i.MX6/7/8/93, TI AM335x, Allwinner SUNXI, TI AM62x and TI AM64x. Please check
+i.MX6/7/8/93, TI AM335x, Allwinner SUNXI, TI AM62x, TI AM64x and Xilinx ZynqMP. Please check
 [supported_socs.yaml](https://github.com/bootlin/snagboot/blob/main/src/snagrecover/supported_socs.yaml) or run `snagrecover
 --list-socs` for a more precise list of supported SoCs.
 
diff --git a/docs/fw_binaries.md b/docs/fw_binaries.md
index 45b7eb6..fb1020f 100644
--- a/docs/fw_binaries.md
+++ b/docs/fw_binaries.md
@@ -258,3 +258,30 @@ DM firmware, U-Boot SPL for A53 and device tree blobs
 configuration:
  * path
 
+## For Xilinx ZynqMP devices
+
+[example](../src/snagrecover/templates/zynqmp-generic.yaml)
+
+Detailed instructions for building the required boot images can be found in the
+[Xilinx documentation](https://xilinx.github.io/Embedded-Design-Tutorials), in
+the "boot-and-configuration" section for ZynqMP SoCs. Please note that the
+first boot image containing only the FSBL and PMUFW should not be required, as
+Snagboot is capable of extracting a working first-stage boot image from the
+full boot image.
+
+The following images are required for all AM6xx SoCs:
+
+**boot:** Xilinx ZynqMP boot image containing the Xilinx FSBL, PMUFW, ATF,
+control DT and U-Boot proper. The FSBL should be compiled with USB support
+enabled.
+
+configuration:
+ * path
+
+**fsbl:** (optional) Xilinx ZynqMP boot image containing the Xilinx FSBL and
+PMUFW. This is optional, and should only be provided if Snagboot fails to
+extract the FSBL and PMUFW from the complete boot image.
+
+configuration:
+ * path
+
diff --git a/src/snagrecover/50-snagboot.rules b/src/snagrecover/50-snagboot.rules
index 58e6df4..d4f4dde 100644
--- a/src/snagrecover/50-snagboot.rules
+++ b/src/snagrecover/50-snagboot.rules
@@ -73,3 +73,6 @@ SUBSYSTEM=="usb", ATTRS{idVendor}=="1f3a", ATTRS{idProduct}=="efe8", MODE="0660"
 SUBSYSTEM=="usb", ATTRS{idVendor}=="0451", ATTRS{idProduct}=="6165", MODE="0660", TAG+="uaccess"
 SUBSYSTEM=="usb", ATTRS{idVendor}=="0451", ATTRS{idProduct}=="6141", MODE="0660", TAG+="uaccess"
 SUBSYSTEM=="usb", ATTRS{idVendor}=="0451", ATTRS{idProduct}=="d022", MODE="0660", TAG+="uaccess"
+
+#Xilinx rules
+SUBSYSTEM=="usb", ATTRS{idVendor}=="03fd", ATTRS{idProduct}=="0050", MODE="0660", TAG+="uaccess"
diff --git a/src/snagrecover/config.py b/src/snagrecover/config.py
index e290da4..db09fb3 100644
--- a/src/snagrecover/config.py
+++ b/src/snagrecover/config.py
@@ -29,6 +29,7 @@
 	"sama5":    "03eb:6124",
 	"sunxi":    "1f3a:efe8",
 	"am6x":     "0451:6165",
+	"zynqmp":   "03fd:0050",
 	"imx": {
 		"imx8qxp": "1fc9:012f",
 		"imx8qm": "1fc9:0129",
diff --git a/src/snagrecover/firmware/firmware.py b/src/snagrecover/firmware/firmware.py
index 1b1767b..4a855e0 100644
--- a/src/snagrecover/firmware/firmware.py
+++ b/src/snagrecover/firmware/firmware.py
@@ -119,6 +119,9 @@ def run_firmware(port, fw_name: str, subfw_name: str = ""):
 		sunxi_run(port, fw_name, fw_blob)
 	elif soc_family == "am6x":
 		am6x_run(port, fw_name, fw_blob)
+	elif soc_family == "zynqmp":
+		from snagrecover.firmware.zynqmp_fw import zynqmp_run
+		zynqmp_run(port, fw_name, fw_blob, subfw_name)
 	else:
 		raise Exception(f"Unsupported SoC family {soc_family}")
 	logger.info(f"Done installing firmware {fw_name}")
diff --git a/src/snagrecover/firmware/zynqmp_fw.py b/src/snagrecover/firmware/zynqmp_fw.py
new file mode 100644
index 0000000..2aff71d
--- /dev/null
+++ b/src/snagrecover/firmware/zynqmp_fw.py
@@ -0,0 +1,175 @@
+"""
+Handle ZynqMP boot images, as generated by the Xilinx bootgen tool.
+"""
+
+import logging
+logger = logging.getLogger("snagrecover")
+from snagrecover.protocols import dfu
+from dataclasses import dataclass, astuple
+import struct
+
+class Header():
+	"""
+	Generic class representing a packed C struct embedded in a bytearray.
+	"""
+
+	fmt = ""
+	size = 0
+
+	@classmethod
+	def read(cls, data, offset=0):
+		obj = cls(*struct.unpack(cls.fmt, data[offset:offset + cls.size]))
+		obj.offset = offset
+
+		return obj
+
+	@classmethod
+	def write(cls, self, data):
+		offset = self.offset
+		data[offset:offset + cls.size] = struct.pack(cls.fmt, *astuple(self))
+
+@dataclass
+class ZynqMPImageTable(Header):
+	version: int
+	part_count: int
+	part_hdr: int
+	image_hdr: int
+	auth_cert: int
+	aux_bootdev: int
+	padding: bytes
+	csum: int
+
+	fmt = "<LLLLLL36sL"
+	size = 0x40
+
+	def update_checksum(self):
+		data = struct.pack(ZynqMPImageTable.fmt, *astuple(self))
+		self.csum = zynqmp_csum(data[:-4])
+
+@dataclass
+class ZynqMPImageHeader(Header):
+	next_image: int
+	part_hdr: int
+	reserved: int
+	part_count: int
+	# Other fields are ignored
+
+	fmt = "<LLLL"
+	size = 16
+
+@dataclass
+class ZynqMPPartHeader(Header):
+	enc_size: int
+	plain_size: int
+	total_size: int
+	next_part: int
+	exec_lo: int
+	exec_hi: int
+	load_lo: int
+	load_hi: int
+	start: int
+	attrs: int
+	sec_count: int
+	csum_word_offset: int
+	image_hdr: int
+	ac_offset: int
+	part_id: int
+	csum: int
+
+	fmt = "<LLLLLLLLLLLLLLLL"
+	size = 0x40
+
+	def update_checksum(self):
+		data = struct.pack(ZynqMPPartHeader.fmt, *astuple(self))
+		self.csum = zynqmp_csum(data[:-4])
+
+def zynqmp_csum(data):
+	s = 0
+
+	for i in range(0, len(data), 4):
+		s = (s + int.from_bytes(data[i:i+4], "little")) % (1 << 32)
+
+	return s ^ 0xffffffff
+
+def find_img_table(boot_bin):
+	boot_hdr = boot_bin[:0xb5]
+	hdr = boot_hdr
+
+	sig = hdr[0x24:0x28]
+	if sig != b"XNLX":
+		raise ValueError("Invalid ZynqMP boot header signature!")
+
+	return int.from_bytes(boot_hdr[0x98:0x9c], "little")
+
+def drop_images(boot_bin: bytearray, keep_images):
+	"""
+	Drop partitions which aren't required for the FSBL stage from the boot
+	image.
+	"""
+
+	if keep_images == 0:
+		raise ValueError("Must keep at least one image!")
+
+	img_table_offset = find_img_table(boot_bin)
+	img_table = ZynqMPImageTable.read(boot_bin, img_table_offset)
+
+	kept_images = []
+	next_image = 4 * img_table.image_hdr
+	for _ in range(keep_images):
+		image = ZynqMPImageHeader.read(boot_bin, next_image)
+		kept_images.append(image)
+		next_image = 4 * image.next_image
+
+	last_kept_image = kept_images[-1]
+	if next_image == 0:
+		return
+
+	# Mark as last image
+	last_kept_image.next_image = 0
+	ZynqMPImageHeader.write(last_kept_image, boot_bin)
+
+	# Find last remaining partition
+	next_part = 4 * last_kept_image.part_hdr
+	for _ in range(last_kept_image.part_count):
+		part = ZynqMPPartHeader.read(boot_bin, next_part)
+		next_part = 4 * part.next_part
+
+	last_kept_part = part
+
+	# Drop unused partition data
+	cutoff = 4 * (last_kept_part.start + last_kept_part.total_size)
+
+	# Mark as last partition and update csum
+	last_kept_part.next_part = 0
+	last_kept_part.update_checksum()
+	ZynqMPPartHeader.write(last_kept_part, boot_bin)
+
+	# Update partition count and csum
+	img_table.part_count = sum([image.part_count for image in kept_images])
+	img_table.update_checksum()
+
+	ZynqMPImageTable.write(img_table, boot_bin)
+
+	return bytes(boot_bin[:cutoff])
+
+def zynqmp_run(dev, fw_name, fw_blob, subfw_name):
+	dfu_cmd = dfu.DFU(dev, stm32=False)
+
+	if fw_name == "fsbl":
+		partid = 0
+	if fw_name == "boot":
+		if subfw_name == "fsbl":
+			partid = 0
+			"""
+			Only the FSBL and PMUFW are required for the first stage.
+			Keep the first two images and drop the rest, so that
+			the firmware blob can fit in SRAM.
+			"""
+			fw_blob = drop_images(bytearray(fw_blob), 2)
+		else:
+			partid = 1
+
+	logger.info("Downloading file...")
+	dfu_cmd.download_and_run(fw_blob, partid, offset=0, size=len(fw_blob))
+	logger.info("Done")
+
diff --git a/src/snagrecover/protocols/dfu.py b/src/snagrecover/protocols/dfu.py
index ed5b076..8509246 100644
--- a/src/snagrecover/protocols/dfu.py
+++ b/src/snagrecover/protocols/dfu.py
@@ -25,6 +25,16 @@
 from errno import EIO,ENODEV,EPIPE
 from snagrecover import utils
 
+def list_partids(dev: usb.core.Device):
+	cfg = dev.get_active_configuration()
+	intfs = cfg.interfaces()
+	partids = []
+
+	for intf in intfs:
+		partids.append(intf.bAlternateSetting)
+
+	return partids
+
 def search_partid(dev: usb.core.Device, partname: str, match_prefix=False) -> int:
 	# search for an altsetting associated with a partition name
 	cfg = dev.get_active_configuration()
diff --git a/src/snagrecover/recoveries/zynqmp.py b/src/snagrecover/recoveries/zynqmp.py
new file mode 100644
index 0000000..3af2867
--- /dev/null
+++ b/src/snagrecover/recoveries/zynqmp.py
@@ -0,0 +1,33 @@
+import usb
+import logging
+logger = logging.getLogger("snagrecover")
+from snagrecover.firmware.firmware import run_firmware
+from snagrecover.utils import get_usb
+from snagrecover.config import recovery_config
+from snagrecover.protocols import dfu
+import time
+
+def altmode1_check(dev: usb.core.Device):
+	try:
+		return 1 in dfu.list_partids(dev)
+	except usb.core.USBError:
+		logger.warning("USB device is unavailable")
+
+	return False
+
+def main():
+	usb_addr = recovery_config["usb_path"]
+	dev = get_usb(usb_addr)
+
+	if "fsbl" not in recovery_config["firmware"]:
+		logger.warning("No FSBL image given, will attempt to extract it from full boot image")
+		run_firmware(dev, "boot", "fsbl")
+	else:
+		run_firmware(dev, "fsbl")
+
+	time.sleep(0.5)
+
+	dev = get_usb(usb_addr, ready_check=altmode1_check)
+
+	run_firmware(dev, "boot")
+
diff --git a/src/snagrecover/supported_socs.yaml b/src/snagrecover/supported_socs.yaml
index fe3fa4d..c36f5d0 100644
--- a/src/snagrecover/supported_socs.yaml
+++ b/src/snagrecover/supported_socs.yaml
@@ -53,6 +53,8 @@ tested:
         family: stm32mp
   stm32mp25:
         family: stm32mp
+  zynqmp:
+        family: zynqmp
 untested:
   a10:
         family: sunxi
diff --git a/src/snagrecover/templates/zynqmp-generic.yaml b/src/snagrecover/templates/zynqmp-generic.yaml
new file mode 100644
index 0000000..d129017
--- /dev/null
+++ b/src/snagrecover/templates/zynqmp-generic.yaml
@@ -0,0 +1,2 @@
+boot:
+  path: boot.bin
diff --git a/src/snagrecover/utils.py b/src/snagrecover/utils.py
index ca8731e..aca0f93 100644
--- a/src/snagrecover/utils.py
+++ b/src/snagrecover/utils.py
@@ -183,6 +183,9 @@ def get_recovery(soc_family: str):
 	elif soc_family == "am6x":
 		from snagrecover.recoveries.am6x import main as am6x_recovery
 		return am6x_recovery
+	elif soc_family == "zynqmp":
+		from snagrecover.recoveries.zynqmp import main as zynqmp_recovery
+		return zynqmp_recovery
 	else:
 		cli_error(f"unsupported board family {soc_family}")