aep_platform/
store.rs

1use std::{
2    collections::BTreeMap,
3    sync::{Arc, Mutex},
4    time::Duration,
5};
6
7use async_trait::async_trait;
8use time::OffsetDateTime;
9
10use crate::{
11    Clock, IdempotencyInput, IdempotencyOperation, IdempotencyResult, IdempotencyStore,
12    IdentityCreation, IdentityListQuery, IdentityListResult, IdentityRecord, IdentityStore,
13    LifecyclePolicy, ManagedAgentStatus, PlatformError, ReplayStore, RequestContext,
14    StoredResponse, validate_identity_record,
15};
16
17type ScopeLock = Arc<futures::lock::Mutex<()>>;
18type IdempotencyKey = (String, String);
19
20#[derive(Default)]
21struct IdentityRecords {
22    by_agent_did: BTreeMap<String, String>,
23    by_agent_did_id: BTreeMap<String, String>,
24    by_scope: BTreeMap<(String, String), String>,
25    records: BTreeMap<String, IdentityRecord>,
26}
27
28#[derive(Default)]
29pub struct MemoryIdentityStore {
30    locks: Mutex<BTreeMap<(String, String), ScopeLock>>,
31    records: Mutex<IdentityRecords>,
32}
33
34impl MemoryIdentityStore {
35    fn scope_lock(&self, key: &(String, String)) -> Result<ScopeLock, PlatformError> {
36        Ok(self
37            .locks
38            .lock()
39            .map_err(lock_error)?
40            .entry(key.clone())
41            .or_default()
42            .clone())
43    }
44
45    fn find_indexed(
46        &self,
47        key: &str,
48        index: impl FnOnce(&IdentityRecords) -> &BTreeMap<String, String>,
49    ) -> Result<Option<IdentityRecord>, PlatformError> {
50        let records = self.records.lock().map_err(lock_error)?;
51        Ok(index(&records)
52            .get(key)
53            .and_then(|identity_id| records.records.get(identity_id))
54            .cloned())
55    }
56
57    fn release_scope_lock(
58        &self,
59        key: &(String, String),
60        scope_lock: &ScopeLock,
61    ) -> Result<(), PlatformError> {
62        let mut locks = self.locks.lock().map_err(lock_error)?;
63        if locks
64            .get(key)
65            .is_some_and(|current| Arc::ptr_eq(current, scope_lock))
66            && Arc::strong_count(scope_lock) == 2
67        {
68            locks.remove(key);
69        }
70        Ok(())
71    }
72}
73
74#[async_trait]
75impl IdentityStore for MemoryIdentityStore {
76    async fn find_or_create(
77        &self,
78        principal: &str,
79        service_did: &str,
80        create: IdentityCreation,
81    ) -> Result<IdentityRecord, PlatformError> {
82        if principal.is_empty() || service_did.is_empty() {
83            return Err(PlatformError::Store(
84                "AEP Platform identity scope must not be empty".to_owned(),
85            ));
86        }
87        let key = (principal.to_owned(), service_did.to_owned());
88        let scope_lock = self.scope_lock(&key)?;
89        let result = async {
90            let _guard = scope_lock.lock().await;
91            let existing = {
92                let records = self.records.lock().map_err(lock_error)?;
93                records
94                    .by_scope
95                    .get(&key)
96                    .and_then(|identity_id| records.records.get(identity_id))
97                    .cloned()
98            };
99            if let Some(existing) = existing {
100                Ok(existing)
101            } else {
102                let identity = create().await?;
103                validate_identity_record(&identity)?;
104                if identity.principal != principal || identity.service_did != service_did {
105                    return Err(PlatformError::Store(
106                        "AEP Platform identity does not match its requested scope".to_owned(),
107                    ));
108                }
109                let mut records = self.records.lock().map_err(lock_error)?;
110                if records.records.contains_key(&identity.agent_identity_id)
111                    || records.by_agent_did.contains_key(&identity.agent_did)
112                    || records.by_agent_did_id.contains_key(&identity.agent_did_id)
113                {
114                    return Err(PlatformError::Store(
115                        "AEP Platform identity material must be unique".to_owned(),
116                    ));
117                }
118                records
119                    .by_scope
120                    .insert(key.clone(), identity.agent_identity_id.clone());
121                records.by_agent_did.insert(
122                    identity.agent_did.clone(),
123                    identity.agent_identity_id.clone(),
124                );
125                records.by_agent_did_id.insert(
126                    identity.agent_did_id.clone(),
127                    identity.agent_identity_id.clone(),
128                );
129                records
130                    .records
131                    .insert(identity.agent_identity_id.clone(), identity.clone());
132                Ok(identity)
133            }
134        }
135        .await;
136        self.release_scope_lock(&key, &scope_lock)?;
137        result
138    }
139
140    async fn find_by_agent_did(
141        &self,
142        agent_did: &str,
143    ) -> Result<Option<IdentityRecord>, PlatformError> {
144        self.find_indexed(agent_did, |records| &records.by_agent_did)
145    }
146
147    async fn find_by_agent_did_id(
148        &self,
149        agent_did_id: &str,
150    ) -> Result<Option<IdentityRecord>, PlatformError> {
151        self.find_indexed(agent_did_id, |records| &records.by_agent_did_id)
152    }
153
154    async fn get(&self, agent_identity_id: &str) -> Result<Option<IdentityRecord>, PlatformError> {
155        Ok(self
156            .records
157            .lock()
158            .map_err(lock_error)?
159            .records
160            .get(agent_identity_id)
161            .cloned())
162    }
163
164    async fn list(
165        &self,
166        principal: &str,
167        query: &IdentityListQuery,
168    ) -> Result<IdentityListResult, PlatformError> {
169        if principal.is_empty() {
170            return Err(PlatformError::Store(
171                "AEP Platform identity-list principal must not be empty".to_owned(),
172            ));
173        }
174        let mut identities = self
175            .records
176            .lock()
177            .map_err(lock_error)?
178            .records
179            .values()
180            .filter(|identity| {
181                identity.principal == principal
182                    && query
183                        .service_did
184                        .as_ref()
185                        .is_none_or(|service_did| &identity.service_did == service_did)
186                    && query.status.is_none_or(|status| identity.status == status)
187            })
188            .cloned()
189            .collect::<Vec<_>>();
190        identities.sort_by(|left, right| {
191            left.created_at
192                .cmp(&right.created_at)
193                .then_with(|| left.agent_identity_id.cmp(&right.agent_identity_id))
194        });
195        if query.descending {
196            identities.reverse();
197        }
198        let total = identities.len();
199        let identities = identities
200            .into_iter()
201            .skip(query.offset)
202            .take(query.limit)
203            .collect();
204        Ok(IdentityListResult { identities, total })
205    }
206
207    async fn update_status(
208        &self,
209        agent_identity_id: &str,
210        status: ManagedAgentStatus,
211        updated_at: OffsetDateTime,
212    ) -> Result<Option<IdentityRecord>, PlatformError> {
213        let mut records = self.records.lock().map_err(lock_error)?;
214        let Some(identity) = records.records.get_mut(agent_identity_id) else {
215            return Ok(None);
216        };
217        identity.status = status;
218        identity.updated_at = updated_at;
219        Ok(Some(identity.clone()))
220    }
221}
222
223#[derive(Clone, Debug)]
224struct IdempotencyRecord {
225    input: IdempotencyInput,
226    response: StoredResponse,
227}
228
229pub struct MemoryIdempotencyStore {
230    clock: Arc<dyn Clock>,
231    locks: Mutex<BTreeMap<IdempotencyKey, ScopeLock>>,
232    records: Mutex<BTreeMap<IdempotencyKey, IdempotencyRecord>>,
233}
234
235impl MemoryIdempotencyStore {
236    pub fn new(clock: Arc<dyn Clock>) -> Self {
237        Self {
238            clock,
239            locks: Mutex::new(BTreeMap::new()),
240            records: Mutex::new(BTreeMap::new()),
241        }
242    }
243
244    fn operation_lock(&self, key: &IdempotencyKey) -> Result<ScopeLock, PlatformError> {
245        Ok(self
246            .locks
247            .lock()
248            .map_err(lock_error)?
249            .entry(key.clone())
250            .or_default()
251            .clone())
252    }
253
254    fn release_operation_lock(
255        &self,
256        key: &IdempotencyKey,
257        operation_lock: &ScopeLock,
258    ) -> Result<(), PlatformError> {
259        let mut locks = self.locks.lock().map_err(lock_error)?;
260        if locks
261            .get(key)
262            .is_some_and(|current| Arc::ptr_eq(current, operation_lock))
263            && Arc::strong_count(operation_lock) == 2
264        {
265            locks.remove(key);
266        }
267        Ok(())
268    }
269}
270
271#[async_trait]
272impl IdempotencyStore for MemoryIdempotencyStore {
273    async fn execute(
274        &self,
275        input: IdempotencyInput,
276        operation: IdempotencyOperation,
277    ) -> Result<IdempotencyResult, PlatformError> {
278        if input.principal.is_empty()
279            || input.idempotency_key.is_empty()
280            || input.request_hash.is_empty()
281        {
282            return Err(PlatformError::Store(
283                "AEP Platform idempotency input is invalid".to_owned(),
284            ));
285        }
286        let key = (input.principal.clone(), input.idempotency_key.clone());
287        let operation_lock = self.operation_lock(&key)?;
288        let result = async {
289            let _guard = operation_lock.lock().await;
290            let now = self.clock.now();
291            let existing = {
292                let mut records = self.records.lock().map_err(lock_error)?;
293                records.retain(|_, record| {
294                    record.response.created_at + Duration::from_secs(3600) > now
295                });
296                records.get(&key).cloned()
297            };
298            if let Some(existing) = existing {
299                if existing.input.operation == input.operation
300                    && existing.input.request_hash == input.request_hash
301                {
302                    Ok(IdempotencyResult::Replayed(existing.response))
303                } else {
304                    Ok(IdempotencyResult::Conflict)
305                }
306            } else {
307                let mut response = operation().await?;
308                if response.created_at == OffsetDateTime::UNIX_EPOCH {
309                    response.created_at = now;
310                }
311                self.records.lock().map_err(lock_error)?.insert(
312                    key.clone(),
313                    IdempotencyRecord {
314                        input,
315                        response: response.clone(),
316                    },
317                );
318                Ok(IdempotencyResult::Created(response))
319            }
320        }
321        .await;
322        self.release_operation_lock(&key, &operation_lock)?;
323        result
324    }
325}
326
327impl Default for MemoryIdempotencyStore {
328    fn default() -> Self {
329        Self::new(Arc::new(SystemClock))
330    }
331}
332
333#[derive(Default)]
334pub struct MemoryReplayStore {
335    records: Mutex<BTreeMap<String, OffsetDateTime>>,
336}
337
338#[async_trait]
339impl ReplayStore for MemoryReplayStore {
340    async fn consume(
341        &self,
342        key: &str,
343        expires_at: OffsetDateTime,
344        now: OffsetDateTime,
345    ) -> Result<bool, PlatformError> {
346        if key.is_empty() || expires_at <= now {
347            return Ok(false);
348        }
349        let mut records = self.records.lock().map_err(lock_error)?;
350        records.retain(|_, expiry| *expiry > now);
351        if records.contains_key(key) {
352            return Ok(false);
353        }
354        records.insert(key.to_owned(), expires_at);
355        Ok(true)
356    }
357}
358
359#[derive(Default)]
360pub struct SystemClock;
361
362impl Clock for SystemClock {
363    fn now(&self) -> OffsetDateTime {
364        OffsetDateTime::now_utc()
365    }
366}
367
368#[derive(Default)]
369pub struct DefaultLifecyclePolicy;
370
371#[async_trait]
372impl LifecyclePolicy for DefaultLifecyclePolicy {
373    async fn can_sign(
374        &self,
375        identity: &IdentityRecord,
376        _context: &RequestContext,
377    ) -> Result<bool, PlatformError> {
378        Ok(identity.status == ManagedAgentStatus::Active)
379    }
380
381    async fn can_transition(
382        &self,
383        _identity: &IdentityRecord,
384        _status: ManagedAgentStatus,
385        _context: &RequestContext,
386    ) -> Result<bool, PlatformError> {
387        Ok(true)
388    }
389
390    async fn can_verify(
391        &self,
392        identity: &IdentityRecord,
393        _context: &RequestContext,
394    ) -> Result<bool, PlatformError> {
395        Ok(identity.status == ManagedAgentStatus::Active)
396    }
397}
398
399fn lock_error<T>(_error: std::sync::PoisonError<T>) -> PlatformError {
400    PlatformError::Store("AEP Platform memory store lock is poisoned".to_owned())
401}
aep_platform/store.rs

aep_platform/
store.rs
