feat: update dependencies and enhance key validation logic

Cargo.lock

@@ -7,7 +7,10 @@ name = "Gemini-Keychecker"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "async-stream",
+ "backon",
  "clap",
+ "futures",
  "regex",
  "reqwest",
  "serde_json",
@@ -95,12 +98,45 @@ version = "1.0.98"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487"
 
+[[package]]
+name = "async-stream"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476"
+dependencies = [
+ "async-stream-impl",
+ "futures-core",
+ "pin-project-lite",
+]
+
+[[package]]
+name = "async-stream-impl"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "atomic-waker"
 version = "1.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"
 
+[[package]]
+name = "backon"
+version = "1.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "302eaff5357a264a2c42f127ecb8bac761cf99749fc3dc95677e2743991f99e7"
+dependencies = [
+ "fastrand",
+ "gloo-timers",
+ "tokio",
+]
+
 [[package]]
 name = "backtrace"
 version = "0.3.75"
@@ -289,6 +325,21 @@ dependencies = [
  "percent-encoding",
 ]
 
+[[package]]
+name = "futures"
+version = "0.3.31"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876"
+dependencies = [
+ "futures-channel",
+ "futures-core",
+ "futures-executor",
+ "futures-io",
+ "futures-sink",
+ "futures-task",
+ "futures-util",
+]
+
 [[package]]
 name = "futures-channel"
 version = "0.3.31"
@@ -296,6 +347,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10"
 dependencies = [
  "futures-core",
+ "futures-sink",
 ]
 
 [[package]]
@@ -304,6 +356,34 @@ version = "0.3.31"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e"
 
+[[package]]
+name = "futures-executor"
+version = "0.3.31"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f"
+dependencies = [
+ "futures-core",
+ "futures-task",
+ "futures-util",
+]
+
+[[package]]
+name = "futures-io"
+version = "0.3.31"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6"
+
+[[package]]
+name = "futures-macro"
+version = "0.3.31"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "futures-sink"
 version = "0.3.31"
@@ -322,10 +402,16 @@ version = "0.3.31"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81"
 dependencies = [
+ "futures-channel",
  "futures-core",
+ "futures-io",
+ "futures-macro",
+ "futures-sink",
  "futures-task",
+ "memchr",
  "pin-project-lite",
  "pin-utils",
+ "slab",
 ]
 
 [[package]]
@@ -357,6 +443,18 @@ version = "0.31.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f"
 
+[[package]]
+name = "gloo-timers"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bbb143cf96099802033e0d4f4963b19fd2e0b728bcf076cd9cf7f6634f092994"
+dependencies = [
+ "futures-channel",
+ "futures-core",
+ "js-sys",
+ "wasm-bindgen",
+]
+
 [[package]]
 name = "h2"
 version = "0.4.11"

Cargo.toml

@@ -5,9 +5,12 @@ edition = "2024"
 
 [dependencies]
 anyhow = "1.0.98"
+backon = "1"
 clap = { version = "4.5.40", features = ["derive"] }
+futures = "0.3"
 regex = "1.11.1"
 reqwest = { version = "0.12.22", features = ["json"] }
 serde_json = "1.0.140"
 tokio = { version = "1.46", features = ["macros", "rt-multi-thread", "time"] }
 url = "2.5.4"
+async-stream = "0.3"

src/main.rs

@@ -1,19 +1,20 @@
 use anyhow::Result;
+use async_stream::stream;
+use backon::{ExponentialBuilder, Retryable};
 use clap::Parser;
+use futures::{pin_mut, stream::StreamExt};
 use regex::Regex;
 use reqwest::{Client, StatusCode};
 use std::{
     collections::HashSet,
     fs,
-    path::PathBuf,
+    io::Write,
-    sync::{Arc, LazyLock},
+    path::{Path, PathBuf},
+    sync::LazyLock,
+    thread::spawn,
     time::Instant,
 };
-use tokio::{
+use tokio::time::Duration;
-    sync::Semaphore,
-    task::JoinSet,
-    time::{Duration, sleep},
-};
 use url::Url;
 static API_KEY_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"^AIzaSy.{33}$").unwrap());
 #[derive(Parser, Debug)]
@@ -40,7 +41,7 @@ enum KeyStatus {
     Invalid,
     Retryable(String),
 }
