diff --git a/Cargo.lock b/Cargo.lock index 45d4250..7ca2a68 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -254,6 +254,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + [[package]] name = "chrono" version = "0.4.42" @@ -364,6 +370,17 @@ dependencies = [ "typenum", ] +[[package]] +name = "ctrlc" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "881c5d0a13b2f1498e2306e82cbada78390e152d4b1378fb28a84f4dcd0dc4f3" +dependencies = [ + "dispatch", + "nix", + "windows-sys 0.61.2", +] + [[package]] name = "digest" version = "0.10.7" @@ -396,6 +413,12 @@ dependencies = [ "winapi", ] +[[package]] +name = "dispatch" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd0c93bb4b0c6d9b77f4435b0ae98c24d17f1c45b2ff844c6151a07256ca923b" + [[package]] name = "displaydoc" version = "0.2.5" @@ -473,6 +496,7 @@ version = "0.1.0" dependencies = [ "anyhow", "clap", + "ctrlc", "fancy-regex", "num-integer", "num_cpus", @@ -884,6 +908,18 @@ version = "1.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "650eef8c711430f1a879fdd01d4745a7deea475becfb90269c06775983bbf086" +[[package]] +name = "nix" +version = "0.30.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6" +dependencies = [ + "bitflags", + "cfg-if", + "cfg_aliases", + "libc", +] + [[package]] name = "nom" version = "7.1.3" diff --git a/Cargo.toml b/Cargo.toml index fa29e32..7acd769 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,7 @@ categories = ["command-line-utilities"] [dependencies] anyhow = "1.0.100" clap = { version = "4.5.50", features = ["derive"] } +ctrlc = "3.5.0" fancy-regex = "0.16.2" num-integer = "0.1.46" num_cpus = "1.17.0" diff --git a/src/lib.rs b/src/lib.rs index b32f94b..535d319 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,6 +28,12 @@ use sequoia_openpgp::{ type SecretKey = Key; +#[allow(clippy::large_enum_variant)] +enum Message { + Key(SecretKey), + Stop, +} + #[derive(Clone, Debug)] pub struct Config { pub regex: Regex, @@ -63,54 +69,71 @@ impl Fingerprunk { } } - pub fn run(mut self) { + pub fn run(mut self) -> anyhow::Result<()> { self.started_instant = Instant::now(); - let (tx, rx) = mpsc::channel(); + let (sender, receiver) = mpsc::channel(); + + { + let sender = sender.clone(); + ctrlc::set_handler(move || { + let _ = sender.send(Message::Stop); + })?; + } thread::scope(|scope| { - const THREAD_SPAWN_EXPECT_MSG: &str = "should be able to spawn thread"; - let ref_self = &self; let status_displayer = if self.config.status_enabled { Some( thread::Builder::new() .name("status_displayer".to_string()) - .spawn_scoped(scope, move || ref_self.status_displayer_thread()) - .expect(THREAD_SPAWN_EXPECT_MSG), + .spawn_scoped(scope, move || ref_self.status_displayer_thread())?, ) } else { None }; for num in 0..num_cpus::get() { - let tx = tx.clone(); + let sender = sender.clone(); thread::Builder::new() .name(format!("worker-{num:03}")) - .spawn_scoped(scope, move || ref_self.worker_thread(tx)) - .expect(THREAD_SPAWN_EXPECT_MSG); + .spawn_scoped(scope, move || ref_self.worker_thread(sender))?; } - let on_stop = || { - // Ask all other threads to stop - self.stop.store(true, Ordering::Relaxed); + let mut stdout = io::stdout().lock(); - // Unpark the status displayer thread - if let Some(status_displayer) = status_displayer { - status_displayer.thread().unpark(); + // Receive and process messages from the workers and the ctrl-c handler + for message in receiver { + match message { + Message::Key(key) => { + let cert = self.key_to_cert(&key)?; + self.serialize_cert(cert, &mut stdout)?; + + // Increase "found" counter and stop if enough matches have been found + let prev = self.counter_found.fetch_add(1, Ordering::Relaxed); + if self.config.stop_after.is_some_and(|s| prev + 1 == s.get()) { + break; + } + } + Message::Stop => break, } - }; + } - thread::Builder::new() - .name("finalizer".to_string()) - .spawn_scoped(scope, move || ref_self.finalizer_thread(rx, on_stop)) - .expect(THREAD_SPAWN_EXPECT_MSG); - }); + // Ask all other threads to stop + self.stop.store(true, Ordering::Relaxed); + + // Unpark the status displayer thread, if existant + if let Some(status_displayer) = status_displayer { + status_displayer.thread().unpark(); + } + + Ok(()) + }) } - fn worker_thread(&self, matches_tx: mpsc::Sender) { + fn worker_thread(&self, sender: mpsc::Sender) { let mut fingerprint_hex = String::with_capacity(20 * 2); while !self.stop.load(Ordering::Relaxed) { @@ -120,8 +143,8 @@ impl Fingerprunk { write!(fingerprint_hex, "{:X}", key.fingerprint()) .expect("should write into string without error"); if self.check_fingerprint(&fingerprint_hex) { - matches_tx - .send(Key::V4(key)) + sender + .send(Message::Key(Key::V4(key))) .expect("should be able to send key"); } self.counter_tried.fetch_add(1, Ordering::Relaxed); @@ -136,27 +159,6 @@ impl Fingerprunk { .expect("should check regex without error") } - fn finalizer_thread(&self, matches_rx: mpsc::Receiver, on_stop: impl FnOnce()) { - let mut stdout = io::stdout().lock(); - - for key in matches_rx { - let cert = self - .key_to_cert(&key) - .expect("should be able to create certificate"); - - self.serialize_cert(cert, &mut stdout) - .expect("should be able to serialize certificate"); - - let prev = self.counter_found.fetch_add(1, Ordering::Relaxed); - - if self.config.stop_after.is_some_and(|s| prev + 1 == s.get()) { - break; - } - } - - on_stop(); - } - fn key_to_cert(&self, key: &SecretKey) -> anyhow::Result { let sig = SignatureBuilder::new(SignatureType::DirectKey) .set_hash_algo(HashAlgorithm::SHA512) diff --git a/src/main.rs b/src/main.rs index 6f575de..5a9c915 100644 --- a/src/main.rs +++ b/src/main.rs @@ -89,7 +89,7 @@ fn main() -> anyhow::Result<()> { password, }; - Fingerprunk::new_from_config(config).run(); + Fingerprunk::new_from_config(config).run()?; Ok(()) }