async_boot_client.rs

use core::{ops::Deref, ptr::NonNull, sync::atomic::{Ordering, compiler_fence}};

use anachro_boot::consts::{POLY_1305_KEY, POLY_TAG_SIZE, PAGE_SIZE_4K, CHUNK_SIZE_256B};
use anachro_client::ManagedArcStr;
use anachro_qspi::{Qspi, ManagedArcSlab, EraseLength, FlashChunk, QSPI_MAPPED_BASE_ADDRESS};
use byte_slab::SlabBox;
use groundhog::RollingTimer;
use groundhog_nrf52::GlobalRollingTimer;
use heapless::spsc::{Consumer, Producer};
use common::{
    apb_icd::{BootloadCommand, BootloadResponse, StartBootload, AbortReason},
    globals::BSLAB,
};
use anachro_485::icd::{TOTAL_SLABS, SLAB_SIZE};
use cassette::yield_now;
use nrf52840_hal::{Spim, pac::SPIM3};
use poly1305::{Block, Key, Poly1305, universal_hash::{NewUniversalHash, UniversalHash}};

use embedded_hal::blocking::spi::Write;

pub struct BootClient<'a> {
    qspi: &'a mut Qspi,
    inc_cmds: Consumer<'a, BootloadCommand<'static, TOTAL_SLABS, SLAB_SIZE>, 8>,
    out_resp: Producer<'a, BootloadResponse<'static, TOTAL_SLABS, SLAB_SIZE>, 8>,
    nxt_metadata: NonNull<u8>,
    nxt_image: NonNull<u8>,
    boot_seq: u32,
    spim: &'a mut Spim<SPIM3>,
}

// This makes me sad.
#[repr(align(4))]
struct WordAligned {
    data: [u8; 256]
}

impl<'a> BootClient<'a> {
    pub fn new(
        qspi: &'a mut Qspi,
        inc_cmds: Consumer<'a, BootloadCommand<'static, TOTAL_SLABS, SLAB_SIZE>, 8>,
        out_resp: Producer<'a, BootloadResponse<'static, TOTAL_SLABS, SLAB_SIZE>, 8>,
        nxt_metadata: NonNull<u8>,
        nxt_image: NonNull<u8>,
        boot_seq: u32,
        spim: &'a mut Spim<SPIM3>,
    ) -> Self {
        // TODO: Send first boot message AND boot success message

        Self {
            qspi,
            inc_cmds,
            out_resp,
            nxt_metadata,
            nxt_image,
            boot_seq,
            spim,
        }
    }

    pub async fn step(&mut self) {
        loop {
            match self.step_inner().await {
                Ok(()) => {
                    defmt::println!("Rebooting!");
                    self.out_resp.enqueue(BootloadResponse::BootloadComplete).ok();
                    return;
                }
                Err(e) => {
                    defmt::println!("Aborting bootload.");
                    self.out_resp.enqueue(BootloadResponse::BootloadAborted(e)).ok();
                }
            }
        }
    }

    async fn verify_page_erase(&mut self, start_addr: usize) -> Result<(), AbortReason> {
        for offset in 0..(PAGE_SIZE_4K / CHUNK_SIZE_256B) {
            let chunk_start = start_addr + (offset * CHUNK_SIZE_256B);
            let mut buf = WordAligned {
                data: [0u8; 256],
            };

            self.qspi
                .read(chunk_start, &mut buf.data)
                .await
                .map_err(|_| AbortReason::ReadFail)?;

            if !buf.data.iter().all(|b| *b == 0xFF) {
                return Err(AbortReason::EraseVerificationQspiFail);
            }
        }

        let sli = self.qspi
            .read_slice(start_addr, PAGE_SIZE_4K)
            .map_err(|n| AbortReason::Other(n))?;

        if !sli.iter().all(|b| *b == 0xFF) {
            return Err(AbortReason::EraseVerificationXipFail);
        }

        Ok(())
    }

