bb_flasher_sd/
flashing.rs

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

bb_flasher_sd/
flashing.rs
