Add uf2 combine command

README.md

@@ -31,7 +31,7 @@ SYNOPSIS:
     picotool reboot [-a] [-u] [-g <partition>] [-c <cpu>] [device-selection]
     picotool otp list|get|set|load|dump|permissions|white-label
     picotool partition info|create
-    picotool uf2 info|convert
+    picotool uf2 info|convert|combine
     picotool version [-s] [<version>]
     picotool coprodis [--quiet] [--verbose] <infile> <outfile>
     picotool help [<cmd>]
@@ -906,6 +906,101 @@ OPTIONS:
             absolute block location (default to 0x10ffff00)
 ```
 
+### combine
+
+This command is used to combine multiple UF2 files (possibly with different family IDs) into a single file with the same family ID. This can be useful for loading a partition table and multiple partitions onto a device using a single file.
+
+```text
+$ picotool help uf2 combine
+UF2 COMBINE:
+    Combine multiple UF2 files.
+
+SYNOPSIS:
+    picotool uf2 combine [--quiet] [--verbose] <infile1> [-t <type>] <infile2> [-t <type>] <outfile> [-t <type>] [--family <family_id>]
+                [--offset <offset>] [--partition <partition>] [[--abs-block] [<abs_block_loc>]]
+
+OPTIONS:
+        --quiet
+            Don't print any output
+        --verbose
+            Print verbose output
+    First file to combine
+        <infile1>
+            The file name
+        -t <type>
+            Specify file type (uf2) explicitly, ignoring file extension
+    Second file to combine
+        <infile2>
+            The file name
+        -t <type>
+            Specify file type (uf2) explicitly, ignoring file extension
+    File to save output to
+        <outfile>
+            The file name
+        -t <type>
+            Specify file type (uf2) explicitly, ignoring file extension
+    UF2 Family options
+        <family_id>
+            family ID for combined UF2 (defaults to first one)
+    Offset options
+        --offset
+            Offset second UF2 by amount
+        <offset>
+            offset amount (default to 0)
+    Partition options
+        --partition
+            Place second UF2 in partition (first UF2 must contain a partition table)
+        <partition>
+            partition number (default to 0)
+    Errata RP2350-E10 Fix
+        --abs-block
+            Add an absolute block
+        <abs_block_loc>
+            absolute block location (default to 0x10ffff00)
+```
+
+The `--partition` argument can be used to place the second file in a partition number, provided that there is a partition table in the first file. For example, take this `pt.json`
+```json
+{
+    "unpartitioned": {
+        "families": ["absolute"],
+        "permissions": {
+            "secure": "rw",
+            "nonsecure": "rw",
+            "bootloader": "rw"
+        }
+    },
+    "partitions": [
+        {
+            "size": "1024K",
+            "families": ["rp2350-arm-s", "rp2350-riscv"],
+            "permissions": {
+                "secure": "rw",
+                "nonsecure": "rw",
+                "bootloader": "rw"
+            }
+        },
+        {
+            "size": "1024K",
+            "families": ["data"],
+            "permissions": {
+                "secure": "rw",
+                "nonsecure": "rw",
+                "bootloader": "rw"
+            }
+        }
+    ]
+}
+```
+
+If you want to load this partition table, and put `code.uf2` and `data.uf2` into the partitions, all using a single UF2, you would run:
+```text
+$ picotool partition create pt.json pt.uf2
+$ picotool uf2 combine pt.uf2 code.uf2 tmp.uf2 --partition 0
+$ picotool uf2 combine tmp.uf2 data.uf2 combined.uf2 --partition 1
+$ picotool load -x combined.uf2
+```
+
 ### info
 
 This command reads the information on a device about why a UF2 download has failed. It will only give information if the most recent download has failed.

elf2uf2/elf2uf2.h

@@ -21,6 +21,7 @@
 #define UF2_PAGE_SIZE (1u << LOG2_PAGE_SIZE)
 
 
+uf2_block gen_abs_block(uint32_t abs_block_loc);
 bool check_abs_block(uf2_block block);
 int bin2uf2(std::shared_ptr<std::iostream> in, std::shared_ptr<std::iostream> out, uint32_t address, uint32_t family_id, model_t model, uint32_t abs_block_loc=0, bool verbose=false);
 int elf2uf2(std::shared_ptr<std::iostream> in, std::shared_ptr<std::iostream> out, uint32_t family_id, model_t model, uint32_t package_addr=0, uint32_t abs_block_loc=0, bool verbose=false);

main.cpp

@@ -597,6 +597,10 @@ struct _settings {
         #else
         uint32_t abs_block_loc = 0;
         #endif
+        uint32_t offset = 0;
+        bool offset_set = false;
+        int partition = -1;
+        bool partition_set = false;
     } uf2;
 };
 _settings settings;
@@ -1370,11 +1374,49 @@ struct uf2_convert_command : public cmd {
     }
 };
 
+struct uf2_combine_command : public cmd {
+    uf2_combine_command() : cmd("combine") {}
+    bool execute(device_map &devices) override;
+    virtual device_support get_device_support() override { return none; }
+
+    group get_cli() override {
+        return (
+                option("--quiet").set(settings.quiet) % "Don't print any output" +
+                option("--verbose").set(settings.verbose) % "Print verbose output" +
+                named_typed_file_selection_x("infile1", 1, "uf2") % "First file to combine" +
+                named_typed_file_selection_x("infile2", 2, "uf2") % "Second file to combine" +
+                named_typed_file_selection_x("outfile", 0, "uf2") % "File to save output to" +
+                (
+                    option("--family") & family_id("family_id").set(settings.family_id) % "family ID for combined UF2 (defaults to first one)"
+                ).force_expand_help(true) % "UF2 Family options" +
+                (
+                    option("--offset").set(settings.uf2.offset_set) % "Offset second UF2 by amount" &
+                        hex("offset").set(settings.uf2.offset) % "offset amount (default to 0)"
+                ).force_expand_help(true) % "Offset options" +
+                (
+                    option("--partition").set(settings.uf2.partition_set) % "Place second UF2 in partition (first UF2 must contain a partition table)" &
+                        integer("partition").min_value(0).max_value(PARTITION_TABLE_MAX_PARTITIONS-1).set(settings.uf2.partition) % "partition number (default to 0)"
+                ).force_expand_help(true) % "Partition options"
+            #if SUPPORT_RP2350_A2
+                + (
+                    option("--abs-block").set(settings.uf2.abs_block) % "Add an absolute block" +
+                        hex("abs_block_loc").set(settings.uf2.abs_block_loc).min(0) % "absolute block location (default to 0x10ffff00)"
+                ).force_expand_help(true).min(0) % "Errata RP2350-E10 Fix"
+            #endif
+        );
+    }
+
+    string get_doc() const override {
+        return "Combine multiple UF2 files.";
+    }
+};
+
 vector<std::shared_ptr<cmd>> uf2_sub_commands {
     #if HAS_LIBUSB
         std::shared_ptr<cmd>(new uf2_info_command()),
     #endif
         std::shared_ptr<cmd>(new uf2_convert_command()),
+        std::shared_ptr<cmd>(new uf2_combine_command()),
 };
 
 struct uf2_command : public multi_cmd {
@@ -3239,6 +3281,25 @@ std::unique_ptr<block> find_last_block(memory_access &raw_access, vector<uint8_t
     return nullptr;
 }
 
+std::vector<std::unique_ptr<block>> find_all_blocks(memory_access &raw_access, vector<uint8_t> &bin) {
+    // todo read the right amount
+    uint32_t read_size = 0x1000;
+    DEBUG_LOG("Reading from %x size %x\n", raw_access.get_binary_start(), read_size);
+    bin = raw_access.read_vector<uint8_t>(raw_access.get_binary_start(), read_size, true);
+
+    std::unique_ptr<block> first_block = find_first_block(bin, raw_access.get_binary_start());
+    if (first_block) {
+        // verify stuff
+        get_more_bin_cb more_cb = [&raw_access](std::vector<uint8_t> &bin, uint32_t offset, uint32_t size) {
+            DEBUG_LOG("Now reading from %x size %x\n", offset, size);
+            bin = raw_access.read_vector<uint8_t>(offset, size, true);
+        };
+        return get_all_blocks(bin, raw_access.get_binary_start(), first_block, more_cb);
+    }
+
+    return std::vector<std::unique_ptr<block>>();
+}
+
 std::shared_ptr<memory_access> get_bi_access(memory_access &raw_access) {
     vector<uint8_t> bin;
     std::unique_ptr<block> best_block = find_best_block(raw_access, bin);
@@ -6775,6 +6836,120 @@ bool uf2_convert_command::execute(device_map &devices) {
     return false;
 }
 
+bool uf2_combine_command::execute(device_map &devices) {
+    if (get_file_type_idx(0) != filetype::uf2 || get_file_type_idx(1) != filetype::uf2 || get_file_type_idx(0) != filetype::uf2) {
+        fail(ERROR_ARGS, "All files must be UF2 files\n");
+    }
+
+    if (settings.uf2.partition_set && settings.uf2.offset_set) {
+        fail(ERROR_ARGS, "Cannot use both partition and offset options together\n");
+    }
+
+    auto out = get_file_idx(ios::out|ios::binary, 0);
+    auto file1 = get_file_idx(ios::in|ios::binary, 1);
+    auto file2 = get_file_idx(ios::in|ios::binary, 2);
+
+    if (settings.uf2.partition_set) {
+        auto access = get_file_memory_access(1);
+
+        vector<uint8_t> bin;
+        auto blocks = find_all_blocks(access, bin);
+        for (auto &block : blocks) {
+            auto partition_table = block->get_item<partition_table_item>();
+            if (partition_table == nullptr) {
+                continue;
+            }
+
+            if (settings.uf2.partition < 0 || settings.uf2.partition >= partition_table->partitions.size()) {
+                fail(ERROR_ARGS, "Partition table only contains partitions 0 -> %d\n", partition_table->partitions.size() - 1);
+            }
+
+            settings.uf2.offset = partition_table->partitions[settings.uf2.partition].first_sector * 4096;
+            settings.uf2.offset_set = true;
+            break;
+        }
+
+        if (!settings.uf2.offset_set) {
+            fail(ERROR_ARGS, "No partition table found in first UF2\n");
+        }
+    }
+
+    out->seekp(0, ios::beg);
+
+    unsigned int num_blocks = 0;
+    for (auto file : {file1, file2}) {
+        uf2_block block;
+        // Seek to 2nd block, in case 1st is abs_block
+        file->seekg(0, ios::beg);
+        file->read((char*)&block, sizeof(uf2_block));
+        if (file->fail()) {
+            fail(ERROR_READ_FAILED, "unexpected end of input file");
+        }
+
+    #if SUPPORT_RP2350_A2
+        if (check_abs_block(block)) {
+            // save abs block address
+            settings.uf2.abs_block = true;
+            settings.uf2.abs_block_loc = block.target_addr;
+            file->read((char*)&block, sizeof(uf2_block));
+        }
+    #endif
+
+        num_blocks += block.num_blocks;
+
+        if (!settings.family_id) {
+            settings.family_id = block.file_size;
+        }
+    }
+
+#if SUPPORT_RP2350_A2
+    if (settings.uf2.abs_block) {
+        uf2_block block = gen_abs_block(settings.uf2.abs_block_loc);
+        out->write((char*)&block, sizeof(uf2_block));
+    }
+#endif
+
+    unsigned int block_no = 0;
+    unsigned int file_no = 0;
+    for (auto file : {file1, file2}) {
+        file->seekg(0, ios::beg);
+        uf2_block block;
+        unsigned int pos = 0;
+        uint32_t next_family_id = 0;
+        do {
+            file->read((char*)&block, sizeof(uf2_block));
+            if (file->fail()) {
+                if (file->eof()) { file->clear(); break; }
+                fail(ERROR_READ_FAILED, "unexpected end of input file");
+            }
+            if (block.magic_start0 == UF2_MAGIC_START0 && block.magic_start1 == UF2_MAGIC_START1 &&
+                block.magic_end == UF2_MAGIC_END) {
+                if (block.flags & UF2_FLAG_FAMILY_ID_PRESENT &&
+                    !(block.flags & UF2_FLAG_NOT_MAIN_FLASH) && block.payload_size == PAGE_SIZE) {
+                    // ignore the absolute block
+                    if (check_abs_block(block)) {
+                        DEBUG_LOG("Ignoring RP2350-E10 absolute block\n");
+                    } else {
+                        block.block_no = block_no; block_no++;
+                        block.num_blocks = num_blocks;
+                        block.file_size = settings.family_id;
+                        if (settings.uf2.offset_set && file_no == 1) {
+                            block.target_addr += settings.uf2.offset;
+                        }
+                        out->write((char*)&block, sizeof(uf2_block));
+                    }
+                }
+            }
+            pos += sizeof(uf2_block);
+        } while (true);
+        file_no++;
+    }
+
+    out->close();
+
+    return false;
+}
+
 
 // Dissassembly helpers
 string gpiodir(int val) {