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