1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109

use crate::connection::Connection;
use crate::elf::{FirmwareImage, RomSegment};
use crate::error::Error;
use crate::flash_target::{begin_command, block_command_with_timeout, FlashTarget};
use crate::flasher::{Command, SpiAttachParams, FLASH_SECTOR_SIZE, FLASH_WRITE_SIZE};
use crate::Chip;
use flate2::write::{ZlibDecoder, ZlibEncoder};
use flate2::Compression;
use indicatif::{ProgressBar, ProgressStyle};
use std::io::Write;

pub struct Esp32Target {
    chip: Chip,
    spi_attach_params: SpiAttachParams,
}

impl Esp32Target {
    pub fn new(chip: Chip, spi_attach_params: SpiAttachParams) -> Self {
        Esp32Target {
            chip,
            spi_attach_params,
        }
    }
}

impl FlashTarget for Esp32Target {
    fn begin(&mut self, connection: &mut Connection, _image: &FirmwareImage) -> Result<(), Error> {
        let spi_params = self.spi_attach_params.encode();
        connection.with_timeout(Command::SpiAttach.timeout(), |connection| {
            connection.command(Command::SpiAttach as u8, spi_params.as_slice(), 0)
        })?;
        Ok(())
    }

    fn write_segment(
        &mut self,
        connection: &mut Connection,
        segment: RomSegment,
    ) -> Result<(), Error> {
        let addr = segment.addr;
        let mut encoder = ZlibEncoder::new(Vec::new(), Compression::best());
        encoder.write_all(&segment.data)?;
        let compressed = encoder.finish()?;
        let block_count = (compressed.len() + FLASH_WRITE_SIZE - 1) / FLASH_WRITE_SIZE;
        let erase_count = (segment.data.len() + FLASH_SECTOR_SIZE - 1) / FLASH_SECTOR_SIZE;

        // round up to sector size
        let erase_size = (erase_count * FLASH_SECTOR_SIZE) as u32;

        begin_command(
            connection,
            Command::FlashDeflateBegin,
            erase_size,
            block_count as u32,
            FLASH_WRITE_SIZE as u32,
            addr,
            self.chip != Chip::Esp32,
        )?;

        let chunks = compressed.chunks(FLASH_WRITE_SIZE);

        let (_, chunk_size) = chunks.size_hint();
        let chunk_size = chunk_size.unwrap_or(0) as u64;
        let pb_chunk = ProgressBar::new(chunk_size);
        pb_chunk.set_style(
            ProgressStyle::default_bar()
                .template("[{elapsed_precise}] {bar:40.cyan/blue} {pos:>7}/{len:7} {msg}")
                .progress_chars("#>-"),
        );

        // decode the chunks to see how much data the device will have to save
        let mut decoder = ZlibDecoder::new(Vec::new());
        let mut decoded_size = 0;

        for (i, block) in chunks.enumerate() {
            decoder.write_all(block)?;
            decoder.flush()?;
            let size = decoder.get_ref().len() - decoded_size;
            decoded_size = decoder.get_ref().len();

            pb_chunk.set_message(format!("segment 0x{:X} writing chunks", addr));
            block_command_with_timeout(
                connection,
                Command::FlashDeflateData,
                block,
                0,
                0xff,
                i as u32,
                Command::FlashDeflateData.timeout_for_size(size as u32),
            )?;
            pb_chunk.inc(1);
        }

        pb_chunk.finish_with_message(format!("segment 0x{:X}", addr));

        Ok(())
    }

    fn finish(&mut self, connection: &mut Connection, reboot: bool) -> Result<(), Error> {
        connection.with_timeout(Command::FlashDeflateEnd.timeout(), |connection| {
            connection.write_command(Command::FlashDeflateEnd as u8, &[1][..], 0)
        })?;
        if reboot {
            connection.reset()
        } else {
            Ok(())
        }
    }
}