From 772280da20b2230dd1f976455cdcfd1f487e3585 Mon Sep 17 00:00:00 2001 From: Yoo1tic <137816438+Yoo1tic@users.noreply.github.com> Date: Sat, 2 Aug 2025 21:42:08 +0800 Subject: [PATCH] refactor: reorganize project structure by moving service modules to validation and utils --- src/config/config.rs | 21 ---- src/config/mod.rs | 2 +- src/error.rs | 16 ++- src/lib.rs | 3 +- src/main.rs | 2 +- src/service/key_tester.rs | 74 -------------- src/service/mod.rs | 7 -- src/{service => utils}/http_client.rs | 0 src/utils/mod.rs | 5 + src/utils/request.rs | 41 ++++++++ src/validation/key_validator.rs | 97 +++++++++++++++++++ src/validation/mod.rs | 80 +++++++++++++++ .../validation_service.rs} | 5 +- 13 files changed, 244 insertions(+), 109 deletions(-) delete mode 100644 src/service/key_tester.rs delete mode 100644 src/service/mod.rs rename src/{service => utils}/http_client.rs (100%) create mode 100644 src/utils/mod.rs create mode 100644 src/utils/request.rs create mode 100644 src/validation/key_validator.rs create mode 100644 src/validation/mod.rs rename src/{service/validation.rs => validation/validation_service.rs} (95%) diff --git a/src/config/config.rs b/src/config/config.rs index d4f4c7c..7672238 100644 --- a/src/config/config.rs +++ b/src/config/config.rs @@ -5,7 +5,6 @@ use figment::{ providers::{Env, Format, Serialized, Toml}, }; use serde::{Deserialize, Serialize}; -use serde_json::Value; use std::fmt::{self, Display}; use std::fs; use std::path::PathBuf; @@ -168,26 +167,6 @@ static DEFAULT_CONFIG: LazyLock = LazyLock::new(|| KeyCheckerC log_level: "info".to_string(), }); -// LazyLock for the test message body used in API key validation -pub static TEST_MESSAGE_BODY: LazyLock = LazyLock::new(|| { - serde_json::json!({ - "contents": [ - { - "parts": [ - { - "text": "Hi" - } - ] - } - ], - "generationConfig": { - "thinkingConfig": { - "thinkingBudget": 0 - } - } - }) -}); - fn default_api_host() -> Url { DEFAULT_CONFIG.api_host.clone() } diff --git a/src/config/mod.rs b/src/config/mod.rs index ee0ad2d..0dfd8ad 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -1,3 +1,3 @@ mod config; -pub use config::{KeyCheckerConfig, TEST_MESSAGE_BODY}; +pub use config::KeyCheckerConfig; diff --git a/src/error.rs b/src/error.rs index 927cb88..522520b 100644 --- a/src/error.rs +++ b/src/error.rs @@ -3,13 +3,13 @@ use thiserror::Error; #[derive(Error, Debug)] pub enum ValidatorError { #[error("HTTP error: {0}")] - ReqwestError(#[from] reqwest::Error), + ReqwestError(Box), #[error("IO error: {0}")] Io(#[from] std::io::Error), #[error("Configuration error: {0}")] - ConfigError(#[from] figment::Error), + ConfigError(Box), #[error("TOML serialization error: {0}")] TomlSer(#[from] toml::ser::Error), @@ -30,4 +30,16 @@ pub enum ValidatorError { KeyFormatInvalid(String), } +impl From for ValidatorError { + fn from(err: reqwest::Error) -> Self { + ValidatorError::ReqwestError(Box::new(err)) + } +} + +impl From for ValidatorError { + fn from(err: figment::Error) -> Self { + ValidatorError::ConfigError(Box::new(err)) + } +} + pub type Result = std::result::Result; diff --git a/src/lib.rs b/src/lib.rs index 9fc54f4..9f6e8ea 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,7 +2,8 @@ pub mod adapters; pub mod config; pub mod error; pub mod types; -pub mod service; +pub mod validation; +pub mod utils; // ASCII art for Gemini pub const BANNER: &str = r#" diff --git a/src/main.rs b/src/main.rs index 372bc6f..5fbcf75 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,5 @@ use gemini_keychecker::error::ValidatorError; -use gemini_keychecker::{BANNER, config::KeyCheckerConfig, service::start_validation}; +use gemini_keychecker::{BANNER, config::KeyCheckerConfig, validation::start_validation}; use mimalloc::MiMalloc; use tracing::info; diff --git a/src/service/key_tester.rs b/src/service/key_tester.rs deleted file mode 100644 index 0bf10c7..0000000 --- a/src/service/key_tester.rs +++ /dev/null @@ -1,74 +0,0 @@ -use backon::{ExponentialBuilder, Retryable}; -use reqwest::{Client, IntoUrl, StatusCode}; -use tokio::time::Duration; -use tracing::{debug, error, info, warn}; -use url::Url; - -use crate::config::{KeyCheckerConfig, TEST_MESSAGE_BODY}; -use crate::error::ValidatorError; -use crate::types::GeminiKey; - -pub async fn validate_key( - client: Client, - api_endpoint: impl IntoUrl, - api_key: GeminiKey, - config: KeyCheckerConfig, -) -> Result { - let api_endpoint = api_endpoint.into_url()?; - - 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) - } - }, - } -} - -async fn send_test_request( - client: Client, - api_endpoint: &Url, - key: GeminiKey, - max_retries: usize, -) -> Result { - let retry_policy = ExponentialBuilder::default() - .with_max_times(max_retries) - .with_min_delay(Duration::from_secs(1)) - .with_max_delay(Duration::from_secs(2)); - - (async || { - let response = client - .post(api_endpoint.clone()) - .header("Content-Type", "application/json") - .header("X-goog-api-key", key.as_ref()) - .json(&*TEST_MESSAGE_BODY) - .send() - .await?; - debug!("Response for key {}: {:?}", key.as_ref(), response.status()); - response.error_for_status() - }) - .retry(&retry_policy) - .when(|e: &reqwest::Error| { - !matches!( - e.status(), - Some(StatusCode::FORBIDDEN | StatusCode::UNAUTHORIZED) - ) - }) - .await -} diff --git a/src/service/mod.rs b/src/service/mod.rs deleted file mode 100644 index 424ef96..0000000 --- a/src/service/mod.rs +++ /dev/null @@ -1,7 +0,0 @@ -pub mod http_client; -pub mod key_tester; -pub mod validation; - -pub use http_client::client_builder; -pub use key_tester::validate_key; -pub use validation::{ValidationService, start_validation}; diff --git a/src/service/http_client.rs b/src/utils/http_client.rs similarity index 100% rename from src/service/http_client.rs rename to src/utils/http_client.rs diff --git a/src/utils/mod.rs b/src/utils/mod.rs new file mode 100644 index 0000000..53cd21a --- /dev/null +++ b/src/utils/mod.rs @@ -0,0 +1,5 @@ +pub mod http_client; +pub mod request; + +pub use http_client::client_builder; +pub use request::send_request; \ No newline at end of file diff --git a/src/utils/request.rs b/src/utils/request.rs new file mode 100644 index 0000000..8bd3bf4 --- /dev/null +++ b/src/utils/request.rs @@ -0,0 +1,41 @@ +use backon::{ExponentialBuilder, Retryable}; +use reqwest::{Client, StatusCode}; +use serde_json::Value; +use tokio::time::Duration; +use tracing::debug; +use url::Url; + +use crate::types::GeminiKey; + +pub async fn send_request( + client: Client, + api_endpoint: &Url, + key: GeminiKey, + payload: &Value, + max_retries: usize, +) -> Result { + let retry_policy = ExponentialBuilder::default() + .with_max_times(max_retries) + .with_min_delay(Duration::from_secs(1)) + .with_max_delay(Duration::from_secs(2)); + + (async || { + let response = client + .post(api_endpoint.clone()) + .header("Content-Type", "application/json") + .header("X-goog-api-key", key.as_ref()) + .json(payload) + .send() + .await?; + debug!("Response for key {}: {:?}", key.as_ref(), response.status()); + response.error_for_status() + }) + .retry(&retry_policy) + .when(|e: &reqwest::Error| { + !matches!( + e.status(), + Some(StatusCode::FORBIDDEN | StatusCode::UNAUTHORIZED) + ) + }) + .await +} diff --git a/src/validation/key_validator.rs b/src/validation/key_validator.rs new file mode 100644 index 0000000..6426e8f --- /dev/null +++ b/src/validation/key_validator.rs @@ -0,0 +1,97 @@ +use reqwest::{Client, IntoUrl, StatusCode}; +use tracing::{error, info, warn}; + +use super::{CACHE_CONTENT_TEST_BODY, GENERATE_CONTENT_TEST_BODY}; +use crate::config::KeyCheckerConfig; +use crate::error::ValidatorError; +use crate::types::GeminiKey; +use crate::utils::send_request; + +pub async fn test_generate_content_api( + client: Client, + api_endpoint: impl IntoUrl, + api_key: GeminiKey, + config: KeyCheckerConfig, +) -> Result { + let api_endpoint = api_endpoint.into_url()?; + + match send_request( + client, + &api_endpoint, + api_key.clone(), + &GENERATE_CONTENT_TEST_BODY, + 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::from(e); + error!("ERROR- {}... - {}", &api_key.as_ref()[..10], req_error); + Err(req_error) + } + }, + } +} + +pub async fn test_cache_content_api( + client: Client, + api_endpoint: impl IntoUrl, + api_key: GeminiKey, + config: KeyCheckerConfig, +) -> Result { + let api_endpoint = api_endpoint.into_url()?; + + match send_request( + client, + &api_endpoint, + api_key.clone(), + &CACHE_CONTENT_TEST_BODY, + config.max_retries, + ) + .await + { + Ok(_) => { + info!( + "CACHE SUCCESS - {}... - Valid key for cache API", + &api_key.as_ref()[..10] + ); + Ok(api_key) + } + Err(e) => match e.status() { + Some( + StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN | StatusCode::TOO_MANY_REQUESTS, + ) => { + warn!( + "CACHE INVALID - {}... - {}", + &api_key.as_ref()[..10], + ValidatorError::KeyInvalid + ); + Err(ValidatorError::KeyInvalid) + } + _ => { + let req_error = ValidatorError::from(e); + error!( + "CACHE ERROR - {}... - {}", + &api_key.as_ref()[..10], + req_error + ); + Err(req_error) + } + }, + } +} diff --git a/src/validation/mod.rs b/src/validation/mod.rs new file mode 100644 index 0000000..6f8ea1d --- /dev/null +++ b/src/validation/mod.rs @@ -0,0 +1,80 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::sync::LazyLock; +pub mod key_validator; +pub mod validation_service; + +pub use key_validator::{test_cache_content_api, test_generate_content_api}; +pub use validation_service::{ValidationService, start_validation}; + +#[derive(Serialize, Deserialize, Debug)] +pub struct TextPart { + pub text: String, +} + +#[derive(Serialize, Deserialize, Debug)] +pub struct ContentPart { + pub parts: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub role: Option, +} + + +#[derive(Serialize, Deserialize, Debug)] +pub struct ThinkingConfig { + #[serde(rename = "thinkingBudget")] + pub thinking_budget: u32, +} + +#[derive(Serialize, Deserialize, Debug)] +pub struct GenerationConfig { + #[serde(skip_serializing_if = "Option::is_none")] + #[serde(rename = "thinkingConfig")] + pub thinking_config: Option, +} + +#[derive(Serialize, Deserialize, Debug)] +pub struct GeminiRequest { + #[serde(skip_serializing_if = "Option::is_none")] + pub model: Option, + pub contents: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + #[serde(rename = "generationConfig")] + pub generation_config: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub ttl: Option, +} + +// LazyLock for the test message body used in API key validation +pub static GENERATE_CONTENT_TEST_BODY: LazyLock = LazyLock::new(|| { + let generate_request = GeminiRequest { + model: None, + contents: vec![ContentPart { + parts: vec![TextPart { + text: "Hi".to_string(), + }], + role: None, + }], + generation_config: Some(GenerationConfig { + thinking_config: Some(ThinkingConfig { thinking_budget: 0 }), + }), + ttl: None, + }; + serde_json::to_value(generate_request).unwrap() +}); + +// LazyLock for the cached content test body used in cache API validation +pub static CACHE_CONTENT_TEST_BODY: LazyLock = LazyLock::new(|| { + // Generate random text content to meet the minimum 1024 tokens requirement for cache API + let long_text = "You are an expert at analyzing transcripts.".repeat(50); + let cache_request = GeminiRequest { + model: Some("models/gemini-2.5-flash-lite".to_string()), + contents: vec![ContentPart { + parts: vec![TextPart { text: long_text }], + role: Some("user".to_string()), + }], + generation_config: None, + ttl: Some("300s".to_string()), + }; + serde_json::to_value(cache_request).unwrap() +}); diff --git a/src/service/validation.rs b/src/validation/validation_service.rs similarity index 95% rename from src/service/validation.rs rename to src/validation/validation_service.rs index 5273e60..ec27f42 100644 --- a/src/service/validation.rs +++ b/src/validation/validation_service.rs @@ -5,10 +5,11 @@ use reqwest::Client; use std::time::Instant; use tokio::{fs, io::AsyncWriteExt, sync::mpsc}; -use super::{http_client::client_builder, key_tester::validate_key}; +use super::key_validator::test_generate_content_api; use crate::adapters::{load_keys, write_keys_txt_file}; use crate::config::KeyCheckerConfig; use crate::types::GeminiKey; +use crate::utils::client_builder; pub struct ValidationService { config: KeyCheckerConfig, @@ -49,7 +50,7 @@ impl ValidationService { // Create stream to validate keys concurrently let valid_keys_stream = stream .map(|key| { - validate_key( + test_generate_content_api( self.client.clone(), self.full_url.clone(), key,