Skip to content

Commit

Permalink
fix: using shared state to signal other plugin workers about unreacha…
Browse files Browse the repository at this point in the history
…ble endpoints
  • Loading branch information
evilsocket committed Nov 9, 2023
1 parent 8473773 commit 6438347
Showing 1 changed file with 10 additions and 6 deletions.
16 changes: 10 additions & 6 deletions src/plugins/manager.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use std::collections::BTreeMap;
use std::sync::Mutex;
use std::sync::{Mutex, RwLock};
use std::time;

use ahash::HashSet;
Expand Down Expand Up @@ -70,10 +70,11 @@ pub(crate) async fn run(
let single = plugin.single_credential();
let override_payload = plugin.override_payload();
let combinations = session.combinations(override_payload, single)?;
let unreachables: Arc<RwLock<HashSet<String>>> = Arc::new(RwLock::new(HashSet::default()));

// spawn worker threads
for _ in 0..session.options.concurrency {
task::spawn(worker(plugin, session.clone()));
task::spawn(worker(plugin, unreachables.clone(), session.clone()));
}

if !session.options.quiet {
Expand All @@ -95,12 +96,15 @@ pub(crate) async fn run(
Ok(())
}

async fn worker(plugin: &dyn Plugin, session: Arc<Session>) {
async fn worker(
plugin: &dyn Plugin,
unreachables: Arc<RwLock<HashSet<String>>>,
session: Arc<Session>,
) {
log::debug!("worker started");

let timeout = time::Duration::from_millis(session.options.timeout);
let retry_time: time::Duration = time::Duration::from_millis(session.options.retry_time);
let mut unreachables: HashSet<String> = HashSet::default();

while let Ok(creds) = session.recv_credentials().await {
if session.is_stop() {
Expand All @@ -125,7 +129,7 @@ async fn worker(plugin: &dyn Plugin, session: Arc<Session>) {
attempt += 1;

// skip attempt if we had enough failures from this specific target
if !unreachables.contains(&creds.target) {
if !unreachables.read().unwrap().contains(&creds.target) {
match plugin.attempt(&creds, timeout).await {
Err(err) => {
errors += 1;
Expand All @@ -142,7 +146,7 @@ async fn worker(plugin: &dyn Plugin, session: Arc<Session>) {
} else {
// add this target to the list of unreachable in order to avoi
// pointless attempts
unreachables.insert(creds.target.clone());
unreachables.write().unwrap().insert(creds.target.clone());

log::error!(
"[{}] attempt {}/{}: {}",
Expand Down

0 comments on commit 6438347

Please sign in to comment.