Add Wi-Fi firmware partition support for Pico 2 W (#1969)

Author: will-v-pi

src/common/boot_uf2_headers/include/boot/uf2.h

@@ -26,13 +26,21 @@
 #define UF2_FLAG_MD5_PRESENT             0x00004000u
 #define UF2_FLAG_EXTENSION_FLAGS_PRESENT 0x00008000u
 
+// Extra family IDs
+#define CYW43_FIRMWARE_FAMILY_ID    0xe48bff55u
+
+// Bootrom supported family IDs
 #define RP2040_FAMILY_ID            0xe48bff56u
 #define ABSOLUTE_FAMILY_ID          0xe48bff57u
 #define DATA_FAMILY_ID              0xe48bff58u
 #define RP2350_ARM_S_FAMILY_ID      0xe48bff59u
 #define RP2350_RISCV_FAMILY_ID      0xe48bff5au
 #define RP2350_ARM_NS_FAMILY_ID     0xe48bff5bu
-#define FAMILY_ID_MAX               0xe48bff5bu
+#define BOOTROM_FAMILY_ID_MIN       RP2040_FAMILY_ID
+#define BOOTROM_FAMILY_ID_MAX       RP2350_ARM_NS_FAMILY_ID
+
+// Defined for backwards compatibility
+#define FAMILY_ID_MAX               BOOTROM_FAMILY_ID_MAX
 
 // 04 e3 57 99
 #define UF2_EXTENSION_RP2_IGNORE_BLOCK 0x9957e304

src/rp2_common/pico_cyw43_driver/BUILD.bazel

@@ -83,3 +83,13 @@ pico_generate_pio_header(
     name = "cyw43_bus_pio",
     srcs = ["cyw43_bus_pio_spi.pio"],
 )
+
+# TODO: https://github.com/raspberrypi/pico-sdk/issues/2055 Support storing
+# Wi-Fi firmware in a separate partition
+filegroup(
+    name = "pico_use_partition_firmware",
+    srcs = [
+        "wifi_firmware.S",
+        "include/cyw43_partition_firmware.h",
+    ]
+)

src/rp2_common/pico_cyw43_driver/CMakeLists.txt

@@ -165,5 +165,89 @@ if (EXISTS ${PICO_CYW43_DRIVER_PATH}/${CYW43_DRIVER_TEST_FILE})
                     )
     endfunction()
 
+    set(PICO_CYW43_DRIVER_CURRENT_PATH ${CMAKE_CURRENT_LIST_DIR})
+    pico_register_common_scope_var(PICO_CYW43_DRIVER_CURRENT_PATH)
+
     pico_promote_common_scope_vars()
