bb_flasher_sd/
flashing.rs

1use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write};
2use std::path::Path;
3use std::sync::Weak;
4
5use futures::channel::mpsc;
6use sha2::{Digest, Sha256};
7
8use crate::customization::Customization;
9use crate::helpers::{chan_send, check_arc, progress, sha256_reader_progress};
10use crate::{Error, Result, Status};
11
12fn read_aligned(mut img: impl Read, buf: &mut [u8]) -> Result<usize> {
13    let mut pos = 0;
14
15    while pos != buf.len() {
16        let count = img.read(&mut buf[pos..])?;
17        // Need to align to 512 byte boundary for writing to sd card
18        if count == 0 {
19            buf[pos..].fill(0);
20            return Ok(pos);
21        }
22
23        pos += count;
24    }
25
26    Ok(pos)
27}
28
29fn write_sd(
30    mut img: impl Read,
31    img_size: u64,
32    mut sd: impl Write,
33    mut hasher: Option<&mut Sha256>,
34    mut chan: Option<&mut mpsc::Sender<Status>>,
35    cancel: Option<&Weak<()>>,
36) -> Result<()> {
37    let mut buf = [0u8; 512];
38    let mut pos = 0;
39
40    chan_send(
41        chan.as_mut().map_or(None, |p| Some(p)),
42        Status::Flashing(0.0),
43    );
44    loop {
45        let count = read_aligned(&mut img, &mut buf)?;
46        if count == 0 {
47            break;
48        }
49
50        // Since buf is 512, just write the whole thing since read_aligned will always fill it
51        // fully.
52        sd.write_all(&buf)?;
53        if let Some(ref mut h) = hasher {
54            h.update(&buf[..count]);
55        }
56
57        pos += count;
58        chan_send(
59            chan.as_mut().map_or(None, |p| Some(p)),
60            Status::Flashing(progress(pos, img_size)),
61        );
62        check_arc(cancel)?;
63    }
64
65    sd.flush().map_err(Into::into)
66}
67
68/// Flash OS image to SD card.
69///
70/// # Customization
71///
72/// Support post flashing customization. Currently only sysconf is supported, which is used by
73/// [BeagleBoard.org].
74///
75/// # Image
76///
77/// Using a resolver function for image and image size. This is to allow downloading the image, or
78/// some kind of lazy loading after SD card permissions have be acquired. This is useful in GUIs
79/// since the user would expect a password prompt at the start of flashing.
80///
81/// Many users might switch task after starting the flashing process, which would make it
82/// frustrating if the prompt occured after downloading.
83///
84/// # Progress
85///
86/// Progress lies between 0 and 1.
87///
88/// # Aborting
89///
90/// The process can be aborted by dropping all strong references to the [`Arc`] that owns the
91/// [`Weak`] passed as `cancel`.
92///
93/// [`Arc`]: std::sync::Arc
94/// [`Weak`]: std::sync::Weak
95/// [BeagleBoard.org]: https://www.beagleboard.org/
96pub fn flash<R: Read>(
97    img_resolver: impl FnOnce() -> std::io::Result<(R, u64)>,
98    dst: &Path,
99    verify: bool,
100    mut chan: Option<mpsc::Sender<Status>>,
101    customization: Option<Customization>,
102    cancel: Option<Weak<()>>,
103) -> Result<()> {
104    let mut sd = crate::pal::open(dst)?;
105    let (img, img_size) = img_resolver()?;
106
107    chan_send(chan.as_mut(), Status::Flashing(0.0));
108    if verify {
109        let mut hasher = Sha256::new();
110        write_sd(
111            img,
112            img_size,
113            BufWriter::new(&mut sd),
114            Some(&mut hasher),
115            chan.as_mut(),
116            cancel.as_ref(),
117        )?;
118
119        let img_hash: [u8; 32] = hasher
120            .finalize()
121            .as_slice()
122            .try_into()
123            .expect("SHA-256 is 32 bytes");
124
125        // Start verification
126        chan_send(chan.as_mut(), Status::Verifying(0.0));
127        sd.seek(std::io::SeekFrom::Start(0))?;
128        let hash =
129            sha256_reader_progress(BufReader::new((&mut sd).take(img_size)), img_size, chan)?;
130
131        check_arc(cancel.as_ref())?;
132        if hash != img_hash {
133            tracing::debug!("Sd SHA256: {:?}", hash);
134            return Err(Error::Sha256Verification);
135        }
136    } else {
137        write_sd(
138            img,
139            img_size,
140            BufWriter::new(&mut sd),
141            None,
142            chan.as_mut(),
143            cancel.as_ref(),
144        )?;
145    }
146
147    check_arc(cancel.as_ref())?;
148
149    if let Some(c) = customization {
150        sd.seek(SeekFrom::Start(0))?;
151        c.customize(sd)?;
152    }
153
154    Ok(())
155}
bb_flasher_sd/flashing.rs

bb_flasher_sd/
flashing.rs
