From 058442ac53177d72c81c5420056ee18c99faa98c Mon Sep 17 00:00:00 2001 From: puffaboo Date: Sun, 3 Jul 2022 22:46:15 +0100 Subject: [PATCH] fixed browser creating " directory, added watcher for chromedriver process, cleanup --- src/browser.rs | 24 +++++++++---- src/chromedriver.rs | 87 +++++++++++++++++++++++++++++++-------------- src/main.rs | 72 +++++++++++++++++++++++++++---------- 3 files changed, 130 insertions(+), 53 deletions(-) diff --git a/src/browser.rs b/src/browser.rs index 6be2553..ff7920e 100644 --- a/src/browser.rs +++ b/src/browser.rs @@ -11,10 +11,14 @@ const DATA_DIR_CHROMIUM: &str = "/.config/chromium"; #[cfg(target_os = "linux")] const DATA_DIR_CHROME: &str = "~/.config/google-chrome"; #[cfg(target_os = "linux")] +const DATA_DIR_CHROME_BETA: &str = "~/.config/google-chrome-beta"; +#[cfg(target_os = "linux")] const HOME_VARIABLE: &str = "HOME"; #[cfg(target_os = "windows")] const DATA_DIR_CHROME: &str = "\\Google\\Chrome\\User Data"; #[cfg(target_os = "windows")] +const DATA_DIR_CHROME_BETA: &str = "\\Google\\Chrome Beta\\User Data"; +#[cfg(target_os = "windows")] const DATA_DIR_CHROMIUM: &str = "\\Google\\Chromium\\User Data"; #[cfg(target_os = "windows")] const HOME_VARIABLE: &str = "LOCALAPPDATA"; @@ -34,11 +38,13 @@ fn get_data_dir() -> String { if home.is_empty() { panic!("cannot find home env variable"); } - if Path::new(format!("{}{}", home, DATA_DIR_CHROMIUM).as_str()).exists() { - println!("{}{}", home, DATA_DIR_CHROMIUM); - return format!("{}{}", home, DATA_DIR_CHROMIUM); + let by_preference = vec![DATA_DIR_CHROMIUM, DATA_DIR_CHROME, DATA_DIR_CHROME_BETA]; + for data_dir in by_preference { + if Path::new(format!("{}{}", home, data_dir).as_str()).exists() { + return format!("{}{}", home, data_dir); + } } - println!("{}{}", home, DATA_DIR_CHROME); + // Default to chrome anyway I guess return format!("{}{}", home, DATA_DIR_CHROME); } @@ -47,7 +53,13 @@ impl BrowserSession { // TODO: error if chrome is already running let mut caps = DesiredCapabilities::chrome(); // Use the normal data directory (for extensions and logins) - caps.add_chrome_arg(format!("{}={}", DATA_DIR_ARG, get_data_dir()).as_str())?; + let data_dir = get_data_dir(); + println!("starting chrome using data directory: {}", data_dir); + + #[cfg(target_os = "windows")] + let data_dir = format!("\"{}\"", data_dir); + + caps.add_chrome_arg(format!("{}={}", DATA_DIR_ARG, data_dir).as_str())?; let driver = WebDriver::new("http://localhost:6444", caps).await?; println!("started browser"); @@ -69,12 +81,10 @@ impl BrowserSession { impl BrowserHandle { pub async fn start(&self, start_url: &str) -> Result<(), anyhow::Error> { self.handle.fullscreen_window().await?; - println!("!"); self.handle .get(format!("http://{}", start_url)) .await .unwrap_or(()); // Ignore errors ig? idfk it errors out in some weirdo way but it works - println!("?"); Ok(()) } diff --git a/src/chromedriver.rs b/src/chromedriver.rs index bd288fa..bdd7f93 100644 --- a/src/chromedriver.rs +++ b/src/chromedriver.rs @@ -1,19 +1,12 @@ +use futures::channel::oneshot::Sender; use std::{ - fmt, - io::{BufRead, BufReader}, - process::{Child, Command, Stdio}, + io::ErrorKind, + process::{Child, Command, Output, Stdio}, + thread::{self, JoinHandle}, + time::Duration, }; use sysinfo::{Pid, ProcessRefreshKind, RefreshKind, System, SystemExt}; -#[derive(Debug, Clone)] -struct AddrInUse; - -impl fmt::Display for AddrInUse { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "address is in use") - } -} - pub struct Driver { child: Child, } @@ -22,27 +15,17 @@ impl Driver { pub fn new() -> Result { let mut child = Command::new("chromedriver") .stdout(Stdio::piped()) + .stderr(Stdio::piped()) .arg("--port=6444") .spawn()?; let child_id = child.id(); - let mut child_out = BufReader::new(child.stdout.as_mut().unwrap()); - let mut line = String::new(); - let mut sys_info = System::new_with_specifics( RefreshKind::new().with_processes(ProcessRefreshKind::new()), ); // Wait for it to say Chromedriver started successfully + print!("chromedriver health check..."); loop { - child_out.read_line(&mut line).unwrap(); - if line.contains("ChromeDriver was started successfully.") { - break; - } - if line.contains("bind() failed: Address already in use") || line.contains("Exiting...") - { - return Err(anyhow::anyhow!(line)); - } - // Check if chromedriver has exited sys_info.refresh_all(); if let None = sys_info.process(Pid::from(to_pid(child_id.clone()))) { @@ -52,16 +35,66 @@ impl Driver { "chromedriver exited with status code {}", status_code.code().unwrap_or(1337), )); + } else { + break; } } + println!(" done."); + + // Check if chrome/chromium is already running. Die if so. + #[cfg(target_os = "linux")] + let chrome_names = ("chrome", "chromium"); + #[cfg(target_os = "windows")] + let chrome_names = ("chrome.exe", "chromium.exe"); + if sys_info.processes_by_exact_name(chrome_names.0).count() > 0 + || sys_info.processes_by_exact_name(chrome_names.1).count() > 0 + { + return Err(anyhow::anyhow!("chrome/chromium is already running")); + } println!("chromedriver started"); Ok(Self { child }) } - pub fn exit(&mut self) -> Result<(), anyhow::Error> { - self.child.kill()?; - Ok(()) + pub fn exit_waiter(&mut self, notify: Sender<()>) -> JoinHandle<()> { + // Check if chromedriver has exited + let mut sys_info = System::new_with_specifics( + RefreshKind::new().with_processes(ProcessRefreshKind::new()), + ); + let pid = Pid::from(to_pid(self.child.id().clone())); + thread::spawn(move || loop { + sys_info.refresh_all(); + let proc = sys_info.process(pid); + // Consider two cases of chromedriver being dead: + // 1) the process itself is gone (might be on windows) + // 2) the process is alive but not reaped (zombie) on linux + // in the second case, we check for if it has zero subtasks. + // For our purposes, that means it's dead. + if !proc.is_none() { + #[cfg(target_os = "linux")] + if proc.unwrap().tasks.len() != 0 { + thread::sleep(Duration::from_secs(3)); + continue; + } + } + println!("chromedriver has exited"); + notify.send(()).unwrap(); + break; + }) + } + + pub fn exit_with_output(mut self) -> Result { + match self.child.kill() { + Err(e) => { + if e.kind() != ErrorKind::InvalidInput { + return Err(anyhow::anyhow!(e)); + } else { + println!("killing child: {}", e); + } + } + _ => {} + } + Ok(self.child.wait_with_output()?) } } diff --git a/src/main.rs b/src/main.rs index e42b533..bcbbb82 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,23 +2,32 @@ mod botdriver; mod browser; mod chromedriver; mod webserve; +use botdriver::BotDriver; +use browser::BrowserSession; +use chromedriver::Driver; +use futures::channel::oneshot; +use serde::{Deserialize, Serialize}; use std::{ io::{self, Write}, sync::{ atomic::{AtomicBool, Ordering}, Arc, }, + time::Duration, }; - -use botdriver::BotDriver; -use browser::BrowserSession; -use chromedriver::Driver; -use serde::{Deserialize, Serialize}; +use tokio::time; const STATIC_ADDR: &str = "127.0.0.1:8010"; #[tokio::main] -async fn main() -> Result<(), anyhow::Error> { +async fn main() { + if let Err(e) = main_but_with_result().await { + println!("frc exited with error: {}", e); + std::process::exit(1); + } + std::process::exit(0); +} +async fn main_but_with_result() -> Result<(), anyhow::Error> { let mut cfg: FRCConfig = confy::load_path(CFG_NAME)?; if cfg.username == None || cfg.token == None { cfg = create_config()?; @@ -36,23 +45,48 @@ async fn main() -> Result<(), anyhow::Error> { let mut bot = BotDriver::new(cfg.token.unwrap(), username); browser.start(STATIC_ADDR).await?; - println!("################"); - println!("# starting frc #"); - println!("################"); - tokio::spawn(async move { - let term = Arc::new(AtomicBool::new(false)); + println!(); + println!("#####################"); + println!("# launching frc #"); + println!("# #"); + println!("# frc can be closed #"); + println!("# by an interrupt #"); + println!("# (Ctl-C) #"); + println!("#####################"); + println!(); + + tokio::spawn(async move { + let should_exit = Arc::new(AtomicBool::new(false)); + + let (chr_sender, mut chr_receiver) = oneshot::channel(); + chr.exit_waiter(chr_sender); + + signal_hook::flag::register(signal_hook::consts::SIGINT, Arc::clone(&should_exit)).unwrap(); + signal_hook::flag::register(signal_hook::consts::SIGTERM, Arc::clone(&should_exit)) + .unwrap(); - signal_hook::flag::register(signal_hook::consts::SIGINT, Arc::clone(&term)).unwrap(); - signal_hook::flag::register(signal_hook::consts::SIGTERM, Arc::clone(&term)).unwrap(); loop { - if term.load(Ordering::Relaxed) { - println!(); - println!("closing down"); - static_handle.abort(); - driver.stop().await.expect("failed killing browser"); - chr.exit().expect("failed killing chromedriver"); + if should_exit.load(Ordering::Relaxed) || chr_receiver.try_recv().unwrap().is_some() { break; } + time::sleep(Duration::from_millis(100)).await; + } + + println!(); + println!("closing down"); + static_handle.abort(); + driver.stop().await.expect("failed killing browser"); + let output = chr.exit_with_output().expect("failed killing chromedriver"); + println!("chromedriver exited with status code {}", output.status); + if output.stderr.len() > 0 { + print!("chromedriver stderr output:"); + io::stdout().write(&output.stderr).unwrap(); + println!(); + } + if output.stdout.len() > 0 { + print!("chromedriver stdout output:"); + io::stdout().write(&output.stdout).unwrap(); + println!(); } std::process::exit(0);