feat: add enable_multiplexing option to configuration and update client builder logic

main
Yoo1tic 2025-07-20 18:18:41 +08:00
parent 824c4f96dd
commit ffd93a6b72
5 changed files with 64 additions and 5 deletions

View File

@ -5,3 +5,4 @@ api_host = "https://generativelanguage.googleapis.com/"
timeout_sec = 20
concurrency = 30
proxy = "http://username:password@host:port"
enable_multiplexing = true

View File

@ -5,17 +5,20 @@ use reqwest::Client;
use crate::config::KeyCheckerConfig;
pub fn client_builder(config: &KeyCheckerConfig) -> Result<Client, reqwest::Error> {
// Adjust connection pool size based on concurrency, optimizing for 0.5 second response time
// Set the maximum number of connections per host based on concurrency.
let pool_size = config.concurrency / 2;
let mut builder = Client::builder()
.timeout(Duration::from_secs(config.timeout_sec))
.pool_max_idle_per_host(pool_size);
// .http2_prior_knowledge();
if let Some(ref proxy_url) = config.proxy {
builder = builder.proxy(reqwest::Proxy::all(proxy_url.clone())?);
}
if !config.enable_multiplexing {
builder = builder.http1_only();
}
builder.build()
}

View File

@ -6,6 +6,7 @@ use figment::{
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::fmt::{self, Display};
use std::fs;
use std::path::PathBuf;
use std::sync::LazyLock;
@ -72,6 +73,10 @@ pub struct KeyCheckerConfig {
// Optional proxy URL for HTTP requests (e.g., --proxy http://user:pass@host:port).
#[serde(default)]
pub proxy: Option<Url>,
// Whether to enable HTTP/2 multiplexing for requests.
#[serde(default)]
pub enable_multiplexing: bool,
}
impl Default for KeyCheckerConfig {
@ -110,6 +115,31 @@ impl KeyCheckerConfig {
}
}
impl Display for KeyCheckerConfig {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
writeln!(f, "API Host: {}", self.api_host)?;
let proxy_status = match &self.proxy {
Some(proxy) => proxy.to_string(),
None => "Disabled".to_string(),
};
writeln!(f, "Proxy: {}", proxy_status)?;
let protocol_status = if self.enable_multiplexing {
"HTTP/2 (Multiplexing Enabled)"
} else {
"HTTP/1.1 (Multiplexing Disabled)"
};
writeln!(f, "Protocol: {}", protocol_status)?;
writeln!(f, "Timeout: {}s", self.timeout_sec)?;
writeln!(f, "Concurrency: {}", self.concurrency)?;
writeln!(f, "Input: {}", self.input_path.display())?;
writeln!(f, "Output: {}", self.output_path.display())?;
write!(f, "Backup: {}", self.backup_path.display())
}
}
// Single LazyLock for entire default configuration
static DEFAULT_CONFIG: LazyLock<KeyCheckerConfig> = LazyLock::new(|| KeyCheckerConfig {
input_path: "keys.txt".into(),
@ -119,6 +149,7 @@ static DEFAULT_CONFIG: LazyLock<KeyCheckerConfig> = LazyLock::new(|| KeyCheckerC
timeout_sec: 15,
concurrency: 50,
proxy: None,
enable_multiplexing: true,
});
// LazyLock for the test message body used in API key validation
@ -136,6 +167,25 @@ pub static TEST_MESSAGE_BODY: LazyLock<Value> = LazyLock::new(|| {
})
});
// ASCII art for Gemini
pub const GEMINI_ASCII: &str = r#"
______ _ _
/ ____/___ ____ ___ (_)____ (_)
/ / __ / _ \ / __ `__ \ / // __ \ / /
/ /_/ // __// / / / / // // / / // /
\____/ \___//_/ /_/ /_//_//_/ /_//_/
"#;
// LazyLock for the application banner
pub static BANNER: LazyLock<String> = LazyLock::new(|| {
format!(
"{}\n{}\nKey Checker - Configuration Status\n{}",
GEMINI_ASCII,
"=".repeat(42),
"=".repeat(42)
)
});
fn default_api_host() -> Url {
DEFAULT_CONFIG.api_host.clone()
}

View File

@ -1,5 +1,5 @@
mod basic_config;
mod basic_client;
mod basic_config;
pub use basic_config::{KeyCheckerConfig, TEST_MESSAGE_BODY};
pub use basic_client::client_builder;
pub use basic_config::{BANNER, KeyCheckerConfig, TEST_MESSAGE_BODY};

View File

@ -1,12 +1,17 @@
use anyhow::Result;
use gemini_keychecker::adapters::load_keys;
use gemini_keychecker::config::{KeyCheckerConfig, client_builder};
use gemini_keychecker::config::{BANNER, KeyCheckerConfig, client_builder};
use gemini_keychecker::validation::ValidationService;
/// Main function - orchestrates the key validation process
#[tokio::main]
async fn main() -> Result<()> {
let config = KeyCheckerConfig::load_config().unwrap();
// Display banner and configuration status at startup
println!("{}", *BANNER);
println!("{}", config);
let keys = load_keys(config.input_path.as_path())?;
let client = client_builder(&config)?;