use std::sync::{Arc, RwLock};
use log::{debug, error, info, warn};
use super::{AccessToken, AccessTokenResponse, TokenStatus};
use crate::errors::{NetDiskError, NetDiskResult};
use crate::http::HttpClient;
#[derive(Debug, Clone)]
pub struct TokenProviderConfig {
pub auto_refresh: bool,
pub refresh_ahead_seconds: u64,
pub max_refresh_retries: usize,
}
impl Default for TokenProviderConfig {
fn default() -> Self {
TokenProviderConfig {
auto_refresh: true,
refresh_ahead_seconds: 86400,
max_refresh_retries: 3,
}
}
}
#[derive(Debug, Clone)]
pub struct TokenProvider {
http_client: HttpClient,
app_key: String,
app_secret: String,
access_token: Arc<RwLock<Option<AccessToken>>>,
config: TokenProviderConfig,
}
impl TokenProvider {
pub fn new(
http_client: HttpClient,
app_key: &str,
app_secret: &str,
config: TokenProviderConfig,
) -> Self {
TokenProvider {
http_client,
app_key: app_key.to_string(),
app_secret: app_secret.to_string(),
access_token: Arc::new(RwLock::new(None)),
config,
}
}
pub fn get_access_token(&self) -> NetDiskResult<String> {
let token = self
.access_token
.read()
.map_err(|e| NetDiskError::SyncError(format!("Failed to read access token: {}", e)))?;
token
.as_ref()
.map(|t| t.access_token.clone())
.ok_or_else(|| NetDiskError::auth_error("No access token available"))
}
pub fn get_access_token_full(&self) -> NetDiskResult<Option<AccessToken>> {
let token = self
.access_token
.read()
.map_err(|e| NetDiskError::SyncError(format!("Failed to read access token: {}", e)))?;
Ok(token.clone())
}
pub fn validate_token(&self) -> NetDiskResult<TokenStatus> {
let token = self
.access_token
.read()
.map_err(|e| NetDiskError::SyncError(format!("Failed to read access token: {}", e)))?;
match token.as_ref() {
Some(t) => Ok(t.validate()),
None => Ok(TokenStatus::Expired),
}
}
pub fn set_access_token(&self, token: AccessToken) -> NetDiskResult<()> {
let status = token.validate();
let mut current = self
.access_token
.write()
.map_err(|e| NetDiskError::SyncError(format!("Failed to write access token: {}", e)))?;
*current = Some(token);
match status {
TokenStatus::Valid => {
info!(
"Access token updated successfully (valid for {} seconds)",
current.as_ref().unwrap().remaining_seconds()
);
}
TokenStatus::ExpiringSoon => {
warn!(
"Access token updated but will expire soon ({} seconds remaining)",
current.as_ref().unwrap().remaining_seconds()
);
}
TokenStatus::Expired => {
error!("Access token updated but is already expired! Please refresh or re-authenticate.");
}
}
Ok(())
}
pub fn needs_refresh(&self) -> NetDiskResult<bool> {
let token = self
.access_token
.read()
.map_err(|e| NetDiskError::SyncError(format!("Failed to read access token: {}", e)))?;
match token.as_ref() {
Some(t) => Ok(t.remaining_seconds() <= self.config.refresh_ahead_seconds),
None => Ok(true),
}
}
pub async fn refresh_token(&self) -> NetDiskResult<AccessToken> {
let refresh_token = {
let current_token = self.access_token.read().map_err(|e| {
NetDiskError::SyncError(format!("Failed to read access token: {}", e))
})?;
current_token
.as_ref()
.ok_or_else(|| NetDiskError::auth_error("No access token to refresh"))?
.refresh_token
.clone()
};
info!("Starting token refresh...");
let params = [
("grant_type", "refresh_token"),
("refresh_token", &refresh_token),
("client_id", &self.app_key),
("client_secret", &self.app_secret),
];
let url = format!(
"https://openapi.baidu.com/oauth/2.0/token?{}",
serde_urlencoded::to_string(params).map_err(|e| {
NetDiskError::Unknown {
message: format!("Failed to encode params: {}", e),
}
})?
);
let response: AccessTokenResponse =
self.http_client.get(&url, None).await.map_err(|e| {
error!("Token refresh failed: {}", e);
NetDiskError::auth_error(&format!(
"Token refresh failed, please re-authenticate: {}",
e
))
})?;
let token: AccessToken = response.into();
self.set_access_token(token.clone())?;
info!(
"Token refreshed successfully, valid for {} seconds",
token.remaining_seconds()
);
Ok(token)
}
pub async fn try_refresh_token(&self) -> NetDiskResult<Option<AccessToken>> {
if !self.needs_refresh()? {
debug!("Token doesn't need refresh");
return Ok(None);
}
match self.refresh_token().await {
Ok(token) => Ok(Some(token)),
Err(e) => {
warn!("Failed to refresh token: {}", e);
Ok(None)
}
}
}
pub async fn get_valid_token(&self) -> NetDiskResult<AccessToken> {
let token_opt = self
.access_token
.read()
.map_err(|e| NetDiskError::SyncError(format!("Failed to read access token: {}", e)))?
.clone();
let token =
token_opt.ok_or_else(|| NetDiskError::auth_error("No access token available"))?;
let status = token.validate();
match status {
TokenStatus::Valid => {
debug!("Token is valid, returning directly");
Ok(token)
}
TokenStatus::ExpiringSoon => {
if self.config.auto_refresh {
info!("Token is expiring soon, attempting refresh...");
match self.refresh_token().await {
Ok(new_token) => Ok(new_token),
Err(_) => {
warn!("Token refresh failed but token is still usable, returning existing token");
Ok(token)
}
}
} else {
warn!("Token is expiring soon, auto-refresh is disabled");
Ok(token)
}
}
TokenStatus::Expired => {
if self.config.auto_refresh {
error!("Token is expired, attempting refresh...");
self.refresh_token().await
} else {
Err(NetDiskError::auth_error(
"Token is expired and auto-refresh is disabled, please re-authenticate",
))
}
}
}
}
pub fn clear_token(&self) -> NetDiskResult<()> {
let mut token = self
.access_token
.write()
.map_err(|e| NetDiskError::SyncError(format!("Failed to write access token: {}", e)))?;
*token = None;
info!("Access token cleared");
Ok(())
}
}