bb-imager: Add customization progress

- Make customization a seperate stage of flashing.
- Gives better cancel messages
- Also cleanup CLI since adding new stages makes the stage id difficult
  to maintain.

Signed-off-by: Ayush Singh <ayush@beagleboard.org>

bb-imager/src/common.rs

@@ -28,6 +28,7 @@ pub enum DownloadFlashingStatus {
     FlashingProgress(f32),
     Verifying,
     VerifyingProgress(f32),
+    Customizing,
 }
 
 #[derive(Debug, Hash, PartialEq, Eq, Clone)]
@@ -194,6 +195,8 @@ impl Flasher {
                 .await?;
 
                 sd::flash(img, &mut disk, &self.chan, config.verify).await?;
+
+                let _ = self.chan.try_send(DownloadFlashingStatus::Customizing);
                 disk.seek(SeekFrom::Start(0)).await?;
 
                 let mut std_disk = disk.into_std().await;

bb-imager/src/flasher/sd.rs

@@ -63,8 +63,6 @@ where
             break;
         }
 
-        tracing::debug!("Write: {} bytes", count);
-
         sd.write_all(&buf[..count]).await?;
 
         pos += count;
@@ -81,6 +79,8 @@ where
         let hash = crate::util::sha256_reader_progress(sd.take(size), size, chan).await?;
 
         if hash != sha256 {
+            tracing::debug!("Image SHA256: {}", const_hex::encode(sha256));
+            tracing::debug!("Sd SHA256: {}", const_hex::encode(hash));
             return Err(Error::Sha256VerificationError.into());
         }
     }

bb-imager/src/util.rs

