cror / src /common /utils.rs
xcolorday's picture
Upload 442 files
6b356c4 verified
Raw
History Blame Contribute Delete
34.4 kB
#![allow(clippy::too_many_arguments)]
// mod checksum;
// mod token;
pub mod base62;
mod base64;
pub mod duration_fmt;
pub mod hex;
pub mod string_builder;
use ::base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use ::core::str::FromStr as _;
use ::prost::Message as _;
use ::reqwest::Client;
use ::std::borrow::Cow;
pub use base64::{from_base64, to_base64};
pub use hex::{byte_to_hex, hex_to_byte};
pub use string_builder::StringBuilder;
use super::model::userinfo::{
GetTeamsResponse, ListActiveSessionsResponse, Session, StripeProfile, Team, UserProfile,
};
use crate::{
app::{
lazy::{
aggregated_usage_events_url, chat_models_url, filtered_usage_events_url,
is_on_new_pricing_url, server_config_url, teams_url,
},
model::{ChainUsage, Checksum, DateTime, ExtToken, GcppHost, Hash, RawToken, Token},
},
common::model::userinfo::{MembershipType, SubscriptionStatus},
core::{
aiserver::v1::{
AvailableModelsRequest, AvailableModelsResponse, GetAggregatedUsageEventsRequest,
GetAggregatedUsageEventsResponse, GetFilteredUsageEventsRequest,
GetFilteredUsageEventsResponse, GetServerConfigResponse,
},
config::key_config,
},
};
mod sealed {
pub trait Sealed: Sized {}
impl Sealed for bool {}
impl Sealed for &'static str {}
impl Sealed for usize {}
}
pub trait ParseFromEnv: sealed::Sealed {
type Result = Self;
fn parse_from_env(key: &str, default: Self) -> Self::Result;
}
impl ParseFromEnv for bool {
#[inline]
fn parse_from_env(key: &str, default: Self) -> Self::Result {
::std::env::var(key)
.ok()
.map(|mut val| {
let res = {
val.make_ascii_lowercase();
val.trim()
};
match res {
"true" | "1" => true,
"false" | "0" => false,
_ => default,
}
})
.unwrap_or(default)
}
}
impl ParseFromEnv for &'static str {
type Result = Cow<'static, str>;
#[inline]
fn parse_from_env(key: &str, default: Self) -> Self::Result {
match ::std::env::var(key) {
Ok(mut value) => {
let trimmed = value.trim();
if trimmed.is_empty() {
// 如果 trim 后为空,使用默认值(不分配)
Cow::Borrowed(default)
} else if trimmed.len() == value.len() {
// 不需要 trim,直接使用
Cow::Owned(value)
} else {
// 需要 trim - 就地修改
let trimmed_len = trimmed.len();
let start_offset = trimmed.as_ptr() as usize - value.as_ptr() as usize;
unsafe {
let vec = value.as_mut_vec();
if start_offset > 0 {
vec.copy_within(start_offset..start_offset + trimmed_len, 0);
}
vec.set_len(trimmed_len);
}
Cow::Owned(value)
}
}
Err(_) => Cow::Borrowed(default),
}
}
}
impl ParseFromEnv for usize {
#[inline]
fn parse_from_env(key: &str, default: Self) -> Self::Result {
::std::env::var(key)
.ok()
.and_then(|v| v.trim().parse().ok())
.unwrap_or(default)
}
}
#[inline]
pub fn parse_from_env<T: ParseFromEnv>(key: &str, default: T) -> T::Result {
ParseFromEnv::parse_from_env(key, default)
}
pub fn now_secs() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("system time before Unix epoch")
.as_secs()
}
const LEN: usize = 2;
pub trait TrimNewlines: Sized {
fn trim_leading_newlines(self) -> Self;
}
impl TrimNewlines for &str {
#[inline(always)]
fn trim_leading_newlines(self) -> Self {
let bytes = self.as_bytes();
if bytes.len() >= LEN && bytes[0] == b'\n' && bytes[1] == b'\n' {
return unsafe { self.get_unchecked(LEN..) };
}
self
}
}
impl TrimNewlines for String {
#[inline(always)]
fn trim_leading_newlines(mut self) -> Self {
let bytes = self.as_bytes();
if bytes.len() >= LEN && bytes[0] == b'\n' && bytes[1] == b'\n' {
unsafe {
let vec = self.as_mut_vec();
vec.drain(..LEN);
}
}
self
}
}
/// 获取完整的token配置文件
/// 协调多个数据源,可选择性获取用户信息
#[inline(never)]
pub async fn get_token_profile(
client: Client,
token: &Token,
maybe_token: Option<&Token>,
is_pri: bool,
include_user: bool,
include_sessions: bool,
) -> (
Option<UserProfile>,
Option<StripeProfile>,
Option<Vec<Session>>,
) {
let maybe_token = maybe_token.unwrap_or(token);
let mut buf = [0; 31];
let user_id = maybe_token.raw().subject.id.to_str(&mut buf) as &str;
if include_user {
if include_sessions {
// 并发获取所有数据,user为必需
let (mut stripe, mut user, teams, is_on_new_pricing, sessions) = tokio::join!(
get_stripe_profile(&client, token.as_str(), is_pri),
get_user_profile(&client, user_id, maybe_token.as_str(), is_pri),
get_teams(&client, user_id, maybe_token.as_str(), is_pri),
get_is_on_new_pricing(&client, user_id, maybe_token.as_str(), is_pri),
get_sessions(&client, user_id, maybe_token.as_str(), is_pri)
);
if let Some(stripe) = stripe.as_mut()
&& teams.is_some_and(|teams| {
teams.into_iter().any(|team| {
team.has_billing
&& team.subscription_status.is_some_and(|subscription_status| {
matches!(subscription_status, SubscriptionStatus::Active)
})
})
})
{
stripe.membership_type = MembershipType::Enterprise;
}
if let Some(user) = user.as_mut() {
user.is_on_new_pricing = is_on_new_pricing.unwrap_or(true);
}
// 所有数据都必需成功
(user, stripe, sessions)
} else {
// 并发获取所有数据,user为必需
let (mut stripe, mut user, teams, is_on_new_pricing) = tokio::join!(
get_stripe_profile(&client, token.as_str(), is_pri),
get_user_profile(&client, user_id, maybe_token.as_str(), is_pri),
get_teams(&client, user_id, maybe_token.as_str(), is_pri),
get_is_on_new_pricing(&client, user_id, maybe_token.as_str(), is_pri)
);
if let Some(stripe) = stripe.as_mut()
&& teams.is_some_and(|teams| {
teams.into_iter().any(|team| {
team.has_billing
&& team.subscription_status.is_some_and(|subscription_status| {
matches!(subscription_status, SubscriptionStatus::Active)
})
})
})
{
stripe.membership_type = MembershipType::Enterprise;
}
if let Some(user) = user.as_mut() {
user.is_on_new_pricing = is_on_new_pricing.unwrap_or(true);
}
// 所有数据都必需成功
(user, stripe, None)
}
} else {
// 仅获取stripe数据
let (mut stripe, teams) = tokio::join!(
get_stripe_profile(&client, token.as_str(), is_pri),
get_teams(&client, user_id, maybe_token.as_str(), is_pri),
);
if let Some(stripe) = stripe.as_mut()
&& teams.is_some_and(|teams| {
teams.into_iter().any(|team| {
team.has_billing
&& team.subscription_status.is_some_and(|subscription_status| {
matches!(subscription_status, SubscriptionStatus::Active)
})
})
})
{
stripe.membership_type = MembershipType::Enterprise;
}
(None, stripe, None)
}
}
// /// 获取用户使用情况配置文件
// pub async fn get_usage_profile(client: &Client, user_id: &str, auth_token: &str, is_pri: bool) {
// if !*crate::app::lazy::log::DEBUG {
// return;
// }
// let request = super::client::build_usage_request(client, user_id, auth_token, is_pri);
// let response = match request.send().await {
// Ok(r) => r,
// Err(_) => {
// crate::debug!("<get_usage_profile> send error");
// return;
// }
// };
// crate::debug!("<get_usage_profile> got {}", response.status());
// let usage = response.json::<UsageProfile>().await.ok();
// crate::debug!(
// "<get_usage_profile> got {}",
// __unwrap!(serde_json::to_string_pretty(&usage))
// );
// }
/// 获取Stripe付费配置文件
pub async fn get_stripe_profile(
client: &Client,
auth_token: &str,
is_pri: bool,
) -> Option<StripeProfile> {
let request = super::client::build_profile_request(client, auth_token, is_pri);
let response = request.send().await.ok()?;
crate::debug!("<get_stripe_profile> {}", response.status());
response.json::<StripeProfile>().await.ok()
}
/// 获取用户基础配置文件
pub async fn get_user_profile(
client: &Client,
user_id: &str,
auth_token: &str,
is_pri: bool,
) -> Option<UserProfile> {
let request = super::client::build_userinfo_request(client, user_id, auth_token, is_pri);
// let response = request.send().await.ok()?;
// crate::debug!("get_user_profile \n{response:?}");
// let bytes = response.bytes().await.ok()?;
// crate::debug!("bytes \n{:?}", unsafe { std::str::from_utf8_unchecked(&bytes[..]) });
// serde_json::from_slice::<UserProfile>(&bytes).ok()
let response = request.send().await.ok()?;
crate::debug!("<get_user_profile> {}", response.status());
response.json::<UserProfile>().await.ok()
}
pub async fn get_available_models(
ext_token: ExtToken,
is_pri: bool,
mut request: AvailableModelsRequest,
) -> Option<AvailableModelsResponse> {
request.exclude_max_named_models = true;
let response = {
let client = super::client::build_client_request(super::client::AiServiceRequest {
ext_token,
fs_client_key: None,
url: chat_models_url(is_pri),
is_stream: false,
trace_id: Some(new_uuid_v4()),
is_pri,
cookie: None,
});
client
.body(__unwrap!(encode_message(&request, false)))
.send()
.await
.ok()?
.bytes()
.await
.ok()?
};
let available_models = AvailableModelsResponse::decode(response.as_ref()).ok()?;
Some(available_models)
}
pub async fn get_token_usage(
ext_token: ExtToken,
is_pri: bool,
time: DateTime,
model_id: &'static str,
) -> Option<ChainUsage> {
const POLL_MAX_ATTEMPTS: u8 = 4;
const POLL_INTERVAL: ::core::time::Duration = ::core::time::Duration::from_millis(1500);
let maybe_token = ext_token
.secondary_token
.as_ref()
.unwrap_or(&ext_token.primary_token);
let mut buf = [0; 31];
let user_id = maybe_token.raw().subject.id.to_str(&mut buf) as &str;
let mut token_usage = None;
for _ in 0..POLL_MAX_ATTEMPTS {
tokio::time::sleep(POLL_INTERVAL).await;
let res = get_filtered_usage_events(
&ext_token.get_client(),
user_id,
maybe_token.as_str(),
is_pri,
FilteredUsageArgs {
start: Some(time),
end: None,
model_id: Some(model_id),
size: Some(10),
},
)
.await?;
if let Some(first) = res.usage_events_display.first()
&& let Some(usage) = first.token_usage
{
token_usage = Some(usage);
break;
};
}
unsafe { ::core::mem::transmute(token_usage) }
}
// pub fn validate_token_and_checksum(auth_token: &str) -> Option<(String, Checksum)> {
// // 尝试使用自定义分隔符查找
// let mut delimiter_pos = auth_token.rfind(*TOKEN_DELIMITER);
// // 如果自定义分隔符未找到,并且 USE_COMMA_DELIMITER 为 true,则尝试使用逗号
// if delimiter_pos.is_none() && *USE_COMMA_DELIMITER {
// delimiter_pos = auth_token.rfind(COMMA);
// }
// // 如果最终都没有找到分隔符,则返回 None
// let comma_pos = delimiter_pos?;
// // 使用找到的分隔符位置分割字符串
// let (token_part, checksum) = auth_token.split_at(comma_pos);
// let checksum = &checksum[1..]; // 跳过逗号
// // 解析 token - 为了向前兼容,忽略最后一个:或%3A前的内容
// let colon_pos = token_part.rfind(':');
// let encoded_colon_pos = token_part.rfind("%3A");
// let token = match (colon_pos, encoded_colon_pos) {
// (None, None) => token_part, // 最简单的构成: token,checksum
// (Some(pos1), None) => &token_part[(pos1 + 1)..],
// (None, Some(pos2)) => &token_part[(pos2 + 3)..],
// (Some(pos1), Some(pos2)) => {
// let pos = pos1.max(pos2);
// let start = if pos == pos2 { pos + 3 } else { pos + 1 };
// &token_part[start..]
// }
// };
// // 验证 token 和 checksum 有效性
// if let Ok(chekcsum) = Checksum::from_str(checksum) {
// if validate_token(token) {
// Some((token.to_string(), chekcsum))
// } else {
// None
// }
// } else {
// None
// }
// }
// pub fn extract_token(auth_token: &str) -> Option<&str> {
// // 尝试使用自定义分隔符查找
// let mut delimiter_pos = auth_token.rfind(*TOKEN_DELIMITER);
// // 如果自定义分隔符未找到,并且 USE_COMMA_DELIMITER 为 true,则尝试使用逗号
// if delimiter_pos.is_none() && *USE_COMMA_DELIMITER {
// delimiter_pos = auth_token.rfind(COMMA);
// }
// // 根据是否找到分隔符来确定 token_part
// let token_part = match delimiter_pos {
// Some(pos) => &auth_token[..pos],
// None => auth_token,
// };
// // 向前兼容
// let colon_pos = token_part.rfind(':');
// let encoded_colon_pos = token_part.rfind("%3A");
// let token = match (colon_pos, encoded_colon_pos) {
// (None, None) => token_part,
// (Some(pos1), None) => &token_part[(pos1 + 1)..],
// (None, Some(pos2)) => &token_part[(pos2 + 3)..],
// (Some(pos1), Some(pos2)) => {
// let pos = pos1.max(pos2);
// let start = if pos == pos2 { pos + 3 } else { pos + 1 };
// &token_part[start..]
// }
// };
// // 验证 token 有效性
// if validate_token(token) {
// Some(token)
// } else {
// None
// }
// }
#[inline(always)]
pub fn format_time_ms(seconds: f64) -> f64 { (seconds * 1000.0).round() / 1000.0 }
/// 将 JWT token 转换为 TokenInfo
#[inline]
pub fn token_to_tokeninfo(
token: RawToken,
checksum: Checksum,
client_key: Hash,
config_version: Option<uuid::Uuid>,
session_id: uuid::Uuid,
proxy_name: Option<String>,
timezone: Option<String>,
gcpp_host: Option<i32>,
) -> key_config::TokenInfo {
key_config::TokenInfo {
token: Some(key_config::token_info::Token::from_raw(token)),
checksum: checksum.into_bytes().to_vec(),
client_key: client_key.into_bytes().to_vec(),
config_version: config_version.map(|v| v.into_bytes().to_vec()),
session_id: session_id.into_bytes().to_vec(),
proxy_name,
timezone,
gcpp_host,
}
}
/// 将 TokenInfo 转换为 JWT token
#[inline]
pub fn tokeninfo_to_token(info: key_config::TokenInfo) -> Option<ExtToken> {
let checksum = Checksum::from_bytes(info.checksum.try_into().ok()?);
let client_key = Hash::from_bytes(info.client_key.try_into().ok()?);
let config_version = info
.config_version
.and_then(|v| uuid::Uuid::from_slice(&v).ok());
let session_id = uuid::Uuid::from_slice(&info.session_id).ok()?;
let timezone = info.timezone.and_then(|s| chrono_tz::Tz::from_str(&s).ok());
let gcpp_host = info.gcpp_host.and_then(GcppHost::from_i32);
Some(ExtToken {
primary_token: Token::new(info.token?.into_raw()?, None),
secondary_token: None,
checksum,
client_key,
config_version,
session_id,
proxy: info.proxy_name,
timezone,
gcpp_host,
user: None,
})
}
/// 压缩数据为gzip格式
#[inline]
fn compress_gzip(data: &[u8]) -> Result<Vec<u8>, ::std::io::Error> {
use ::std::io::Write as _;
use flate2::{Compression, write::GzEncoder};
const LEVEL: Compression = Compression::new(6);
// 预分配容量,gzip压缩后通常会变小,但预留一些空间给gzip头部
let estimated_size = data.len() / 2 + 18;
let mut encoder = GzEncoder::new(Vec::with_capacity(estimated_size), LEVEL);
encoder.write_all(data)?;
encoder.finish()
}
#[inline(always)]
pub fn encode_message(
message: &impl ::prost::Message,
maybe_stream: bool,
) -> Result<Vec<u8>, Box<dyn ::std::error::Error + Send + Sync>> {
const COMPRESSION_THRESHOLD: usize = 1024; // 1KB
const LENGTH_OVERFLOW_MSG: &str = "Message length exceeds ~4 GiB";
let estimated_size = message.encoded_len();
if !maybe_stream {
let mut encoded = Vec::with_capacity(estimated_size);
message.encode_raw(&mut encoded);
return Ok(encoded);
}
// 提前检查长度是否溢出
if estimated_size > u32::MAX as usize {
__cold_path!();
return Err(LENGTH_OVERFLOW_MSG.into());
}
use ::core::mem::MaybeUninit;
// 创建未初始化的缓冲区:[头部 5B][消息体]
let mut buffer = Vec::<MaybeUninit<u8>>::with_capacity(5 + estimated_size);
unsafe {
// 设置长度,但内容仍未初始化
buffer.set_len(5 + estimated_size);
// 获取消息体部分的可变引用并编码
let header_ptr: *mut u8 = buffer.as_mut_ptr().cast();
let body_ptr = header_ptr.add(5);
message.encode_raw(&mut ::core::slice::from_raw_parts_mut(
body_ptr,
estimated_size,
));
// 压缩处理:仅当压缩后更小时才使用压缩版本
let (compression_flag, final_len) = if estimated_size >= COMPRESSION_THRESHOLD {
let body_slice = ::core::slice::from_raw_parts(body_ptr, estimated_size);
let compressed = compress_gzip(body_slice)?;
let compressed_len = compressed.len();
if compressed_len < estimated_size {
// 覆盖为压缩数据
::core::ptr::copy_nonoverlapping(compressed.as_ptr(), body_ptr, compressed_len);
// 截断到实际使用的长度
buffer.set_len(5 + compressed_len);
(0x01, compressed_len)
} else {
(0x00, estimated_size)
}
} else {
(0x00, estimated_size)
};
// 初始化头部
*header_ptr = compression_flag;
// 构建协议头:[压缩标志 1B][长度 4B BE]
::core::hint::assert_unchecked(final_len <= u32::MAX as usize);
let len_bytes = (final_len as u32).to_be_bytes();
// 直接拷贝长度字节
::core::ptr::copy_nonoverlapping(len_bytes.as_ptr(), header_ptr.add(1), 4);
// 转换为已初始化的 Vec<u8>
Ok(::core::mem::transmute(buffer))
}
}
/// 生成 PKCE code_verifier 和对应的 code_challenge (S256 method)
#[inline]
fn generate_pkce_pair() -> ([u8; 43], [u8; 43]) {
use ::core::mem::MaybeUninit;
use rand::TryRngCore as _;
use sha2::Digest as _;
// 生成 32 字节随机数作为 verifier
let mut verifier_bytes = MaybeUninit::<[u8; 32]>::uninit();
// 直接填充到未初始化内存
unsafe {
let bytes_ptr = verifier_bytes.as_mut_ptr().cast();
let bytes_slice = ::core::slice::from_raw_parts_mut(bytes_ptr, 32);
rand::rngs::OsRng
.try_fill_bytes(bytes_slice)
.expect("获取系统安全随机数失败,这是一个严重错误!");
// 此时 verifier_bytes 已完全初始化
let verifier_bytes = verifier_bytes.assume_init();
// Base64 编码为 code_verifier
let mut code_verifier = MaybeUninit::<[u8; 43]>::uninit();
let verifier_ptr = code_verifier.as_mut_ptr().cast();
let verifier_slice = ::core::slice::from_raw_parts_mut(verifier_ptr, 43);
__unwrap!(URL_SAFE_NO_PAD.encode_slice(verifier_bytes, verifier_slice));
let code_verifier = code_verifier.assume_init();
// SHA-256 哈希后 Base64 编码为 code_challenge
let hash_result = sha2::Sha256::digest(code_verifier);
let mut code_challenge = MaybeUninit::<[u8; 43]>::uninit();
let challenge_ptr = code_challenge.as_mut_ptr().cast();
let challenge_slice = ::core::slice::from_raw_parts_mut(challenge_ptr, 43);
__unwrap!(URL_SAFE_NO_PAD.encode_slice(hash_result, challenge_slice));
let code_challenge = code_challenge.assume_init();
(code_verifier, code_challenge)
}
}
pub async fn get_new_token(ext_token: &mut ExtToken, is_pri: bool) -> bool {
let is_session = ext_token.primary_token.is_session();
match if is_session {
refresh_token(ext_token, is_pri).await
} else {
upgrade_token(ext_token, is_pri).await
} {
Some((new_token, s)) => {
let tmp = Token::new(new_token, Some(s));
if !is_session && ext_token.secondary_token.is_none() {
let old_token = ::core::mem::replace(&mut ext_token.primary_token, tmp);
ext_token.secondary_token = Some(old_token);
} else {
ext_token.primary_token = tmp;
}
true
}
None => false,
}
}
async fn upgrade_token(ext_token: &ExtToken, is_pri: bool) -> Option<(RawToken, String)> {
const POLL_MAX_ATTEMPTS: u8 = 5;
const POLL_INTERVAL: ::core::time::Duration = ::core::time::Duration::from_secs(1);
#[derive(::serde::Deserialize)]
#[serde(rename_all = "camelCase")]
struct PollResponse {
pub access_token: String,
// pub refresh_token: String,
// pub challenge: String,
// pub auth_id: String,
// pub uuid: String,
}
let (verifier, challenge) = generate_pkce_pair();
let verifier = unsafe { ::core::str::from_utf8_unchecked(&verifier) };
let challenge = unsafe { ::core::str::from_utf8_unchecked(&challenge) };
let mut buf = [0; 36];
let uuid = uuid::Uuid::new_v4().hyphenated().encode_lower(&mut buf) as &str;
let token = ext_token
.secondary_token
.as_ref()
.unwrap_or(&ext_token.primary_token);
let mut buf = [0; 31];
let user_id = token.raw().subject.id.to_str(&mut buf) as &str;
let auth_token = token.as_str();
// 发起刷新请求
let upgrade_response = super::client::build_token_upgrade_request(
&ext_token.get_client(),
uuid,
challenge,
user_id,
auth_token,
is_pri,
)
.send()
.await
.ok()?;
if !upgrade_response.status().is_success() {
return None;
}
// 轮询获取token
for _ in 0..POLL_MAX_ATTEMPTS {
let poll_response = super::client::build_token_poll_request(
&ext_token.get_client(),
uuid,
verifier,
is_pri,
)
.send()
.await
.ok()?;
match poll_response.status() {
reqwest::StatusCode::OK => {
let token = poll_response
.json::<PollResponse>()
.await
.ok()?
.access_token;
return parse_token(token);
}
reqwest::StatusCode::NOT_FOUND => {
tokio::time::sleep(POLL_INTERVAL).await;
}
_ => return None,
}
}
None
}
async fn refresh_token(ext_token: &ExtToken, is_pri: bool) -> Option<(RawToken, String)> {
const CLIENT_ID: &str = "KbZUR41cY7W6zRSdpSUJ7I7mLYBKOCmB";
struct RefreshTokenRequest<'a> {
refresh_token: &'a str,
}
impl ::serde::Serialize for RefreshTokenRequest<'_> {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: ::serde::Serializer,
{
use ::serde::ser::SerializeStruct as _;
let mut state = serializer.serialize_struct("RefreshTokenRequest", 3)?;
state.serialize_field("grant_type", "refresh_token")?;
state.serialize_field("client_id", CLIENT_ID)?;
state.serialize_field("refresh_token", self.refresh_token)?;
state.end()
}
}
#[derive(::serde::Deserialize)]
struct RefreshTokenResponse {
access_token: String,
// id_token: String,
// #[serde(rename = "shouldLogout")]
// should_logout: bool,
}
let refresh_request = RefreshTokenRequest {
refresh_token: ext_token.primary_token.as_str(),
};
let body = serde_json::to_vec(&refresh_request).ok()?;
let response =
super::client::build_token_refresh_request(&ext_token.get_client(), is_pri, body)
.send()
.await
.ok()?;
let token = response
.json::<RefreshTokenResponse>()
.await
.ok()?
.access_token;
parse_token(token)
}
// 提取token解析逻辑
#[inline]
fn parse_token(token_string: String) -> Option<(RawToken, String)> {
let raw_token = token_string.parse().ok()?;
Some((raw_token, token_string))
}
pub async fn get_server_config(ext_token: ExtToken, is_pri: bool) -> Option<uuid::Uuid> {
let response = {
let client = super::client::build_client_request(super::client::AiServiceRequest {
ext_token,
fs_client_key: None,
url: server_config_url(is_pri),
is_stream: false,
trace_id: Some(new_uuid_v4()),
is_pri,
cookie: None,
});
client.send().await.ok()?.bytes().await.ok()?
};
let server_config = GetServerConfigResponse::decode(response.as_ref()).ok()?;
uuid::Uuid::try_parse(&server_config.config_version).ok()
}
// pub async fn get_geo_cpp_backend_url(
// client: Client,
// auth_token: &str,
// checksum: Checksum,
// client_key: Hash,
// timezone: &'static str,
// session_id: Option<uuid::Uuid>,
// is_pri: bool,
// ) -> Option<String> {
// let response = {
// let client = super::client::build_client_request(super::client::AiServiceRequest {
// client,
// auth_token,
// checksum,
// client_key,
// fs_client_key: None,
// url: crate::app::lazy::cpp_config_url(is_pri),
// is_stream: false,
// config_version: None,
// timezone,
// trace_id: Some(new_uuid_v4()),
// session_id,
// is_pri,
// });
// let request = crate::core::aiserver::v1::CppConfigRequest::default();
// client
// .body(__unwrap!(encode_message(&request, false)))
// .send()
// .await
// .ok()?
// .bytes()
// .await
// .ok()?
// };
// crate::core::aiserver::v1::CppConfigResponse::decode(response.as_ref())
// .ok()
// .map(|res| res.geo_cpp_backend_url)
// }
const EMPTY_JSON: bytes::Bytes = bytes::Bytes::from_static(b"{}");
pub async fn get_teams(
client: &Client,
user_id: &str,
auth_token: &str,
is_pri: bool,
) -> Option<Vec<Team>> {
let request = super::client::build_proto_web_request(
client, user_id, auth_token, teams_url, is_pri, EMPTY_JSON,
);
request
.send()
.await
.ok()?
.json::<GetTeamsResponse>()
.await
.ok()
.map(|r| r.teams)
}
pub async fn get_is_on_new_pricing(
client: &Client,
user_id: &str,
auth_token: &str,
is_pri: bool,
) -> Option<bool> {
let request = super::client::build_proto_web_request(
client,
user_id,
auth_token,
is_on_new_pricing_url,
is_pri,
EMPTY_JSON,
);
#[derive(serde::Deserialize)]
struct PricingConfig {
#[serde(rename = "isOnNewPricing")]
is_on_new_pricing: bool,
}
request
.send()
.await
.ok()?
.json::<PricingConfig>()
.await
.ok()
.map(|r| r.is_on_new_pricing)
}
pub async fn get_sessions(
client: &Client,
user_id: &str,
auth_token: &str,
is_pri: bool,
) -> Option<Vec<Session>> {
let request = super::client::build_sessions_request(client, user_id, auth_token, is_pri);
request
.send()
.await
.ok()?
.json::<ListActiveSessionsResponse>()
.await
.ok()
.map(|r| r.sessions)
}
pub async fn get_aggregated_usage_events(
client: &Client,
user_id: &str,
auth_token: &str,
is_pri: bool,
) -> Option<GetAggregatedUsageEventsResponse> {
let request = super::client::build_proto_web_request(
client,
user_id,
auth_token,
aggregated_usage_events_url,
is_pri,
bytes::Bytes::from(__unwrap!(serde_json::to_vec(&{
const DELTA: chrono::TimeDelta = __unwrap!(chrono::TimeDelta::new(2629743, 765840000));
let now = DateTime::utc_now();
let start_date = now - DELTA;
GetAggregatedUsageEventsRequest {
team_id: -1,
start_date: Some(start_date.timestamp_millis()),
end_date: Some(now.timestamp_millis()),
user_id: None,
}
}))),
);
request
.send()
.await
.ok()?
.json::<GetAggregatedUsageEventsResponse>()
.await
.ok()
}
pub struct FilteredUsageArgs {
pub start: Option<DateTime>,
pub end: Option<DateTime>,
pub model_id: Option<&'static str>,
pub size: Option<i32>,
}
impl From<FilteredUsageArgs> for GetFilteredUsageEventsRequest {
#[inline]
fn from(args: FilteredUsageArgs) -> Self {
const TZ: chrono::FixedOffset = __unwrap!(chrono::FixedOffset::west_opt(16 * 3600));
const TIME: chrono::NaiveTime = __unwrap!(chrono::NaiveTime::from_hms_opt(0, 0, 0));
const START: chrono::TimeDelta = chrono::TimeDelta::days(-7);
const END: chrono::TimeDelta = __unwrap!(chrono::TimeDelta::new(86399, 999000000));
let (start_date, end_date) = if let (Some(a), Some(b)) = (args.start, args.end) {
(a.timestamp_millis(), b.timestamp_millis())
} else {
let now = chrono::DateTime::<chrono::FixedOffset>::from_naive_utc_and_offset(
DateTime::naive_now(),
TZ,
)
.date_naive()
.and_time(TIME);
match (args.start, args.end) {
(None, None) => (
(now + START)
.and_local_timezone(TZ)
.unwrap()
.timestamp_millis(),
(now + END)
.and_local_timezone(TZ)
.unwrap()
.timestamp_millis(),
),
(None, Some(b)) => (
(now + START)
.and_local_timezone(TZ)
.unwrap()
.timestamp_millis(),
b.timestamp_millis(),
),
(Some(a), None) => (
a.timestamp_millis(),
(now + END)
.and_local_timezone(TZ)
.unwrap()
.timestamp_millis(),
),
(Some(_), Some(_)) => unsafe { ::core::hint::unreachable_unchecked() },
}
};
Self {
team_id: 0,
start_date: Some(start_date),
end_date: Some(end_date),
user_id: None,
model_id: args.model_id.map(ToString::to_string),
page: Some(1),
page_size: Some(args.size.unwrap_or(100)),
}
}
}
pub async fn get_filtered_usage_events(
client: &Client,
user_id: &str,
auth_token: &str,
is_pri: bool,
args: FilteredUsageArgs,
) -> Option<GetFilteredUsageEventsResponse> {
let request = super::client::build_proto_web_request(
client,
user_id,
auth_token,
filtered_usage_events_url,
is_pri,
bytes::Bytes::from(__unwrap!(serde_json::to_vec(&{
let req: GetFilteredUsageEventsRequest = args.into();
req
}))),
);
let res = request.send().await.ok()?;
crate::debug!("<get_filtered_usage_events> {}", res.status());
let res = res.bytes().await.ok()?;
crate::debug!("<get_filtered_usage_events> {}", unsafe {
::core::str::from_utf8_unchecked(&res[..])
});
serde_json::from_slice(&res[..]).ok()
// .json::<GetFilteredUsageEventsResponse>()
// .await
// .ok()
}
#[inline]
pub fn new_uuid_v4() -> [u8; 36] {
let mut buf = [0; 36];
uuid::Uuid::new_v4().hyphenated().encode_lower(&mut buf);
buf
}