diff --git a/boot_diags/src/app.rs b/boot_diags/src/app.rs index 220e587..6854bf8 100644 --- a/boot_diags/src/app.rs +++ b/boot_diags/src/app.rs @@ -27,7 +27,7 @@ use core::{ cpu::get_cpu_name, drivers, }, - tasks::{Task, Tasks}, + tasks::{TaskType, Tasks}, tui::{Event, Tui}, }; use std::{ @@ -178,7 +178,7 @@ impl App { Mode::InjectDrivers | Mode::InstallDrivers => self.clone.scan_drivers(), Mode::ScanDisks => { if self.tasks.idle() { - self.tasks.add(Task::ScanDisks); + self.tasks.add(TaskType::ScanDisks); } self.action_tx.send(Action::DisplayPopup( popup::Type::Info, @@ -297,7 +297,7 @@ impl App { match self.cur_mode { Mode::BootDiags => { // Check result of background task (if finished) - if let Some(handle) = self.tasks.poll()? { + if let Some(task) = self.tasks.poll()? { // TODO: Impl logic } } diff --git a/core/src/system/boot.rs b/core/src/system/boot.rs index 5dfd4dd..97862ef 100644 --- a/core/src/system/boot.rs +++ b/core/src/system/boot.rs @@ -14,7 +14,7 @@ // along with Deja-Vu. If not, see . // use super::{disk::PartitionTableType, drivers::Driver}; -use crate::tasks::Task; +use crate::tasks::TaskType; use color_eyre::Result; use std::path::PathBuf; @@ -30,11 +30,11 @@ pub fn configure_disk( letter_os: &str, system32: &str, table_type: PartitionTableType, -) -> Vec { +) -> Vec { let mut tasks = Vec::new(); // Create - tasks.push(Task::Command( + tasks.push(TaskType::Command( PathBuf::from(format!("{system32}/bcdboot.exe")), vec![ format!("{letter_os}:\\Windows"), @@ -50,7 +50,7 @@ pub fn configure_disk( // Update boot sector (for legacy setups) if table_type == PartitionTableType::Legacy { - tasks.push(Task::Command( + tasks.push(TaskType::Command( PathBuf::from(format!("{system32}/bootsect.exe")), vec![ String::from("/nt60"), @@ -71,10 +71,10 @@ pub fn configure_disk( tasks } -pub fn inject_driver(driver: &Driver, letter_os: &str, system32: &str) -> Result { +pub fn inject_driver(driver: &Driver, letter_os: &str, system32: &str) -> Result { //if let Some(driver_path_str) = driver.path.to_str() { let driver_path = driver.path.to_str().unwrap(); - Ok(Task::Command( + Ok(TaskType::Command( PathBuf::from(format!("{system32}/dism.exe")), vec![ format!("/image:{letter_os}:\\"), @@ -90,7 +90,7 @@ pub fn set_mode( mode: SafeMode, system32: &str, table_type: &PartitionTableType, -) -> Result { +) -> Result { let bcd_path = match table_type { PartitionTableType::Guid => { format!("{letter_boot}:\\EFI\\Microsoft\\Boot\\BCD") @@ -115,7 +115,7 @@ pub fn set_mode( cmd_args.push(String::from("minimal")); } } - Ok(Task::Command( + Ok(TaskType::Command( PathBuf::from(format!("{system32}/bcdedit.exe")), cmd_args, )) diff --git a/core/src/tasks.rs b/core/src/tasks.rs index 7e885e5..13ba8d5 100644 --- a/core/src/tasks.rs +++ b/core/src/tasks.rs @@ -34,7 +34,7 @@ use crate::{ }; #[derive(Clone, Debug)] -pub enum Task { +pub enum TaskType { Command(PathBuf, Vec), // (command, args) Diskpart(String), // (script_as_string) ScanDisks, @@ -44,16 +44,26 @@ pub enum Task { } #[derive(Debug)] -pub struct TaskHandle { - task: Task, - handle: JoinHandle<()>, +pub struct Task { + pub handle: Option>, + pub task_type: TaskType, +} + +impl Task { + pub fn new(task_type: TaskType) -> Task { + Task { + handle: None, + task_type, + } + } } #[derive(Debug)] pub struct Tasks { action_tx: mpsc::UnboundedSender, disk_list: Arc>>, - handle: Option, + 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 @@ -68,24 +78,25 @@ impl Tasks { Tasks { action_tx, disk_list: disk_list_arc, - handle: None, + cur_handle: None, + cur_task: None, task_list: VecDeque::new(), task_rx, task_tx, } } - pub fn add(&mut self, task: Task) { - info!("Adding task: {:?}", &task); - self.task_list.push_back(task); + pub fn add(&mut self, task_type: TaskType) { + info!("Adding task: {:?}", &task_type); + self.task_list.push_back(Task::new(task_type)); } pub fn idle(&self) -> bool { - self.handle.is_none() + self.cur_handle.is_none() } - pub fn poll(&mut self) -> Result> { - let mut return_handle: Option = None; + 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()); @@ -94,9 +105,13 @@ impl Tasks { // Check status of current task (if one is running). // NOTE: Action::NextScreen is sent once all tasks are complete - if let Some(task_handle) = self.handle.take() { - if task_handle.handle.is_finished() { - return_handle = Some(task_handle); + if let Some(task_handle) = self.cur_handle.take() { + if task_handle.is_finished() { + // Need to return task with handle + if let Some(mut cur_task) = self.cur_task.take() { + cur_task.handle = Some(task_handle); + return_task = Some(cur_task); + } if self.task_list.is_empty() { // No tasks remain self.task_tx.send(Action::NextScreen)?; @@ -105,85 +120,76 @@ impl Tasks { self.start()?; } } else { - // Task not complete, return handle - self.handle = Some(task_handle); + // TaskType not complete, return handle + self.cur_handle.replace(task_handle); } } else if !self.task_list.is_empty() { // No current task but one is available self.start()?; } - Ok(return_handle) + Ok(return_task) } pub fn start(&mut self) -> Result<()> { - if let Some(task) = self.task_list.pop_front() { + self.cur_task = self.task_list.pop_front(); + if let Some(task) = self.cur_task.take() { let task_tx = self.task_tx.clone(); - match task { - Task::Command(cmd_path, cmd_args) => { - self.handle = Some(run_task_command( + match task.task_type { + TaskType::Command(ref cmd_path, ref cmd_args) => { + self.cur_handle = Some(run_task_command( cmd_path.clone(), cmd_args.clone(), task_tx, )); } - Task::Diskpart(script) => { - self.handle = Some(run_task_diskpart(&script, task_tx)); + TaskType::Diskpart(ref script) => { + self.cur_handle = Some(run_task_diskpart(&script, task_tx)); } - Task::ScanDisks => { + TaskType::ScanDisks => { let disk_list_arc = self.disk_list.clone(); // Queue UpdateDiskList for various components - self.add(Task::UpdateDiskList); + self.add(TaskType::UpdateDiskList); - self.handle = Some(TaskHandle { - task: task.clone(), - handle: thread::spawn(move || { - let mut disks = disk_list_arc.lock().unwrap(); - *disks = disk::get_disks() - }), - }); + self.cur_handle = Some(thread::spawn(move || { + let mut disks = disk_list_arc.lock().unwrap(); + *disks = disk::get_disks() + })); } - Task::Sleep => { - self.handle = Some(TaskHandle { - task: task.clone(), - handle: thread::spawn(|| sleep(Duration::from_millis(250))), - }); + TaskType::Sleep => { + self.cur_handle = Some(thread::spawn(|| sleep(Duration::from_millis(250)))); } - Task::UpdateDestDisk(index) => { + TaskType::UpdateDestDisk(index) => { self.action_tx.send(Action::DisplayPopup( popup::Type::Info, String::from("Refreshing disk info"), ))?; // Queue UpdateDiskList for various components - self.add(Task::Sleep); - self.add(Task::UpdateDiskList); + self.add(TaskType::Sleep); + self.add(TaskType::UpdateDiskList); // Update destination disk ~in-place let disk_list_arc = self.disk_list.clone(); - self.handle = Some(TaskHandle { - task: task.clone(), - handle: thread::spawn(move || { - let mut disks = disk_list_arc.lock().unwrap(); - let old_disk = &mut disks[index]; - disks[index] = disk::refresh_disk_info(old_disk); - }), - }); + self.cur_handle = Some(thread::spawn(move || { + let mut disks = disk_list_arc.lock().unwrap(); + let old_disk = &mut disks[index]; + disks[index] = disk::refresh_disk_info(old_disk); + })); } - Task::UpdateDiskList => { + TaskType::UpdateDiskList => { let disks = self.disk_list.lock().unwrap(); let disks_copy = disks.clone(); let action_tx = self.action_tx.clone(); - self.handle = Some(TaskHandle { - handle: thread::spawn(move || { - if let Err(err) = action_tx.send(Action::UpdateDiskList(disks_copy)) { - panic!("Failed to send Action: {err:?}"); - } - }), - task: task.clone(), - }); + self.cur_handle = Some(thread::spawn(move || { + if let Err(err) = action_tx.send(Action::UpdateDiskList(disks_copy)) { + panic!("Failed to send Action: {err:?}"); + } + })); } } + // Done + self.cur_task.replace(task); } Ok(()) } @@ -193,9 +199,8 @@ fn run_task_command( cmd_path: PathBuf, cmd_args: Vec, task_tx: mpsc::UnboundedSender, -) -> TaskHandle { - let task = Task::Command(cmd_path.clone(), cmd_args.clone()); - let handle = if cfg!(windows) { +) -> JoinHandle<()> { + if cfg!(windows) { thread::spawn(move || { let result = Command::new(cmd_path) .args(cmd_args) @@ -227,14 +232,13 @@ fn run_task_command( } else { // Simulate task if not running under Windows thread::spawn(|| sleep(Duration::from_millis(250))) - }; - TaskHandle { handle, task } + } } -fn run_task_diskpart(script: &str, task_tx: mpsc::UnboundedSender) -> TaskHandle { - let task = Task::Diskpart(String::from(script)); +fn run_task_diskpart(script: &str, task_tx: mpsc::UnboundedSender) -> JoinHandle<()> { + let task = TaskType::Diskpart(String::from(script)); let task_str = format!("{:?}", &task); - let handle = if cfg!(windows) { + if cfg!(windows) { let script = String::from(script); thread::spawn(move || { let output = diskpart::run_script_raw(&script); @@ -251,6 +255,5 @@ fn run_task_diskpart(script: &str, task_tx: mpsc::UnboundedSender) -> Ta } else { // Simulate task if not running under Windows thread::spawn(|| sleep(Duration::from_millis(250))) - }; - TaskHandle { handle, task } + } } diff --git a/deja_vu/src/app.rs b/deja_vu/src/app.rs index 463565e..cbbe506 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::{Task, Tasks}, + tasks::{TaskType, Tasks}, tui::{Event, Tui}, }; use std::{ @@ -163,7 +163,7 @@ impl App { Mode::ScanDisks => { self.prev_mode = self.cur_mode; if self.tasks.idle() { - self.tasks.add(Task::ScanDisks); + self.tasks.add(TaskType::ScanDisks); } self.action_tx.send(Action::DisplayPopup( popup::Type::Info, @@ -191,7 +191,7 @@ impl App { ))?; // (Re)Enable volume mounting - self.tasks.add(Task::Command( + self.tasks.add(TaskType::Command( PathBuf::from(format!("{system32}/mountvol.exe")), vec![String::from("/e")], )); @@ -202,7 +202,7 @@ impl App { if let Some(disk) = disk_list.get(disk_index) { let table_type = self.clone.table_type.clone().unwrap(); let diskpart_script = build_dest_format_script(disk.id, &table_type); - self.tasks.add(Task::Diskpart(diskpart_script)); + self.tasks.add(TaskType::Diskpart(diskpart_script)); } } } @@ -211,12 +211,12 @@ impl App { popup::Type::Info, String::from("Running Clone Tool"), ))?; - self.tasks.add(Task::Command( + self.tasks.add(TaskType::Command( self.config.clone_app_path.clone(), Vec::new(), )); if let Some(dest_index) = self.clone.disk_index_dest { - self.tasks.add(Task::UpdateDestDisk(dest_index)); + self.tasks.add(TaskType::UpdateDestDisk(dest_index)); } } Mode::PostClone => { diff --git a/pe_menu/src/app.rs b/pe_menu/src/app.rs index 4922332..c8d4389 100644 --- a/pe_menu/src/app.rs +++ b/pe_menu/src/app.rs @@ -22,7 +22,7 @@ use core::{ config::Config, line::DVLine, state::Mode, - tasks::{Task, Tasks}, + tasks::{TaskType, Tasks}, tui::{Event, Tui}, }; use std::{ @@ -244,7 +244,7 @@ impl App { self.action_tx.send(Action::Select(None, None))?; } Action::OpenTerminal => { - self.tasks.add(Task::Command( + self.tasks.add(TaskType::Command( PathBuf::from("cmd.exe"), vec![String::from("-new_console:n")], )); @@ -254,11 +254,11 @@ impl App { popup::Type::Info, String::from("Restarting..."), ))?; - self.tasks.add(Task::Command( + self.tasks.add(TaskType::Command( PathBuf::from("X:/Windows/System32/sync64.exe"), vec![String::from("-r")], )); - self.tasks.add(Task::Command( + self.tasks.add(TaskType::Command( PathBuf::from("X:/Windows/System32/wpeutil.exe"), vec![String::from("reboot")], )); @@ -269,11 +269,11 @@ impl App { popup::Type::Info, String::from("Powering off..."), ))?; - self.tasks.add(Task::Command( + self.tasks.add(TaskType::Command( PathBuf::from("X:/Windows/System32/sync64.exe"), vec![String::from("-r")], )); - self.tasks.add(Task::Command( + self.tasks.add(TaskType::Command( PathBuf::from("X:/Windows/System32/wpeutil.exe"), vec![String::from("shutdown")], )); @@ -370,7 +370,7 @@ fn get_chunks(r: Rect) -> Vec { chunks } -pub fn build_command(app: &App, tool: &Tool) -> Task { +pub fn build_command(app: &App, tool: &Tool) -> TaskType { let cmd_path: PathBuf; let mut cmd_args: Vec = Vec::new(); let start_index: usize; @@ -390,7 +390,7 @@ pub fn build_command(app: &App, tool: &Tool) -> Task { }); } } - Task::Command(cmd_path, cmd_args) + TaskType::Command(cmd_path, cmd_args) } fn build_left_items(app: &App) -> Action {