Use custom partition table to prevent firmware getting overwritten

firmware/.vscode/launch.json

@@ -10,6 +10,14 @@
       "coreConfigs": [
         {
           "programBinary": "target/xtensa-esp32s3-none-elf/debug/acid-firmware",
+          "rttEnabled": true,
+          "rttChannelFormats": [
+            {
+              "channelNumber": 0,
+              "dataFormat": "String",
+              "mode": "BlockIfFull"
+            }
+          ]
         },
       ],
     },
@@ -22,6 +30,14 @@
       "coreConfigs": [
         {
           "programBinary": "target/xtensa-esp32s3-none-elf/release/acid-firmware",
+          "rttEnabled": true,
+          "rttChannelFormats": [
+            {
+              "channelNumber": 0,
+              "dataFormat": "String",
+              "mode": "BlockIfFull"
+            }
+          ]
         },
       ],
     },
@@ -32,12 +48,23 @@
       "request": "launch",
       "flashingConfig": {
         "flashingEnabled": true,
+        "formatOptions": {
+          "idf_partition_table": "partition-table.csv"
+        }
       },
       "probe": "303a:1001",
       "chip": "esp32s3",
       "coreConfigs": [
         {
           "programBinary": "target/xtensa-esp32s3-none-elf/debug/acid-firmware",
+          "rttEnabled": true,
+          "rttChannelFormats": [
+            {
+              "channelNumber": 0,
+              "dataFormat": "String",
+              "mode": "BlockIfFull"
+            }
+          ]
         },
       ],
     },
@@ -48,12 +75,23 @@
       "request": "launch",
       "flashingConfig": {
         "flashingEnabled": true,
+        "formatOptions": {
+          "idf_partition_table": "partition-table.csv"
+        }
       },
       "probe": "303a:1001",
       "chip": "esp32s3",
       "coreConfigs": [
         {
           "programBinary": "target/xtensa-esp32s3-none-elf/release/acid-firmware",
+          "rttEnabled": true,
+          "rttChannelFormats": [
+            {
+              "channelNumber": 0,
+              "dataFormat": "String",
+              "mode": "BlockIfFull"
+            }
+          ]
         },
       ],
     },

firmware/.vscode/tasks.json

