1use aep_core::{
2 AUTHORIZATION_SCHEME, AgentStatus, BuiltInGrantResponse, Command, EnrollRequest,
3 EnrollResponse, GrantRequest, GrantType, HttpRequest, MEDIA_TYPE, PROBLEM_MEDIA_TYPE,
4 RevokeRequest, RevokeResponse, StatusResponse, StringBoolean, parse_built_in_grant_response,
5 parse_enroll_response, parse_grant_request, parse_problem_details, parse_revoke_response,
6 parse_status_response, validate_enroll_request, validate_revoke_request,
7};
8use http::{HeaderMap, HeaderValue, Method, header};
9use serde::Serialize;
10use serde_json::Value;
11use time::{OffsetDateTime, format_description::well_known::Rfc3339};
12
13use crate::http_response::media_type_matches;
14use crate::{
15 AgentError, AgentIdentity, AssertionSigner, CommandResult, CredentialRecord, EnrollOptions,
16 GrantOptions, GrantResult, Inspection, OperationKey, RevokeOptions, Session, WaitOptions,
17 assertion_operation,
18};
19
20impl Session {
21 pub async fn enroll(
22 &self,
23 options: EnrollOptions,
24 ) -> Result<CommandResult<EnrollResponse>, AgentError> {
25 let (inspection, identity, signer) = self.command_context(&Command::Enroll, true).await?;
26 let required = inspection
27 .document
28 .claims
29 .as_ref()
30 .map_or(&[][..], |claims| claims.required.as_slice());
31 let missing = aep_core::missing_required_claim_names(required, options.claims.as_ref());
32 if !missing.is_empty() {
33 return Err(AgentError::claims(missing));
34 }
35 let key = self
36 .idempotency_key(
37 &inspection,
38 &Command::Enroll,
39 None,
40 None,
41 options.idempotency_key,
42 )
43 .await?;
44 let body = EnrollRequest {
45 agent_did: identity.agent_did.clone(),
46 claims: options.claims,
47 idempotency_key: Some(key.clone()),
48 additional: Default::default(),
49 };
50 validate_enroll_request(&body).map_err(aep_core::CoreError::from)?;
51 self.execute_command(
52 CommandRequest {
53 inspection: &inspection,
54 identity: &identity,
55 signer: signer.as_ref(),
56 command: &Command::Enroll,
57 method: Method::POST,
58 body: Some(&body),
59 idempotency_key: Some(&key),
60 },
61 parse_enroll_response,
62 )
63 .await
64 }
65
66 pub async fn status(&self) -> Result<CommandResult<StatusResponse>, AgentError> {
67 let (inspection, identity, signer) = self.command_context(&Command::Status, true).await?;
68 self.execute_command::<Value, StatusResponse>(
69 CommandRequest {
70 inspection: &inspection,
71 identity: &identity,
72 signer: signer.as_ref(),
73 command: &Command::Status,
74 method: Method::GET,
75 body: None,
76 idempotency_key: None,
77 },
78 parse_status_response,
79 )
80 .await
81 }
82
83 pub async fn wait_for_active(
84 &self,
85 options: WaitOptions,
86 ) -> Result<CommandResult<StatusResponse>, AgentError> {
87 if options.interval.is_zero() || options.timeout.is_zero() {
88 return Err(AgentError::InvalidConfiguration(
89 "AEP Status polling interval and timeout must be positive".to_owned(),
90 ));
91 }
92 let started = self.client.clock.now();
93 let mut remaining_attempts = options
98 .timeout
99 .as_nanos()
100 .checked_div(options.interval.as_nanos())
101 .unwrap_or_default()
102 .saturating_add(1);
103 loop {
104 let result = self.status().await?;
105 if result.body.status == AgentStatus::Active {
106 return Ok(result);
107 }
108 if matches!(
109 result.body.status,
110 AgentStatus::Rejected | AgentStatus::Suspended | AgentStatus::Terminated
111 ) {
112 return Err(AgentError::EnrollmentState {
113 status: result.body.status,
114 });
115 }
116 let elapsed = self.client.clock.now() - started;
117 let timeout = time::Duration::try_from(options.timeout).map_err(|_| {
118 AgentError::InvalidConfiguration(
119 "AEP Status polling timeout is too large".to_owned(),
120 )
121 })?;
122 remaining_attempts = remaining_attempts.saturating_sub(1);
123 if elapsed >= timeout || remaining_attempts == 0 {
124 return Err(AgentError::PollingTimeout);
125 }
126 self.client.delay.sleep(options.interval).await;
127 }
128 }
129
130 pub async fn grant(
131 &self,
132 options: GrantOptions,
133 ) -> Result<CommandResult<GrantResult>, AgentError> {
134 let (inspection, identity, signer) = self.command_context(&Command::Grant, false).await?;
135 let grant_type = select_grant_type(
136 &inspection,
137 options.grant_type.as_ref(),
138 &options.preferred_grant_types,
139 )?;
140 if !inspection
141 .document
142 .commands
143 .supported
144 .contains(&Command::Status)
145 {
146 return Err(AgentError::CommandNotAdvertised(
147 "status required by grant".to_owned(),
148 ));
149 }
150 let status = self
151 .execute_command::<Value, StatusResponse>(
152 CommandRequest {
153 inspection: &inspection,
154 identity: &identity,
155 signer: signer.as_ref(),
156 command: &Command::Status,
157 method: Method::GET,
158 body: None,
159 idempotency_key: None,
160 },
161 parse_status_response,
162 )
163 .await?;
164 if status.body.status != AgentStatus::Active {
165 return Err(AgentError::Command {
169 status: 401,
170 problem: Some(Box::new(aep_core::new_problem_details(
171 aep_core::ErrorCode::NotRecognized,
172 "Grant requires an active enrollment",
173 401,
174 ))),
175 });
176 }
177 let key = self
178 .idempotency_key(
179 &inspection,
180 &Command::Grant,
181 Some(grant_type.clone()),
182 None,
183 options.idempotency_key,
184 )
185 .await?;
186 let body = GrantRequest {
187 grant_type: grant_type.clone(),
188 requested_scopes: options.requested_scopes,
189 additional: Default::default(),
190 };
191 parse_grant_request(&serde_json::to_vec(&body)?)?;
192 let raw = self
193 .execute_raw(CommandRequest {
194 inspection: &inspection,
195 identity: &identity,
196 signer: signer.as_ref(),
197 command: &Command::Grant,
198 method: Method::POST,
199 body: Some(&body),
200 idempotency_key: Some(&key),
201 })
202 .await?;
203 let value: Value = serde_json::from_slice(&raw.body)?;
204 if !value.is_object() {
205 return Err(AgentError::Credential(
206 "AEP Grant response must be a JSON object".to_owned(),
207 ));
208 }
209 if value
210 .get("credential_id")
211 .and_then(Value::as_str)
212 .is_none_or(str::is_empty)
213 {
214 return Err(AgentError::Credential(
215 "AEP Grant response must contain a non-empty credential_id".to_owned(),
216 ));
217 }
218 let credential = match grant_type {
219 GrantType::OAuthBearer | GrantType::ApiKey | GrantType::Basic => {
220 Some(parse_built_in_grant_response(&grant_type, &raw.body)?)
221 }
222 GrantType::Other(_) => None,
223 };
224 if let Some(credential) = credential.as_ref() {
225 self.client
226 .credential_store
227 .save(credential_record(
228 credential,
229 value.clone(),
230 &inspection,
231 self.client.clock.now(),
232 )?)
233 .await?;
234 }
235 Ok(CommandResult {
236 body: GrantResult {
237 credential,
238 grant_type,
239 raw: value,
240 },
241 status: raw.status,
242 url: raw.url,
243 })
244 }
245
246 pub async fn revoke(
247 &self,
248 options: RevokeOptions,
249 ) -> Result<CommandResult<RevokeResponse>, AgentError> {
250 if options.all_grant_types
251 && (options.grant_type.is_some() || options.credential_id.is_some())
252 {
253 return Err(AgentError::InvalidConfiguration(
254 "AEP all-grant-types Revoke cannot include a grant type or credential ID"
255 .to_owned(),
256 ));
257 }
258 let (inspection, identity, signer) = self.command_context(&Command::Revoke, true).await?;
259 if let Some(grant_type) = options.grant_type.as_ref()
261 && !inspection
262 .document
263 .commands
264 .grant_types
265 .contains(grant_type)
266 {
267 return Err(AgentError::NoCompatibleGrantType);
268 }
269 let body = RevokeRequest {
270 grant_type: options.grant_type.clone(),
271 credential_id: options.credential_id.clone(),
272 all_grant_types: options.all_grant_types.then_some(StringBoolean::True),
273 additional: Default::default(),
274 };
275 validate_revoke_request(&body).map_err(aep_core::CoreError::from)?;
276 if body.credential_id.is_some() {
277 let advertised = body
278 .grant_type
279 .as_ref()
280 .and_then(|grant_type| {
281 inspection
282 .document
283 .commands
284 .grant_types_config
285 .get(grant_type.as_str())
286 })
287 .and_then(|config| config.supports_per_credential_revoke)
288 == Some(StringBoolean::True);
289 if !advertised {
290 return Err(AgentError::Credential(
291 "AEP Service does not advertise per-credential Revoke".to_owned(),
292 ));
293 }
294 }
295 let key = self
296 .idempotency_key(
297 &inspection,
298 &Command::Revoke,
299 body.grant_type.clone(),
300 body.credential_id.clone(),
301 options.idempotency_key,
302 )
303 .await?;
304 let result = self
305 .execute_command(
306 CommandRequest {
307 inspection: &inspection,
308 identity: &identity,
309 signer: signer.as_ref(),
310 command: &Command::Revoke,
311 method: Method::POST,
312 body: Some(&body),
313 idempotency_key: Some(&key),
314 },
315 parse_revoke_response,
316 )
317 .await?;
318 self.delete_revoked_credentials(&inspection.document.service.did, &body)
319 .await?;
320 Ok(result)
321 }
322
323 async fn command_context(
324 &self,
325 command: &Command,
326 create_identity: bool,
327 ) -> Result<
328 (
329 Inspection,
330 AgentIdentity,
331 std::sync::Arc<dyn AssertionSigner>,
332 ),
333 AgentError,
334 > {
335 let inspection = self.inspect().await?;
336 if !inspection.document.commands.supported.contains(command) {
337 return Err(AgentError::CommandNotAdvertised(
338 command.as_str().to_owned(),
339 ));
340 }
341 let identity = self.resolve_identity(&inspection, create_identity).await?;
342 let signer = self.client.identity_provider.signer_for(&identity).await?;
343 Ok((inspection, identity, signer))
344 }
345
346 async fn execute_command<B: Serialize + ?Sized, T>(
347 &self,
348 request: CommandRequest<'_, B>,
349 parser: fn(&[u8]) -> Result<T, aep_core::ParseError>,
350 ) -> Result<CommandResult<T>, AgentError> {
351 let raw = self.execute_raw(request).await?;
352 Ok(CommandResult {
353 body: parser(&raw.body)?,
354 status: raw.status,
355 url: raw.url,
356 })
357 }
358
359 async fn execute_raw<B: Serialize + ?Sized>(
360 &self,
361 request: CommandRequest<'_, B>,
362 ) -> Result<RawCommandResult, AgentError> {
363 let url = request.inspection.command_url(request.command)?;
364 let operation = assertion_operation(request.command)
365 .ok_or_else(|| AgentError::CommandNotAdvertised(request.command.as_str().to_owned()))?;
366 let assertion = self
367 .client
368 .sign_assertion(
369 request.inspection,
370 request.identity,
371 request.signer,
372 operation,
373 None,
374 )
375 .await?;
376 let mut headers = HeaderMap::new();
377 headers.insert(header::ACCEPT, HeaderValue::from_static(MEDIA_TYPE));
378 headers.insert(
379 header::AUTHORIZATION,
380 HeaderValue::from_str(&format!("{AUTHORIZATION_SCHEME} {assertion}")).map_err(
381 |_| {
382 AgentError::Identity("AEP assertion is not a valid HTTP field value".to_owned())
383 },
384 )?,
385 );
386 let encoded = request
387 .body
388 .map(serde_json::to_vec)
389 .transpose()?
390 .unwrap_or_default();
391 if request.body.is_some() {
392 headers.insert(header::CONTENT_TYPE, HeaderValue::from_static(MEDIA_TYPE));
393 }
394 if let Some(key) = request.idempotency_key {
395 headers.insert(
396 "idempotency-key",
397 HeaderValue::from_str(key).map_err(|_| {
398 AgentError::InvalidConfiguration(
399 "AEP idempotency key is not a valid HTTP field value".to_owned(),
400 )
401 })?,
402 );
403 }
404 let response = self
405 .client
406 .command_transport
407 .send(HttpRequest {
408 method: request.method,
409 url: url.clone(),
410 headers,
411 body: encoded,
412 })
413 .await
414 .map_err(|error| AgentError::Transport(error.to_string()))?;
415 if response.final_url != url {
416 return Err(AgentError::Transport(
417 "AEP command redirects are not allowed".to_owned(),
418 ));
419 }
420 if response.body.len() > self.client.maximum_response_bytes {
421 return Err(AgentError::Transport(
422 "AEP command response exceeds the configured limit".to_owned(),
423 ));
424 }
425 if !response.status.is_success() {
426 let problem = media_type_matches(&response.headers, PROBLEM_MEDIA_TYPE)
430 .then(|| parse_problem_details(&response.body).ok())
431 .flatten()
432 .filter(|problem| problem.status == i64::from(response.status.as_u16()))
433 .map(Box::new);
434 return Err(AgentError::Command {
435 status: response.status.as_u16(),
436 problem,
437 });
438 }
439 if !media_type_matches(&response.headers, MEDIA_TYPE) {
440 return Err(AgentError::Transport(
441 "AEP command response media type is invalid".to_owned(),
442 ));
443 }
444 Ok(RawCommandResult {
445 body: response.body,
446 status: response.status.as_u16(),
447 url,
448 })
449 }
450
451 async fn idempotency_key(
452 &self,
453 inspection: &Inspection,
454 command: &Command,
455 grant_type: Option<GrantType>,
456 credential_id: Option<String>,
457 provided: Option<String>,
458 ) -> Result<String, AgentError> {
459 if let Some(value) = provided {
460 if value.trim().is_empty() {
463 return Err(AgentError::InvalidConfiguration(
464 "AEP idempotency key must not be empty".to_owned(),
465 ));
466 }
467 return Ok(value);
468 }
469 let value = self
470 .client
471 .idempotency_keys
472 .create_key(&OperationKey {
473 command: command.clone(),
474 credential_id,
475 grant_type,
476 service_did: inspection.document.service.did.clone(),
477 service_url: self.service_url.clone(),
478 })
479 .await?;
480 if value.trim().is_empty() {
481 return Err(AgentError::InvalidConfiguration(
482 "AEP idempotency key provider returned an empty key".to_owned(),
483 ));
484 }
485 Ok(value)
486 }
487
488 async fn delete_revoked_credentials(
489 &self,
490 service_did: &str,
491 selector: &RevokeRequest,
492 ) -> Result<(), AgentError> {
493 for record in self.client.credential_store.list(service_did).await? {
494 let matches = selector.all_grant_types == Some(StringBoolean::True)
495 || selector.credential_id.as_deref() == Some(record.credential_id.as_str())
496 || (selector.credential_id.is_none()
497 && selector.grant_type.as_ref() == Some(&record.grant_type));
498 if matches {
499 self.client
500 .credential_store
501 .delete(service_did, &record.credential_id)
502 .await?;
503 }
504 }
505 Ok(())
506 }
507}
508
509struct RawCommandResult {
510 body: Vec<u8>,
511 status: u16,
512 url: url::Url,
513}
514
515struct CommandRequest<'a, B: ?Sized> {
516 inspection: &'a Inspection,
517 identity: &'a AgentIdentity,
518 signer: &'a dyn AssertionSigner,
519 command: &'a Command,
520 method: Method,
521 body: Option<&'a B>,
522 idempotency_key: Option<&'a str>,
523}
524
525fn select_grant_type(
526 inspection: &Inspection,
527 selected: Option<&GrantType>,
528 preferred: &[GrantType],
529) -> Result<GrantType, AgentError> {
530 let advertised = &inspection.document.commands.grant_types;
531 if let Some(selected) = selected {
532 return advertised
533 .contains(selected)
534 .then(|| selected.clone())
535 .ok_or(AgentError::NoCompatibleGrantType);
536 }
537 let candidates = if preferred.is_empty() {
538 advertised
539 } else {
540 preferred
541 };
542 candidates
543 .iter()
544 .find(|candidate| advertised.contains(candidate))
545 .cloned()
546 .ok_or(AgentError::NoCompatibleGrantType)
547}
548
549fn credential_record(
550 credential: &BuiltInGrantResponse,
551 payload: Value,
552 inspection: &Inspection,
553 issued_at: OffsetDateTime,
554) -> Result<CredentialRecord, AgentError> {
555 let (credential_id, expires_at) = match credential {
556 BuiltInGrantResponse::OAuthBearer(value) => (&value.credential_id, &value.expires_at),
557 BuiltInGrantResponse::ApiKey(value) => (&value.credential_id, &value.expires_at),
558 BuiltInGrantResponse::Basic(value) => (&value.credential_id, &value.expires_at),
559 };
560 let expires_at = OffsetDateTime::parse(expires_at, &Rfc3339).map_err(|_| {
561 AgentError::Credential("AEP credential expiration is not RFC 3339".to_owned())
562 })?;
563 Ok(CredentialRecord {
564 credential_id: credential_id.clone(),
565 expires_at,
566 grant_type: credential.grant_type(),
567 issued_at,
568 payload,
569 service_did: inspection.document.service.did.clone(),
570 service_url: inspection.service_url.clone(),
571 })
572}