From 5f1b30772b44c14e33e19e3c12d0005f6e93e348 Mon Sep 17 00:00:00 2001
From: Yoo1tic <137816438+Yoo1tic@users.noreply.github.com>
Date: Thu, 31 Jul 2025 20:57:19 +0800
Subject: [PATCH] feat: add configurable retry mechanism for API requests

---
 Config.toml.example       |  1 +
 src/config/config.rs      |  9 +++++++++
 src/service/key_tester.rs | 40 ++++++++++++++++++++++-----------------
 src/service/validation.rs | 10 ++++++++--
 4 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/Config.toml.example b/Config.toml.example
index 805ef00..7bf9e41 100644
--- a/Config.toml.example
+++ b/Config.toml.example
@@ -3,6 +3,7 @@ output_path = "output_keys.txt"
 backup_path = "backup_keys.txt"
 api_host = "https://generativelanguage.googleapis.com/"
 timeout_sec = 20
+max_retries = 2
 concurrency = 30
 proxy = "http://username:password@host:port"
 enable_multiplexing = true
diff --git a/src/config/config.rs b/src/config/config.rs
index d1109a7..d4f4c7c 100644
--- a/src/config/config.rs
+++ b/src/config/config.rs
@@ -39,6 +39,10 @@ struct Cli {
     #[serde(skip_serializing_if = "Option::is_none")]
     concurrency: Option<usize>,
 
+    #[arg(short = 'r', long)]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    max_retries: Option<usize>,
+
     #[arg(short = 'x', long)]
     #[serde(skip_serializing_if = "Option::is_none")]
     proxy: Option<Url>,
@@ -66,6 +70,10 @@ pub struct KeyCheckerConfig {
     #[serde(default)]
     pub timeout_sec: u64,
 
+    // Maximum number of retries for failed requests.
+    #[serde(default)]
+    pub max_retries: usize,
+
     // Maximum number of concurrent requests.
     #[serde(default)]
     pub concurrency: usize,
@@ -154,6 +162,7 @@ static DEFAULT_CONFIG: LazyLock<KeyCheckerConfig> = LazyLock::new(|| KeyCheckerC
     api_host: Url::parse("https://generativelanguage.googleapis.com/").unwrap(),
     timeout_sec: 15,
     concurrency: 50,
+    max_retries: 2,
     proxy: None,
     enable_multiplexing: true,
     log_level: "info".to_string(),
diff --git a/src/service/key_tester.rs b/src/service/key_tester.rs
index 984c54b..0bf10c7 100644
--- a/src/service/key_tester.rs
+++ b/src/service/key_tester.rs
@@ -4,7 +4,7 @@ use tokio::time::Duration;
 use tracing::{debug, error, info, warn};
 use url::Url;
 
-use crate::config::TEST_MESSAGE_BODY;
+use crate::config::{KeyCheckerConfig, TEST_MESSAGE_BODY};
 use crate::error::ValidatorError;
 use crate::types::GeminiKey;
 
@@ -12,27 +12,32 @@ pub async fn validate_key(
     client: Client,
     api_endpoint: impl IntoUrl,
     api_key: GeminiKey,
+    config: KeyCheckerConfig,
 ) -> Result<GeminiKey, ValidatorError> {
     let api_endpoint = api_endpoint.into_url()?;
 
-    match send_test_request(client, &api_endpoint, api_key.clone()).await {
+    match send_test_request(client, &api_endpoint, api_key.clone(), config.max_retries).await {
         Ok(_) => {
             info!("SUCCESS - {}... - Valid key found", &api_key.as_ref()[..10]);
             Ok(api_key)
         }
-        Err(e) => {
-            match e.status() {
-                Some(StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN | StatusCode::TOO_MANY_REQUESTS) => {
-                    warn!("INVALID - {}... - {}", &api_key.as_ref()[..10], ValidatorError::KeyInvalid);
-                    Err(ValidatorError::KeyInvalid)
-                }
-                _ => {
-                    let req_error = ValidatorError::ReqwestError(e);
-                    error!("ERROR-  {}... - {}", &api_key.as_ref()[..10], req_error);
-                    Err(req_error)
-                }
+        Err(e) => match e.status() {
+            Some(
+                StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN | StatusCode::TOO_MANY_REQUESTS,
+            ) => {
+                warn!(
+                    "INVALID - {}... - {}",
+                    &api_key.as_ref()[..10],
+                    ValidatorError::KeyInvalid
+                );
+                Err(ValidatorError::KeyInvalid)
             }
-        }
+            _ => {
+                let req_error = ValidatorError::ReqwestError(e);
+                error!("ERROR-  {}... - {}", &api_key.as_ref()[..10], req_error);
+                Err(req_error)
+            }
+        },
     }
 }
 
@@ -40,11 +45,12 @@ async fn send_test_request(
     client: Client,
     api_endpoint: &Url,
     key: GeminiKey,
+    max_retries: usize,
 ) -> Result<reqwest::Response, reqwest::Error> {
     let retry_policy = ExponentialBuilder::default()
-        .with_max_times(3)
-        .with_min_delay(Duration::from_secs(3))
-        .with_max_delay(Duration::from_secs(5));
+        .with_max_times(max_retries)
+        .with_min_delay(Duration::from_secs(1))
+        .with_max_delay(Duration::from_secs(2));
 
     (async || {
         let response = client
diff --git a/src/service/validation.rs b/src/service/validation.rs
index 7fb8132..5273e60 100644
--- a/src/service/validation.rs
+++ b/src/service/validation.rs
@@ -48,7 +48,14 @@ impl ValidationService {
 
         // Create stream to validate keys concurrently
         let valid_keys_stream = stream
-            .map(|key| validate_key(self.client.clone(), self.full_url.clone(), key))
+            .map(|key| {
+                validate_key(
+                    self.client.clone(),
+                    self.full_url.clone(),
+                    key,
+                    self.config.clone(),
+                )
+            })
             .buffer_unordered(self.config.concurrency)
             .filter_map(|result| async { result.ok() });
         pin_mut!(valid_keys_stream);
@@ -72,7 +79,6 @@ impl ValidationService {
     }
 }
 
-
 pub async fn start_validation() -> Result<(), ValidatorError> {
     let config = KeyCheckerConfig::load_config()?;