Skip to content

Instantly share code, notes, and snippets.

@jamesmunns
Last active Dec 29, 2021
Embed
What would you like to do?
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
}
}
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment