From 742c2e1c5c32ffb66fe9ce92f229c5e22a39bcfb Mon Sep 17 00:00:00 2001 From: wisdgod Date: Sat, 18 Jan 2025 03:50:38 +0800 Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E7=8E=AF=E5=A2=83=E5=8F=98?= =?UTF-8?q?=E9=87=8F=20`TOKEN=5FDELIMITER`=20=E6=94=AF=E6=8C=81=E8=87=AA?= =?UTF-8?q?=E5=AE=9A=E4=B9=89=E5=88=86=E9=9A=94=E7=AC=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .env.example | 3 +++ src/app/lazy.rs | 9 +++++++-- src/chat/service.rs | 21 --------------------- src/common/utils.rs | 28 +++++++++++++++++++++++----- 4 files changed, 33 insertions(+), 28 deletions(-) diff --git a/.env.example b/.env.example index 58b77da..a87a2ec 100644 --- a/.env.example +++ b/.env.example @@ -47,3 +47,6 @@ REVERSE_PROXY_HOST= # 请求体大小限制（单位为MB） # 默认为2MB (2,097,152 字节) REQUEST_BODY_LIMIT_MB=2 + +# OpenAI 请求时，token 和 checksum 的分隔符 +TOKEN_DELIMITER=, \ No newline at end of file diff --git a/src/app/lazy.rs b/src/app/lazy.rs index 1bfa317..66c1756 100644 --- a/src/app/lazy.rs +++ b/src/app/lazy.rs @@ -3,7 +3,7 @@ use crate::{ CURSOR_API2_HOST, CURSOR_HOST, DEFAULT_TOKEN_FILE_NAME, DEFAULT_TOKEN_LIST_FILE_NAME, EMPTY_STRING, }, - common::utils::parse_string_from_env, + common::utils::{parse_char_from_env, parse_string_from_env}, }; use std::sync::LazyLock; @@ -49,12 +49,17 @@ pub fn get_start_time() -> chrono::DateTime { def_pub_static!(DEFAULT_INSTRUCTIONS, env: "DEFAULT_INSTRUCTIONS", default: "Respond in Chinese by default"); -def_pub_static!(REVERSE_PROXY_HOST, env: "REVERSE_PROXY_HOST", default: ""); +def_pub_static!(REVERSE_PROXY_HOST, env: "REVERSE_PROXY_HOST", default: EMPTY_STRING); def_pub_static!(SHARED_AUTH_TOKEN, env: "SHARED_AUTH_TOKEN", default: EMPTY_STRING); pub static USE_SHARE: LazyLock = LazyLock::new(|| !SHARED_AUTH_TOKEN.is_empty()); +pub static TOKEN_DELIMITER: LazyLock = + LazyLock::new(|| parse_char_from_env("TOKEN_DELIMITER", ',')); + +pub static TOKEN_DELIMITER_LEN: LazyLock = LazyLock::new(|| TOKEN_DELIMITER.len_utf8()); + pub static USE_PROXY: LazyLock = LazyLock::new(|| !REVERSE_PROXY_HOST.is_empty()); pub static CURSOR_API2_CHAT_URL: LazyLock = LazyLock::new(|| { diff --git a/src/chat/service.rs b/src/chat/service.rs index 2fe809f..9fc2959 100644 --- a/src/chat/service.rs +++ b/src/chat/service.rs @@ -243,27 +243,6 @@ pub async fn handle_chat( }; // 构建请求客户端 - // let client_key = match generate_client_key(&checksum) { - // Some(key) => key, - // None => { - // let mut state = state.lock().await; - // if let Some(log) = state - // .request_logs - // .iter_mut() - // .rev() - // .find(|log| log.id == current_id) - // { - // log.status = STATUS_FAILED; - // log.error = Some(ERR_CHECKSUM_NO_GOOD.to_string()); - // } - // state.active_requests -= 1; - // state.error_requests += 1; - // return Err(( - // StatusCode::BAD_REQUEST, - // Json(ChatError::RequestFailed(ERR_CHECKSUM_NO_GOOD.to_string()).to_json()), - // )); - // } - // }; let client = build_client(&auth_token, &checksum); let response = client.body(hex_data).send().await; diff --git a/src/common/utils.rs b/src/common/utils.rs index ef24055..d14b4cc 100644 --- a/src/common/utils.rs +++ b/src/common/utils.rs @@ -4,7 +4,10 @@ mod tokens; pub use tokens::*; use super::models::userinfo::{StripeProfile, TokenProfile, UsageProfile, UserProfile}; -use crate::app::constant::{FALSE, TRUE}; +use crate::app::{ + constant::{FALSE, TRUE}, + lazy::{TOKEN_DELIMITER, TOKEN_DELIMITER_LEN}, +}; pub fn parse_bool_from_env(key: &str, default: bool) -> bool { std::env::var(key) @@ -21,6 +24,20 @@ pub fn parse_string_from_env(key: &str, default: &str) -> String { std::env::var(key).unwrap_or_else(|_| default.to_string()) } +pub fn parse_char_from_env(key: &str, default: char) -> char { + std::env::var(key) + .ok() + .and_then(|v| { + let chars: Vec = v.chars().collect(); + if chars.len() == 1 { + Some(chars[0]) + } else { + None + } + }) + .unwrap_or(default) +} + pub fn parse_usize_from_env(key: &str, default: usize) -> usize { std::env::var(key) .ok() @@ -86,9 +103,9 @@ pub async fn get_user_profile(auth_token: &str) -> Option { pub fn validate_token_and_checksum(auth_token: &str) -> Option<(String, String)> { // 找最后一个逗号 - let comma_pos = auth_token.rfind(',')?; + let comma_pos = auth_token.rfind(*TOKEN_DELIMITER)?; let (token_part, checksum) = auth_token.split_at(comma_pos); - let checksum = &checksum[1..]; // 跳过逗号 + let checksum = &checksum[*TOKEN_DELIMITER_LEN..]; // 跳过逗号 // 解析 token - 为了向前兼容,忽略最后一个:或%3A前的内容 let colon_pos = token_part.rfind(':'); @@ -115,11 +132,12 @@ pub fn validate_token_and_checksum(auth_token: &str) -> Option<(String, String)> pub fn extract_token(auth_token: &str) -> Option { // 解析 token - let token_part = match auth_token.rfind(',') { + let token_part = match auth_token.rfind(*TOKEN_DELIMITER) { Some(pos) => &auth_token[..pos], - None => auth_token + None => auth_token, }; + // 向前兼容 let colon_pos = token_part.rfind(':'); let encoded_colon_pos = token_part.rfind("%3A");