1use std::future::Future;
2use std::sync::Arc;
3
4use http::{Method, Request as HttpRequest};
5use reqwest::{Request, Response};
6
7use crate::client::injector::InteractionCallback;
8use crate::client::reqwest::deferred::{AgentDeferredOptions, poll_deferred_with};
9use crate::client::reqwest::send::SignedSend;
10use crate::error::{AAuthError, Result};
11use crate::headers::parse_aauth_requirement;
12use crate::types::{
13 AAuthProtocolError, PersonServerMetadata, RequirementLevel, TokenExchangeRequest,
14 TokenResponseBody,
15};
16
17const PREFER_WAIT: u64 = 45;
18
19#[derive(Debug, Clone)]
20pub struct TokenExchangeResult {
21 pub auth_token: String,
22 pub expires_in: u64,
23}
24
25#[derive(Clone)]
26pub struct TokenExchangeOptions {
27 pub(crate) person_server_url: String,
28 pub(crate) person_server_metadata: Option<PersonServerMetadata>,
29 pub(crate) on_metadata: Option<Arc<dyn Fn(PersonServerMetadata) + Send + Sync>>,
30 pub(crate) resource_token: String,
31 pub(crate) justification: Option<String>,
32 pub(crate) localhost_callback: Option<String>,
33 pub(crate) login_hint: Option<String>,
34 pub(crate) tenant: Option<String>,
35 pub(crate) domain_hint: Option<String>,
36 pub(crate) capabilities: Option<Vec<String>>,
37 pub(crate) prompt: Option<String>,
38 pub(crate) on_interaction: Option<InteractionCallback>,
39 pub(crate) on_clarification: Option<crate::client::injector::ClarificationCallback>,
40 pub(crate) max_poll_duration_secs: Option<u64>,
41}
42
43#[derive(Clone)]
44pub struct TokenExchangeOptionsBuilder {
45 person_server_url: String,
46 resource_token: String,
47 person_server_metadata: Option<PersonServerMetadata>,
48 on_metadata: Option<Arc<dyn Fn(PersonServerMetadata) + Send + Sync>>,
49 justification: Option<String>,
50 localhost_callback: Option<String>,
51 login_hint: Option<String>,
52 tenant: Option<String>,
53 domain_hint: Option<String>,
54 capabilities: Option<Vec<String>>,
55 prompt: Option<String>,
56 on_interaction: Option<InteractionCallback>,
57 on_clarification: Option<crate::client::injector::ClarificationCallback>,
58 max_poll_duration_secs: Option<u64>,
59}
60
61impl TokenExchangeOptions {
62 pub fn builder(
63 person_server_url: impl Into<String>,
64 resource_token: impl Into<String>,
65 ) -> TokenExchangeOptionsBuilder {
66 TokenExchangeOptionsBuilder::new(person_server_url, resource_token)
67 }
68}
69
70impl TokenExchangeOptionsBuilder {
71 pub fn new(
72 person_server_url: impl Into<String>,
73 resource_token: impl Into<String>,
74 ) -> Self {
75 Self {
76 person_server_url: person_server_url.into(),
77 resource_token: resource_token.into(),
78 person_server_metadata: None,
79 on_metadata: None,
80 justification: None,
81 localhost_callback: None,
82 login_hint: None,
83 tenant: None,
84 domain_hint: None,
85 capabilities: None,
86 prompt: None,
87 on_interaction: None,
88 on_clarification: None,
89 max_poll_duration_secs: None,
90 }
91 }
92
93 pub fn person_server_metadata(mut self, metadata: PersonServerMetadata) -> Self {
94 self.person_server_metadata = Some(metadata);
95 self
96 }
97
98 pub fn on_metadata(
99 mut self,
100 callback: Arc<dyn Fn(PersonServerMetadata) + Send + Sync>,
101 ) -> Self {
102 self.on_metadata = Some(callback);
103 self
104 }
105
106 pub fn justification(mut self, justification: impl Into<String>) -> Self {
107 self.justification = Some(justification.into());
108 self
109 }
110
111 pub fn localhost_callback(mut self, url: impl Into<String>) -> Self {
112 self.localhost_callback = Some(url.into());
113 self
114 }
115
116 pub fn login_hint(mut self, login_hint: impl Into<String>) -> Self {
117 self.login_hint = Some(login_hint.into());
118 self
119 }
120
121 pub fn tenant(mut self, tenant: impl Into<String>) -> Self {
122 self.tenant = Some(tenant.into());
123 self
124 }
125
126 pub fn domain_hint(mut self, domain_hint: impl Into<String>) -> Self {
127 self.domain_hint = Some(domain_hint.into());
128 self
129 }
130
131 pub fn capabilities(mut self, capabilities: Vec<String>) -> Self {
132 self.capabilities = Some(capabilities);
133 self
134 }
135
136 pub fn prompt(mut self, prompt: impl Into<String>) -> Self {
137 self.prompt = Some(prompt.into());
138 self
139 }
140
141 pub fn on_interaction(mut self, callback: InteractionCallback) -> Self {
142 self.on_interaction = Some(callback);
143 self
144 }
145
146 pub fn on_clarification(
147 mut self,
148 callback: crate::client::injector::ClarificationCallback,
149 ) -> Self {
150 self.on_clarification = Some(callback);
151 self
152 }
153
154 pub fn max_poll_duration_secs(mut self, secs: u64) -> Self {
155 self.max_poll_duration_secs = Some(secs);
156 self
157 }
158
159 pub fn build(self) -> TokenExchangeOptions {
160 TokenExchangeOptions {
161 person_server_url: self.person_server_url,
162 resource_token: self.resource_token,
163 person_server_metadata: self.person_server_metadata,
164 on_metadata: self.on_metadata,
165 justification: self.justification,
166 localhost_callback: self.localhost_callback,
167 login_hint: self.login_hint,
168 tenant: self.tenant,
169 domain_hint: self.domain_hint,
170 capabilities: self.capabilities,
171 prompt: self.prompt,
172 on_interaction: self.on_interaction,
173 on_clarification: self.on_clarification,
174 max_poll_duration_secs: self.max_poll_duration_secs,
175 }
176 }
177}
178
179#[derive(Debug, Clone)]
180pub struct TokenExchangeError {
181 pub status: u16,
182 pub aauth_error: Option<AAuthProtocolError>,
183}
184
185impl std::fmt::Display for TokenExchangeError {
186 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
187 if let Some(err) = &self.aauth_error {
188 write!(
189 f,
190 "{}",
191 err.error_description.as_deref().unwrap_or(&err.error)
192 )
193 } else {
194 write!(f, "Token exchange failed with status {}", self.status)
195 }
196 }
197}
198
199impl std::error::Error for TokenExchangeError {}
200
201pub async fn exchange_token<F, Fut>(
202 options: TokenExchangeOptions,
203 send: F,
204) -> Result<TokenExchangeResult>
205where
206 F: FnMut(Request) -> Fut + Send,
207 Fut: Future<Output = Result<Response>> + Send,
208{
209 struct Adapter<F>(F);
210
211 #[async_trait::async_trait]
212 impl<F, Fut> SignedSend for Adapter<F>
213 where
214 F: FnMut(Request) -> Fut + Send,
215 Fut: Future<Output = Result<Response>> + Send,
216 {
217 async fn send(&mut self, req: Request) -> Result<Response> {
218 (self.0)(req).await
219 }
220 }
221
222 exchange_token_with(options, &mut Adapter(send)).await
223}
224
225pub(crate) async fn exchange_token_with<S: SignedSend>(
226 options: TokenExchangeOptions,
227 send: &mut S,
228) -> Result<TokenExchangeResult> {
229 let metadata = if let Some(metadata) = options.person_server_metadata.clone() {
230 metadata
231 } else {
232 let metadata = fetch_metadata(&options.person_server_url, send).await?;
233 if let Some(on_metadata) = &options.on_metadata {
234 on_metadata(metadata.clone());
235 }
236 metadata
237 };
238
239 let body = TokenExchangeRequest {
240 resource_token: options.resource_token,
241 justification: options.justification,
242 localhost_callback: options.localhost_callback,
243 login_hint: options.login_hint,
244 tenant: options.tenant,
245 domain_hint: options.domain_hint,
246 capabilities: options.capabilities,
247 prompt: options.prompt,
248 };
249
250 let token_body =
251 serde_json::to_string(&body).map_err(|e| AAuthError::Message(e.to_string()))?;
252
253 let http_req = HttpRequest::builder()
254 .method(Method::POST)
255 .uri(&metadata.token_endpoint)
256 .header("content-type", "application/json")
257 .header("prefer", format!("wait={PREFER_WAIT}"))
258 .body(token_body.into_bytes())
259 .expect("valid http request");
260 let response = send
261 .send(Request::try_from(http_req).expect("valid reqwest request"))
262 .await?;
263
264 if response.status().as_u16() == 202 {
265 let location = response
266 .headers()
267 .get("location")
268 .and_then(|v| v.to_str().ok())
269 .ok_or_else(|| AAuthError::Message("202 response missing Location header".into()))?
270 .to_string();
271
272 let mut deferred = AgentDeferredOptions::builder(resolve_url(
273 &options.person_server_url,
274 &location,
275 ));
276 if let Some(header) = response
277 .headers()
278 .get("aauth-requirement")
279 .and_then(|v| v.to_str().ok())
280 {
281 if let Ok(challenge) = parse_aauth_requirement(header) {
282 if challenge.requirement == RequirementLevel::Interaction {
283 if let (Some(url), Some(code)) = (challenge.url, challenge.code) {
284 deferred = deferred.interaction(url, code);
285 }
286 }
287 }
288 }
289 if let Some(cb) = options.on_interaction {
290 deferred = deferred.on_interaction(cb);
291 }
292 if let Some(cb) = options.on_clarification {
293 deferred = deferred.on_clarification(cb);
294 }
295 if let Some(secs) = options.max_poll_duration_secs {
296 deferred = deferred.max_poll_duration_secs(secs);
297 }
298
299 let result = poll_deferred_with(deferred.build(), send).await?;
300
301 if result.response.status().is_success() {
302 let parsed: TokenResponseBody = result
303 .response
304 .json()
305 .await
306 .map_err(|e| AAuthError::Message(e.to_string()))?;
307 return Ok(TokenExchangeResult {
308 auth_token: parsed.auth_token,
309 expires_in: parsed.expires_in,
310 });
311 }
312
313 return Err(AAuthError::Message(
314 TokenExchangeError {
315 status: result.response.status().as_u16(),
316 aauth_error: result.error,
317 }
318 .to_string(),
319 ));
320 }
321
322 if response.status().is_success() {
323 let parsed: TokenResponseBody = response
324 .json()
325 .await
326 .map_err(|e| AAuthError::Message(e.to_string()))?;
327 return Ok(TokenExchangeResult {
328 auth_token: parsed.auth_token,
329 expires_in: parsed.expires_in,
330 });
331 }
332
333 Err(AAuthError::Message(
334 TokenExchangeError {
335 status: response.status().as_u16(),
336 aauth_error: None,
337 }
338 .to_string(),
339 ))
340}
341
342async fn fetch_metadata<S: SignedSend>(
343 person_server_url: &str,
344 send: &mut S,
345) -> Result<PersonServerMetadata> {
346 let metadata_url = format!(
347 "{}/.well-known/aauth-person.json",
348 person_server_url.trim_end_matches('/')
349 );
350 let http_req = HttpRequest::builder()
351 .method(Method::GET)
352 .uri(&metadata_url)
353 .body(Vec::new())
354 .expect("valid http request");
355 let response = send
356 .send(Request::try_from(http_req).expect("valid reqwest request"))
357 .await?;
358
359 if !response.status().is_success() {
360 return Err(AAuthError::Message(format!(
361 "Failed to fetch person server metadata: {}",
362 response.status()
363 )));
364 }
365
366 let metadata: PersonServerMetadata = response
367 .json()
368 .await
369 .map_err(|e| AAuthError::Message(e.to_string()))?;
370 metadata.validate().map_err(AAuthError::Message)?;
371 Ok(metadata)
372}
373
374fn resolve_url(base: &str, url: &str) -> String {
375 if url.starts_with("http://") || url.starts_with("https://") {
376 url.to_string()
377 } else {
378 url::Url::parse(base)
379 .and_then(|b| b.join(url))
380 .map(|u| u.to_string())
381 .unwrap_or_else(|_| url.to_string())
382 }
383}