File size: 8,514 Bytes
3c2af29 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 |
mod checksum;
use ::base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
pub use checksum::*;
mod token;
pub use token::*;
mod base64;
pub use base64::*;
use super::model::{token::TokenPayload, userinfo::{StripeProfile, TokenProfile, UsageProfile, UserProfile}};
use crate::app::{
constant::{COMMA, FALSE, TRUE},
lazy::{TOKEN_DELIMITER, USE_COMMA_DELIMITER},
};
pub fn parse_bool_from_env(key: &str, default: bool) -> bool {
std::env::var(key)
.ok()
.map(|v| match v.to_lowercase().as_str() {
TRUE | "1" => true,
FALSE | "0" => false,
_ => default,
})
.unwrap_or(default)
}
pub fn parse_string_from_env(key: &str, default: &str) -> String {
std::env::var(key).unwrap_or_else(|_| default.to_string())
}
pub fn parse_ascii_char_from_env(key: &str, default: char) -> char {
std::env::var(key)
.ok()
.and_then(|v| {
let chars: Vec<char> = v.chars().collect();
if chars.len() == 1 && chars[0].is_ascii() {
Some(chars[0])
} else {
None
}
})
.unwrap_or(default)
}
pub fn parse_usize_from_env(key: &str, default: usize) -> usize {
std::env::var(key)
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(default)
}
pub trait TrimNewlines {
fn trim_leading_newlines(self) -> Self;
}
impl TrimNewlines for String {
#[inline(always)]
fn trim_leading_newlines(mut self) -> Self {
let bytes = self.as_bytes();
if bytes.len() >= 2 && bytes[0] == b'\n' && bytes[1] == b'\n' {
unsafe {
let start_ptr = self.as_mut_ptr();
let new_len = self.len() - 2;
std::ptr::copy(start_ptr.add(2), start_ptr, new_len);
self.as_mut_vec().set_len(new_len);
}
}
self
}
}
pub async fn get_token_profile(auth_token: &str) -> Option<TokenProfile> {
let user_id = extract_user_id(auth_token)?;
// 构建请求客户端
let client = super::client::build_usage_client(&user_id, auth_token);
// 发送请求并获取响应
// let response = client.send().await.ok()?;
// let bytes = response.bytes().await?;
// println!("Raw response bytes: {:?}", bytes);
// let usage = serde_json::from_str::<UsageProfile>(&text).ok()?;
let usage = client
.send()
.await
.ok()?
.json::<UsageProfile>()
.await
.ok()?;
let user = get_user_profile(auth_token).await?;
// 从 Stripe 获取用户资料
let stripe = get_stripe_profile(auth_token).await?;
// 映射响应数据到 TokenProfile
Some(TokenProfile {
usage,
user,
stripe,
})
}
pub async fn get_stripe_profile(auth_token: &str) -> Option<StripeProfile> {
let client = super::client::build_profile_client(auth_token);
let response = client
.send()
.await
.ok()?
.json::<StripeProfile>()
.await
.ok()?;
Some(response)
}
pub async fn get_user_profile(auth_token: &str) -> Option<UserProfile> {
let user_id = extract_user_id(auth_token)?;
// 构建请求客户端
let client = super::client::build_userinfo_client(&user_id, auth_token);
// 发送请求并获取响应
let user_profile = client.send().await.ok()?.json::<UserProfile>().await.ok()?;
Some(user_profile)
}
pub fn validate_token_and_checksum(auth_token: &str) -> Option<(String, String)> {
// 尝试使用自定义分隔符查找
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 validate_token(token) && validate_checksum(checksum) {
Some((token.to_string(), generate_checksum_with_repair(checksum)))
} else {
None
}
}
pub fn extract_token(auth_token: &str) -> Option<String> {
// 尝试使用自定义分隔符查找
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.to_string())
} else {
None
}
}
pub fn format_time_ms(seconds: f64) -> f64 {
(seconds * 1000.0).round() / 1000.0
}
use crate::chat::config::key_config;
/// 将 JWT token 转换为 TokenInfo
pub fn token_to_tokeninfo(auth_token: &str) -> Option<key_config::TokenInfo> {
let (token, checksum) = validate_token_and_checksum(auth_token)?;
// JWT token 由3部分组成,用 . 分隔
let parts: Vec<&str> = token.split('.').collect();
if parts.len() != 3 {
return None;
}
// 解码 payload (第二部分)
let payload = match URL_SAFE_NO_PAD.decode(parts[1]) {
Ok(decoded) => decoded,
Err(_) => return None,
};
// 将 payload 转换为字符串
let payload_str = match String::from_utf8(payload) {
Ok(s) => s,
Err(_) => return None,
};
// 解析为 TokenPayload
let payload: TokenPayload = match serde_json::from_str(&payload_str) {
Ok(p) => p,
Err(_) => return None,
};
let (machine_id_hash, mac_id_hash) = extract_hashes(&checksum)?;
// 构建 TokenInfo
Some(key_config::TokenInfo {
sub: payload.sub,
exp: payload.exp,
randomness: payload.randomness,
signature: parts[2].to_string(),
machine_id: machine_id_hash,
mac_id: mac_id_hash,
})
}
/// 将 TokenInfo 转换为 JWT token
pub fn tokeninfo_to_token(info: &key_config::TokenInfo) -> Option<(String, String)> {
// 构建 payload
let payload = TokenPayload {
sub: info.sub.clone(),
exp: info.exp,
randomness: info.randomness.clone(),
time: (info.exp - 2592000000).to_string(), // exp - 30000天
iss: ISSUER.to_string(),
scope: SCOPE.to_string(),
aud: AUDIENCE.to_string(),
};
let payload_str = match serde_json::to_string(&payload) {
Ok(s) => s,
Err(_) => return None,
};
let payload_b64 = URL_SAFE_NO_PAD.encode(payload_str.as_bytes());
// 从 TokenInfo 中获取 machine_id 和 mac_id 的 hex 字符串
let device_id = hex::encode(&info.machine_id);
let mac_addr = if !info.mac_id.is_empty() {
Some(hex::encode(&info.mac_id))
} else {
None
};
// 组合 token
Some((format!("{}.{}.{}", HEADER_B64, payload_b64, info.signature), generate_checksum(&device_id, mac_addr.as_deref())))
}
|