@@ -5,7 +5,10 @@
 			"label": "rust: cargo build",
 			"type": "cargo",
 			"options": {
-				"cwd": "${workspaceFolder}/acid-firmware"
+				"cwd": "${workspaceFolder}/acid-firmware",
+				"env": {
+					"ESP_LOG": "info"
+				}
 			},
 			"command": "build",
 			"args": [
@@ -23,7 +26,10 @@
 			"label": "rust: cargo build --release",
 			"type": "cargo",
 			"options": {
-				"cwd": "${workspaceFolder}/acid-firmware"
+				"cwd": "${workspaceFolder}/acid-firmware",
+				"env": {
+					"ESP_LOG": "info"
+				}
 			},
 			"command": "build",
 			"args": [

firmware/acid-firmware/.cargo/config.toml

@@ -1,5 +1,5 @@
 [target.'cfg(all(any(target_arch = "riscv32", target_arch = "xtensa"), target_os = "none"))']
-runner = "espflash flash --monitor"
+runner = "espflash flash --partition-table partition-table.csv --monitor"
 # runner = "probe-rs run --chip esp32s3 --preverify"
 
 [build]

firmware/acid-firmware/Cargo.toml

@@ -87,6 +87,7 @@ chrono = { version = "0.4.43", default-features = false, features = ["alloc", "s
 tinyvec = { version = "1.10.0", default-features = false, features = ["alloc"] }
 esp-metadata-generated = { version = "0.3.0", features = ["esp32s3"] }
 hex = { version = "0.4.3", default-features = false, features = ["alloc"] }
+indoc = "2.0.7"
 
 # A fork of slint with patches for `allocator_api` support.
 # Don't forget to change `slint-build` in build dependencies, if this is changed.

firmware/acid-firmware/src/main.rs

@@ -18,12 +18,15 @@ extern crate alloc;
 
 use core::alloc::Layout;
 use core::cell::RefCell;
+use core::fmt::Write;
 use core::sync::atomic::{AtomicBool, Ordering};
 
 use alloc::boxed::Box;
 use alloc::collections::vec_deque::VecDeque;
+use alloc::string::String;
 use alloc::sync::Arc;
 use alloc::vec;
+use alloc::vec::Vec;
 use embassy_embedded_hal::adapter::BlockingAsync;
 use embassy_embedded_hal::flash::partition::Partition;
 use embassy_executor::Spawner;
@@ -33,6 +36,7 @@ use embassy_sync::mutex::Mutex;
 use embassy_sync::signal::Signal;
 use embassy_time::{Duration, Timer};
 use esp_alloc::{HeapRegion, MemoryCapability};
+use esp_bootloader_esp_idf::partitions::PartitionTable;
 use esp_hal::Blocking;
 use esp_hal::clock::CpuClock;
 use esp_hal::dma::{BurstConfig, DmaDescriptor, DmaTxBuf, ExternalBurstConfig};
@@ -50,6 +54,7 @@ use esp_hal::system::Stack;
 use esp_hal::timer::timg::TimerGroup;
 use esp_rtos::embassy::Executor;
 use esp_storage::FlashStorage;
+use indoc::writedoc;
 use log::{error, info, warn};
 use rmk::channel::{CONTROLLER_CHANNEL, ControllerSub};
 use rmk::config::{BehaviorConfig, PositionalConfig, RmkConfig, StorageConfig, VialConfig};
@@ -242,10 +247,20 @@ async fn main(_spawner: Spawner) {
     };
 
     // Initialize the flash
-    static FLASH: StaticCell<Mutex<CriticalSectionRawMutex, BlockingAsync<FlashStorage>>> =
+    static PARTITION_TABLE_BUFFER: StaticCell<Vec<u8, &'static esp_alloc::EspHeap>> =
         StaticCell::new();
-    let flash = FLASH.init_with(|| {
+    let partition_table_buffer = PARTITION_TABLE_BUFFER.init_with(|| {
-        let flash = FlashStorage::new(peripherals.FLASH)
+        let mut buffer = Vec::<u8, _>::new_in(&PSRAM_ALLOCATOR);
+        buffer.resize(1024, 0_u8);
+        buffer
+    });
+
+    static FLASH: StaticCell<(
+        Mutex<CriticalSectionRawMutex, BlockingAsync<FlashStorage>>,
+        PartitionTable<'static>,
+    )> = StaticCell::new();
+    let (flash, partition_table) = FLASH.init_with(|| {
+        let mut flash = FlashStorage::new(peripherals.FLASH)
             // Flash memory may not be written to while another core is executing from it.
             // By default, `FlashStorage` is configured to abort the operation and log an error message.
             // However, it can also be configured to auto-park the other core, such that writing to
@@ -254,19 +269,68 @@ async fn main(_spawner: Spawner) {
             // to avoid having to park the other core, which could result in better performance.
             // Invalid configuration would then present itself as freezing/UB.
             .multicore_auto_park();
+        let partition_table = {
+            esp_bootloader_esp_idf::partitions::read_partition_table(
+                &mut flash,
+                partition_table_buffer,
+            )
+            .expect("Failed to read the partition table.")
+        };
 
-        Mutex::<CriticalSectionRawMutex, _>::new(async_flash_wrapper(flash))
+        (
+            Mutex::<CriticalSectionRawMutex, _>::new(async_flash_wrapper(flash)),
+            partition_table,
+        )
     });
-    const FLASH_SIZE_TOTAL: u32 = 16 * 1024 * 1024;
+
-    const FLASH_PART_FIRMWARE_OFFSET: u32 = 0;
+    {
-    const FLASH_PART_FIRMWARE_SIZE: u32 = 0x3f0000;
+        let mut buffer = String::new();
-    const FLASH_PART_RMK_OFFSET: u32 = FLASH_PART_FIRMWARE_OFFSET + FLASH_PART_FIRMWARE_SIZE;
+
-    const FLASH_PART_RMK_SIZE_IN_SECTORS: u32 = 16;
+        writeln!(buffer, "Partition table:").unwrap();
-    const FLASH_PART_RMK_SIZE: u32 = FLASH_PART_RMK_SIZE_IN_SECTORS * FlashStorage::SECTOR_SIZE;
+
-    const FLASH_PART_ACID_OFFSET: u32 = FLASH_PART_RMK_OFFSET + FLASH_PART_RMK_SIZE;
+        for (index, partition) in partition_table.iter().enumerate() {
-    const FLASH_PART_ACID_SIZE: u32 = FLASH_SIZE_TOTAL - FLASH_PART_ACID_OFFSET;
+            writedoc!(
-    let flash_part_rmk = Partition::new(flash, FLASH_PART_RMK_OFFSET, FLASH_PART_RMK_SIZE);
+                buffer,
-    let flash_part_acid = Partition::new(flash, FLASH_PART_ACID_OFFSET, FLASH_PART_ACID_SIZE);
+                "
+                Partition #{index} {label:?}:
+                    offset: 0x{offset:x}
+                    length: 0x{len:x}
+                    type: 0x{type:?}
+                    read only: {read_only}
+                    encrypted: {encrypted}
+                    magic: {magic}
+                ",
+                label = partition.label_as_str(),
+                offset = partition.offset(),
+                len = partition.len(),
+                type = partition.partition_type(),
+                read_only = partition.is_read_only(),
+                encrypted = partition.is_encrypted(),
+                magic = partition.magic(),
+            )
+            .unwrap();
+        }
+
+        info!("{}", buffer);
+    }
+    let flash_part_info_rmk = partition_table
+        .iter()
+        .find(|partition| partition.label_as_str() == "rmk")
+        .expect("No \"rmk\" partition found. Make sure to use the custom partition-table.csv when flashing.");
+    let flash_part_info_acid = partition_table
+        .iter()
+        .find(|partition| partition.label_as_str() == "acid")
+        .expect("No \"acid\" partition found. Make sure to use the custom partition-table.csv when flashing.");
+    let flash_part_rmk = Partition::new(
+        flash,
+        flash_part_info_rmk.offset(),
+        flash_part_info_rmk.len(),
+    );
+    let flash_part_acid = Partition::new(
+        flash,
+        flash_part_info_acid.offset(),
+        flash_part_info_acid.len(),
+    );
 
     info!("Flash memory configured!");
 
@@ -313,7 +377,15 @@ async fn main(_spawner: Spawner) {
     let vial_config = VialConfig::new(VIAL_KEYBOARD_ID, VIAL_KEYBOARD_DEF, &[(0, 0), (1, 1)]);
     let storage_config = StorageConfig {
         start_addr: 0,
-        num_sectors: FLASH_PART_RMK_SIZE_IN_SECTORS as u8,
+        num_sectors: {
+            assert!(
+                flash_part_info_rmk.len() as u32 % FlashStorage::SECTOR_SIZE == 0,
+                "The size of the RMK partition must be a multiple of {} bytes. Current size: {}",
+                FlashStorage::SECTOR_SIZE,
+                flash_part_info_rmk.len()
+            );
+            (flash_part_info_rmk.len() as u32 / FlashStorage::SECTOR_SIZE) as u8
+        },
         ..Default::default()
     };
     let rmk_config = RmkConfig {