    async fn verify_page_write(&mut self, start_addr: usize, data: &[u8]) -> Result<(), AbortReason> {
        let mut buf = WordAligned {
            data: [0u8; 256],
        };

        core::sync::atomic::compiler_fence(Ordering::SeqCst);

        self.qspi
            .read(start_addr, &mut buf.data)
            .await
            .map_err(|_| AbortReason::ReadFail)?;

        core::sync::atomic::compiler_fence(Ordering::SeqCst);

        Write::write(self.spim, &20u32.to_le_bytes()).ok();
        Write::write(self.spim, &start_addr.to_le_bytes()).ok();
        Write::write(self.spim, &0u32.to_le_bytes()).ok();
        Write::write(self.spim, &buf.data).ok();

        if &buf.data != data {
            return Err(AbortReason::WriteVerificationQspiFail);
        }

        core::sync::atomic::compiler_fence(Ordering::SeqCst);

        let sli = self.qspi
            .read_slice(start_addr, CHUNK_SIZE_256B)
            .map_err(|n| AbortReason::Other(n))?;

        core::sync::atomic::compiler_fence(Ordering::SeqCst);

        Write::write(self.spim, &20u32.to_le_bytes()).ok();
        Write::write(self.spim, &start_addr.to_le_bytes()).ok();
        Write::write(self.spim, &0u32.to_le_bytes()).ok();
        Write::write(self.spim, sli).ok();

        if !(sli == data) {
            return Err(AbortReason::WriteVerificationXipFail);
        }

        core::sync::atomic::compiler_fence(Ordering::SeqCst);

        Ok(())
    }

    async fn erase_metadata(&mut self) -> Result<(), AbortReason> {
        let metadata_usize = self.nxt_metadata.as_ptr() as usize - QSPI_MAPPED_BASE_ADDRESS;

        #[cfg(not(feature = "no-bootloader"))]
        self.qspi
            .erase(metadata_usize, EraseLength::_4KB)
            .await
            .map_err(|_| AbortReason::EraseFail)?;

        #[cfg(not(feature = "no-bootloader"))]
        self.verify_page_erase(metadata_usize).await?;

        Ok(())
    }

    async fn erase_image(&mut self, chunks_256_ct: usize) ->Result<(), AbortReason> {
        // Round up to the number of sectors that need to be erased
        let erase_sz_4k = ((chunks_256_ct as usize * CHUNK_SIZE_256B) + (PAGE_SIZE_4K - 1)) / PAGE_SIZE_4K;
        let start = self.nxt_image.as_ptr() as usize - QSPI_MAPPED_BASE_ADDRESS;

        defmt::println!(
            "Erasing {=usize} 4K sectors, Starting at 0x{=usize:08X}...",
            erase_sz_4k,
            start
        );

        #[cfg(not(feature = "no-bootloader"))]
        for sect_idx in 0..erase_sz_4k {
            let sect_start = start + (PAGE_SIZE_4K * sect_idx);

            self.qspi
                .erase(sect_start, EraseLength::_4KB)
                .await
                .map_err(|_| AbortReason::EraseFail)?;

            self.verify_page_erase(sect_start).await?;
        }

        compiler_fence(Ordering::SeqCst);

        Ok(())
    }

