1  #![no_std]
  2  #![no_main]
  3  
  4  use defmt::info;
  5  use embassy_executor::Spawner;
  6  use embassy_net::{Runner, StackResources, tcp::TcpSocket};
  7  use embassy_time::{Duration, Timer, with_timeout};
  8  use esp_hal::{
  9      clock::CpuClock,
 10      interrupt::software::SoftwareInterruptControl,
 11      rng::Rng,
 12      system::software_reset,
 13      timer::timg::TimerGroup,
 14  };
 15  use esp_hal_ota::Ota;
 16  use esp_radio::wifi::{Config, ControllerConfig, Interface, WifiController, sta::StationConfig};
 17  use esp_storage::FlashStorage;
 18  use panic_rtt_target as _;
 19  use static_cell::StaticCell;
 20  
 21  const WIFI_SSID: &str = env!("WIFI_SSID");
 22  const WIFI_PASSWORD: &str = env!("WIFI_PSK");
 23  const OTA_LISTEN_PORT: u16 = 3232;
 24  
 25  const RX_BUFFER_SIZE: usize = 16384;
 26  const TX_BUFFER_SIZE: usize = 16384;
 27  const OTA_CHUNK_SIZE: usize = 8192;
 28  
 29  esp_bootloader_esp_idf::esp_app_desc!();
 30  
 31  macro_rules! mk_static {
 32      ($type:ty, $value:expr) => {{
 33          static STATIC_CELL: StaticCell<$type> = StaticCell::new();
 34          STATIC_CELL.uninit().write($value)
 35      }};
 36  }
 37  
 38  fn random_seed(random_number_generator: &mut Rng) -> u64 {
 39      (u64::from(random_number_generator.random()) << 32)
 40          | u64::from(random_number_generator.random())
 41  }
 42  
 43  #[embassy_executor::task]
 44  async fn wifi_connection_task(mut wifi_controller: WifiController<'static>) {
 45      loop {
 46          info!("attempting Wi-Fi STA connection");
 47  
 48          match wifi_controller.connect_async().await {
 49              Ok(_connected_info) => {
 50                  info!("Wi-Fi STA connected");
 51                  let _ = wifi_controller.wait_for_disconnect_async().await;
 52                  info!("Wi-Fi STA disconnected");
 53              }
 54              Err(error) => {
 55                  info!("Wi-Fi STA connect failed: {:?}", error);
 56                  Timer::after(Duration::from_secs(5)).await;
 57              }
 58          }
 59      }
 60  }
 61  
 62  #[embassy_executor::task]
 63  async fn wifi_net_task(mut runner: Runner<'static, Interface<'static>>) {
 64      runner.run().await;
 65  }
 66  
 67  async fn read_exact(socket: &mut TcpSocket<'_>, buf: &mut [u8]) -> Result<(), embassy_net::tcp::Error> {
 68      let mut remaining = buf;
 69      while !remaining.is_empty() {
 70          match socket.read(remaining).await {
 71              Ok(0) => return Err(embassy_net::tcp::Error::ConnectionReset),
 72              Ok(n) => {
 73                  let (done, rest) = remaining.split_at_mut(n);
 74                  remaining = rest;
 75                  if done.len() != n {
 76                      break;
 77                  }
 78              }
 79              Err(e) => return Err(e),
 80          }
 81      }
 82      Ok(())
 83  }
 84  
 85  #[esp_rtos::main]
 86  async fn main(spawner: Spawner) -> ! {
 87      rtt_target::rtt_init_defmt!();
 88  
 89      let config = esp_hal::Config::default().with_cpu_clock(CpuClock::max());
 90      let peripherals = esp_hal::init(config);
 91  
 92      esp_alloc::heap_allocator!(#[esp_hal::ram(reclaimed)] size: 73744);
 93      esp_alloc::heap_allocator!(size: 64 * 1024);
 94  
 95      let timg0 = TimerGroup::new(peripherals.TIMG0);
 96      let sw_ints = SoftwareInterruptControl::new(peripherals.SW_INTERRUPT);
 97      esp_rtos::start(timg0.timer0, sw_ints.software_interrupt0);
 98  
 99      info!("Embassy initialized!");
100  
101      {
102          let mut ota = Ota::new(FlashStorage::new()).expect("Cannot create OTA");
103          if let Err(error) = ota.ota_mark_app_valid() {
104              info!("ota_mark_app_valid failed (may be on factory partition): {:?}", error);
105          } else {
106              info!("marked current OTA slot as valid");
107          }
108      }
109  
110      let station_config = Config::Station(
111          StationConfig::default()
112              .with_ssid(WIFI_SSID)
113              .with_password(WIFI_PASSWORD.into()),
114      );
115  
116      info!("starting Wi-Fi in STA mode");
117      let (wifi_controller, interfaces) = esp_radio::wifi::new(
118          peripherals.WIFI,
119          ControllerConfig::default().with_initial_config(station_config),
120      )
121      .expect("Failed to initialize Wi-Fi controller");
122  
123      info!("attempting STA connection to '{}'", WIFI_SSID);
124  
125      let mut random_number_generator = Rng::new();
126      let seed = random_seed(&mut random_number_generator);
127  
128      let network_config = embassy_net::Config::dhcpv4(Default::default());
129      let (stack, runner) = embassy_net::new(
130          interfaces.station,
131          network_config,
132          mk_static!(StackResources<3>, StackResources::<3>::new()),
133          seed,
134      );
135  
136      spawner.spawn(wifi_connection_task(wifi_controller).unwrap());
137      spawner.spawn(wifi_net_task(runner).unwrap());
138  
139      info!("waiting for STA DHCP configuration...");
140  
141      match with_timeout(Duration::from_secs(30), stack.wait_config_up()).await {
142          Ok(()) => {
143              if let Some(ip_config) = stack.config_v4() {
144                  info!("STA connected with IP: {}", ip_config.address);
145              } else {
146                  info!("STA DHCP completed but no IPv4 config available");
147              }
148          }
149          Err(_) => {
150              info!("STA DHCP timed out after 30s");
151          }
152      }
153  
154      loop {
155          if stack.is_link_up() {
156              info!("network link is up");
157              break;
158          }
159          Timer::after(Duration::from_millis(500)).await;
160      }
161  
162      info!("OTA receiver listening on TCP port {}", OTA_LISTEN_PORT);
163  
164      loop {
165          static mut RX_BUFF: [u8; RX_BUFFER_SIZE] = [0; RX_BUFFER_SIZE];
166          static mut TX_BUFF: [u8; TX_BUFFER_SIZE] = [0; TX_BUFFER_SIZE];
167  
168          let mut socket = unsafe {
169              TcpSocket::new(
170                  stack,
171                  &mut *core::ptr::addr_of_mut!(RX_BUFF),
172                  &mut *core::ptr::addr_of_mut!(TX_BUFF),
173              )
174          };
175  
176          socket.set_timeout(Some(Duration::from_secs(10)));
177  
178          match socket.accept(OTA_LISTEN_PORT).await {
179              Ok(()) => {
180                  if let Some(remote) = socket.remote_endpoint() {
181                      info!("host connected from {}", remote);
182                  } else {
183                      info!("host connected, remote endpoint unavailable");
184                  }
185  
186                  info!("receiving firmware header");
187                  let mut header_buffer = [0u8; 8];
188                  if let Err(error) = read_exact(&mut socket, &mut header_buffer).await {
189                      info!("failed to read firmware header: {:?}", error);
190                      Timer::after(Duration::from_secs(5)).await;
191                      continue;
192                  }
193  
194                  let firmware_size = u32::from_le_bytes(header_buffer[..4].try_into().unwrap());
195                  let target_crc = u32::from_le_bytes(header_buffer[4..8].try_into().unwrap());
196  
197                  info!("firmware size: {} bytes, target CRC: {:#010x}", firmware_size, target_crc);
198  
199                  let mut ota = Ota::new(FlashStorage::new()).expect("Cannot create OTA");
200                  ota.ota_begin(firmware_size, target_crc).expect("ota_begin failed");
201  
202                  let mut ota_chunk_buffer = [0u8; OTA_CHUNK_SIZE];
203                  let mut bytes_received_total: usize = 0;
204                  let mut last_reported_percent: u32 = 0;
205  
206                  loop {
207                      let bytes_to_read = OTA_CHUNK_SIZE.min(firmware_size as usize - bytes_received_total);
208                      if bytes_to_read == 0 {
209                          break;
210                      }
211  
212                      if let Err(error) = read_exact(&mut socket, &mut ota_chunk_buffer[..bytes_to_read]).await {
213                          info!("failed to read firmware chunk: {:?}", error);
214                          Timer::after(Duration::from_secs(5)).await;
215                          break;
216                      }
217  
218                      let write_complete = ota
219                          .ota_write_chunk(&ota_chunk_buffer[..bytes_to_read])
220                          .expect("ota_write_chunk failed");
221  
222                      bytes_received_total += bytes_to_read;
223  
224                      let progress = (bytes_received_total as u32 * 100) / firmware_size;
225                      if progress >= last_reported_percent + 5 || progress == 100 {
226                          info!(
227                              "OTA progress: {}% ({}/{} bytes)",
228                              progress,
229                              bytes_received_total,
230                              firmware_size
231                          );
232                          last_reported_percent = progress;
233                      }
234  
235                      socket.write(&[0]).await.ok();
236  
237                      if write_complete {
238                          break;
239                      }
240                  }
241  
242                  info!("firmware received, flushing OTA");
243  
244                  match ota.ota_flush(false, true) {
245                      Ok(()) => {
246                          info!("OTA complete! Rebooting into new firmware");
247                          Timer::after(Duration::from_millis(1000)).await;
248                          software_reset();
249                      }
250                      Err(error) => {
251                          info!("OTA flush failed: {:?}", error);
252                          Timer::after(Duration::from_secs(5)).await;
253                      }
254                  }
255                  software_reset();
256              }
257              Err(error) => {
258                  info!("failed to accept OTA connection, retrying in 5s: {:?}", error);
259                  Timer::after(Duration::from_secs(5)).await;
260              }
261          }
262      }
263  }