Cancel thread if --captcha-string was found in body

src/cli.rs

@@ -6,7 +6,7 @@ pub struct Args {
     pub base_uri: String,
     pub uri_file: String,
     pub user_agent: String,
-
+    pub captcha_string: String,
     pub verbose: bool,
     pub bypass: bool,
 }
@@ -45,7 +45,15 @@ pub fn get_args() -> Args {
                 .help("User-Agent to use")
                 .takes_value(true),
         )
-        .arg(Arg::with_name("bypass").short("c").long("--bypass").help(
+        .arg(
+            Arg::with_name("captcha-string")
+                .short("c")
+                .long("--captcha-string")
+                .value_name("STRING")
+                .help("Stop processing when STRING was found in body")
+                .takes_value(true),
+        )
+        .arg(Arg::with_name("bypass").short("p").long("--bypass").help(
             "Set cacheupdate cookie to bypass cache",
         ))
         .arg(Arg::with_name("verbose").short("v").long("verbose").help(
@@ -62,6 +70,7 @@ pub fn get_args() -> Args {
         verbose: args.is_present("verbose"),
         bypass: args.is_present("bypass"),
         base_uri: args.value_of("base-uri").unwrap_or("").to_string(),
+        captcha_string: args.value_of("captcha-string").unwrap_or("").to_string(),
         uri_file: args.value_of("uri-file").unwrap().to_string(),
         user_agent: args.value_of("user-agent")
             .unwrap_or("Googlebot (cache warmer)")

src/main.rs

@@ -60,11 +60,12 @@ fn main() {
         // Clone values before move
         let uris = uris.clone();
         let user_agent = user_agent.clone();
+        let captcha_string = args.captcha_string.clone();
         let verbose = args.verbose;
         let bypass = args.bypass;
 
         workers.push(thread::spawn(move || {
-            worker::spawn(uris, user_agent, verbose, bypass);
+            worker::spawn(uris, user_agent, &captcha_string, verbose, bypass);
         }));
     }
 

src/worker.rs

@@ -9,7 +9,13 @@ use tokio_core::reactor::Core;
 // Make custom X-Cache-Status header known
 header! { (XCacheStatus, "X-Cache-Status") => [String] }
 
-pub fn spawn(uris: Arc<Mutex<Vec<Uri>>>, user_agent: UserAgent, verbose: bool, bypass: bool) {
+pub fn spawn(
+    uris: Arc<Mutex<Vec<Uri>>>,
+    user_agent: UserAgent,
+    captcha_string: &str,
+    verbose: bool,
+    bypass: bool,
+) {
     let mut core = Core::new().unwrap();
     let handle = core.handle();
 
@@ -45,9 +51,17 @@ pub fn spawn(uris: Arc<Mutex<Vec<Uri>>>, user_agent: UserAgent, verbose: bool, b
                 );
             }
 
-            // We need to read out the full body, so the connection can be closed.
-            // TODO: Is there a more efficient way of consuming the body?
-            res.body().for_each(|_| Ok(()))
+            res.body().concat2().and_then(move |body| {
+                // body is a &[8], so from_utf8_lossy() is required here
+                let html = String::from_utf8_lossy(body.as_ref());
+                if captcha_string.len() > 0 && html.contains(captcha_string) {
+                    println!(
+                        "Found '{}' in response body. Stopping thread.",
+                        captcha_string
+                    );
+                }
+                Ok(())
+            })
         });
 
         core.run(work).unwrap();