    async fn process_chunk(&mut self, chunk_idx: usize) -> Result<(), AbortReason> {
        let start = self.nxt_image.as_ptr() as usize - QSPI_MAPPED_BASE_ADDRESS;
        let chunk_addr = (start + (chunk_idx * CHUNK_SIZE_256B)) as u32;

        core::sync::atomic::compiler_fence(Ordering::SeqCst);

        // Note: We have already verified the individual page signature here
        let chunk = self
            .wait_for_chunk(chunk_idx as u32, chunk_addr)
            .await?;

        core::sync::atomic::compiler_fence(Ordering::SeqCst);

        let data = chunk.deref();
        let mut adata = WordAligned { data: [0u8; 256] };
        adata.data.copy_from_slice(data);

        Write::write(self.spim, &10u32.to_le_bytes()).ok();
        Write::write(self.spim, &chunk_addr.to_le_bytes()).ok();
        Write::write(self.spim, &(adata.data.as_ptr() as u32).to_le_bytes()).ok();
        Write::write(self.spim, &adata.data).ok();

        core::sync::atomic::compiler_fence(Ordering::SeqCst);

        #[cfg(not(feature = "no-bootloader"))]
        self.qspi
            .write(FlashChunk {
                addr: chunk_addr as usize,
                data: ManagedArcSlab::<TOTAL_SLABS, SLAB_SIZE>::Borrowed(&adata.data)
            })
            .await
            .map_err(|_| AbortReason::WriteFail)?;

        core::sync::atomic::compiler_fence(Ordering::SeqCst);

        #[cfg(not(feature = "no-bootloader"))]
        self.verify_page_write(chunk_addr as usize, data).await?;

        core::sync::atomic::compiler_fence(Ordering::SeqCst);

        Ok(())
    }

    async fn write_metadata(
        &mut self,
        new_fw_id: &[u8; 16],
        ttl_sig: &[u8; 16],
        chunk_ct: u32,
    ) -> Result<(), AbortReason> {
        let mut payload = WordAligned {
            data: [0xFFu8; 256],
        };

        let metadata_usize = self.nxt_metadata.as_ptr() as usize;

        // Fill to 0xFF's, to be nice to the flash
        payload.data.iter_mut().for_each(|b| *b = 0xFF);
        payload.data[..16].copy_from_slice(new_fw_id);
        payload.data[16..32].copy_from_slice(ttl_sig);
        payload.data[32..36].copy_from_slice(&chunk_ct.to_le_bytes());

        payload.data[128..132].copy_from_slice(&(self.boot_seq + 1).to_le_bytes());
        // Leave the other tagwords all 0xFFs

        let mas = ManagedArcSlab::<TOTAL_SLABS, SLAB_SIZE>::Borrowed(&payload.data);

        #[cfg(not(feature = "no-bootloader"))]
        self.qspi
            .write(FlashChunk { addr: metadata_usize, data: mas })
            .await
            .map_err(|_| AbortReason::WriteFail)?;

        Ok(())
    }

    async fn step_inner(&mut self) -> Result<(), AbortReason> {
        let start_msg = self.wait_for_start().await;

        // Copy relevant data from the incoming message
        let chunks_256_ct = start_msg.chunks_256_ct as usize;
        let mut new_fw_id = [0u8; 16];
        let mut ttl_sig = [0u8; 16];
        new_fw_id.copy_from_slice(start_msg.firmware_id.deref());
        ttl_sig.copy_from_slice(start_msg.ttl_signature.deref());

        // Drop initial message to make sure the page is freed
        drop(start_msg);

        defmt::println!("Erasing metadata...");
        self.erase_metadata().await?;

        defmt::println!("Erasing (old) image...");
        self.erase_image(chunks_256_ct).await?;

        defmt::println!("Processing chunks...");
        for pg_idx in 0..chunks_256_ct {
            self.process_chunk(pg_idx).await?;
        }

        self.write_metadata(&new_fw_id, &ttl_sig, chunks_256_ct as u32).await?;

        // Good Luck, starfox.
        Ok(())
    }

