eculocate

ota.rs (6004B)

---

 use nb::block;
 
 use embassy_time::{Timer};
 
 use ed25519_dalek::{VerifyingKey, Verifier, Signature};
 
 use esp_hal_ota::Ota;
 use esp_storage::FlashStorage;
 use embedded_io_async::{Read,ReadExactError};
 
 use esp_hal::sha;
 use core::convert::Infallible;
 
 use log::{warn, info};
 
 use alloc::string::String;
 
 
 const ERASE_BLOCK_BYTES : usize  = 4096;
 const ED25519_SIGNATURE_BYTES : usize = 64;
 const SHA256_BYTES : usize = 32;
 
 pub struct Flasher<'a> {
     peripheral: esp_hal::peripherals::FLASH<'a>,
     sha_engine : sha::Sha<'a>
 }
 
 #[derive(Debug)]
 pub enum FlashError {
     UnexpectedEof,
     OtaFailed,
     InvalidSignature,
     Infallible
 }
 
 impl<E> From<ReadExactError<E>> for FlashError {
     fn from(_e : ReadExactError<E>) -> Self {
 	Self::UnexpectedEof
     }
 }
 
 impl From<esp_hal_ota::OtaError> for FlashError {
     fn from(_e :  esp_hal_ota::OtaError) -> Self {
 	Self::OtaFailed
     }
 }
 
 impl From<Infallible> for FlashError {
     fn from(_e :  Infallible) -> Self {
 	Self::Infallible
     }
 }
 
 
 impl Flasher<'_> {
     pub fn new(flash_peripheral: esp_hal::peripherals::FLASH<'static>,
                sha_peripheral: esp_hal::peripherals::SHA<'static>) -> Self {
 	let sha_engine = sha::Sha::new(sha_peripheral);
 
 	Self { peripheral: flash_peripheral, sha_engine }
     }
 
     pub async fn write(&mut self, sock: &mut embassy_net::tcp::TcpSocket<'_>) -> Result<(), FlashError> {
 	let storage = FlashStorage::new(self.peripheral.reborrow());
 	let mut ota = Ota::new(storage).unwrap();
 	write_ota(&mut ota, &mut self.sha_engine, sock).await
     }
 
     pub fn var(&mut self, name: &str) -> Option<String> {
 	let storage = FlashStorage::new(self.peripheral.reborrow());
 	let mut nvs =
             esp_nvs::Nvs::new(0x9000, 0x4000, storage).expect("failed to create nvs");
 	let ns = const { esp_nvs::Key::from_str("eculocate") };
 	match nvs.get(&ns, &esp_nvs::Key::from_str(name)) {
 	    Ok(s) => Some(s),
 	    Err(e) => None
 	}
     }
 
     pub fn set_var(&mut self, name: &str, value: Option<&str>) ->  Result<(), esp_nvs::error::Error> {
 	let storage = FlashStorage::new(self.peripheral.reborrow());
 	let mut nvs =
             esp_nvs::Nvs::new(0x9000, 0x4000, storage).expect("failed to create nvs");
 	let ns = const { esp_nvs::Key::from_str("eculocate") };
 	match value {
 	    Some(v) => {
 		nvs.set(&ns, &esp_nvs::Key::from_str(name), v)
 	    }
 	    None => {
 		nvs.delete(&ns, &esp_nvs::Key::from_str(name))
 	    }
 	}
     }
 }
 
 
 async fn write_ota(ota : &mut Ota<FlashStorage<'_>>,
                    sha_engine : &mut sha::Sha<'_>,
                    sock: &mut embassy_net::tcp::TcpSocket<'_>) -> Result<(), FlashError> {
     let mut buf = [0u8; SHA256_BYTES + ERASE_BLOCK_BYTES];
     let mut buf_offset = 0;
     let mut signed;
     let mut signature = [0u8; ED25519_SIGNATURE_BYTES ];
     let mut next_block_hash = [0u8; SHA256_BYTES];
 
     // ed25519 signature of first block
     // first block is sha256 of second block, plus length as u32
     // second+ block is sha256 of successor block, plus data
 
     let mut block0 = [0u8; SHA256_BYTES + 4];
 
     sock.read_exact(&mut signature).await?;
     sock.read_exact(&mut block0[..SHA256_BYTES + 4]).await?;
 
     let verifying_key_bytes = include_bytes!("../../pubkey.bin");
     let verifying_key: VerifyingKey = VerifyingKey::from_bytes(verifying_key_bytes).expect("verify");
     signed = verifying_key.verify(&mut block0,
                                   &Signature::try_from(&signature).expect("sig")).is_ok();
 
     info!("initial sig checks out? {:?}", signed);
 
     let (sha, len_bytes) = block0.split_at(SHA256_BYTES);
     let len = u32::from_be_bytes(len_bytes.try_into().expect("4 bytes"));
     next_block_hash.clone_from_slice(sha);
 
     info!("ota length {}", len);
 
     ota.ota_begin(len, 0)?;
 
     // total payload to expect is image size plus per-block hashes
     let mut remaining = len + (len/(ERASE_BLOCK_BYTES as u32)+ 1) * (SHA256_BYTES as u32);
 
     while remaining > 0 {
         // info!("remaining {}, offset {}", remaining, buf_offset);
         match sock.read(&mut buf[buf_offset..]).await {
             Ok(0) => {
                 info!("unexpected eof on read");
                 break;
             }
             Ok(chunksize) => {
                 buf_offset += chunksize;
                 remaining -= chunksize as u32;
             }
             Err(x) => {
                 warn!("error reading ota image: {:?}", x);
                 break;
             }
         }
 
         if remaining <= 0 || buf_offset == SHA256_BYTES + ERASE_BLOCK_BYTES {
             let mut actual_hash = [0u8; SHA256_BYTES];
             let mut hasher = sha_engine.start::<sha::Sha256>();
             let mut p = &buf[..buf_offset];
             while !p.is_empty() {
                 match  block!(hasher.update(p)) {
                     Ok(remaining) => { p = remaining; }
                     _ => ()
                 }
             }
             block!(hasher.finish(&mut actual_hash))?;
 
             signed = signed && next_block_hash == actual_hash;
 
             if signed {
                 if let Err(e) =  ota.ota_write_chunk(&buf[SHA256_BYTES..buf_offset]) {
 		    warn!("error writing chunk {:?}", e);
                     break;
                 }
             } else {
                 // accept the rest of the update, just don't flash it.
                 warn!("signature invalid with {remaining} left");
             }
 
             buf_offset = 0;
             next_block_hash.copy_from_slice(&buf[0..SHA256_BYTES]);
         }
     }
     if remaining == 0 {
         info!("wrote all bytes");
         info!("write {:?}",  sock.write(b"BYE\n").await);
         let _ = sock.flush().await;
         sock.close();
         if signed {
 	    ota.ota_flush(false, false)?;
             Timer::after_millis(1000).await;
             esp_hal::system::software_reset();
 	} else {
 	    return Err(FlashError::InvalidSignature)
 	}
     } else {
         warn!("stream ended early or flash failed with {} bytes left", remaining);
 	Err(FlashError::UnexpectedEof)
     }
 }