+
+    # pico_use_wifi_firmware_partition(TARGET [NO_EMBEDDED_PT])
+    # \brief\ Use a partition for the Wi-Fi firmware
+    #
+    # This will read the CYW43 firmware from a partition with the ID 0x776966696669726d,
+    # instead of embedding the firmware blob in the binary. By default it will also embed
+    # a compatible partition table in the binary, but this can be disabled by passing the
+    # NO_EMBEDDED_PT argument (for example, if you need to chain into the binary, it
+    # can't contain a partition table).
+    #
+    # \param\ NO_EMBEDDED_PT If set, will not embed a partition table in the binary
+    function(pico_use_wifi_firmware_partition TARGET)
+        set(options NO_EMBEDDED_PT)
+        cmake_parse_arguments(PARSE_ARGV 1 OPTS "${options}" "" "")
+        if (PICO_PLATFORM STREQUAL "rp2040")
+                message(FATAL_ERROR "RP2040 does not support storing wi-fi firmware in partitions")
+        endif()
+        target_compile_definitions(${TARGET} PRIVATE CYW43_USE_FIRMWARE_PARTITION=1)
+
+        if (NOT OPTS_NO_EMBEDDED_PT)
+            get_target_property(picotool_embed_pt ${TARGET} PICOTOOL_EMBED_PT)
+            if (NOT picotool_embed_pt)
+                pico_embed_pt_in_binary(${TARGET} ${PICO_CYW43_DRIVER_CURRENT_PATH}/wifi_pt.json)
+            endif()
+        endif()
+
+        find_package (Python3 REQUIRED COMPONENTS Interpreter)
+
+        # CYW43 firmware blob
+        add_custom_target(${TARGET}_firmware_blob DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/firmware_blob.S)
+        add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/firmware_blob.S
+                COMMAND ${Python3_EXECUTABLE} ${PICO_CYW43_DRIVER_CURRENT_PATH}/cyw43_firmware.py ${PICO_CYW43_DRIVER_PATH}/firmware/wb43439A0_7_95_49_00_combined.h ${CMAKE_CURRENT_BINARY_DIR}/firmware_blob.S
+                )
+
+        # Create UF2s for regular and TBYB firmwares
+        add_executable(${TARGET}_firmware
+                ${PICO_CYW43_DRIVER_CURRENT_PATH}/wifi_firmware.S)
+        add_executable(${TARGET}_firmware_tbyb
+                ${PICO_CYW43_DRIVER_CURRENT_PATH}/wifi_firmware.S)
+
+        add_dependencies(${TARGET}_firmware ${TARGET}_firmware_blob)
+        add_dependencies(${TARGET}_firmware_tbyb ${TARGET}_firmware_blob)
+
+        target_include_directories(${TARGET}_firmware PRIVATE ${CMAKE_CURRENT_BINARY_DIR})
+        target_include_directories(${TARGET}_firmware_tbyb PRIVATE ${CMAKE_CURRENT_BINARY_DIR})
+
+        target_compile_definitions(${TARGET}_firmware PRIVATE
+                NO_PICO_PLATFORM=1
+                )
+        target_compile_definitions(${TARGET}_firmware_tbyb PRIVATE
+                NO_PICO_PLATFORM=1
+                PICO_CRT0_IMAGE_TYPE_TBYB=1
+                )
+
+        target_link_options(${TARGET}_firmware PRIVATE -nostartfiles -nodefaultlibs LINKER:--script=${PICO_CYW43_DRIVER_CURRENT_PATH}/wifi_firmware.ld)
+        target_link_options(${TARGET}_firmware_tbyb PRIVATE -nostartfiles -nodefaultlibs LINKER:--script=${PICO_CYW43_DRIVER_CURRENT_PATH}/wifi_firmware.ld)
+
+        target_link_libraries(${TARGET}_firmware boot_picobin_headers)
+        target_link_libraries(${TARGET}_firmware_tbyb boot_picobin_headers)
+
+        get_target_property(hasSigfile ${TARGET} PICOTOOL_SIGFILE)
+        if (hasSigfile)
+                pico_sign_binary(${TARGET}_firmware ${sigfile})
+                pico_sign_binary(${TARGET}_firmware_tbyb ${sigfile})
+        endif()
+
+        pico_hash_binary(${TARGET}_firmware)
+        pico_hash_binary(${TARGET}_firmware_tbyb)
+
+        pico_set_uf2_family(${TARGET}_firmware cyw43-firmware)
+        pico_set_uf2_family(${TARGET}_firmware_tbyb cyw43-firmware)
+
+        pico_package_uf2_output(${TARGET}_firmware 0x10000000)
+        pico_package_uf2_output(${TARGET}_firmware_tbyb 0x10000000)
+
+        pico_add_extra_outputs(${TARGET}_firmware)
+        pico_add_extra_outputs(${TARGET}_firmware_tbyb)
+
+        add_dependencies(${TARGET}
+                ${TARGET}_firmware ${TARGET}_firmware_tbyb)
+    endfunction()
 endif()

src/rp2_common/pico_cyw43_driver/cyw43_driver.c

@@ -19,7 +19,18 @@
 #define CYW43_SLEEP_CHECK_MS 50
 #endif
 
-static async_context_t *cyw43_async_context;
+static async_context_t *cyw43_async_context = NULL;
+
+#if CYW43_USE_FIRMWARE_PARTITION
+    #include "pico/bootrom.h"
+    #include "hardware/flash.h"
+    #include "boot/picobin.h"
+    #include <stdlib.h>
+
+    int32_t cyw43_wifi_fw_len;
+    int32_t cyw43_clm_len;
+    uintptr_t fw_data;
+#endif
 
 static void cyw43_sleep_timeout_reached(async_context_t *context, async_at_time_worker_t *worker);
 static void cyw43_do_poll(async_context_t *context, async_when_pending_worker_t *worker);
