pub(crate) mod token_store;
mod headers;
pub use token_store::{
EncryptedFileTokenStore, FileTokenStore, KeychainTokenStore, TokenSet, TokenStore,
};
pub use headers::{create_basic_header, create_bearer_header};
use crate::error::AuthError;
use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD as BASE64_URL};
use chrono::{DateTime, Duration, Utc};
use parking_lot::RwLock;
use ring::rand::SecureRandom;
use ring::{digest, rand as ring_rand};
use secrecy::SecretString;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::Mutex;
use tracing::{debug, info, warn};
#[cfg(feature = "callback-server")]
use axum::{Router, extract::Query, response::Html, routing::get};
pub type AuthResult<T> = std::result::Result<T, AuthError>;
#[derive(Debug, Clone)]
pub enum TokenNotification {
AccessTokenExpiring { seconds_remaining: i64 },
RefreshTokenExpiring { hours_remaining: i64 },
RefreshTokenExpired,
TokenFileCorrupted,
SessionRecreated,
}
type Refused = (DateTime<Utc>, String);
#[derive(Serialize, Deserialize)]
struct PendingState {
state_sha256: String,
created_at: i64,
}
fn state_digest(state: &str) -> String {
digest::digest(&digest::SHA256, state.as_bytes())
.as_ref()
.iter()
.map(|b| format!("{b:02x}"))
.collect()
}
const AUTHORIZATION_TTL: std::time::Duration = std::time::Duration::from_secs(300);
#[cfg(feature = "callback-server")]
#[derive(Debug, Deserialize)]
struct CallbackParams {
code: Option<String>,
state: Option<String>,
error: Option<String>,
}
#[cfg(feature = "callback-server")]
fn callback_tls_identity(
config: &OAuthConfig,
) -> AuthResult<(
Vec<rustls::pki_types::CertificateDer<'static>>,
rustls::pki_types::PrivateKeyDer<'static>,
)> {
use rustls::pki_types::pem::PemObject;
use rustls::pki_types::{CertificateDer, PrivateKeyDer};
if let (Some(cert_file), Some(key_file)) =
(&config.callback_cert_file, &config.callback_key_file)
{
let unreadable = |file: &std::path::Path, e: &dyn std::fmt::Display| {
AuthError::OAuthFlow(format!("Cannot read {}: {e}", file.display()))
};
let chain = CertificateDer::pem_file_iter(cert_file)
.and_then(Iterator::collect::<Result<Vec<_>, _>>)
.map_err(|e| unreadable(cert_file, &e))?;
if chain.is_empty() {
return Err(unreadable(cert_file, &"no PEM certificate in it"));
}
let key = PrivateKeyDer::from_pem_file(key_file).map_err(|e| unreadable(key_file, &e))?;
info!(
"Callback server uses the certificate in {}",
cert_file.display()
);
return Ok((chain, key));
}
let cert =
rcgen::generate_simple_self_signed(vec!["127.0.0.1".to_string(), "localhost".to_string()])
.map_err(|e| {
AuthError::OAuthFlow(format!("Failed to generate TLS certificate: {e}"))
})?;
info!(
"Callback server uses a self-signed certificate — your browser will warn once; \
choose 'Advanced → Proceed' to complete the flow, or set SCHWAB_CALLBACK_CERT_FILE \
and SCHWAB_CALLBACK_KEY_FILE to a trusted one"
);
Ok((
vec![CertificateDer::from(cert.cert.der().to_vec())],
PrivateKeyDer::Pkcs8(cert.key_pair.serialize_der().into()),
))
}
#[derive(Clone, Serialize, Deserialize)]
pub struct OAuthConfig {
pub app_key: String,
pub app_secret: String,
pub callback_url: String,
pub auth_url: String,
pub token_url: String,
pub tokens_file: PathBuf,
pub capture_callback: bool,
pub auto_refresh: bool,
pub refresh_buffer_seconds: i64,
#[serde(default)]
pub pkce_enabled: bool,
#[serde(default)]
pub token_store_kind: TokenStoreKind,
#[serde(default)]
pub allow_external_callback: bool,
#[serde(default)]
pub callback_cert_file: Option<PathBuf>,
#[serde(default)]
pub callback_key_file: Option<PathBuf>,
#[serde(skip)]
pub on_token_notification: Option<Arc<dyn Fn(TokenNotification) + Send + Sync>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TokenStoreKind {
File,
EncryptedFile,
Keychain,
}
#[cfg(target_os = "macos")]
impl Default for TokenStoreKind {
fn default() -> Self {
tracing::debug!("Defaulting to Keychain storage on macOS");
TokenStoreKind::Keychain
}
}
#[cfg(not(target_os = "macos"))]
impl Default for TokenStoreKind {
fn default() -> Self {
tracing::debug!("Defaulting to EncryptedFile storage");
TokenStoreKind::EncryptedFile
}
}
impl std::fmt::Debug for OAuthConfig {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("OAuthConfig")
.field("app_key", &"***")
.field("app_secret", &"***")
.field("callback_url", &self.callback_url)
.field("auth_url", &self.auth_url)
.field("token_url", &self.token_url)
.field("tokens_file", &self.tokens_file)
.field("capture_callback", &self.capture_callback)
.field("auto_refresh", &self.auto_refresh)
.field("refresh_buffer_seconds", &self.refresh_buffer_seconds)
.field("callback_cert_file", &self.callback_cert_file)
.field("callback_key_file", &self.callback_key_file)
.field(
"on_token_notification",
&self.on_token_notification.is_some(),
)
.finish()
}
}
impl Default for OAuthConfig {
fn default() -> Self {
Self {
app_key: String::new(),
app_secret: String::new(),
callback_url: "https://127.0.0.1:8080".to_string(),
auth_url: "https://api.schwabapi.com/v1/oauth/authorize".to_string(),
token_url: "https://api.schwabapi.com/v1/oauth/token".to_string(),
tokens_file: PathBuf::from("schwab_tokens.json"),
capture_callback: false,
auto_refresh: true,
refresh_buffer_seconds: 61, pkce_enabled: false, token_store_kind: TokenStoreKind::default(),
allow_external_callback: false, callback_cert_file: None,
callback_key_file: None,
on_token_notification: None,
}
}
}
const REFRESH_TOKEN_MIN_REMAINING_SECS: i64 = 3630;
const OAUTH_TOKEN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
const OAUTH_TOKEN_RETRY_ATTEMPTS: u32 = 3;
const OAUTH_TOKEN_RETRY_BASE: std::time::Duration = std::time::Duration::from_millis(250);
#[derive(Debug, Deserialize)]
struct TokenResponse {
access_token: SecretString,
refresh_token: Option<SecretString>,
id_token: Option<SecretString>,
token_type: String,
expires_in: u64,
scope: Option<String>,
}
impl TokenResponse {
fn into_token_set(self, issued_at: DateTime<Utc>) -> TokenSet {
TokenSet::from_secrets(
self.access_token,
self.refresh_token.unwrap_or_else(|| SecretString::from("")),
self.id_token.unwrap_or_else(|| SecretString::from("")),
self.token_type,
i64::try_from(self.expires_in).unwrap_or(i64::MAX),
self.scope.unwrap_or_default(),
issued_at,
)
}
}
async fn read_token_response(mut response: reqwest::Response) -> Result<TokenResponse, String> {
let limit = token_store::TOKEN_JSON_CAPACITY;
let mut body = zeroize::Zeroizing::new(Vec::with_capacity(limit));
while let Some(chunk) = response.chunk().await.map_err(|e| e.to_string())? {
if body.len() + chunk.len() > limit {
return Err(format!("token response is over {limit} bytes"));
}
body.extend_from_slice(&chunk);
}
serde_json::from_slice(&body).map_err(|e| {
format!(
"{:?} error at line {} column {}",
e.classify(),
e.line(),
e.column()
)
})
}
#[derive(Clone)]
pub struct AuthManager {
config: OAuthConfig,
tokens: Arc<RwLock<Option<TokenSet>>>,
client: Arc<RwLock<reqwest::Client>>,
refresh_lock: Arc<Mutex<()>>,
refresh_task: Arc<RwLock<Option<tokio::task::JoinHandle<()>>>>,
pkce_verifier: Arc<RwLock<Option<String>>>,
oauth_state: Arc<RwLock<Option<(String, std::time::Instant)>>>,
token_store: Arc<dyn TokenStore + Send + Sync>,
last_notified_rt_hour: Arc<RwLock<Option<i64>>>,
reauthorization: Arc<RwLock<Option<Refused>>>,
}
impl std::fmt::Debug for AuthManager {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AuthManager")
.field("config", &self.config)
.field("tokens", &"<tokens>")
.field("client", &"<client>")
.field("refresh_lock", &"<mutex>")
.field("refresh_task", &self.refresh_task.read().is_some())
.field("pkce", &self.config.pkce_enabled)
.finish()
}
}
#[derive(Debug, Clone, Copy)]
enum OAuthGrant {
AuthCode,
Refresh,
}
impl OAuthGrant {
fn grant_type(self) -> &'static str {
match self {
OAuthGrant::AuthCode => "authorization_code",
OAuthGrant::Refresh => "refresh_token",
}
}
fn wrap(self, msg: impl Into<String>) -> AuthError {
match self {
OAuthGrant::AuthCode => AuthError::OAuthFlow(msg.into()),
OAuthGrant::Refresh => AuthError::RefreshFailed(msg.into()),
}
}
}
async fn post_oauth_token(
client: &reqwest::Client,
config: &OAuthConfig,
form: &[(&str, String)],
grant: OAuthGrant,
) -> AuthResult<TokenResponse> {
let result = post_oauth_token_with_retry(client, config, form, grant).await;
metrics::counter!(
crate::metrics::TOKEN_REFRESH_TOTAL,
"outcome" => if result.is_ok() { "success" } else { "failure" },
"grant" => grant.grant_type(),
)
.increment(1);
result
}
async fn post_oauth_token_with_retry(
client: &reqwest::Client,
config: &OAuthConfig,
form: &[(&str, String)],
grant: OAuthGrant,
) -> AuthResult<TokenResponse> {
let auth_header = create_basic_header(&config.app_key, &config.app_secret);
let mut last_err: Option<AuthError> = None;
for attempt in 1..=OAUTH_TOKEN_RETRY_ATTEMPTS {
let result = client
.post(&config.token_url)
.header("Authorization", &auth_header)
.header("Content-Type", "application/x-www-form-urlencoded")
.form(form)
.send()
.await;
let response = match result {
Ok(r) => r,
Err(e) => {
let err = grant.wrap(format!("oauth request error: {e}"));
last_err = Some(err);
if attempt < OAUTH_TOKEN_RETRY_ATTEMPTS {
backoff(attempt).await;
continue;
}
break;
}
};
let status = response.status();
if status.is_success() {
return read_token_response(response)
.await
.map_err(|e| grant.wrap(format!("failed to parse token response: {e}")));
}
if status.is_server_error() {
warn!(
"oauth token endpoint returned {} on attempt {}/{}",
status, attempt, OAUTH_TOKEN_RETRY_ATTEMPTS
);
last_err = Some(grant.wrap(format!("oauth token endpoint returned {status}")));
if attempt < OAUTH_TOKEN_RETRY_ATTEMPTS {
backoff(attempt).await;
continue;
}
break;
}
if matches!(grant, OAuthGrant::Refresh)
&& matches!(
status,
reqwest::StatusCode::BAD_REQUEST | reqwest::StatusCode::UNAUTHORIZED
)
{
return Err(AuthError::ReauthorizationRequired(format!(
"Schwab refused the refresh token ({status}); it was revoked, replaced \
by a newer authorization, or expired"
)));
}
return Err(grant.wrap(format!("oauth token endpoint returned {status}")));
}
Err(last_err.unwrap_or_else(|| grant.wrap("oauth token request failed")))
}
async fn backoff(attempt: u32) {
let base = OAUTH_TOKEN_RETRY_BASE * 2u32.saturating_pow(attempt - 1);
let jitter_ms: u64 = (rand::random::<u64>() % 100) + 1;
tokio::time::sleep(base + std::time::Duration::from_millis(jitter_ms)).await;
}
pub fn open_token_store(config: &OAuthConfig) -> Arc<dyn TokenStore + Send + Sync> {
match config.token_store_kind {
TokenStoreKind::File => {
tracing::info!("Using FileTokenStore (basic security: 0600 permissions only)");
Arc::new(FileTokenStore::new(config.tokens_file.clone()))
}
TokenStoreKind::EncryptedFile => {
tracing::info!(
"Using EncryptedFileTokenStore (ChaCha20Poly1305 encryption + 0600 permissions)"
);
Arc::new(EncryptedFileTokenStore::new(config.tokens_file.clone()))
}
TokenStoreKind::Keychain => {
tracing::info!("Using KeychainTokenStore (OS-native credential storage)");
Arc::new(KeychainTokenStore::new(
"rustyschwab".to_string(),
config.app_key.clone(),
))
}
}
}
impl AuthManager {
pub fn new(config: OAuthConfig) -> AuthResult<Self> {
if config.app_key.is_empty() {
return Err(AuthError::MissingConfig(
"App key cannot be empty".to_string(),
));
}
if config.app_secret.is_empty() {
return Err(AuthError::MissingConfig(
"App secret cannot be empty".to_string(),
));
}
for (url, field) in [
(&config.token_url, "token_url"),
(&config.auth_url, "auth_url"),
] {
crate::config::SchwabConfig::validate_url_domain(url, field)
.map_err(|e| AuthError::MissingConfig(e.to_string()))?;
}
let key_len = config.app_key.len();
let secret_len = config.app_secret.len();
if !key_len.is_multiple_of(2) || !secret_len.is_multiple_of(2) || key_len + secret_len < 32
{
return Err(AuthError::MissingConfig(format!(
"App key or app secret likely invalid: lengths must be even and total at \
least 32 chars (got key={key_len}, secret={secret_len})"
)));
}
if !config.callback_url.starts_with("https") {
return Err(AuthError::MissingConfig(
"Callback URL must be HTTPS".to_string(),
));
}
if config.callback_url.ends_with('/') {
return Err(AuthError::MissingConfig(
"Callback URL cannot be path (ends with '/')".to_string(),
));
}
if config.callback_cert_file.is_some() != config.callback_key_file.is_some() {
return Err(AuthError::MissingConfig(
"Set both callback_cert_file and callback_key_file, or neither".to_string(),
));
}
let manager = Self {
config: config.clone(),
tokens: Arc::new(RwLock::new(None)),
client: Arc::new(RwLock::new(
reqwest::Client::builder()
.timeout(OAUTH_TOKEN_TIMEOUT)
.redirect(reqwest::redirect::Policy::none())
.retry(reqwest::retry::never())
.build()
.map_err(|e| AuthError::OAuthFlow(e.to_string()))?,
)),
refresh_lock: Arc::new(Mutex::new(())),
refresh_task: Arc::new(RwLock::new(None)),
pkce_verifier: Arc::new(RwLock::new(None)),
oauth_state: Arc::new(RwLock::new(None)),
last_notified_rt_hour: Arc::new(RwLock::new(None)),
reauthorization: Arc::new(RwLock::new(None)),
token_store: open_token_store(&config),
};
match manager.load_tokens() {
Ok(tokens) => {
*manager.tokens.write() = Some(tokens);
}
Err(e) if token_store::is_absent(&e) => {
info!("No stored tokens; authorization is required");
}
Err(e) => {
warn!(error = %e, "Could not load stored tokens; authorization may be required");
if let Some(ref notify) = config.on_token_notification {
notify(TokenNotification::TokenFileCorrupted);
}
}
}
Ok(manager)
}
pub fn config(&self) -> &OAuthConfig {
&self.config
}
pub async fn start(&self) -> AuthResult<()> {
if let Ok(tokens) = self.store_load().await {
*self.tokens.write() = Some(tokens);
}
if self.config.auto_refresh {
let already_running = self
.refresh_task
.read()
.as_ref()
.map(|h| !h.is_finished())
.unwrap_or(false);
if !already_running {
self.start_refresh_task().await;
}
}
Ok(())
}
pub async fn get_access_token(&self) -> AuthResult<String> {
{
let token_guard = self.tokens.read();
if let Some(tokens) = token_guard.as_ref()
&& tokens.is_valid()
&& !tokens.needs_refresh(self.config.refresh_buffer_seconds)
{
return tokens.access_token();
}
}
self.ensure_valid_tokens().await?;
self.tokens
.read()
.as_ref()
.ok_or(AuthError::TokenExpired)?
.access_token()
}
pub async fn ensure_valid_tokens(&self) -> AuthResult<()> {
self.refresh_if_needed(None).await
}
pub async fn refresh_rejected_token(&self, rejected: &str) -> AuthResult<()> {
self.refresh_if_needed(Some(rejected)).await
}
async fn refresh_if_needed(&self, rejected: Option<&str>) -> AuthResult<()> {
let _lock = self.refresh_lock.lock().await;
let _process_lock = token_store::refresh_lock(&self.config.tokens_file).await?;
if let Ok(stored) = self.store_load().await {
let mut current = self.tokens.write();
if current
.as_ref()
.is_none_or(|t| stored.issued_at >= t.issued_at)
{
*current = Some(stored);
}
}
let old = self.tokens.read().clone().ok_or_else(|| {
AuthError::ReauthorizationRequired("no tokens are stored; authorize first".into())
})?;
if let Some((issued_at, reason)) = self.reauthorization.read().as_ref()
&& *issued_at == old.issued_at
{
return Err(AuthError::ReauthorizationRequired(reason.clone()));
}
if let Some(rejected) = rejected {
if old.access_token()? != rejected && old.is_valid() {
return Ok(());
}
} else if old.is_valid() && !old.needs_refresh(self.config.refresh_buffer_seconds) {
return Ok(());
}
let remaining = (old
.refresh_token_expires_at
.unwrap_or(old.issued_at + Duration::days(7))
- Utc::now())
.num_seconds();
if remaining <= 0 {
return Err(AuthError::ReauthorizationRequired(
"the refresh token expired (seven days after authorization)".into(),
));
}
if remaining <= REFRESH_TOKEN_MIN_REMAINING_SECS {
return Err(AuthError::RefreshTokenExpiring {
seconds_remaining: remaining,
});
}
let refresh_token = old.refresh_token()?;
if refresh_token.is_empty() {
return Err(AuthError::MissingRefreshToken);
}
let mut new_tokens = match self.refresh_tokens(&refresh_token).await {
Ok(tokens) => tokens,
Err(AuthError::ReauthorizationRequired(reason)) => {
*self.reauthorization.write() = Some((old.issued_at, reason.clone()));
return Err(AuthError::ReauthorizationRequired(reason));
}
Err(error) => return Err(error),
};
new_tokens.refresh_token_expires_at = Some(
old.refresh_token_expires_at
.unwrap_or(old.issued_at + Duration::days(7)),
);
self.store_save(new_tokens.clone()).await?;
*self.tokens.write() = Some(new_tokens);
Ok(())
}
pub async fn authorize(&self) -> AuthResult<(String, String)> {
self.authorize_with(|url| info!("Open this URL in your browser:\n{url}"))
.await
}
pub async fn authorize_with(
&self,
on_url: impl FnOnce(&str) + Send,
) -> AuthResult<(String, String)> {
let state = uuid::Uuid::new_v4().to_string();
*self.oauth_state.write() = Some((state.clone(), std::time::Instant::now()));
let pending = serde_json::to_vec(&PendingState {
state_sha256: state_digest(&state),
created_at: Utc::now().timestamp_millis(),
})
.map_err(|e| AuthError::OAuthFlow(e.to_string()))?;
let path = self.pending_state_path();
tokio::task::spawn_blocking(move || crate::security::secure_file_write(&path, &pending))
.await
.map_err(|e| AuthError::TokenStorage(std::io::Error::other(e)))?
.map_err(|e| AuthError::OAuthFlow(format!("Cannot record the OAuth state: {e}")))?;
debug!("Generated OAuth state parameter for CSRF protection");
let mut url = url::Url::parse(&self.config.auth_url)
.map_err(|_| AuthError::OAuthFlow("Invalid authorization URL".into()))?;
url.query_pairs_mut()
.append_pair("client_id", &self.config.app_key)
.append_pair("redirect_uri", &self.config.callback_url)
.append_pair("state", &state);
if self.config.pkce_enabled {
let verifier = Self::generate_code_verifier()?;
let challenge = Self::code_challenge_s256(&verifier);
*self.pkce_verifier.write() = Some(verifier);
url.query_pairs_mut()
.append_pair("code_challenge", &challenge)
.append_pair("code_challenge_method", "S256");
}
let auth_url = url.to_string();
on_url(&auth_url);
if self.config.capture_callback {
let auth_code = self.capture_callback_code().await?;
let tokens = self.exchange_code_for_tokens(&auth_code).await?;
*self.tokens.write() = Some(tokens.clone());
self.store_save(tokens).await?;
Ok((auth_url, auth_code))
} else {
Ok((auth_url, String::new()))
}
}
pub async fn exchange_code(&self, code: String) -> AuthResult<TokenSet> {
if code.trim().is_empty() {
return Err(AuthError::OAuthFlow(
"Authorization code is empty".to_string(),
));
}
let code = code.trim();
let extracted_code = if code.starts_with("https://") || code.starts_with("http://") {
self.extract_code_from_callback_url(code)?
} else {
urlencoding::decode(code)
.map_err(|_| {
AuthError::OAuthFlow("Authorization code is not valid UTF-8".to_string())
})?
.into_owned()
};
let code_without_at = extracted_code.trim_end_matches('@');
if code_without_at
.chars()
.any(|c| c.is_whitespace() || c.is_control() || "&=?#/".contains(c))
{
return Err(AuthError::OAuthFlow(
"Invalid authorization code format: paste the whole redirect URL or only the \
value of its `code` parameter"
.to_string(),
));
}
if code_without_at.len() < 10 {
return Err(AuthError::OAuthFlow(
"Authorization code too short - may be truncated or invalid".to_string(),
));
}
self.exchange_code_for_tokens(&extracted_code).await
}
fn consume_callback_state(&self, supplied: Option<&str>) -> AuthResult<()> {
let refused = || {
AuthError::OAuthFlow(
"Missing, invalid, expired, or already consumed OAuth state".into(),
)
};
let supplied = supplied.ok_or_else(refused)?;
let path = self.pending_state_path();
{
let mut pending = self.oauth_state.write();
if let Some((expected, created)) = pending.as_ref()
&& supplied == expected
&& created.elapsed() < AUTHORIZATION_TTL
{
*pending = None;
let _ = std::fs::remove_file(&path);
return Ok(());
}
}
let stored: PendingState = crate::security::read_private_file(&path)
.ok()
.and_then(|bytes| serde_json::from_slice(&bytes).ok())
.ok_or_else(refused)?;
let age = Utc::now().timestamp_millis() - stored.created_at;
if stored.state_sha256 != state_digest(supplied)
|| !(0..AUTHORIZATION_TTL.as_millis() as i64).contains(&age)
{
return Err(refused());
}
std::fs::remove_file(&path).map_err(|_| refused())?;
Ok(())
}
fn pending_state_path(&self) -> PathBuf {
self.config.tokens_file.with_extension("oauth-state")
}
fn extract_code_from_callback_url(&self, url: &str) -> AuthResult<String> {
let parsed = url::Url::parse(url)
.map_err(|e| AuthError::InvalidCallbackUrl(format!("Invalid callback URL: {}", e)))?;
let callback_state = parsed
.query_pairs()
.find(|(key, _)| key == "state")
.map(|(_, value)| value.into_owned());
self.consume_callback_state(callback_state.as_deref())?;
let code = parsed
.query_pairs()
.find(|(key, _)| key == "code")
.map(|(_, value)| value.into_owned())
.ok_or_else(|| {
AuthError::OAuthFlow("Callback URL missing 'code' parameter".to_string())
})?;
let decoded = if code.ends_with('@') {
code
} else {
format!("{}@", code)
};
Ok(decoded)
}
async fn exchange_code_for_tokens(&self, code: &str) -> AuthResult<TokenSet> {
let code_with_at = if !code.ends_with('@') {
format!("{}@", code)
} else {
code.to_string()
};
debug!("Exchanging code for tokens");
debug!("Token URL: {}", self.config.token_url);
debug!("Redirect URI: {}", self.config.callback_url);
debug!(
"Code length: {}, ends with @: {}",
code_with_at.len(),
code_with_at.ends_with('@')
);
debug!(
"App Key length: {}, Secret length: {}",
self.config.app_key.len(),
self.config.app_secret.len()
);
let mut form_params: Vec<(&str, String)> = vec![
("grant_type", OAuthGrant::AuthCode.grant_type().to_string()),
("code", code_with_at.clone()),
("redirect_uri", self.config.callback_url.clone()),
];
if self.config.pkce_enabled {
let verifier = self.pkce_verifier.write().take().ok_or_else(|| {
AuthError::OAuthFlow(
"PKCE enabled but no code verifier found. Call authorize() first.".to_string(),
)
})?;
form_params.push(("code_verifier", verifier));
}
let client = self.client.read().clone();
let token_response =
post_oauth_token(&client, &self.config, &form_params, OAuthGrant::AuthCode).await?;
let tokens = token_response.into_token_set(Utc::now());
let _file_lock = token_store::refresh_lock(&self.config.tokens_file).await?;
self.store_save(tokens.clone()).await?;
*self.tokens.write() = Some(tokens.clone());
Ok(tokens)
}
async fn refresh_tokens(&self, refresh_token: &str) -> AuthResult<TokenSet> {
debug!("Refreshing access token");
let client = self.client.read().clone();
let mut token_response = post_oauth_token(
&client,
&self.config,
&[
("grant_type", OAuthGrant::Refresh.grant_type().to_string()),
("refresh_token", refresh_token.to_string()),
],
OAuthGrant::Refresh,
)
.await?;
if token_response.refresh_token.is_none() {
token_response.refresh_token = Some(SecretString::from(refresh_token));
}
let tokens = token_response.into_token_set(Utc::now());
info!(
"Token refreshed successfully, expires at: {}",
tokens.expires_at
);
Ok(tokens)
}
async fn start_refresh_task(&self) {
let mut manager = self.clone();
manager.refresh_task = Arc::new(RwLock::new(None));
let task = tokio::spawn(async move {
let mut interval = tokio::time::interval(std::time::Duration::from_secs(30));
let mut failing: Option<String> = None;
loop {
interval.tick().await;
let remaining = manager
.tokens
.read()
.as_ref()
.and_then(|t| t.refresh_token_expires_at)
.map(|t| (t - Utc::now()).num_seconds());
if let Some(remaining) = remaining {
let hour = if remaining <= 0 { -1 } else { remaining / 3600 };
let notify = {
let mut last = manager.last_notified_rt_hour.write();
let changed = remaining <= 43200 && *last != Some(hour);
if changed {
*last = Some(hour);
}
changed
};
if notify && let Some(ref callback) = manager.config.on_token_notification {
callback(if remaining <= 0 {
TokenNotification::RefreshTokenExpired
} else {
TokenNotification::RefreshTokenExpiring {
hours_remaining: hour.max(0),
}
});
}
}
match manager.ensure_valid_tokens().await {
Ok(()) => failing = None,
Err(error) => {
let message = error.to_string();
if failing.as_deref() != Some(message.as_str()) {
if error.requires_reauthorization() {
tracing::error!(%error, "Automatic token refresh needs a new authorization");
} else {
warn!(%error, "Automatic token refresh failed");
}
failing = Some(message);
}
}
}
}
});
*self.refresh_task.write() = Some(task);
}
#[cfg(feature = "callback-server")]
async fn capture_callback_code(&self) -> AuthResult<String> {
use std::sync::Arc;
use tokio::sync::oneshot;
let (tx, rx) = oneshot::channel();
let tx = Arc::new(Mutex::new(Some(tx)));
let callback_handler = {
let tx = tx.clone();
let manager = self.clone();
move |Query(params): Query<CallbackParams>| {
let manager = manager.clone();
async move {
if manager
.consume_callback_state(params.state.as_deref())
.is_err()
{
return Html("Invalid or expired authorization callback".to_string());
}
let outcome = match (params.code, params.error) {
(Some(code), None) if !code.is_empty() => Ok(if code.ends_with('@') {
code
} else {
format!("{code}@")
}),
(_, Some(error)) if error == "access_denied" => {
Err(AuthError::AuthorizationDenied)
}
(_, Some(error)) => Err(AuthError::OAuthFlow(format!(
"Authorization failed: {}",
error.chars().take(64).collect::<String>()
))),
_ => Err(AuthError::OAuthFlow(
"Authorization callback carried no code".into(),
)),
};
let succeeded = outcome.is_ok();
if let Some(tx) = tx.lock().await.take() {
let _ = tx.send(outcome);
}
if !succeeded {
return Html(
"Authorization was not granted; return to your application."
.to_string(),
);
}
Html(
r#"<!DOCTYPE html>
<html>
<head>
<title>Authorization Successful</title>
<style>
body {
font-family: system-ui, -apple-system, sans-serif;
display: flex;
justify-content: center;
align-items: center;
height: 100vh;
margin: 0;
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
}
.container {
background: white;
padding: 2rem;
border-radius: 10px;
box-shadow: 0 10px 40px rgba(0,0,0,0.2);
text-align: center;
}
.success { color: #48bb78; font-size: 3rem; }
</style>
</head>
<body>
<div class="container">
<div class="success">✓</div>
<h1>Authorization Successful!</h1>
<p>You can close this window and return to your application.</p>
</div>
</body>
</html>"#
.to_string(),
)
}
}
};
let callback_path = url::Url::parse(&self.config.callback_url)
.map_err(|_| AuthError::InvalidCallbackUrl("Invalid URL".into()))?;
let app = Router::new().route(callback_path.path(), get(callback_handler));
let callback_url = url::Url::parse(&self.config.callback_url)
.map_err(|e| AuthError::InvalidCallbackUrl(e.to_string()))?;
let port = callback_url.port_or_known_default().unwrap_or(443);
let addr = if self.config.allow_external_callback {
warn!(
"⚠️ SECURITY: OAuth callback server binding to 0.0.0.0:{} \
(all network interfaces). This exposes the callback to your LAN/WAN — \
only enable when terminating TLS in front of this process \
(ngrok / cloudflared / reverse proxy).",
port
);
std::net::SocketAddr::from(([0, 0, 0, 0], port))
} else {
info!(
"OAuth callback server binding to 127.0.0.1:{} (localhost only)",
port
);
std::net::SocketAddr::from(([127, 0, 0, 1], port))
};
if self.config.allow_external_callback {
let listener = tokio::net::TcpListener::bind(addr).await.map_err(|e| {
AuthError::OAuthFlow(format!("Failed to bind to port {}: {}", port, e))
})?;
let server = axum::serve(listener, app);
let wait = async {
tokio::select! {
code = rx => code.map_err(|_| {
AuthError::OAuthFlow("Failed to receive authorization code".to_string())
})?,
result = server => Err(AuthError::OAuthFlow(format!(
"Callback server stopped unexpectedly: {result:?}"
))),
}
};
tokio::time::timeout(AUTHORIZATION_TTL, wait)
.await
.unwrap_or_else(|_| {
Err(AuthError::OAuthFlow(
"Timed out waiting for the authorization callback".into(),
))
})
} else {
let (chain, key) = callback_tls_identity(&self.config)?;
let provider = Arc::new(rustls::crypto::ring::default_provider());
let server_config = rustls::ServerConfig::builder_with_provider(provider)
.with_safe_default_protocol_versions()
.map_err(|e| AuthError::OAuthFlow(format!("TLS config error: {e}")))?
.with_no_client_auth()
.with_single_cert(chain, key)
.map_err(|e| AuthError::OAuthFlow(format!("TLS config error: {e}")))?;
let rustls_config =
axum_server::tls_rustls::RustlsConfig::from_config(Arc::new(server_config));
let server =
axum_server::bind_rustls(addr, rustls_config).serve(app.into_make_service());
let wait = async {
tokio::select! {
code = rx => code.map_err(|_| {
AuthError::OAuthFlow("Failed to receive authorization code".to_string())
})?,
result = server => Err(AuthError::OAuthFlow(format!(
"Callback server stopped unexpectedly: {result:?}"
))),
}
};
tokio::time::timeout(AUTHORIZATION_TTL, wait)
.await
.unwrap_or_else(|_| {
Err(AuthError::OAuthFlow(
"Timed out waiting for the authorization callback".into(),
))
})
}
}
#[cfg(not(feature = "callback-server"))]
async fn capture_callback_code(&self) -> AuthResult<String> {
Err(AuthError::OAuthFlow(
"Callback server not enabled. Enable the 'callback-server' feature \
or manually copy the authorization code."
.to_string(),
))
}
async fn store_load(&self) -> AuthResult<TokenSet> {
let store = self.token_store.clone();
tokio::task::spawn_blocking(move || store.load())
.await
.map_err(|e| AuthError::TokenStorage(std::io::Error::other(e)))?
}
async fn store_save(&self, tokens: TokenSet) -> AuthResult<()> {
let store = self.token_store.clone();
tokio::task::spawn_blocking(move || store.save(&tokens))
.await
.map_err(|e| AuthError::TokenStorage(std::io::Error::other(e)))??;
debug!("Tokens saved");
Ok(())
}
fn load_tokens(&self) -> AuthResult<TokenSet> {
let tokens = self.token_store.load()?;
debug!("Tokens loaded");
Ok(tokens)
}
pub fn has_valid_tokens(&self) -> bool {
self.tokens
.read()
.as_ref()
.map(|t| t.is_valid())
.unwrap_or(false)
}
pub fn get_http_client(&self) -> Arc<RwLock<reqwest::Client>> {
self.client.clone()
}
fn generate_code_verifier() -> AuthResult<String> {
let rng = ring_rand::SystemRandom::new();
let mut bytes = [0u8; 32];
rng.fill(&mut bytes).map_err(|_| {
AuthError::OAuthFlow(
"Failed to generate PKCE verifier - cryptographic RNG unavailable".to_string(),
)
})?;
Ok(BASE64_URL.encode(bytes))
}
fn code_challenge_s256(verifier: &str) -> String {
let digest = digest::digest(&digest::SHA256, verifier.as_bytes());
BASE64_URL.encode(digest.as_ref())
}
}
impl Drop for AuthManager {
fn drop(&mut self) {
if Arc::strong_count(&self.refresh_task) == 1
&& let Some(task) = self.refresh_task.write().take()
{
task.abort();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn config_with_creds(key_len: usize, secret_len: usize) -> OAuthConfig {
OAuthConfig {
app_key: "k".repeat(key_len),
app_secret: "s".repeat(secret_len),
callback_url: "https://127.0.0.1:8080".to_string(),
token_store_kind: TokenStoreKind::File,
tokens_file: std::env::temp_dir().join("schwab_auth_validation_test.json"),
..Default::default()
}
}
#[cfg(feature = "callback-server")]
#[test]
fn callback_server_uses_a_configured_certificate() {
let dir = tempfile::tempdir().unwrap();
let issued = rcgen::generate_simple_self_signed(vec!["127.0.0.1".to_string()]).unwrap();
let (cert_file, key_file) = (dir.path().join("cert.pem"), dir.path().join("key.pem"));
std::fs::write(&cert_file, issued.cert.pem()).unwrap();
std::fs::write(&key_file, issued.key_pair.serialize_pem()).unwrap();
let mut config = config_with_creds(32, 16);
let (chain, _) = callback_tls_identity(&config).unwrap();
assert_ne!(
chain[0].as_ref(),
issued.cert.der().as_ref(),
"generated when unset"
);
config.callback_cert_file = Some(cert_file.clone());
assert!(
AuthManager::new(config.clone()).is_err(),
"a cert without a key"
);
config.callback_key_file = Some(key_file);
assert!(AuthManager::new(config.clone()).is_ok());
let (chain, _) = callback_tls_identity(&config).unwrap();
assert_eq!(chain.len(), 1);
assert_eq!(chain[0].as_ref(), issued.cert.der().as_ref());
std::fs::write(&cert_file, "").unwrap();
let error = callback_tls_identity(&config).unwrap_err().to_string();
assert!(error.contains("cert.pem"), "{error}");
}
#[test]
fn credential_validation_matches_upstream_rule() {
assert!(AuthManager::new(config_with_creds(32, 16)).is_ok());
assert!(AuthManager::new(config_with_creds(48, 64)).is_ok());
assert!(AuthManager::new(config_with_creds(28, 8)).is_ok());
assert!(AuthManager::new(config_with_creds(31, 16)).is_err());
assert!(AuthManager::new(config_with_creds(32, 15)).is_err());
assert!(AuthManager::new(config_with_creds(16, 14)).is_err());
}
}
#[cfg(test)]
mod callback_regressions {
use super::*;
#[tokio::test]
async fn callback_state_is_required_expires_and_is_single_use() {
let dir = tempfile::tempdir().unwrap();
let manager = AuthManager::new(OAuthConfig {
app_key: "k".repeat(32),
app_secret: "s".repeat(16),
token_store_kind: TokenStoreKind::File,
tokens_file: dir.path().join("tokens.json"),
capture_callback: false,
auto_refresh: false,
..Default::default()
})
.unwrap();
let (url, _) = manager.authorize().await.unwrap();
let state = url::Url::parse(&url)
.unwrap()
.query_pairs()
.find(|(k, _)| k == "state")
.unwrap()
.1
.into_owned();
assert!(manager.consume_callback_state(None).is_err());
assert!(manager.consume_callback_state(Some("wrong")).is_err());
manager.consume_callback_state(Some(&state)).unwrap();
assert!(manager.consume_callback_state(Some(&state)).is_err());
*manager.oauth_state.write() = Some((
state.clone(),
std::time::Instant::now() - std::time::Duration::from_secs(301),
));
assert!(manager.consume_callback_state(Some(&state)).is_err());
}
}