1use std::{
2 collections::BTreeMap,
3 sync::{Arc, Mutex},
4};
5
6use aep_core::{
7 AUTHORIZATION_HEADER, AuthenticationMethod, AuthorizationCarrier, BuiltInGrantResponse,
8 CredentialScheme, GrantRequest, GrantType, GrantTypeConfig, RevokeRequest, StringBoolean,
9 parse_protected_resource_authorization, validate_built_in_grant_response,
10};
11use async_trait::async_trait;
12use base64::{Engine as _, engine::general_purpose::STANDARD};
13use http::{HeaderMap, HeaderName, header};
14use serde::Serialize;
15use sha2::{Digest as _, Sha256};
16use subtle::ConstantTimeEq as _;
17use time::{OffsetDateTime, format_description::well_known::Rfc3339};
18
19use crate::{
20 AuthenticatedPrincipal, AuthenticationKind, CredentialAuthenticationInput, CredentialMatch,
21 GrantContext, GrantTypeDefinition, GrantTypeHandler, ServiceCredentialRecord,
22 ServiceCredentialStore, ServiceError, StoredApiKeyGrantTypeOptions,
23 StoredBasicGrantTypeOptions, StoredCredentialGrantTypeOptions,
24 StoredOAuthBearerGrantTypeOptions,
25};
26
27pub fn stored_oauth_bearer_grant_type(
28 options: StoredOAuthBearerGrantTypeOptions,
29) -> Result<GrantTypeDefinition, ServiceError> {
30 stored_credential_grant_type(GrantType::OAuthBearer, options, |credential| {
31 BuiltInGrantResponse::OAuthBearer(credential)
32 })
33}
34
35pub fn stored_api_key_grant_type(
36 options: StoredApiKeyGrantTypeOptions,
37) -> Result<GrantTypeDefinition, ServiceError> {
38 configured_api_key_headers(&options.config)?;
39 stored_credential_grant_type(GrantType::ApiKey, options, |credential| {
40 BuiltInGrantResponse::ApiKey(credential)
41 })
42}
43
44pub fn stored_basic_grant_type(
45 options: StoredBasicGrantTypeOptions,
46) -> Result<GrantTypeDefinition, ServiceError> {
47 stored_credential_grant_type(GrantType::Basic, options, |credential| {
48 BuiltInGrantResponse::Basic(credential)
49 })
50}
51
52fn stored_credential_grant_type<T>(
53 grant_type: GrantType,
54 options: StoredCredentialGrantTypeOptions<T>,
55 into_credential: fn(T) -> BuiltInGrantResponse,
56) -> Result<GrantTypeDefinition, ServiceError>
57where
58 T: Serialize + Send + 'static,
59{
60 let mut config = options.config;
61 config.supports_per_credential_revoke = Some(StringBoolean::True);
62 Ok(GrantTypeDefinition {
63 config: Some(config.clone()),
64 grant_type: grant_type.clone(),
65 handler: Some(Arc::new(StoredCredentialHandler {
66 config,
67 grant_type,
68 into_credential,
69 issue: options.issue,
70 store: options.store,
71 })),
72 })
73}
74
75struct StoredCredentialHandler<T> {
76 config: GrantTypeConfig,
77 grant_type: GrantType,
78 into_credential: fn(T) -> BuiltInGrantResponse,
79 issue: crate::BuiltInCredentialIssuer<T>,
80 store: Arc<dyn ServiceCredentialStore>,
81}
82
83#[async_trait]
84impl<T> GrantTypeHandler for StoredCredentialHandler<T>
85where
86 T: Serialize + Send + 'static,
87{
88 async fn grant(
89 &self,
90 request: &GrantRequest,
91 context: &GrantContext,
92 ) -> Result<serde_json::Value, ServiceError> {
93 let issued = (self.issue)(request.clone(), context.clone()).await?;
94 let response = serde_json::to_value(&issued)?;
95 let credential = (self.into_credential)(issued);
96 validate_built_in_grant_response(&self.grant_type, &credential)
97 .map_err(aep_core::CoreError::from)?;
98 validate_issued_credential_config(&credential, &self.config)?;
99 self.store
100 .save(service_credential_record(
101 context.agent_did.clone(),
102 credential.clone(),
103 context.now,
104 )?)
105 .await?;
106 Ok(response)
107 }
108
109 async fn revoke(
110 &self,
111 request: &RevokeRequest,
112 context: &GrantContext,
113 ) -> Result<(), ServiceError> {
114 if let Some(credential_id) = &request.credential_id {
115 self.store
116 .revoke_credential(
117 &context.agent_did,
118 &self.grant_type,
119 credential_id,
120 context.now,
121 )
122 .await
123 } else {
124 self.store
125 .revoke_grant_type(&context.agent_did, &self.grant_type, context.now)
126 .await
127 }
128 }
129
130 async fn authenticate(
131 &self,
132 input: &CredentialAuthenticationInput,
133 ) -> Result<Option<AuthenticatedPrincipal>, ServiceError> {
134 let Some(matched) = self.store.authenticate(&self.grant_type, input).await? else {
135 return Ok(None);
136 };
137 if matched.agent_did.is_empty()
138 || matched.credential_id.is_empty()
139 || matched.expires_at <= input.now
140 || matched.grant_type != self.grant_type
141 {
142 return Err(ServiceError::Store(
143 "AEP credential store returned an invalid match".to_owned(),
144 ));
145 }
146 Ok(Some(AuthenticatedPrincipal {
147 agent_did: matched.agent_did,
148 authentication_kind: AuthenticationKind::SessionCredential,
149 authentication_method: authentication_method(&self.grant_type)?,
150 credential_id: Some(matched.credential_id),
151 grant_type: Some(matched.grant_type),
152 scopes: matched.scopes,
153 }))
154 }
155
156 async fn has_presentation(
157 &self,
158 input: &CredentialAuthenticationInput,
159 ) -> Result<bool, ServiceError> {
160 self.store.has_presentation(&self.grant_type, input).await
161 }
162}
163
164fn authentication_method(grant_type: &GrantType) -> Result<AuthenticationMethod, ServiceError> {
165 match grant_type {
166 GrantType::OAuthBearer => Ok(AuthenticationMethod::OAuthBearer),
167 GrantType::ApiKey => Ok(AuthenticationMethod::ApiKey),
168 GrantType::Basic => Ok(AuthenticationMethod::Basic),
169 GrantType::Other(_) => Err(ServiceError::InvalidConfiguration(
170 "stored credentials require a built-in AEP grant type".to_owned(),
171 )),
172 }
173}
174
175fn validate_issued_credential_config(
176 credential: &BuiltInGrantResponse,
177 config: &GrantTypeConfig,
178) -> Result<(), ServiceError> {
179 let BuiltInGrantResponse::ApiKey(value) = credential else {
180 return Ok(());
181 };
182 let Some(headers) = configured_api_key_headers(config)? else {
183 return Ok(());
184 };
185 if headers
186 .iter()
187 .any(|header| header == &canonical_header(&value.header))
188 {
189 Ok(())
190 } else {
191 Err(ServiceError::Handler(
192 "issued API-key header is not advertised by the Service".to_owned(),
193 ))
194 }
195}
196
197fn configured_api_key_headers(
198 config: &GrantTypeConfig,
199) -> Result<Option<Vec<String>>, ServiceError> {
200 let Some(value) = config.additional.get("header_names") else {
201 return Ok(None);
202 };
203 let headers = serde_json::from_value::<Vec<String>>(value.clone()).map_err(|_| {
204 ServiceError::InvalidConfiguration(
205 "API-key header_names must be an array of HTTP field names".to_owned(),
206 )
207 })?;
208 let mut canonical = Vec::with_capacity(headers.len());
209 for value in headers {
210 let name = HeaderName::try_from(value.as_str()).map_err(|_| {
211 ServiceError::InvalidConfiguration(
212 "API-key header_names contains an invalid HTTP field name".to_owned(),
213 )
214 })?;
215 let name = name.as_str().to_owned();
216 if canonical.contains(&name) {
217 return Err(ServiceError::InvalidConfiguration(
218 "API-key header_names contains a duplicate HTTP field name".to_owned(),
219 ));
220 }
221 canonical.push(name);
222 }
223 Ok(Some(canonical))
224}
225
226fn service_credential_record(
227 agent_did: String,
228 credential: BuiltInGrantResponse,
229 created_at: OffsetDateTime,
230) -> Result<ServiceCredentialRecord, ServiceError> {
231 if agent_did.is_empty() {
232 return Err(ServiceError::Handler(
233 "issued credential requires an Agent DID".to_owned(),
234 ));
235 }
236 let (credential_id, expires_at, _) = credential_metadata(&credential);
237 let expires_at = OffsetDateTime::parse(expires_at, &Rfc3339).map_err(|_| {
238 ServiceError::Handler("issued credential has an invalid expiration".to_owned())
239 })?;
240 if expires_at <= created_at {
241 return Err(ServiceError::Handler(
242 "issued credential must expire after issuance".to_owned(),
243 ));
244 }
245 Ok(ServiceCredentialRecord {
246 agent_did,
247 created_at,
248 credential_id: credential_id.to_owned(),
249 expires_at,
250 grant_type: credential.grant_type(),
251 credential,
252 })
253}
254
255#[derive(Clone)]
256struct MemoryCredentialRecord {
257 agent_did: String,
258 credential_id: String,
259 expires_at: OffsetDateTime,
260 grant_type: GrantType,
261 header: String,
262 revoked_at: Option<OffsetDateTime>,
263 scopes: Vec<String>,
264 verifier: [u8; 32],
265}
266
267#[derive(Default)]
268pub struct MemoryServiceCredentialStore {
269 records: Mutex<BTreeMap<String, MemoryCredentialRecord>>,
270}
271
272impl MemoryServiceCredentialStore {
273 pub fn new(
274 records: impl IntoIterator<Item = ServiceCredentialRecord>,
275 ) -> Result<Self, ServiceError> {
276 let store = Self::default();
277 {
278 let mut destination = store.records.lock().map_err(lock_error)?;
279 for record in records {
280 save_record(&mut destination, memory_credential_record(record)?)?;
281 }
282 }
283 Ok(store)
284 }
285}
286
287#[async_trait]
288impl ServiceCredentialStore for MemoryServiceCredentialStore {
289 async fn authenticate(
290 &self,
291 grant_type: &GrantType,
292 input: &CredentialAuthenticationInput,
293 ) -> Result<Option<CredentialMatch>, ServiceError> {
294 let records = self.records.lock().map_err(lock_error)?;
295 let presentations = presentations(grant_type, &input.headers, records.values());
296 if presentations.len() != 1 {
297 return Ok(None);
298 }
299 let presentation = &presentations[0];
300 let candidate: [u8; 32] = Sha256::digest(presentation.value.as_bytes()).into();
301 for record in records.values() {
302 if &record.grant_type == grant_type
303 && record.header == presentation.header
304 && record.revoked_at.is_none()
305 && record.expires_at > input.now
306 && bool::from(record.verifier.ct_eq(&candidate))
307 {
308 return Ok(Some(CredentialMatch {
309 agent_did: record.agent_did.clone(),
310 credential_id: record.credential_id.clone(),
311 expires_at: record.expires_at,
312 grant_type: record.grant_type.clone(),
313 scopes: record.scopes.clone(),
314 }));
315 }
316 }
317 Ok(None)
318 }
319
320 async fn has_presentation(
321 &self,
322 grant_type: &GrantType,
323 input: &CredentialAuthenticationInput,
324 ) -> Result<bool, ServiceError> {
325 let records = self.records.lock().map_err(lock_error)?;
326 Ok(!presentations(grant_type, &input.headers, records.values()).is_empty())
327 }
328
329 async fn revoke_credential(
330 &self,
331 agent_did: &str,
332 grant_type: &GrantType,
333 credential_id: &str,
334 revoked_at: OffsetDateTime,
335 ) -> Result<(), ServiceError> {
336 let mut records = self.records.lock().map_err(lock_error)?;
337 if let Some(record) = records.get_mut(credential_id)
338 && record.agent_did == agent_did
339 && &record.grant_type == grant_type
340 {
341 record.revoked_at = Some(revoked_at);
342 }
343 Ok(())
344 }
345
346 async fn revoke_grant_type(
347 &self,
348 agent_did: &str,
349 grant_type: &GrantType,
350 revoked_at: OffsetDateTime,
351 ) -> Result<(), ServiceError> {
352 for record in self.records.lock().map_err(lock_error)?.values_mut() {
353 if record.agent_did == agent_did && &record.grant_type == grant_type {
354 record.revoked_at = Some(revoked_at);
355 }
356 }
357 Ok(())
358 }
359
360 async fn save(&self, record: ServiceCredentialRecord) -> Result<(), ServiceError> {
361 let record = memory_credential_record(record)?;
362 let mut records = self.records.lock().map_err(lock_error)?;
363 save_record(&mut records, record)
364 }
365}
366
367fn save_record(
368 records: &mut BTreeMap<String, MemoryCredentialRecord>,
369 record: MemoryCredentialRecord,
370) -> Result<(), ServiceError> {
371 if records.contains_key(&record.credential_id) {
372 return Err(ServiceError::Store(
373 "AEP credential identifier has already been issued".to_owned(),
374 ));
375 }
376 if records.values().any(|existing| {
377 existing.grant_type == record.grant_type
378 && existing.header == record.header
379 && bool::from(existing.verifier.ct_eq(&record.verifier))
380 }) {
381 return Err(ServiceError::Store(
382 "AEP credential secret has already been issued".to_owned(),
383 ));
384 }
385 records.insert(record.credential_id.clone(), record);
386 Ok(())
387}
388
389fn memory_credential_record(
390 record: ServiceCredentialRecord,
391) -> Result<MemoryCredentialRecord, ServiceError> {
392 validate_built_in_grant_response(&record.grant_type, &record.credential)
393 .map_err(aep_core::CoreError::from)?;
394 let (credential_id, expires_at, scopes) = credential_metadata(&record.credential);
395 let parsed_expiry = OffsetDateTime::parse(expires_at, &Rfc3339).map_err(|_| {
396 ServiceError::Store("AEP credential record has an invalid expiration".to_owned())
397 })?;
398 if record.agent_did.is_empty()
399 || credential_id != record.credential_id
400 || parsed_expiry != record.expires_at
401 || record.expires_at <= record.created_at
402 || record.grant_type != record.credential.grant_type()
403 {
404 return Err(ServiceError::Store(
405 "AEP credential store received an invalid record".to_owned(),
406 ));
407 }
408 let (header, secret) = match &record.credential {
409 BuiltInGrantResponse::OAuthBearer(value) => (
410 canonical_header(header::AUTHORIZATION.as_str()),
411 value.access_token.clone(),
412 ),
413 BuiltInGrantResponse::ApiKey(value) => {
414 (canonical_header(&value.header), value.api_key.clone())
415 }
416 BuiltInGrantResponse::Basic(value) => (
417 canonical_header(header::AUTHORIZATION.as_str()),
418 STANDARD.encode(format!("{}:{}", value.username, value.password)),
419 ),
420 };
421 Ok(MemoryCredentialRecord {
422 agent_did: record.agent_did,
423 credential_id: record.credential_id,
424 expires_at: record.expires_at,
425 grant_type: record.grant_type,
426 header,
427 revoked_at: None,
428 scopes: scopes.to_vec(),
429 verifier: Sha256::digest(secret.as_bytes()).into(),
430 })
431}
432
433fn credential_metadata(credential: &BuiltInGrantResponse) -> (&str, &str, &[String]) {
434 match credential {
435 BuiltInGrantResponse::OAuthBearer(value) => {
436 (&value.credential_id, &value.expires_at, &value.scopes)
437 }
438 BuiltInGrantResponse::ApiKey(value) => {
439 (&value.credential_id, &value.expires_at, &value.scopes)
440 }
441 BuiltInGrantResponse::Basic(value) => {
442 (&value.credential_id, &value.expires_at, &value.scopes)
443 }
444 }
445}
446
447struct CredentialPresentation {
448 header: String,
449 value: String,
450}
451
452fn presentations<'a>(
453 grant_type: &GrantType,
454 headers: &HeaderMap,
455 records: impl Iterator<Item = &'a MemoryCredentialRecord>,
456) -> Vec<CredentialPresentation> {
457 if grant_type == &GrantType::ApiKey {
458 let unique_headers = records
459 .filter(|record| record.grant_type == GrantType::ApiKey)
460 .map(|record| record.header.clone())
461 .collect::<std::collections::BTreeSet<_>>();
462 return unique_headers
463 .iter()
464 .flat_map(|name| header_values(headers, name))
465 .map(|(header, value)| CredentialPresentation { header, value })
466 .collect();
467 }
468 let expected = if grant_type == &GrantType::OAuthBearer {
469 CredentialScheme::Bearer
470 } else if grant_type == &GrantType::Basic {
471 CredentialScheme::Basic
472 } else {
473 return Vec::new();
474 };
475 [
476 (
477 header::AUTHORIZATION.as_str(),
478 AuthorizationCarrier::Standard,
479 ),
480 (AUTHORIZATION_HEADER, AuthorizationCarrier::Dedicated),
481 ]
482 .into_iter()
483 .flat_map(|(name, carrier)| {
484 header_values(headers, name).filter_map(move |(_, value)| {
485 parse_protected_resource_authorization(&value, carrier)
486 .ok()
487 .filter(|parsed| parsed.scheme == expected)
488 .map(|parsed| CredentialPresentation {
489 header: canonical_header(header::AUTHORIZATION.as_str()),
490 value: parsed.credentials,
491 })
492 })
493 })
494 .collect()
495}
496
497fn header_values<'a>(
498 headers: &'a HeaderMap,
499 name: &'a str,
500) -> impl Iterator<Item = (String, String)> + 'a {
501 headers
502 .get_all(name)
503 .iter()
504 .filter_map(|value| value.to_str().ok())
505 .map(move |value| (canonical_header(name), value.to_owned()))
506}
507
508fn canonical_header(value: &str) -> String {
509 value.to_ascii_lowercase()
510}
511
512fn lock_error<T>(_error: std::sync::PoisonError<T>) -> ServiceError {
513 ServiceError::Store("AEP Service credential memory store lock is poisoned".to_owned())
514}