@@ -104,6 +115,87 @@ static void cyw43_sleep_timeout_reached(async_context_t *context, __unused async
 }
 
 bool cyw43_driver_init(async_context_t *context) {
+#if CYW43_USE_FIRMWARE_PARTITION
+    uint32_t buffer[(16 * 4) + 1] = {}; // maximum of 16 partitions, each with maximum of 4 words returned, plus 1
+    int ret = rom_get_partition_table_info(buffer, count_of(buffer), PT_INFO_PARTITION_LOCATION_AND_FLAGS | PT_INFO_PARTITION_ID);
+
+    assert(buffer[0] == (PT_INFO_PARTITION_LOCATION_AND_FLAGS | PT_INFO_PARTITION_ID));
+
+    if (ret > 0) {
+        int i = 1;
+        int p = 0;
+        int picked_p = -1;
+        while (i < ret) {
+            i++;
+            uint32_t flags_and_permissions = buffer[i++];
+            bool has_id = (flags_and_permissions & PICOBIN_PARTITION_FLAGS_HAS_ID_BITS);
+            if (has_id) {
+                uint64_t id = 0;
+                id |= buffer[i++];
+                id |= ((uint64_t)(buffer[i++]) << 32ull);
+                if (id == CYW43_FIRMWARE_PARTITION_ID) {
+                    picked_p = p;
+                    break;
+                }
+            }
+
+            p++;
+        }
+
+        if (picked_p >= 0) {
+            #ifdef __riscv
+                // Increased bootrom stack is required for this function
+                bootrom_stack_t stack = {
+                    .base = malloc(0x400),
+                    .size = 0x400
+                };
+                rom_set_bootrom_stack(&stack);
+            #endif
+            uint32_t* workarea = malloc(0x1000);
+            picked_p = rom_pick_ab_update_partition(workarea, 0x1000, picked_p);
+            free(workarea);
+            #ifdef __riscv
+                // Reset bootrom stack
+                rom_set_bootrom_stack(&stack);
+                free(stack.base);
+            #endif
+
+            if (picked_p < 0) {
+                if (picked_p == BOOTROM_ERROR_NOT_FOUND) {
+                    CYW43_DEBUG("Chosen CYW43 firmware partition was not verified\n");
+                } else if (picked_p == BOOTROM_ERROR_NOT_PERMITTED) {
+                    CYW43_DEBUG("Too many update boots going on at once\n");
+                }
+                return false;
+            }
+
+            CYW43_DEBUG("Chosen CYW43 firmware in partition %d\n", picked_p);
+            int ret = rom_get_partition_table_info(buffer, count_of(buffer), PT_INFO_PARTITION_LOCATION_AND_FLAGS | PT_INFO_SINGLE_PARTITION | (picked_p << 24));
+            assert(buffer[0] == (PT_INFO_PARTITION_LOCATION_AND_FLAGS | PT_INFO_SINGLE_PARTITION));
+            assert(ret == 3);
+
+            uint32_t location_and_permissions = buffer[1];
+            uint32_t saddr = ((location_and_permissions >> PICOBIN_PARTITION_LOCATION_FIRST_SECTOR_LSB) & 0x1fffu) * FLASH_SECTOR_SIZE;
+            uint32_t eaddr = (((location_and_permissions >> PICOBIN_PARTITION_LOCATION_LAST_SECTOR_LSB) & 0x1fffu) + 1) * FLASH_SECTOR_SIZE;
+            // Starts with metadata block
+            while(saddr < eaddr && *(uint32_t*)(XIP_NOCACHE_NOALLOC_NOTRANSLATE_BASE + saddr) != PICOBIN_BLOCK_MARKER_END) {
+                saddr += 4;
+            }
+            saddr += 4;
+            // Now onto data
+            cyw43_wifi_fw_len = *(uint32_t*)(XIP_NOCACHE_NOALLOC_NOTRANSLATE_BASE + saddr);
+            cyw43_clm_len = *(uint32_t*)(XIP_NOCACHE_NOALLOC_NOTRANSLATE_BASE + saddr + 4);
+            fw_data = XIP_NOCACHE_NOALLOC_NOTRANSLATE_BASE + saddr + 8;
+        } else {
+            CYW43_DEBUG("No CYW43 firmware partition found, so cannot get firmware from partition\n");
+            return false;
+        }
+    } else {
+        CYW43_DEBUG("No partition table, so cannot get firmware from partition - get_partition_table_info returned %d\n", ret);
+        return false;
+    }
+
+#endif
     cyw43_init(&cyw43_state);
     cyw43_async_context = context;
     // we need the IRQ to be on the same core as the context, because we need to be able to enable/disable the IRQ
@@ -114,13 +206,15 @@ bool cyw43_driver_init(async_context_t *context) {
 }
 
 void cyw43_driver_deinit(async_context_t *context) {
-    assert(context == cyw43_async_context);
-    async_context_remove_at_time_worker(context, &sleep_timeout_worker);
-    async_context_remove_when_pending_worker(context, &cyw43_poll_worker);
-    // the IRQ IS on the same core as the context, so must be de-initialized there
-    async_context_execute_sync(context, cyw43_irq_deinit, NULL);
-    cyw43_deinit(&cyw43_state);
-    cyw43_async_context = NULL;
+    if (cyw43_async_context != NULL) {
+        assert(context == cyw43_async_context);
+        async_context_remove_at_time_worker(context, &sleep_timeout_worker);
+        async_context_remove_when_pending_worker(context, &cyw43_poll_worker);
+        // the IRQ IS on the same core as the context, so must be de-initialized there
+        async_context_execute_sync(context, cyw43_irq_deinit, NULL);
+        cyw43_deinit(&cyw43_state);
+        cyw43_async_context = NULL;
+    }
 }
 
 // todo maybe add an #ifdef in cyw43_driver

src/rp2_common/pico_cyw43_driver/cyw43_firmware.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2024 Raspberry Pi (Trading) Ltd.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+import sys
+import re
+
+assert len(sys.argv) == 3
+
+cyw43_wifi_fw_len = -1
+cyw43_clm_len = -1
+
+with open(sys.argv[1], "r") as f:
+    data = f.read()
+    statements = data.split(";")
+    for line in statements[1].split("\n"):
+        if "#define CYW43_WIFI_FW_LEN" in line:
+            matches = re.search(r"#define\s+\S+\s+\((\S+)\)", line)
+            cyw43_wifi_fw_len = int(matches[1])
+        elif "#define CYW43_CLM_LEN" in line:
+            matches = re.search(r"#define\s+\S+\s+\((\S+)\)", line)
+            cyw43_clm_len = int(matches[1])
+        if cyw43_wifi_fw_len > 0 and cyw43_clm_len > 0:
+            break
+    data = statements[0]
+    bits = data.split(",")
+    bits[0] = bits[0].split("{")[-1]
+    bits[-1] = bits[-1].split("}")[0]
+    for i in range(len(bits)):
+        bits[i] = bits[i].strip()
+        bits[i] = bits[i].strip(',')
+        bits[i] = int(bits[i], base=0)
+
+data = (
+    cyw43_wifi_fw_len.to_bytes(4, 'little', signed=True) +
+    cyw43_clm_len.to_bytes(4, 'little', signed=True) +
+    bytearray(bits)
+)
+
+with open(sys.argv[2], "w") as f:
+    for b in data:
+        f.write(f".byte 0x{b:02x}\n")

src/rp2_common/pico_cyw43_driver/include/cyw43_configport.h

@@ -63,12 +63,16 @@ extern "C" {
 #endif
 
 #ifndef CYW43_CHIPSET_FIRMWARE_INCLUDE_FILE
+#if CYW43_USE_FIRMWARE_PARTITION
+#define CYW43_CHIPSET_FIRMWARE_INCLUDE_FILE "cyw43_partition_firmware.h"
+#else
 #if CYW43_ENABLE_BLUETOOTH
 #define CYW43_CHIPSET_FIRMWARE_INCLUDE_FILE "wb43439A0_7_95_49_00_combined.h"
 #else
 #define CYW43_CHIPSET_FIRMWARE_INCLUDE_FILE "w43439A0_7_95_49_00_combined.h"
 #endif
 #endif
+#endif
 
 #ifndef CYW43_WIFI_NVRAM_INCLUDE_FILE
 #define CYW43_WIFI_NVRAM_INCLUDE_FILE "wifi_nvram_43439.h"

src/rp2_common/pico_cyw43_driver/include/cyw43_partition_firmware.h

@@ -0,0 +1,15 @@
+/*
+ * Copyright (c) 2024 Raspberry Pi (Trading) Ltd.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+extern int cyw43_wifi_fw_len;
+extern int cyw43_clm_len;
+
+#define CYW43_WIFI_FW_LEN (cyw43_wifi_fw_len)
+#define CYW43_CLM_LEN (cyw43_clm_len)
+extern uintptr_t fw_data;
+
+#include "boot/picobin.h"
+#include "pico/bootrom.h"

src/rp2_common/pico_cyw43_driver/include/pico/cyw43_driver.h

@@ -20,6 +20,14 @@
 #include "cyw43_configport.h"
 #endif
 
+#if CYW43_USE_FIRMWARE_PARTITION
+// PICO_CONFIG: CYW43_FIRMWARE_PARTITION_ID, ID of Wi-Fi firmware partition which must match the ID used in the partition table JSON, type=int, default=0x776966696669726d, group=pico_cyw43_driver
+#ifndef CYW43_FIRMWARE_PARTITION_ID
+// The default 0x776966696669726d value is the ASCII encoding of "wififirm"
+#define CYW43_FIRMWARE_PARTITION_ID 0x776966696669726d
+#endif
+#endif
+
 #ifdef __cplusplus
 extern "C" {
 #endif