eculocate/ecu-esp32/src/ota.rs

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, error};

const ERASE_BLOCK_BYTES : usize  = 4096;
const ED25519_SIGNATURE_BYTES : usize = 64;
const SHA256_BYTES : usize = 32;

pub struct Flasher<'a> {
    ota : Ota<FlashStorage<'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 flash = FlashStorage::new(flash_peripheral);
	let sha_engine = sha::Sha::new(sha_peripheral);
	let ota = Ota::new(flash).unwrap();

	Self { ota, sha_engine }
    }

    pub async fn write(&mut self, sock: &mut embassy_net::tcp::TcpSocket<'_>) -> Result<(), FlashError> {
	write_ota(&mut self.ota, &mut self.sha_engine, sock).await
    }
}


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(_) =  ota.ota_write_chunk(&buf[ED25519_SIGNATURE_BYTES..buf_offset]) {
                    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)
    }
}