diff --git a/core/src/tasks.rs b/core/src/tasks.rs index 13ba8d5..296ea6a 100644 --- a/core/src/tasks.rs +++ b/core/src/tasks.rs @@ -33,6 +33,12 @@ use crate::{ system::{disk, diskpart}, }; +#[derive(Clone, Debug)] +pub enum TaskResult { + Error(String), + Output(String, String, bool), // stdout, stderr, success +} + #[derive(Clone, Debug)] pub enum TaskType { Command(PathBuf, Vec), // (command, args) @@ -46,6 +52,7 @@ pub enum TaskType { #[derive(Debug)] pub struct Task { pub handle: Option>, + pub result: Option, pub task_type: TaskType, } @@ -53,6 +60,7 @@ impl Task { pub fn new(task_type: TaskType) -> Task { Task { handle: None, + result: None, task_type, } } @@ -65,8 +73,8 @@ pub struct Tasks { cur_handle: Option>, cur_task: Option, task_list: VecDeque, - task_rx: mpsc::UnboundedReceiver, // Used to forward Actions from Tasks to App - task_tx: mpsc::UnboundedSender, // Used to forward Actions from Tasks to App + task_rx: mpsc::UnboundedReceiver, + task_tx: mpsc::UnboundedSender, } impl Tasks { @@ -97,10 +105,12 @@ impl Tasks { pub fn poll(&mut self) -> Result> { let mut return_task: Option = None; - // Forward any actions to main app - if let Ok(action) = self.task_rx.try_recv() { - let result = self.action_tx.send(action.clone()); - assert!(result.is_ok(), "Failed to send Action: {action:?}"); + // Handle task channel item(s) + if let Ok(result) = self.task_rx.try_recv() { + if let Some(mut task) = self.cur_task.take() { + task.result.replace(result); + self.cur_task.replace(task); + } } // Check status of current task (if one is running). @@ -114,7 +124,7 @@ impl Tasks { } if self.task_list.is_empty() { // No tasks remain - self.task_tx.send(Action::NextScreen)?; + self.action_tx.send(Action::NextScreen)?; } else { // Start next task self.start()?; @@ -195,10 +205,17 @@ impl Tasks { } } +fn parse_bytes_as_str(bytes: Vec) -> String { + match String::from_utf8(bytes) { + Ok(s) => s.trim().to_string(), + Err(_) => String::from("Failed to parse bytes as UTF-8 text"), + } +} + fn run_task_command( cmd_path: PathBuf, cmd_args: Vec, - task_tx: mpsc::UnboundedSender, + task_tx: mpsc::UnboundedSender, ) -> JoinHandle<()> { if cfg!(windows) { thread::spawn(move || { @@ -206,27 +223,19 @@ fn run_task_command( .args(cmd_args) .stdout(Stdio::piped()) .output(); - if let Some(action) = match result { - Ok(output) => { - if output.status.success() { - None - } else { - // Command returned an error status - let mut msg = String::new(); - if let Ok(stdout) = String::from_utf8(output.stdout) { - msg = String::from(stdout.trim()); - } - if msg.is_empty() { - msg = String::from("Generic error"); - } - Some(Action::Error(format!("Command failed: {msg}",))) - } + match result { + Err(e) => { + task_tx + .send(TaskResult::Error(format!("{:?}", &e))) + .expect("Failed to propegate error?"); + } + Ok(output) => { + let stderr = parse_bytes_as_str(output.stderr.to_owned()); + let stdout = parse_bytes_as_str(output.stdout.to_owned()); + let task_result = TaskResult::Output(stdout, stderr, output.status.success()); + let err_str = format!("Failed to send TaskResult: {:?}", &task_result); + task_tx.send(task_result).expect(err_str.as_str()); } - Err(err) => Some(Action::Error(format!("Failed to run command: {err:?}"))), - } { - let msg = format!("{:?}", &action); - let result = task_tx.send(action); - assert!(result.is_ok(), "Failed to send Action: {msg}"); } }) } else { @@ -235,22 +244,16 @@ fn run_task_command( } } -fn run_task_diskpart(script: &str, task_tx: mpsc::UnboundedSender) -> JoinHandle<()> { - let task = TaskType::Diskpart(String::from(script)); - let task_str = format!("{:?}", &task); +fn run_task_diskpart(script: &str, task_tx: mpsc::UnboundedSender) -> JoinHandle<()> { if cfg!(windows) { - let script = String::from(script); + let script = script.to_owned(); thread::spawn(move || { let output = diskpart::run_script_raw(&script); - if !output.status.success() - && task_tx - .send(Action::Error(String::from( - "Diskpart script returned an error", - ))) - .is_err() - { - panic!("Failed to send Action: {task_str:?}"); - } + let stderr = parse_bytes_as_str(output.stderr.to_owned()); + let stdout = parse_bytes_as_str(output.stdout.to_owned()); + let task_result = TaskResult::Output(stdout, stderr, output.status.success()); + let err_str = format!("Failed to send TaskResult: {:?}", &task_result); + task_tx.send(task_result).expect(err_str.as_str()); }) } else { // Simulate task if not running under Windows diff --git a/deja_vu/src/app.rs b/deja_vu/src/app.rs index cbbe506..7ee37dd 100644 --- a/deja_vu/src/app.rs +++ b/deja_vu/src/app.rs @@ -26,7 +26,7 @@ use core::{ boot, cpu::get_cpu_name, disk::PartitionTableType, diskpart::build_dest_format_script, drivers, }, - tasks::{TaskType, Tasks}, + tasks::{TaskResult, TaskType, Tasks}, tui::{Event, Tui}, }; use std::{ @@ -375,8 +375,37 @@ impl App { self.last_tick_key_events.drain(..); match self.cur_mode { Mode::ScanDisks | Mode::PreClone | Mode::Clone | Mode::PostClone => { - // Check background task - self.tasks.poll()?; // Once all are complete Action::NextScreen is sent + // Check background task (Action::NextScreen is sent when task(s) are done) + if let Some(task) = self.tasks.poll()? { + match task.task_type { + TaskType::Command(_, _) | TaskType::Diskpart(_) => { + if let Some(result) = &task.result { + match result { + TaskResult::Error(msg) => { + self.action_tx.send(Action::Error(format!( + "{task:?} Failed: {msg}" + )))?; + } + TaskResult::Output(stdout, stderr, success) => { + if !success { + let msg = if !stdout.is_empty() { + stdout.clone() + } else if !stderr.is_empty() { + stderr.clone() + } else { + String::from("Unknown Error") + }; + self.action_tx.send(Action::Error( + format!("{task:?} Failed: {msg}"), + ))?; + } + } + } + } + } + _ => {} + } + } } _ => {} }