From db94d6d86c51a42415471844dfb61e84268b3204 Mon Sep 17 00:00:00 2001
From: Yoo1tic <137816438+Yoo1tic@users.noreply.github.com>
Date: Sat, 19 Jul 2025 23:09:15 +0800
Subject: [PATCH] feat: enhance validation service with full API URL and
 refactor key validation logic

---
 src/config/basic_config.rs | 26 ++++++++++++++++++++++++--
 src/config/mod.rs          |  2 +-
 src/key_validator.rs       | 26 ++++++--------------------
 src/validation.rs          | 10 ++++++++--
 4 files changed, 39 insertions(+), 25 deletions(-)

diff --git a/src/config/basic_config.rs b/src/config/basic_config.rs
index 74c56b8..12d20f7 100644
--- a/src/config/basic_config.rs
+++ b/src/config/basic_config.rs
@@ -5,6 +5,7 @@ use figment::{
     providers::{Env, Format, Serialized, Toml},
 };
 use serde::{Deserialize, Serialize};
+use serde_json::Value;
 use std::fs;
 use std::path::PathBuf;
 use std::sync::LazyLock;
@@ -98,10 +99,15 @@ impl KeyCheckerConfig {
             .merge(Serialized::defaults(Cli::parse()))
             .extract()?;
 
-        dbg!(&config);
-
         Ok(config)
     }
+
+    /// Returns the complete Gemini API URL for generateContent endpoint
+    pub fn gemini_api_url(&self) -> Url {
+        self.api_host
+            .join("v1beta/models/gemini-2.0-flash-exp:generateContent")
+            .expect("Failed to join API URL")
+    }
 }
 
 // Single LazyLock for entire default configuration
@@ -114,6 +120,22 @@ static DEFAULT_CONFIG: LazyLock<KeyCheckerConfig> = LazyLock::new(|| KeyCheckerC
     concurrency: 50,
     proxy: None,
 });
+
+// LazyLock for the test message body used in API key validation
+pub static TEST_MESSAGE_BODY: LazyLock<Value> = LazyLock::new(|| {
+    serde_json::json!({
+        "contents": [
+            {
+                "parts": [
+                    {
+                        "text": "Hi"
+                    }
+                ]
+            }
+        ]
+    })
+});
+
 fn default_api_host() -> Url {
     DEFAULT_CONFIG.api_host.clone()
 }
diff --git a/src/config/mod.rs b/src/config/mod.rs
index 4580cbf..373507c 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -1,5 +1,5 @@
 mod basic_config;
 mod basic_client;
 
-pub use basic_config::KeyCheckerConfig;
+pub use basic_config::{KeyCheckerConfig, TEST_MESSAGE_BODY};
 pub use basic_client::client_builder;
diff --git a/src/key_validator.rs b/src/key_validator.rs
index 8304531..a0f7a42 100644
--- a/src/key_validator.rs
+++ b/src/key_validator.rs
@@ -1,15 +1,15 @@
 use anyhow::Result;
 use backon::{ExponentialBuilder, Retryable};
 use reqwest::{Client, StatusCode};
-use serde_json;
 use tokio::time::Duration;
 use url::Url;
 
+use crate::config::TEST_MESSAGE_BODY;
 use crate::types::{GeminiKey, KeyStatus};
 
 pub async fn validate_key_with_retry(
     client: Client,
-    api_host: Url,
+    full_url: Url,
     key: GeminiKey,
 ) -> Option<GeminiKey> {
     let retry_policy = ExponentialBuilder::default()
@@ -17,7 +17,7 @@ pub async fn validate_key_with_retry(
         .with_min_delay(Duration::from_secs(3))
         .with_max_delay(Duration::from_secs(5));
 
-    let result = (async || match keytest(client.to_owned(), &api_host, &key).await {
+    let result = (async || match keytest(client.to_owned(), &full_url, &key).await {
         Ok(KeyStatus::Valid) => {
             println!("Key: {}... -> SUCCESS", &key.as_ref()[..10]);
             Ok(Some(key.clone()))
@@ -58,26 +58,12 @@ pub async fn validate_key_with_retry(
     }
 }
 
-async fn keytest(client: Client, api_host: &Url, key: &GeminiKey) -> Result<KeyStatus> {
-    const API_PATH: &str = "v1beta/models/gemini-2.0-flash-exp:generateContent";
-    let full_url = api_host.join(API_PATH)?;
-    let request_body = serde_json::json!({
-        "contents": [
-            {
-                "parts": [
-                    {
-                        "text": "Hi"
-                    }
-                ]
-            }
-        ]
-    });
-
+async fn keytest(client: Client, full_url: &Url, key: &GeminiKey) -> Result<KeyStatus> {
     let response = client
-        .post(full_url)
+        .post(full_url.clone())
         .header("Content-Type", "application/json")
         .header("X-goog-api-key", key.as_ref())
-        .json(&request_body)
+        .json(&*TEST_MESSAGE_BODY)
         .send()
         .await?;
 
diff --git a/src/validation.rs b/src/validation.rs
index f3e7d7f..9ad79c3 100644
--- a/src/validation.rs
+++ b/src/validation.rs
@@ -13,11 +13,17 @@ use crate::types::GeminiKey;
 pub struct ValidationService {
     config: KeyCheckerConfig,
     client: Client,
+    full_url: url::Url,
 }
 
 impl ValidationService {
     pub fn new(config: KeyCheckerConfig, client: Client) -> Self {
-        Self { config, client }
+        let full_url = config.gemini_api_url();
+        Self {
+            config,
+            client,
+            full_url,
+        }
     }
 
     pub async fn validate_keys(&self, keys: Vec<GeminiKey>) -> Result<()> {
@@ -42,7 +48,7 @@ impl ValidationService {
 
         // Create stream to validate keys concurrently
         let valid_keys_stream = stream
-            .map(|key| validate_key_with_retry(self.client.to_owned(), self.config.api_host.clone(), key))
+            .map(|key| validate_key_with_retry(self.client.to_owned(), self.full_url.clone(), key))
             .buffer_unordered(self.config.concurrency)
             .filter_map(|r| async { r });
         pin_mut!(valid_keys_stream);