main.rs (8621B)
1 #![no_std] 2 #![no_main] 3 #![feature(type_alias_impl_trait)] 4 5 extern crate alloc; 6 use core::str::FromStr; 7 use embassy_executor::Spawner; 8 use embassy_net::{Config, Runner, Stack, StackResources, tcp::TcpSocket}; 9 use embassy_time::{Duration, Timer}; 10 use esp_backtrace as _; 11 use esp_hal::timer::timg::TimerGroup; 12 use esp_hal_ota::Ota; 13 use esp_radio::{ 14 Controller, 15 wifi::{ClientConfig, ModeConfig, WifiController, WifiDevice, WifiEvent, WifiStaState}, 16 }; 17 use esp_storage::FlashStorage; 18 19 const WIFI_SSID: &'static str = env!("SSID"); 20 const WIFI_PSK: &'static str = env!("PSK"); 21 const OTA_SERVER_IP: &'static str = env!("OTA_IP"); 22 23 const RX_BUFFER_SIZE: usize = 16384; 24 const TX_BUFFER_SIZE: usize = 16384; 25 static mut TX_BUFF: [u8; TX_BUFFER_SIZE] = [0; TX_BUFFER_SIZE]; 26 static mut RX_BUFF: [u8; RX_BUFFER_SIZE] = [0; RX_BUFFER_SIZE]; 27 28 macro_rules! mk_static { 29 ($t:ty,$val:expr) => {{ 30 static STATIC_CELL: static_cell::StaticCell<$t> = static_cell::StaticCell::new(); 31 #[deny(unused_attributes)] 32 let x = STATIC_CELL.uninit().write(($val)); 33 x 34 }}; 35 } 36 37 esp_bootloader_esp_idf::esp_app_desc!(); 38 39 #[esp_rtos::main] 40 async fn main(spawner: Spawner) { 41 #[cfg(not(feature = "esp32"))] 42 { 43 esp_alloc::heap_allocator!(size: 150 * 1024); 44 } 45 46 #[cfg(feature = "esp32")] 47 { 48 static mut HEAP: core::mem::MaybeUninit<[u8; 30 * 1024]> = core::mem::MaybeUninit::uninit(); 49 50 #[unsafe(link_section = ".dram2_uninit")] 51 static mut HEAP2: core::mem::MaybeUninit<[u8; 64 * 1024]> = 52 core::mem::MaybeUninit::uninit(); 53 54 unsafe { 55 #[allow(static_mut_refs)] 56 esp_alloc::HEAP.add_region(esp_alloc::HeapRegion::new( 57 HEAP.as_mut_ptr() as *mut u8, 58 core::mem::size_of_val(&*core::ptr::addr_of!(HEAP)), 59 esp_alloc::MemoryCapability::Internal.into(), 60 )); 61 62 #[allow(static_mut_refs)] 63 esp_alloc::HEAP.add_region(esp_alloc::HeapRegion::new( 64 HEAP2.as_mut_ptr() as *mut u8, 65 core::mem::size_of_val(&*core::ptr::addr_of!(HEAP2)), 66 esp_alloc::MemoryCapability::Internal.into(), 67 )); 68 } 69 } 70 71 let peripherals = esp_hal::init(esp_hal::Config::default()); 72 //let io = Io::new(peripherals.GPIO, peripherals.IO_MUX); 73 74 esp_println::logger::init_logger_from_env(); 75 //log::set_max_level(log::LevelFilter::Info); // only for esp32s3?? 76 77 let timg0 = TimerGroup::new(peripherals.TIMG0); 78 79 #[cfg(any(feature = "esp32s3", feature = "esp32"))] 80 esp_rtos::start(timg0.timer0); 81 82 #[cfg(feature = "esp32c3")] 83 let software_interrupt = 84 esp_hal::interrupt::software::SoftwareInterruptControl::new(peripherals.SW_INTERRUPT); 85 86 #[cfg(feature = "esp32c3")] 87 esp_rtos::start(timg0.timer0, software_interrupt.software_interrupt0); 88 89 let init = &*mk_static!(Controller<'static>, esp_radio::init().unwrap()); 90 91 let wifi = peripherals.WIFI; 92 let (controller, interfaces) = esp_radio::wifi::new(&init, wifi, Default::default()).unwrap(); 93 94 let config = Config::dhcpv4(Default::default()); 95 let seed = 69420; 96 97 let (stack, runner) = embassy_net::new( 98 interfaces.sta, 99 config, 100 { 101 static STATIC_CELL: static_cell::StaticCell<StackResources<3>> = 102 static_cell::StaticCell::new(); 103 STATIC_CELL.uninit().write(StackResources::<3>::new()) 104 }, 105 seed, 106 ); 107 108 spawner 109 .spawn(connection(controller, stack)) 110 .expect("connection spawn"); 111 spawner.spawn(net_task(runner)).expect("net task spawn"); 112 113 // mark ota partition valid 114 { 115 let mut ota = Ota::new(FlashStorage::new(unsafe { 116 peripherals.FLASH.clone_unchecked() 117 })) 118 .expect("Cannot create ota"); 119 _ = ota.ota_mark_app_valid(); //do not unwrap here if using factory/test partition 120 } 121 122 loop { 123 log::info!("Wait for wifi!"); 124 Timer::after(Duration::from_secs(1)).await; 125 126 if let Some(config) = stack.config_v4() { 127 log::info!("Got IP: {}", config.address); 128 break; 129 } 130 } 131 132 let mut socket = unsafe { 133 TcpSocket::new( 134 stack, 135 &mut *core::ptr::addr_of_mut!(RX_BUFF), 136 &mut *core::ptr::addr_of_mut!(TX_BUFF), 137 ) 138 }; 139 140 let ip = embassy_net::IpEndpoint::from_str(OTA_SERVER_IP).expect("Wrong ip addr"); 141 socket.connect(ip).await.expect("Cannot connect!"); 142 let mut ota_buff = [0; 4096 * 2]; 143 socket 144 .read(&mut ota_buff[..4]) 145 .await 146 .expect("Cannot read firmware size!"); 147 let flash_size = u32::from_le_bytes(ota_buff[..4].try_into().unwrap()); 148 149 socket 150 .read(&mut ota_buff[..4]) 151 .await 152 .expect("Cannot read target crc!"); 153 let target_crc = u32::from_le_bytes(ota_buff[..4].try_into().unwrap()); 154 155 log::info!("flash_size: {flash_size}"); 156 log::info!("target_crc: {target_crc}"); 157 158 let mut ota = Ota::new(FlashStorage::new(unsafe { 159 peripherals.FLASH.clone_unchecked() 160 })) 161 .expect("Cannot create ota"); 162 ota.ota_begin(flash_size, target_crc).unwrap(); 163 164 let mut bytes_read: usize = 0; 165 loop { 166 let bytes_to_read = 8192.min(flash_size - bytes_read as u32); 167 let res = read_exact(&mut socket, &mut ota_buff[..bytes_to_read as usize]).await; 168 if let Ok(_) = res { 169 bytes_read += bytes_to_read as usize; 170 if bytes_to_read == 0 { 171 break; 172 } 173 174 let res = ota.ota_write_chunk(&ota_buff[..bytes_to_read as usize]); 175 _ = socket.write(&[0]).await; 176 177 match res { 178 Ok(true) => { 179 let res = ota.ota_flush(false, true); 180 if let Err(e) = res { 181 log::error!("Ota flush error: {e:?}"); 182 break; 183 } 184 185 log::info!("Ota OK! Rebooting in 1s!"); 186 Timer::after_millis(1000).await; 187 esp_hal::system::software_reset(); 188 } 189 Err(e) => { 190 log::error!("Ota write error: {e:?}"); 191 break; 192 } 193 _ => {} 194 } 195 } 196 197 log::info!("Progress: {}%", (ota.get_ota_progress() * 100.0) as u8); 198 } 199 200 loop { 201 log::info!("bump"); 202 Timer::after_millis(15000).await; 203 } 204 } 205 206 #[embassy_executor::task] 207 async fn connection(mut controller: WifiController<'static>, stack: Stack<'static>) { 208 log::info!("start connection task"); 209 log::info!("Device capabilities: {:?}", controller.capabilities()); 210 loop { 211 if esp_radio::wifi::sta_state() == WifiStaState::Connected { 212 // wait until we're no longer connected 213 controller.wait_for_event(WifiEvent::StaDisconnected).await; 214 Timer::after(Duration::from_millis(5000)).await 215 } 216 if !matches!(controller.is_started(), Ok(true)) { 217 let client_config = ModeConfig::Client( 218 ClientConfig::default() 219 .with_ssid(WIFI_SSID.try_into().expect("Wifi ssid parse")) 220 .with_password(WIFI_PSK.try_into().expect("Wifi psk parse")), 221 ); 222 controller.set_config(&client_config).unwrap(); 223 log::info!("Starting wifi"); 224 controller.start_async().await.unwrap(); 225 log::info!("Wifi started!"); 226 } 227 log::info!("About to connect..."); 228 229 match controller.connect_async().await { 230 Ok(_) => { 231 log::info!("Wifi connected!"); 232 233 loop { 234 if stack.is_link_up() { 235 break; 236 } 237 Timer::after(Duration::from_millis(500)).await; 238 } 239 } 240 Err(e) => { 241 log::info!("Failed to connect to wifi: {e:?}"); 242 Timer::after(Duration::from_millis(5000)).await 243 } 244 } 245 } 246 } 247 248 #[embassy_executor::task] 249 async fn net_task(mut runner: Runner<'static, WifiDevice<'static>>) { 250 runner.run().await 251 } 252 253 pub async fn read_exact( 254 socket: &mut TcpSocket<'_>, 255 mut buf: &mut [u8], 256 ) -> Result<(), embassy_net::tcp::Error> { 257 while !buf.is_empty() { 258 match socket.read(buf).await { 259 Ok(0) => return Err(embassy_net::tcp::Error::ConnectionReset), 260 Ok(n) => { 261 buf = &mut buf[n..]; 262 } 263 Err(e) => return Err(e), 264 } 265 } 266 Ok(()) 267 }