1 // Copyright 2025 Google LLC 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 // 15 //////////////////////////////////////////////////////////////////////////////// 16 17 //! In-memory implementation of the storage trait for the purpose of unit-tests 18 // TODO: remove this 19 #![allow(dead_code)] 20 #![allow(unused_variables)] 21 use alloc::sync::Arc; 22 use authgraph_core::key::{InstanceIdentifier, Policy}; 23 use authmgr_be::data_structures::{ClientId, PersistentClientContext, PersistentInstanceContext}; 24 use authmgr_be::traits::{PersistentStorage, Storage}; 25 use authmgr_be::{ 26 am_err, 27 error::{Error, ErrorCode}, 28 }; 29 use std::collections::{hash_map::Entry, HashMap}; 30 31 /// Instance sequence number is part of the fully qualified path of a client's storage 32 #[derive(Clone, Debug, Eq, Hash, PartialEq)] 33 pub struct InstanceSeqNumber(i32); 34 35 #[derive(Clone, Debug, Eq, Hash, PartialEq)] 36 struct FullyQualifiedClientId(InstanceSeqNumber, Arc<ClientId>); 37 38 /// In-memory implementation for AuthMgr persistent storage 39 #[derive(Default)] 40 pub struct MockPersistentStorage { 41 global_seq_num: i32, 42 instances: HashMap<Arc<InstanceIdentifier>, PersistentInstanceContext>, 43 clients: HashMap<FullyQualifiedClientId, PersistentClientContext>, 44 } 45 46 impl MockPersistentStorage { 47 /// Constructor for `MockPersistentStorage` new() -> Self48 pub fn new() -> Self { 49 Self { global_seq_num: 0, instances: HashMap::new(), clients: HashMap::new() } 50 } 51 } 52 53 impl Storage for MockPersistentStorage { 54 type InstanceContext = PersistentInstanceContext; 55 type ClientContext = PersistentClientContext; 56 update_instance_policy_in_storage( &mut self, instance_id: &Arc<InstanceIdentifier>, latest_dice_policy: &Arc<Policy>, ) -> Result<(), Error>57 fn update_instance_policy_in_storage( 58 &mut self, 59 instance_id: &Arc<InstanceIdentifier>, 60 latest_dice_policy: &Arc<Policy>, 61 ) -> Result<(), Error> { 62 let instance = self 63 .instances 64 .get_mut(instance_id) 65 .ok_or(am_err!(InternalError, "instance does not exist."))?; 66 instance.dice_policy = Arc::clone(latest_dice_policy); 67 Ok(()) 68 } 69 update_client_policy_in_storage( &mut self, instance_seq_number: i32, client_id: &Arc<ClientId>, latest_dice_policy: &Arc<Policy>, ) -> Result<(), Error>70 fn update_client_policy_in_storage( 71 &mut self, 72 instance_seq_number: i32, 73 client_id: &Arc<ClientId>, 74 latest_dice_policy: &Arc<Policy>, 75 ) -> Result<(), Error> { 76 let client = self 77 .clients 78 .get_mut(&FullyQualifiedClientId( 79 InstanceSeqNumber(instance_seq_number), 80 Arc::clone(client_id), 81 )) 82 .ok_or(am_err!(InternalError, "client does not exist"))?; 83 client.dice_policy = Arc::clone(latest_dice_policy); 84 Ok(()) 85 } 86 read_instance_context( &self, instance_id: &Arc<InstanceIdentifier>, ) -> Result<Option<Self::InstanceContext>, Error>87 fn read_instance_context( 88 &self, 89 instance_id: &Arc<InstanceIdentifier>, 90 ) -> Result<Option<Self::InstanceContext>, Error> { 91 Ok(self.instances.get(instance_id).cloned()) 92 } 93 create_instance_context( &mut self, instance_id: &Arc<InstanceIdentifier>, instance_info: Self::InstanceContext, ) -> Result<(), Error>94 fn create_instance_context( 95 &mut self, 96 instance_id: &Arc<InstanceIdentifier>, 97 instance_info: Self::InstanceContext, 98 ) -> Result<(), Error> { 99 if let Entry::Vacant(e) = self.instances.entry(Arc::clone(instance_id)) { 100 e.insert(instance_info); 101 Ok(()) 102 } else { 103 Err(am_err!(InternalError, "instance already exists")) 104 } 105 } 106 read_client_context( &self, instance_seq_number: i32, client_id: &Arc<ClientId>, ) -> Result<Option<Self::ClientContext>, Error>107 fn read_client_context( 108 &self, 109 instance_seq_number: i32, 110 client_id: &Arc<ClientId>, 111 ) -> Result<Option<Self::ClientContext>, Error> { 112 Ok(self 113 .clients 114 .get(&FullyQualifiedClientId( 115 InstanceSeqNumber(instance_seq_number), 116 Arc::clone(client_id), 117 )) 118 .cloned()) 119 } 120 create_client_context( &mut self, instance_seq_number: i32, client_id: &Arc<ClientId>, client_info: Self::ClientContext, ) -> Result<(), Error>121 fn create_client_context( 122 &mut self, 123 instance_seq_number: i32, 124 client_id: &Arc<ClientId>, 125 client_info: Self::ClientContext, 126 ) -> Result<(), Error> { 127 if let Entry::Vacant(e) = self.clients.entry(FullyQualifiedClientId( 128 InstanceSeqNumber(instance_seq_number), 129 Arc::clone(client_id), 130 )) { 131 e.insert(client_info); 132 Ok(()) 133 } else { 134 Err(am_err!(InternalError, "client already exists")) 135 } 136 } 137 } 138 139 impl PersistentStorage for MockPersistentStorage { get_or_create_global_sequence_number(&mut self) -> Result<i32, Error>140 fn get_or_create_global_sequence_number(&mut self) -> Result<i32, Error> { 141 Ok(self.global_seq_num) 142 } 143 increment_global_sequence_number(&mut self) -> Result<i32, Error>144 fn increment_global_sequence_number(&mut self) -> Result<i32, Error> { 145 self.global_seq_num += 1; 146 Ok(self.global_seq_num) 147 } 148 } 149