    async fn wait_for_start(&mut self) -> StartBootload<'static, TOTAL_SLABS, SLAB_SIZE> {
        loop {
            match self.inc_cmds.dequeue() {
                Some(BootloadCommand::StartBootload(sbm)) => {
                    let good_fwid_len = sbm.firmware_id.len() == 16;
                    let good_sig_len = sbm.ttl_signature.len() == 16;
                    let good_pg_ct = (sbm.chunks_256_ct >= 1) && (sbm.chunks_256_ct <= 3072);

                    let uuid_all_zeroes = sbm.firmware_id.iter().all(|b| *b == 0);
                    let uuid_all_effffs = sbm.firmware_id.iter().all(|b| *b == 0xFF);

                    let good_uuid = !(uuid_all_zeroes || uuid_all_effffs);

                    if !(good_fwid_len && good_sig_len && good_pg_ct && good_uuid) {
                        defmt::println!("{=bool}, {=bool}, {=bool}, {=bool}", good_fwid_len, good_sig_len, good_pg_ct, good_uuid);
                        defmt::println!("Wrong start message data!");
                    } else {
                        return sbm
                    }
                }
                Some(_) => {
                    defmt::println!("Unexpected bootloader command! Waiting for start...");
                }
                None => yield_now().await,
            }
        }
    }

    async fn request_chunk(&mut self, chunk_id: u32, chunk_addr: u32) {
        loop {
            match self.out_resp.enqueue(BootloadResponse::RequestChunk { chunk_idx: chunk_id, dest_addr: chunk_addr }) {
                Ok(()) => break,
                Err(_) => yield_now().await,
            }
        }
    }

    async fn wait_for_chunk(&mut self, chunk_id: u32, chunk_addr: u32) -> Result<ManagedArcSlab<'static, TOTAL_SLABS, SLAB_SIZE>, AbortReason> {
        // Make sure we have successfully sent a request
        // TODO: Do this with some kind of retry logic?
        self.request_chunk(chunk_id, chunk_addr).await;
        let timer = GlobalRollingTimer::default();

        let mut last_request = timer.get_ticks();
        let mut num_re_requests = 0;
        let mut bad_chunk_ct = 0;

        loop {
            match self.inc_cmds.dequeue() {
                Some(BootloadCommand::ChunkData(data)) => {
                    let good_data_len = data.chunk_data.len() == 256;
                    let good_sig_len = data.chunk_signature.len() == 16;
                    let right_chunk = chunk_id == data.chunk_idx;

                    let good_poly = if !(good_data_len && good_sig_len && right_chunk) {
                        defmt::println!("Bad chunk data (len)");
                        false
                    } else {
                        let key = Key::from_slice(POLY_1305_KEY);
                        let mut poly = Poly1305::new(key);

                        for chunk in data.chunk_data.chunks_exact(POLY_TAG_SIZE) {
                            poly.update(Block::from_slice(chunk));
                        }

                        let calc_poly: [u8; POLY_TAG_SIZE] = poly.finalize().into_bytes().into();

                        &calc_poly == data.chunk_signature.deref()
                    };

                    if good_poly {
                        return Ok(data.chunk_data);
                    } else {
                        if bad_chunk_ct < 3 {
                            bad_chunk_ct += 1;
                            defmt::println!("Bad chunk data! Retrying...");
                            self.request_chunk(chunk_id, chunk_addr).await;
                        } else {
                            defmt::println!("Bad chunk data three times in a row!");
                            return Err(AbortReason::InvalidChunk);
                        }
                    }
                }
                Some(BootloadCommand::AbortBootload) => {
                    defmt::println!("Aborting bootload on request!");
                    return Err(AbortReason::AbortRequested);
                }
                Some(_) => {
                    defmt::println!("Unexpected bootloader command! Waiting for chunk {=u32}", chunk_id);
                }
                None => {
                    // Re-request the chunk if it has been a while
                    if timer.millis_since(last_request) >= 1000 {
                        if num_re_requests < 10 {
                            self.request_chunk(chunk_id, chunk_addr).await;
                            last_request = timer.get_ticks();
                            num_re_requests += 1;
                        } else {
                            defmt::println!("Chunk request timeout! Aborting.");
                            return Err(AbortReason::ChunkTimeout);
                        }
                    }
                    yield_now().await
                }
            }
        }
    }
}