-fn load_keys(path: &PathBuf) -> Result<Vec<String>> {
+fn load_keys(path: &Path) -> Result<Vec<String>> {
     let keys_txt = fs::read_to_string(path)?;
     let unique_keys_set: HashSet<&str> = keys_txt
         .lines()
@@ -52,19 +53,46 @@ fn load_keys(path: &PathBuf) -> Result<Vec<String>> {
     Ok(keys)
 }
 
-fn output_file_txt(keys: &[String], output_path: &PathBuf) -> Result<()> {
+async fn validate_key_with_retry(client: &Client, api_host: &Url, key: String) -> Option<String> {
-    let content = keys.join("\n");
+    let retry_policy = ExponentialBuilder::default()
-    fs::write(output_path, content)?;
+        .with_max_times(3)
-    println!(
+        .with_min_delay(Duration::from_secs(3))
-        "Successfully wrote {} keys to {:?}",
+        .with_max_delay(Duration::from_secs(8));
-        keys.len(),
+
-        output_path
+    let result = (|| async {
-    );
+        match keytest(&client, &api_host, &key).await {
-    Ok(())
+            Ok(KeyStatus::Valid) => {
+                println!("Key: {}... -> SUCCESS", &key[..10]);
+                Ok(Some(key.clone()))
+            }
+            Ok(KeyStatus::Invalid) => {
+                println!("Key: {}... -> INVALID (Forbidden)", &key[..10]);
+                Ok(None)
+            }
+            Ok(KeyStatus::Retryable(reason)) => {
+                eprintln!("Key: {}... -> RETRYABLE (Reason: {})", &key[..10], reason);
+                Err(anyhow::anyhow!("Retryable error: {}", reason))
+            }
+            Err(e) => {
+                eprintln!("Key: {}... -> NETWORK ERROR (Reason: {})", &key[..10], e);
+                Err(e)
+            }
+        }
+    })
+    .retry(retry_policy)
+    .await;
+
+    match result {
+        Ok(key_result) => key_result,
+        Err(_) => {
+            eprintln!("Key: {}... -> FAILED after all retries.", &key[..10]);
+            None
+        }
+    }
 }
 
 async fn keytest(client: &Client, api_host: &Url, keys: &str) -> Result<KeyStatus> {
-    const API_PATH: &str = "v1beta/models/gemini-2.0-flash-exp:generateContent";
+    const API_PATH: &str = "v1beta/models/gemini-2.0-flash:generateContent";
     let full_url = api_host.join(API_PATH)?;
     let request_body = serde_json::json!({
         "contents": [
@@ -110,65 +138,40 @@ async fn main() -> Result<()> {
         .timeout(Duration::from_millis(config.timeout_ms))
         .build()?;
 
-    let semaphore = Arc::new(Semaphore::new(config.concurrency));
+    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<String>();
-    let mut set = JoinSet::new();
+    let stream = stream! {
+        while let Some(item) = rx.recv().await {
+            yield item;
+        }
+    };
 
-    for key in keys {
+    spawn(move || {
-        let client_clone = client.clone();
+        for key in keys {
-        let api_host_clone = config.api_host.clone();
+            if API_KEY_REGEX.is_match(&key) {
-        let semaphore_clone = Arc::clone(&semaphore);
+                if let Err(e) = tx.send(key) {
-
+                    eprintln!("Failed to send key: {}", e);
-        set.spawn(async move {
-            const MAX_RETRIES: u32 = 3; 
-            let _permit = semaphore_clone.acquire().await.unwrap();
-            for attempt in 0..MAX_RETRIES {
-                match keytest(&client_clone, &api_host_clone, &key).await {
-                    Ok(KeyStatus::Valid) => {
-                        println!("Key: {}... -> SUCCESS", &key[..10]);
-                        return Some(key);
-                    }
-                    Ok(KeyStatus::Invalid) => {
-                        println!("Key: {}... -> INVALID (Forbidden)", &key[..10]);
-                        return None;
-                    }
-                    Ok(KeyStatus::Retryable(reason)) => {
-                        eprintln!(
-                            "Key: {}... -> RETRYABLE (Attempt {}/{}, Reason: {})",
-                            &key[..10],
-                            attempt + 1,
-                            MAX_RETRIES,
-                            reason
-                        );
-                        if attempt < MAX_RETRIES - 1 {
-                            sleep(Duration::from_secs(2_u64.pow(attempt))).await;
-                        }
-                    }
-                    Err(e) => {
-                        eprintln!(
-                            "Key: {}... -> NETWORK ERROR (Attempt {}/{}, Reason: {})",
-                            &key[..10],
-                            attempt + 1,
-                            MAX_RETRIES,
-                            e
-                        );
-                        if attempt < MAX_RETRIES - 1 {
-                            sleep(Duration::from_secs(2_u64.pow(attempt))).await;
-                        }
-                    }
                 }
+            } else {
+                eprintln!("Invalid key format: {}", key);
             }
+        }
+    });
 
-            eprintln!("Key: {}... -> FAILED after all retries.", &key[..10]);
+    let valid_keys_stream = stream
-            None
+        .map(|key| validate_key_with_retry(&client, &config.api_host, key))
-        });
+        .buffer_unordered(config.concurrency)
-    }
+        .filter_map(|r| async { r });
-    let mut valid_keys = Vec::new();
+    pin_mut!(valid_keys_stream);
-    while let Some(res) = set.join_next().await {
+    // open output file
-        if let Ok(Some(key)) = res {
+    let mut output_file = fs::File::create(&config.output_path)?;
-            valid_keys.push(key);
+    // Write valid keys to output file
+    while let Some(valid_key) = valid_keys_stream.next().await {
+        // Collect valid keys
+        println!("Valid key found: {}", valid_key);
+        if let Err(e) = writeln!(output_file, "{}", valid_key) {
+            eprintln!("Failed to write key to output file: {}", e);
         }
     }
-    output_file_txt(&valid_keys, &config.output_path)?;
+    println!("Total Elapsed Time: {:?}", start_time.elapsed());
-    println!("Total cost time:{:?}", start_time.elapsed());
     Ok(())
 }