@@ -28,10 +28,10 @@ pub(crate) async fn sha256_reader_progress<R: tokio::io::AsyncReadExt + Unpin>(
 
     loop {
         let count = reader.read(&mut buffer).await?;
-
         if count == 0 {
             break;
         }
+
         hasher.update(&buffer[..count]);
 
         pos += count;

cli/src/main.rs

@@ -66,7 +66,11 @@ async fn main() {
                     println!("| {: <12} | {: <12} |", "Sd Card", "Size (in G)");
                     println!("|--------------|--------------|");
                     for d in dsts {
-                        println!("| {: <12} | {: <12} |", d.path().to_str().unwrap(), d.size() / (1024 * 1024 * 1024))
+                        println!(
+                            "| {: <12} | {: <12} |",
+                            d.path().to_str().unwrap(),
+                            d.size() / (1024 * 1024 * 1024)
+                        )
                     }
                 }
                 FlashTarget::Bcf | FlashTarget::Msp430 => {
@@ -102,7 +106,7 @@ async fn flash(img: PathBuf, dst: String, target: FlashTarget, quite: bool, veri
                         static PREPARING: Once = Once::new();
 
                         PREPARING.call_once(|| {
-                            println!("[1/3] Preparing");
+                            println!("Preparing");
                         });
                     }
                     DownloadFlashingStatus::DownloadingProgress(_) => {
@@ -113,7 +117,7 @@ async fn flash(img: PathBuf, dst: String, target: FlashTarget, quite: bool, veri
                             let bar = bars.add(indicatif::ProgressBar::new(100));
                             bar.set_style(
                                 indicatif::ProgressStyle::with_template(
-                                    "[2/3] {msg}  [{wide_bar}] [{percent} %]",
+                                    "{msg}  [{wide_bar}] [{percent} %]",
                                 )
                                 .expect("Failed to create progress bar"),
                             );
@@ -132,7 +136,7 @@ async fn flash(img: PathBuf, dst: String, target: FlashTarget, quite: bool, veri
                             }
                         }
 
-                        VERIFYING.call_once(|| println!("[3/3] Verifying"));
+                        VERIFYING.call_once(|| println!("Verifying"));
                     }
                     DownloadFlashingStatus::VerifyingProgress(p) => {
                         if let Some(x) = FLASHING.get() {
@@ -145,7 +149,7 @@ async fn flash(img: PathBuf, dst: String, target: FlashTarget, quite: bool, veri
                             let bar = bars.add(indicatif::ProgressBar::new(100));
                             bar.set_style(
                                 indicatif::ProgressStyle::with_template(
-                                    "[3/3] {msg} [{wide_bar}] [{percent} %]",
+                                    "{msg} [{wide_bar}] [{percent} %]",
                                 )
                                 .expect("Failed to create progress bar"),
                             );
@@ -155,6 +159,20 @@ async fn flash(img: PathBuf, dst: String, target: FlashTarget, quite: bool, veri
 
                         bar.set_position((p * 100.0) as u64);
                     }
+                    DownloadFlashingStatus::Customizing => {
+                        static CUSTOMIZING: Once = Once::new();
+
+                        // Finish verifying progress if not already done
+                        if let Some(x) = VERIFYING.get() {
+                            if !x.is_finished() {
+                                x.finish()
+                            }
+                        }
+
+                        CUSTOMIZING.call_once(|| {
+                            println!("Customizing");
+                        });
+                    }
                 };
             }
         });

gui/src/constants.rs

-pub const BB_IMAGER_ORIGINAL_CONFIG: &str = "https://dl.beagle.cc/distros.json";
+pub const BB_IMAGER_ORIGINAL_CONFIG: &str = "https://www.beagleboard.org/distros.json";
 pub const DEFAULT_CONFIG: &[u8] = include_bytes!("../config.json");
 pub const APP_NAME: &str = "BeagleBoard Imager";
 

gui/src/helpers.rs

@@ -3,6 +3,7 @@ use std::{
     collections::{HashMap, HashSet},
 };
 
+use bb_imager::DownloadFlashingStatus;
 use iced::{
     widget::{self, button, text},
     Element,
@@ -43,22 +44,22 @@ pub struct ProgressBarState {
     label: Cow<'static, str>,
     progress: f32,
     state: ProgressBarStatus,
-    inner_state: Option<bb_imager::DownloadFlashingStatus>,
+    inner_state: Option<DownloadFlashingStatus>,
 }
 
 impl ProgressBarState {
     pub const FLASHING_SUCCESS: Self =
         Self::const_new("Flashing Successful", 1.0, ProgressBarStatus::Success, None);
-    pub const PREPARING: Self =
-        Self::loading("Preparing...", bb_imager::DownloadFlashingStatus::Preparing);
-    pub const VERIFYING: Self =
-        Self::loading("Verifying...", bb_imager::DownloadFlashingStatus::Verifying);
+    pub const PREPARING: Self = Self::loading("Preparing...", DownloadFlashingStatus::Preparing);
+    pub const VERIFYING: Self = Self::loading("Verifying...", DownloadFlashingStatus::Verifying);
+    pub const CUSTOMIZING: Self =
+        Self::loading("Customizing...", DownloadFlashingStatus::Customizing);
 
     const fn const_new(
         label: &'static str,
         progress: f32,
         state: ProgressBarStatus,
-        inner_state: Option<bb_imager::DownloadFlashingStatus>,
+        inner_state: Option<DownloadFlashingStatus>,
     ) -> Self {
         Self {
             label: Cow::Borrowed(label),
@@ -76,7 +77,7 @@ impl ProgressBarState {
         label: impl Into<Cow<'static, str>>,
         progress: f32,
         state: ProgressBarStatus,
-        inner_state: Option<bb_imager::DownloadFlashingStatus>,
+        inner_state: Option<DownloadFlashingStatus>,
     ) -> Self {
         Self {
             label: label.into(),
@@ -87,11 +88,7 @@ impl ProgressBarState {
     }
 
     /// Progress should be between 0 to 1.0
-    fn progress(
-        prefix: &'static str,
-        progress: f32,
-        inner_state: bb_imager::DownloadFlashingStatus,
-    ) -> Self {
+    fn progress(prefix: &'static str, progress: f32, inner_state: DownloadFlashingStatus) -> Self {
         Self::new(
             format!("{prefix}... {}%", (progress * 100.0).round() as usize),
             progress,
@@ -100,7 +97,7 @@ impl ProgressBarState {
         )
     }
 
-    const fn loading(label: &'static str, inner_state: bb_imager::DownloadFlashingStatus) -> Self {
+    const fn loading(label: &'static str, inner_state: DownloadFlashingStatus) -> Self {
         Self::const_new(label, 0.5, ProgressBarStatus::Loading, Some(inner_state))
     }
 
@@ -126,42 +123,38 @@ impl ProgressBarState {
     }
 
     pub fn cancel(&self) -> Option<Self> {
-        match self.inner_state? {
-            bb_imager::DownloadFlashingStatus::Preparing => {
-                Some(Self::fail("Preparation cancelled by user"))
-            }
-            bb_imager::DownloadFlashingStatus::DownloadingProgress(_) => {
-                Some(Self::fail("Downloading cancelled by user"))
-            }
-            bb_imager::DownloadFlashingStatus::FlashingProgress(_) => {
-                Some(Self::fail("Flashing cancelled by user"))
+        let x = match self.inner_state? {
+            DownloadFlashingStatus::Preparing => Self::fail("Preparation cancelled by user"),
+            DownloadFlashingStatus::DownloadingProgress(_) => {
+                Self::fail("Downloading cancelled by user")
             }
-            bb_imager::DownloadFlashingStatus::Verifying
-            | bb_imager::DownloadFlashingStatus::VerifyingProgress(_) => {
-                Some(Self::fail("Verification cancelled by user"))
+            DownloadFlashingStatus::FlashingProgress(_) => Self::fail("Flashing cancelled by user"),
+            DownloadFlashingStatus::Verifying | DownloadFlashingStatus::VerifyingProgress(_) => {
+                Self::fail("Verification cancelled by user")
             }
-        }
+            DownloadFlashingStatus::Customizing => Self::fail("Customization cancelled by user"),
+        };
+        Some(x)
     }
 }
 
-impl From<bb_imager::DownloadFlashingStatus> for ProgressBarState {
-    fn from(value: bb_imager::DownloadFlashingStatus) -> Self {
+impl From<DownloadFlashingStatus> for ProgressBarState {
+    fn from(value: DownloadFlashingStatus) -> Self {
         match value {
-            bb_imager::DownloadFlashingStatus::Preparing => Self::PREPARING,
-            bb_imager::DownloadFlashingStatus::DownloadingProgress(p) => Self::progress(
+            DownloadFlashingStatus::Preparing => Self::PREPARING,
+            DownloadFlashingStatus::DownloadingProgress(p) => Self::progress(
                 "Downloading Image",
                 p,
-                bb_imager::DownloadFlashingStatus::DownloadingProgress(0.0),
-            ),
-            bb_imager::DownloadFlashingStatus::FlashingProgress(p) => Self::progress(
-                "Flashing",
-                p,
-                bb_imager::DownloadFlashingStatus::FlashingProgress(0.0),
+                DownloadFlashingStatus::DownloadingProgress(0.0),
             ),
-            bb_imager::DownloadFlashingStatus::Verifying => Self::VERIFYING,
-            bb_imager::DownloadFlashingStatus::VerifyingProgress(p) => {
-                Self::progress("Verifying", p, bb_imager::DownloadFlashingStatus::Verifying)
+            DownloadFlashingStatus::FlashingProgress(p) => {
+                Self::progress("Flashing", p, DownloadFlashingStatus::FlashingProgress(0.0))
+            }
+            DownloadFlashingStatus::Verifying => Self::VERIFYING,
+            DownloadFlashingStatus::VerifyingProgress(p) => {
+                Self::progress("Verifying", p, DownloadFlashingStatus::Verifying)
             }
+            DownloadFlashingStatus::Customizing => Self::CUSTOMIZING,
         }
     }
 }