Skip to main content

aep_agent/
commands.rs

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        // The elapsed-time bound below is the one the caller asked for. It relies on the supplied
94        // clock advancing as the supplied delay sleeps, so polling is also bounded by the number
95        // of intervals that fit in the timeout; a clock that does not advance then stops here
96        // instead of polling forever.
97        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            // Grant never implicitly enrolls, so an identity the Service has not activated is
166            // refused before the request is sent. CORE-ERR-009 pairs `not_recognized` with 401,
167            // which is the answer the Service would give.
168            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        // CORE-RVK-002: Revoke may only name a grant type the Service advertises.
260        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            // CORE-ERR-001 serves an AEP error as Problem Details under its own media type, and
427            // RFC 9457 carries the response status inside the document. A body that disagrees is
428            // reported as the bare status rather than as a structured AEP problem.
429            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            // CORE-IDP-003 treats an empty field value as a missing key, and a value that is only
461            // whitespace renders as one.
462            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}