use crate::error::{BotError, Result};
use serde::{Deserialize, Serialize};
use std::fmt;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::Mutex;
#[derive(Clone, Serialize, Deserialize)]
pub struct Token {
app_id: String,
secret: String,
#[serde(skip)]
access_token: Option<String>,
#[serde(skip)]
expires_at: Option<u64>,
#[serde(skip)]
refresh_mutex: Arc<Mutex<()>>,
}
impl Token {
pub fn new(app_id: impl Into<String>, secret: impl Into<String>) -> Self {
Self {
app_id: app_id.into(),
secret: secret.into(),
access_token: None,
expires_at: None,
refresh_mutex: Arc::new(Mutex::new(())),
}
}
pub fn app_id(&self) -> &str {
&self.app_id
}
pub fn secret(&self) -> &str {
&self.secret
}
pub async fn authorization_header(&self) -> Result<String> {
self.ensure_valid_token().await?;
if let Some(access_token) = &self.access_token {
Ok(format!("QQBot {access_token}"))
} else {
Err(BotError::auth("No valid access token available"))
}
}
pub async fn bot_token(&self) -> Result<String> {
self.authorization_header().await
}
async fn ensure_valid_token(&self) -> Result<()> {
let current_time = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_err(|_| BotError::internal("Failed to get current time"))?
.as_secs();
if self.access_token.is_none() || self.expires_at.is_none_or(|exp| current_time >= exp) {
self.refresh_access_token().await?;
}
Ok(())
}
async fn refresh_access_token(&self) -> Result<()> {
let _guard = self.refresh_mutex.lock().await;
let current_time = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_err(|_| BotError::internal("Failed to get current time"))?
.as_secs();
if let Some(expires_at) = self.expires_at {
if current_time < expires_at && self.access_token.is_some() {
return Ok(());
}
}
let client = reqwest::Client::new();
let request_body = serde_json::json!({
"appId": self.app_id,
"clientSecret": self.secret
});
let response = client
.post("https://bots.qq.com/app/getAppAccessToken")
.json(&request_body)
.timeout(std::time::Duration::from_secs(20))
.send()
.await
.map_err(|e| BotError::connection(format!("Failed to request access token: {e}")))?;
if !response.status().is_success() {
return Err(BotError::api(
response.status().as_u16() as u32,
format!(
"Token request failed: {}",
response.text().await.unwrap_or_default()
),
));
}
let token_response: serde_json::Value = response.json().await.map_err(BotError::Http)?;
let access_token = token_response
.get("access_token")
.and_then(|v| v.as_str())
.ok_or_else(|| BotError::auth("No access_token in response"))?;
let expires_in = token_response
.get("expires_in")
.and_then(|v| v.as_str())
.and_then(|s| s.parse::<u64>().ok())
.ok_or_else(|| BotError::auth("No expires_in in response"))?;
unsafe {
let self_mut = self as *const Self as *mut Self;
(*self_mut).access_token = Some(access_token.to_string());
(*self_mut).expires_at = Some(current_time + expires_in);
}
Ok(())
}
pub fn validate(&self) -> Result<()> {
if self.app_id.is_empty() {
return Err(BotError::auth("App ID cannot be empty"));
}
if self.secret.is_empty() {
return Err(BotError::auth("Secret cannot be empty"));
}
Ok(())
}
pub fn from_env() -> Result<Self> {
let app_id = std::env::var("QQ_BOT_APP_ID")
.map_err(|_| BotError::config("QQ_BOT_APP_ID environment variable not found"))?;
let secret = std::env::var("QQ_BOT_SECRET")
.map_err(|_| BotError::config("QQ_BOT_SECRET environment variable not found"))?;
let token = Self::new(app_id, secret);
token.validate()?;
Ok(token)
}
pub fn safe_display(&self) -> String {
let masked_secret = if self.secret.len() > 8 {
format!(
"{}****{}",
&self.secret[..4],
&self.secret[self.secret.len() - 4..]
)
} else {
"****".to_string()
};
format!(
"Token {{ app_id: {}, secret: {} }}",
self.app_id, masked_secret
)
}
}
impl fmt::Display for Token {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.safe_display())
}
}
impl PartialEq for Token {
fn eq(&self, other: &Self) -> bool {
self.app_id == other.app_id && self.secret == other.secret
}
}
impl Eq for Token {}
impl fmt::Debug for Token {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Token")
.field("app_id", &self.app_id)
.field("secret", &"[REDACTED]")
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_token_creation() {
let token = Token::new("123456", "secret123");
assert_eq!(token.app_id(), "123456");
assert_eq!(token.secret(), "secret123");
}
#[tokio::test]
async fn test_authorization_header() {
let token = Token::new("test", "secret");
let result = token.authorization_header().await;
assert!(
result.is_err(),
"Expected authorization_header to fail with invalid credentials"
);
}
#[tokio::test]
async fn test_bot_token() {
let token = Token::new("test", "secret");
let bot_token_result = token.bot_token().await;
let auth_header_result = token.authorization_header().await;
assert!(bot_token_result.is_err());
assert!(auth_header_result.is_err());
}
#[test]
fn test_validation() {
let valid_token = Token::new("123", "secret");
assert!(valid_token.validate().is_ok());
let empty_app_id = Token::new("", "secret");
assert!(empty_app_id.validate().is_err());
let empty_secret = Token::new("123", "");
assert!(empty_secret.validate().is_err());
}
#[test]
fn test_safe_display() {
let token = Token::new("123456", "verylongsecret123");
let display = token.safe_display();
assert!(display.contains("123456"));
assert!(display.contains("very"));
assert!(display.contains("123"));
assert!(display.contains("****"));
assert!(!display.contains("longsecret"));
let short_token = Token::new("123", "short");
let short_display = short_token.safe_display();
assert!(short_display.contains("****"));
assert!(!short_display.contains("short"));
}
#[test]
fn test_debug_format() {
let token = Token::new("123456", "secret123");
let debug_str = format!("{:?}", token);
assert!(debug_str.contains("123456"));
assert!(debug_str.contains("[REDACTED]"));
assert!(!debug_str.contains("secret123"));
}
}