diff --git a/Cargo.lock b/Cargo.lock index bf0251a..0aa4f07 100644 --- a/Cargo.lock +++ b/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" diff --git a/Cargo.toml b/Cargo.toml index c63b0c9..0f556e4 100644 --- a/Cargo.toml +++ b/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" \ No newline at end of file +url = "2.5.4" +async-stream = "0.3" \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index d330b0b..adefc77 100644 --- a/src/main.rs +++ b/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, - sync::{Arc, LazyLock}, + io::Write, + path::{Path, PathBuf}, + sync::LazyLock, + thread::spawn, time::Instant, }; -use tokio::{ - sync::Semaphore, - task::JoinSet, - time::{Duration, sleep}, -}; +use tokio::time::Duration; use url::Url; static API_KEY_REGEX: LazyLock = 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> { +fn load_keys(path: &Path) -> Result> { 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> { Ok(keys) } -fn output_file_txt(keys: &[String], output_path: &PathBuf) -> Result<()> { - let content = keys.join("\n"); - fs::write(output_path, content)?; - println!( - "Successfully wrote {} keys to {:?}", - keys.len(), - output_path - ); - Ok(()) +async fn validate_key_with_retry(client: &Client, api_host: &Url, key: String) -> Option { + let retry_policy = ExponentialBuilder::default() + .with_max_times(3) + .with_min_delay(Duration::from_secs(3)) + .with_max_delay(Duration::from_secs(8)); + + let result = (|| async { + match keytest(&client, &api_host, &key).await { + 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 { - 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 mut set = JoinSet::new(); + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::(); + let stream = stream! { + while let Some(item) = rx.recv().await { + yield item; + } + }; - for key in keys { - let client_clone = client.clone(); - let api_host_clone = config.api_host.clone(); - let semaphore_clone = Arc::clone(&semaphore); - - 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; - } - } + spawn(move || { + for key in keys { + if API_KEY_REGEX.is_match(&key) { + if let Err(e) = tx.send(key) { + eprintln!("Failed to send key: {}", e); } + } else { + eprintln!("Invalid key format: {}", key); } + } + }); - eprintln!("Key: {}... -> FAILED after all retries.", &key[..10]); - None - }); - } - let mut valid_keys = Vec::new(); - while let Some(res) = set.join_next().await { - if let Ok(Some(key)) = res { - valid_keys.push(key); + let valid_keys_stream = stream + .map(|key| validate_key_with_retry(&client, &config.api_host, key)) + .buffer_unordered(config.concurrency) + .filter_map(|r| async { r }); + pin_mut!(valid_keys_stream); + // open output file + let mut output_file = fs::File::create(&config.output_path)?; + // 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 cost time:{:?}", start_time.elapsed()); + println!("Total Elapsed Time: {:?}", start_time.elapsed()); Ok(()) -} +} \ No newline at end of file