use crate::chip::Chip;
use crate::Error;
use crate::{connection::Connection, elf::RomSegment};
use indicatif::{HumanBytes, ProgressBar, ProgressStyle};
use serial::{BaudRate, SerialPort};
use sha2::{Digest, Sha256};
use std::{
io::{Cursor, Read, Write},
time::{Duration, Instant},
};
use std::{ops::Range, thread::sleep};
fn get_bar(len: u64) -> ProgressBar {
let bar = ProgressBar::new(len);
bar.set_style(
ProgressStyle::default_bar()
.template(" {wide_bar} {bytes}/{total_bytes} {bytes_per_sec} {eta} ")
.progress_chars("#>-"),
);
bar
}
pub struct Flasher {
connection: Connection,
boot_info: protocol::BootInfo,
chip: Box<dyn Chip>,
flash_speed: BaudRate,
}
impl Flasher {
pub fn connect(
chip: impl Chip + 'static,
serial: impl SerialPort + 'static,
initial_speed: BaudRate,
flash_speed: BaudRate,
) -> Result<Self, Error> {
let mut flasher = Flasher {
connection: Connection::new(serial),
boot_info: protocol::BootInfo::default(),
chip: Box::new(chip),
flash_speed,
};
flasher.connection.set_baud(initial_speed)?;
flasher.start_connection()?;
flasher.connection.set_timeout(Duration::from_secs(10))?;
flasher.boot_info = flasher.boot_rom().get_boot_info()?;
Ok(flasher)
}
pub fn into_inner(self) -> Connection {
self.connection
}
pub fn boot_info(&self) -> &protocol::BootInfo {
&self.boot_info
}
pub fn load_segments<'a>(
&'a mut self,
force: bool,
segments: impl Iterator<Item = RomSegment<'a>>,
) -> Result<(), Error> {
self.load_eflash_loader()?;
for segment in segments {
let local_hash = Sha256::digest(&segment.data[0..segment.size() as usize]);
if !force {
let sha256 = self
.eflash_loader()
.sha256_read(segment.addr, segment.size())?;
if sha256 == &local_hash[..] {
log::info!(
"Skip segment addr: {:x} size: {} sha256 matches",
segment.addr,
segment.size()
);
continue;
}
}
log::info!(
"Erase flash addr: {:x} size: {}",
segment.addr,
segment.size()
);
self.eflash_loader()
.flash_erase(segment.addr, segment.addr + segment.size())?;
let mut reader = Cursor::new(&segment.data);
let mut cur = segment.addr;
let start = Instant::now();
log::info!("Program flash... {:x}", local_hash);
let pb = get_bar(segment.size() as u64);
loop {
let size = self.eflash_loader().flash_program(cur, &mut reader)?;
cur += size;
pb.inc(size as u64);
if size == 0 {
break;
}
}
pb.finish_and_clear();
let elapsed = start.elapsed();
log::info!(
"Program done {:?} {}/s",
elapsed,
HumanBytes((segment.size() as f64 / elapsed.as_millis() as f64 * 1000.0) as u64)
);
let sha256 = self
.eflash_loader()
.sha256_read(segment.addr, segment.size())?;
if sha256 != &local_hash[..] {
log::warn!(
"sha256 not match: {} != {}",
hex::encode(sha256),
hex::encode(local_hash)
);
}
}
Ok(())
}
pub fn check_segments<'a>(
&'a mut self,
segments: impl Iterator<Item = RomSegment<'a>>,
) -> Result<(), Error> {
self.load_eflash_loader()?;
for segment in segments {
let local_hash = Sha256::digest(&segment.data[0..segment.size() as usize]);
let sha256 = self
.eflash_loader()
.sha256_read(segment.addr, segment.size())?;
if sha256 != &local_hash[..] {
log::warn!(
"{:x} sha256 not match: {} != {}",
segment.addr,
hex::encode(sha256),
hex::encode(local_hash)
);
} else {
log::info!("{:x} sha256 match", segment.addr);
}
}
Ok(())
}
pub fn dump_flash(&mut self, range: Range<u32>, mut writer: impl Write) -> Result<(), Error> {
self.load_eflash_loader()?;
const BLOCK_SIZE: usize = 4096;
let mut cur = range.start;
let pb = get_bar(range.len() as u64);
while cur < range.end {
let data = self
.eflash_loader()
.flash_read(cur, (range.end - cur).min(BLOCK_SIZE as u32))?;
writer.write_all(&data)?;
cur += data.len() as u32;
pb.inc(data.len() as u64);
}
pb.finish_and_clear();
Ok(())
}
pub fn load_eflash_loader(&mut self) -> Result<(), Error> {
let input = self.chip.get_eflash_loader().to_vec();
let len = input.len();
let mut reader = Cursor::new(input);
self.boot_rom().load_boot_header(&mut reader)?;
self.boot_rom().load_segment_header(&mut reader)?;
let start = Instant::now();
log::info!("Sending eflash_loader...");
let pb = get_bar(len as u64);
loop {
let size = self.boot_rom().load_segment_data(&mut reader)?;
pb.inc(size as u64);
if size == 0 {
break;
}
}
pb.finish_and_clear();
let elapsed = start.elapsed();
log::info!(
"Finished {:?} {}/s",
elapsed,
HumanBytes((len as f64 / elapsed.as_millis() as f64 * 1000.0) as u64)
);
self.boot_rom().check_image()?;
self.boot_rom().run_image()?;
sleep(Duration::from_millis(500));
self.connection.set_baud(self.flash_speed)?;
self.handshake()?;
log::info!("Entered eflash_loader");
Ok(())
}
pub fn reset(&mut self) -> Result<(), Error> {
Ok(self.connection.reset()?)
}
fn boot_rom(&mut self) -> BootRom {
BootRom(&mut self.connection)
}
fn eflash_loader(&mut self) -> EflashLoader {
EflashLoader(&mut self.connection)
}
fn handshake(&mut self) -> Result<(), Error> {
self.connection
.with_timeout(Duration::from_millis(200), |connection| {
let len = connection.calc_duration_length(Duration::from_millis(5));
log::trace!("5ms send count {}", len);
let data: Vec<u8> = std::iter::repeat(0x55u8).take(len).collect();
let start = Instant::now();
connection.write_all(&data)?;
connection.flush()?;
log::trace!("handshake sent elapsed {:?}", start.elapsed());
sleep(Duration::from_millis(200));
for _ in 0..5 {
if connection.read_response(0).is_ok() {
return Ok(());
}
}
Err(Error::Timeout)
})
}
fn start_connection(&mut self) -> Result<(), Error> {
log::info!("Start connection...");
self.connection.reset_to_flash()?;
for i in 1..=10 {
self.connection.flush()?;
if self.handshake().is_ok() {
log::info!("Connection Succeed");
return Ok(());
} else {
log::debug!("Retry {}", i);
}
}
Err(Error::ConnectionFailed)
}
}
pub struct BootRom<'a>(&'a mut Connection);
impl<'a> BootRom<'a> {
pub fn run_image(&mut self) -> Result<(), Error> {
self.0.command(protocol::RunImage {})?;
Ok(())
}
pub fn check_image(&mut self) -> Result<(), Error> {
self.0.command(protocol::CheckImage {})?;
Ok(())
}
pub fn load_boot_header(&mut self, reader: &mut impl Read) -> Result<(), Error> {
let mut boot_header = vec![0u8; protocol::LOAD_BOOT_HEADER_LEN];
reader.read_exact(&mut boot_header)?;
self.0.command(protocol::LoadBootHeader { boot_header })?;
Ok(())
}
pub fn load_segment_header(&mut self, reader: &mut impl Read) -> Result<(), Error> {
let mut segment_header = vec![0u8; protocol::LOAD_SEGMENT_HEADER_LEN];
reader.read_exact(&mut segment_header)?;
let resp = self.0.command(protocol::LoadSegmentHeaderReq {
segment_header: segment_header.clone(),
})?;
if resp.data != segment_header {
log::warn!(
"Segment header not match req:{:x?} != resp:{:x?}",
segment_header,
resp.data
)
}
Ok(())
}
pub fn load_segment_data(&mut self, reader: &mut impl Read) -> Result<u32, Error> {
let mut segment_data = vec![0u8; 4000];
let size = reader.read(&mut segment_data)?;
if size == 0 {
return Ok(0);
}
segment_data.truncate(size);
self.0.command(protocol::LoadSegmentData { segment_data })?;
Ok(size as u32)
}
pub fn get_boot_info(&mut self) -> Result<protocol::BootInfo, Error> {
self.0.command(protocol::BootInfoReq {})
}
}
pub struct EflashLoader<'a>(&'a mut Connection);
impl<'a> EflashLoader<'a> {
pub fn sha256_read(&mut self, addr: u32, len: u32) -> Result<[u8; 32], Error> {
Ok(self.0.command(protocol::Sha256Read { addr, len })?.digest)
}
pub fn flash_read(&mut self, addr: u32, size: u32) -> Result<Vec<u8>, Error> {
Ok(self.0.command(protocol::FlashRead { addr, size })?.data)
}
pub fn flash_program(&mut self, addr: u32, reader: &mut impl Read) -> Result<u32, Error> {
let mut data = vec![0u8; 4000];
let size = reader.read(&mut data)?;
if size == 0 {
return Ok(0);
}
data.truncate(size);
self.0.command(protocol::FlashProgram { addr, data })?;
Ok(size as u32)
}
pub fn flash_erase(&mut self, start: u32, end: u32) -> Result<(), Error> {
self.0.command(protocol::FlashErase { start, end })?;
Ok(())
}
}
mod protocol {
use crate::connection::{Command, Response};
use deku::prelude::*;
pub const LOAD_BOOT_HEADER_LEN: usize = 176;
pub const LOAD_SEGMENT_HEADER_LEN: usize = 16;
#[derive(Debug, DekuWrite, Default)]
pub struct CheckImage {}
impl_command!(0x19, CheckImage);
#[derive(Debug, DekuWrite, Default)]
pub struct RunImage {}
impl_command!(0x1a, RunImage);
#[derive(Debug, DekuWrite, Default)]
pub struct BootInfoReq {}
#[derive(Debug, DekuRead, Default)]
pub struct BootInfo {
pub len: u16,
pub bootrom_version: u32,
pub otp_info: [u8; 16],
}
impl_command!(0x10, BootInfoReq, BootInfo);
#[derive(Debug, DekuWrite, Default)]
pub struct LoadBootHeader {
pub boot_header: Vec<u8>,
}
impl_command!(0x11, LoadBootHeader);
#[derive(Debug, DekuWrite, Default)]
pub struct LoadSegmentHeaderReq {
pub segment_header: Vec<u8>,
}
#[derive(Debug, DekuRead)]
pub struct LoadSegmentHeader {
pub len: u16,
#[deku(count = "len")]
pub data: Vec<u8>,
}
impl_command!(0x17, LoadSegmentHeaderReq, LoadSegmentHeader);
#[derive(Debug, DekuWrite, Default)]
pub struct LoadSegmentData {
pub segment_data: Vec<u8>,
}
impl_command!(0x18, LoadSegmentData);
#[derive(Debug, DekuWrite, Default)]
pub struct FlashErase {
pub start: u32,
pub end: u32,
}
impl_command!(0x30, FlashErase);
#[derive(Debug, DekuWrite, Default)]
pub struct FlashProgram {
pub addr: u32,
pub data: Vec<u8>,
}
impl_command!(0x31, FlashProgram);
#[derive(Debug, DekuWrite, Default)]
pub struct FlashRead {
pub addr: u32,
pub size: u32,
}
#[derive(Debug, DekuRead)]
pub struct FlashReadResp {
pub len: u16,
#[deku(count = "len")]
pub data: Vec<u8>,
}
impl_command!(0x32, FlashRead, FlashReadResp);
#[derive(Debug, DekuWrite, Default)]
pub struct Sha256Read {
pub addr: u32,
pub len: u32,
}
#[derive(Debug, DekuRead)]
#[deku(magic = b"\x20\x00")]
pub struct Sha256ReadResp {
pub digest: [u8; 32],
}
impl_command!(0x3d, Sha256Read, Sha256ReadResp);
}