ota.rs
1 use alloc::boxed::Box; 2 use core::sync::atomic::{AtomicBool, Ordering}; 3 use defmt::info; 4 use embassy_net::{Stack, tcp::TcpSocket}; 5 use embassy_time::{Duration, Timer}; 6 use embedded_storage::Storage; 7 use esp_bootloader_esp_idf::ota::OtaImageState; 8 use esp_bootloader_esp_idf::ota_updater::OtaUpdater; 9 use esp_bootloader_esp_idf::partitions::PARTITION_TABLE_MAX_LEN; 10 use esp_hal::system::software_reset; 11 use esp_storage::FlashStorage; 12 13 use crate::networking::tcp::read_exact; 14 15 static FIRMWARE_UPGRADE_IN_PROGRESS: AtomicBool = AtomicBool::new(false); 16 17 const OTA_STATUS_READY: u8 = 0xA5; 18 const OTA_STATUS_BEGIN_FAILED: u8 = 0xE1; 19 20 #[embassy_executor::task] 21 pub async fn task(stack: Stack<'static>, mut flash: FlashStorage<'static>) { 22 info!("OTA receiver listening on TCP port {}", crate::config::app::ota::PORT); 23 24 loop { 25 static mut RX_BUFFER: [u8; crate::config::app::ota::RX_BUF_SIZE] = [0; crate::config::app::ota::RX_BUF_SIZE]; 26 static mut TX_BUFFER: [u8; crate::config::app::ota::TX_BUF_SIZE] = [0; crate::config::app::ota::TX_BUF_SIZE]; 27 28 let mut socket = unsafe { 29 TcpSocket::new( 30 stack, 31 &mut *core::ptr::addr_of_mut!(RX_BUFFER), 32 &mut *core::ptr::addr_of_mut!(TX_BUFFER), 33 ) 34 }; 35 36 socket.set_timeout(Some(Duration::from_secs(10))); 37 38 match socket.accept(crate::config::app::ota::PORT).await { 39 Ok(()) => { 40 if let Some(remote) = socket.remote_endpoint() { 41 info!("OTA host connected from {}", remote); 42 } else { 43 info!("OTA host connected (remote endpoint unavailable)"); 44 } 45 46 let mut header_buffer = [0u8; 8]; 47 if let Err(error) = read_exact(&mut socket, &mut header_buffer).await { 48 info!("failed to read OTA header: {:?}", error); 49 Timer::after(Duration::from_secs(2)).await; 50 continue; 51 } 52 53 let firmware_size = u32::from_le_bytes( 54 header_buffer[..4] 55 .try_into() 56 .expect("header_buffer[..4] is statically 4 bytes"), 57 ); 58 let _target_crc = u32::from_le_bytes( 59 header_buffer[4..8] 60 .try_into() 61 .expect("header_buffer[4..8] is statically 4 bytes"), 62 ); 63 64 info!("OTA header received: size={} bytes", firmware_size); 65 66 let mut pt_buffer = Box::new([0u8; PARTITION_TABLE_MAX_LEN]); 67 let mut ota = match OtaUpdater::new(&mut flash, &mut pt_buffer) { 68 Ok(ota) => ota, 69 Err(error) => { 70 info!("failed to create OTA updater: {:?}", error); 71 let _ = socket.write(&[OTA_STATUS_BEGIN_FAILED]).await; 72 Timer::after(Duration::from_secs(2)).await; 73 continue; 74 } 75 }; 76 77 let (mut target_partition, target_slot) = match ota.next_partition() { 78 Ok(result) => result, 79 Err(error) => { 80 info!("failed to get next OTA partition: {:?}", error); 81 let _ = socket.write(&[OTA_STATUS_BEGIN_FAILED]).await; 82 Timer::after(Duration::from_secs(2)).await; 83 continue; 84 } 85 }; 86 87 info!("OTA target partition: {:?}", target_slot); 88 89 FIRMWARE_UPGRADE_IN_PROGRESS.store(true, Ordering::Release); 90 91 if socket.write(&[OTA_STATUS_READY]).await.is_err() { 92 info!("failed to send OTA ready status"); 93 FIRMWARE_UPGRADE_IN_PROGRESS.store(false, Ordering::Release); 94 Timer::after(Duration::from_secs(2)).await; 95 continue; 96 } 97 98 let mut chunk_buffer = [0u8; crate::config::app::ota::CHUNK_SIZE]; 99 let mut bytes_written: u32 = 0; 100 let mut last_reported_percent: u32 = 0; 101 102 let write_result: Result<(), ()> = loop { 103 let bytes_remaining = (firmware_size).saturating_sub(bytes_written); 104 if bytes_remaining == 0 { 105 break Ok(()); 106 } 107 108 let bytes_to_read = (bytes_remaining as usize).min(crate::config::app::ota::CHUNK_SIZE); 109 if let Err(error) = 110 read_exact(&mut socket, &mut chunk_buffer[..bytes_to_read]).await 111 { 112 info!("failed to read OTA chunk: {:?}", error); 113 break Err(()); 114 } 115 116 if let Err(error) = target_partition.write(bytes_written, &chunk_buffer[..bytes_to_read]) { 117 info!("failed to write OTA chunk at offset {}: {:?}", bytes_written, error); 118 break Err(()); 119 } 120 121 bytes_written += bytes_to_read as u32; 122 123 let progress = if firmware_size > 0 { 124 (bytes_written as u64 * 100 / firmware_size as u64) as u32 125 } else { 126 0 127 }; 128 if progress >= last_reported_percent + 5 || progress == 100 { 129 info!( 130 "OTA progress: {}% ({}/{} bytes)", 131 progress, bytes_written, firmware_size 132 ); 133 last_reported_percent = progress; 134 } 135 136 if socket.write(&[0]).await.is_err() { 137 info!("failed to ACK OTA chunk"); 138 break Err(()); 139 } 140 }; 141 142 FIRMWARE_UPGRADE_IN_PROGRESS.store(false, Ordering::Release); 143 144 if write_result.is_err() { 145 Timer::after(Duration::from_secs(2)).await; 146 continue; 147 } 148 149 info!("OTA payload received ({} bytes), activating partition", bytes_written); 150 151 if let Err(error) = ota.activate_next_partition() { 152 info!("failed to activate next partition: {:?}", error); 153 Timer::after(Duration::from_secs(2)).await; 154 continue; 155 } 156 157 if let Err(error) = ota.set_current_ota_state(OtaImageState::New) { 158 info!("failed to set OTA state to New: {:?}", error); 159 } 160 161 info!("OTA complete, rebooting into new firmware"); 162 Timer::after(Duration::from_millis(1000)).await; 163 software_reset(); 164 } 165 Err(error) => { 166 info!("OTA accept failed: {:?}", error); 167 Timer::after(Duration::from_secs(2)).await; 168 } 169 } 170 } 171 } 172 173 pub fn spawn(spawner: &embassy_executor::Spawner, stack: Stack<'static>, flash: FlashStorage<'static>) { 174 spawner.spawn(task(stack, flash).unwrap()); 175 }