use std::future::Future;
use std::sync::Arc;
use http::{Method, Request as HttpRequest};
use reqwest::{Request, Response};
use crate::client::injector::InteractionCallback;
use crate::client::reqwest::deferred::{AgentDeferredOptions, poll_deferred_with};
use crate::client::reqwest::send::SignedSend;
use crate::error::{AAuthError, Result};
use crate::headers::parse_aauth_requirement;
use crate::types::{
AAuthProtocolError, PersonServerMetadata, RequirementLevel, TokenExchangeRequest,
TokenResponseBody,
};
const PREFER_WAIT: u64 = 45;
#[derive(Debug, Clone)]
pub struct TokenExchangeResult {
pub auth_token: String,
pub expires_in: u64,
}
#[derive(Clone)]
pub struct TokenExchangeOptions {
pub(crate) person_server_url: String,
pub(crate) person_server_metadata: Option<PersonServerMetadata>,
pub(crate) on_metadata: Option<Arc<dyn Fn(PersonServerMetadata) + Send + Sync>>,
pub(crate) resource_token: String,
pub(crate) justification: Option<String>,
pub(crate) localhost_callback: Option<String>,
pub(crate) login_hint: Option<String>,
pub(crate) tenant: Option<String>,
pub(crate) domain_hint: Option<String>,
pub(crate) capabilities: Option<Vec<String>>,
pub(crate) prompt: Option<String>,
pub(crate) on_interaction: Option<InteractionCallback>,
pub(crate) on_clarification: Option<crate::client::injector::ClarificationCallback>,
pub(crate) max_poll_duration_secs: Option<u64>,
}
#[derive(Clone)]
pub struct TokenExchangeOptionsBuilder {
person_server_url: String,
resource_token: String,
person_server_metadata: Option<PersonServerMetadata>,
on_metadata: Option<Arc<dyn Fn(PersonServerMetadata) + Send + Sync>>,
justification: Option<String>,
localhost_callback: Option<String>,
login_hint: Option<String>,
tenant: Option<String>,
domain_hint: Option<String>,
capabilities: Option<Vec<String>>,
prompt: Option<String>,
on_interaction: Option<InteractionCallback>,
on_clarification: Option<crate::client::injector::ClarificationCallback>,
max_poll_duration_secs: Option<u64>,
}
impl TokenExchangeOptions {
pub fn builder(
person_server_url: impl Into<String>,
resource_token: impl Into<String>,
) -> TokenExchangeOptionsBuilder {
TokenExchangeOptionsBuilder::new(person_server_url, resource_token)
}
}
impl TokenExchangeOptionsBuilder {
pub fn new(
person_server_url: impl Into<String>,
resource_token: impl Into<String>,
) -> Self {
Self {
person_server_url: person_server_url.into(),
resource_token: resource_token.into(),
person_server_metadata: None,
on_metadata: None,
justification: None,
localhost_callback: None,
login_hint: None,
tenant: None,
domain_hint: None,
capabilities: None,
prompt: None,
on_interaction: None,
on_clarification: None,
max_poll_duration_secs: None,
}
}
pub fn person_server_metadata(mut self, metadata: PersonServerMetadata) -> Self {
self.person_server_metadata = Some(metadata);
self
}
pub fn on_metadata(
mut self,
callback: Arc<dyn Fn(PersonServerMetadata) + Send + Sync>,
) -> Self {
self.on_metadata = Some(callback);
self
}
pub fn justification(mut self, justification: impl Into<String>) -> Self {
self.justification = Some(justification.into());
self
}
pub fn localhost_callback(mut self, url: impl Into<String>) -> Self {
self.localhost_callback = Some(url.into());
self
}
pub fn login_hint(mut self, login_hint: impl Into<String>) -> Self {
self.login_hint = Some(login_hint.into());
self
}
pub fn tenant(mut self, tenant: impl Into<String>) -> Self {
self.tenant = Some(tenant.into());
self
}
pub fn domain_hint(mut self, domain_hint: impl Into<String>) -> Self {
self.domain_hint = Some(domain_hint.into());
self
}
pub fn capabilities(mut self, capabilities: Vec<String>) -> Self {
self.capabilities = Some(capabilities);
self
}
pub fn prompt(mut self, prompt: impl Into<String>) -> Self {
self.prompt = Some(prompt.into());
self
}
pub fn on_interaction(mut self, callback: InteractionCallback) -> Self {
self.on_interaction = Some(callback);
self
}
pub fn on_clarification(
mut self,
callback: crate::client::injector::ClarificationCallback,
) -> Self {
self.on_clarification = Some(callback);
self
}
pub fn max_poll_duration_secs(mut self, secs: u64) -> Self {
self.max_poll_duration_secs = Some(secs);
self
}
pub fn build(self) -> TokenExchangeOptions {
TokenExchangeOptions {
person_server_url: self.person_server_url,
resource_token: self.resource_token,
person_server_metadata: self.person_server_metadata,
on_metadata: self.on_metadata,
justification: self.justification,
localhost_callback: self.localhost_callback,
login_hint: self.login_hint,
tenant: self.tenant,
domain_hint: self.domain_hint,
capabilities: self.capabilities,
prompt: self.prompt,
on_interaction: self.on_interaction,
on_clarification: self.on_clarification,
max_poll_duration_secs: self.max_poll_duration_secs,
}
}
}
#[derive(Debug, Clone)]
pub struct TokenExchangeError {
pub status: u16,
pub aauth_error: Option<AAuthProtocolError>,
}
impl std::fmt::Display for TokenExchangeError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
if let Some(err) = &self.aauth_error {
write!(
f,
"{}",
err.error_description.as_deref().unwrap_or(&err.error)
)
} else {
write!(f, "Token exchange failed with status {}", self.status)
}
}
}
impl std::error::Error for TokenExchangeError {}
pub async fn exchange_token<F, Fut>(
options: TokenExchangeOptions,
send: F,
) -> Result<TokenExchangeResult>
where
F: FnMut(Request) -> Fut + Send,
Fut: Future<Output = Result<Response>> + Send,
{
struct Adapter<F>(F);
#[async_trait::async_trait]
impl<F, Fut> SignedSend for Adapter<F>
where
F: FnMut(Request) -> Fut + Send,
Fut: Future<Output = Result<Response>> + Send,
{
async fn send(&mut self, req: Request) -> Result<Response> {
(self.0)(req).await
}
}
exchange_token_with(options, &mut Adapter(send)).await
}
pub(crate) async fn exchange_token_with<S: SignedSend>(
options: TokenExchangeOptions,
send: &mut S,
) -> Result<TokenExchangeResult> {
let metadata = if let Some(metadata) = options.person_server_metadata.clone() {
metadata
} else {
let metadata = fetch_metadata(&options.person_server_url, send).await?;
if let Some(on_metadata) = &options.on_metadata {
on_metadata(metadata.clone());
}
metadata
};
let body = TokenExchangeRequest {
resource_token: options.resource_token,
justification: options.justification,
localhost_callback: options.localhost_callback,
login_hint: options.login_hint,
tenant: options.tenant,
domain_hint: options.domain_hint,
capabilities: options.capabilities,
prompt: options.prompt,
};
let token_body =
serde_json::to_string(&body).map_err(|e| AAuthError::Message(e.to_string()))?;
let http_req = HttpRequest::builder()
.method(Method::POST)
.uri(&metadata.token_endpoint)
.header("content-type", "application/json")
.header("prefer", format!("wait={PREFER_WAIT}"))
.body(token_body.into_bytes())
.expect("valid http request");
let response = send
.send(Request::try_from(http_req).expect("valid reqwest request"))
.await?;
if response.status().as_u16() == 202 {
let location = response
.headers()
.get("location")
.and_then(|v| v.to_str().ok())
.ok_or_else(|| AAuthError::Message("202 response missing Location header".into()))?
.to_string();
let mut deferred = AgentDeferredOptions::builder(resolve_url(
&options.person_server_url,
&location,
));
if let Some(header) = response
.headers()
.get("aauth-requirement")
.and_then(|v| v.to_str().ok())
{
if let Ok(challenge) = parse_aauth_requirement(header) {
if challenge.requirement == RequirementLevel::Interaction {
if let (Some(url), Some(code)) = (challenge.url, challenge.code) {
deferred = deferred.interaction(url, code);
}
}
}
}
if let Some(cb) = options.on_interaction {
deferred = deferred.on_interaction(cb);
}
if let Some(cb) = options.on_clarification {
deferred = deferred.on_clarification(cb);
}
if let Some(secs) = options.max_poll_duration_secs {
deferred = deferred.max_poll_duration_secs(secs);
}
let result = poll_deferred_with(deferred.build(), send).await?;
if result.response.status().is_success() {
let parsed: TokenResponseBody = result
.response
.json()
.await
.map_err(|e| AAuthError::Message(e.to_string()))?;
return Ok(TokenExchangeResult {
auth_token: parsed.auth_token,
expires_in: parsed.expires_in,
});
}
return Err(AAuthError::Message(
TokenExchangeError {
status: result.response.status().as_u16(),
aauth_error: result.error,
}
.to_string(),
));
}
if response.status().is_success() {
let parsed: TokenResponseBody = response
.json()
.await
.map_err(|e| AAuthError::Message(e.to_string()))?;
return Ok(TokenExchangeResult {
auth_token: parsed.auth_token,
expires_in: parsed.expires_in,
});
}
Err(AAuthError::Message(
TokenExchangeError {
status: response.status().as_u16(),
aauth_error: None,
}
.to_string(),
))
}
async fn fetch_metadata<S: SignedSend>(
person_server_url: &str,
send: &mut S,
) -> Result<PersonServerMetadata> {
let metadata_url = format!(
"{}/.well-known/aauth-person.json",
person_server_url.trim_end_matches('/')
);
let http_req = HttpRequest::builder()
.method(Method::GET)
.uri(&metadata_url)
.body(Vec::new())
.expect("valid http request");
let response = send
.send(Request::try_from(http_req).expect("valid reqwest request"))
.await?;
if !response.status().is_success() {
return Err(AAuthError::Message(format!(
"Failed to fetch person server metadata: {}",
response.status()
)));
}
let metadata: PersonServerMetadata = response
.json()
.await
.map_err(|e| AAuthError::Message(e.to_string()))?;
metadata.validate().map_err(AAuthError::Message)?;
Ok(metadata)
}
fn resolve_url(base: &str, url: &str) -> String {
if url.starts_with("http://") || url.starts_with("https://") {
url.to_string()
} else {
url::Url::parse(base)
.and_then(|b| b.join(url))
.map(|u| u.to_string())
.unwrap_or_else(|_| url.to_string())
}
}