1use std::{collections::BTreeMap, sync::Arc};
2
3use autonomi::ScratchpadAddress;
4use tokio::sync::RwLock;
5
6use crate::{
7 error::Error,
8 events::{GetCallback, PurgeCallback, SyncCallback},
9 index::{
10 master_index::{IndexEntry, MasterIndex, StorageStats},
11 PadInfo,
12 },
13 network::{Network, NetworkChoice, DEV_TESTNET_PRIVATE_KEY_HEX},
14 ops::Data,
15};
16
17use mutant_protocol::{
18 HealthCheckCallback, HealthCheckResult, PurgeResult, PutCallback, StorageMode, SyncResult,
19};
20
21#[derive(Clone)]
26pub struct MutAnt {
27 index: Arc<RwLock<MasterIndex>>,
28 data: Arc<RwLock<Data>>,
29}
30
31impl MutAnt {
32 async fn init_all(private_key_hex: &str, network_choice: NetworkChoice) -> Result<Self, Error> {
33 let network = Arc::new(Network::new(private_key_hex, network_choice)?);
34 let index = Arc::new(RwLock::new(MasterIndex::new(network_choice)));
35 let data = Arc::new(RwLock::new(Data::new(network.clone(), index.clone())));
36
37 Ok(Self { index, data })
38 }
39 pub async fn init(private_key_hex: &str) -> Result<Self, Error> {
40 Self::init_all(private_key_hex, NetworkChoice::Mainnet).await
41 }
42
43 pub async fn init_public() -> Result<Self, Error> {
44 Self::init_all(DEV_TESTNET_PRIVATE_KEY_HEX, NetworkChoice::Mainnet).await
45 }
46
47 pub async fn init_local() -> Result<Self, Error> {
48 Self::init_all(DEV_TESTNET_PRIVATE_KEY_HEX, NetworkChoice::Devnet).await
49 }
50
51 pub async fn init_public_local() -> Result<Self, Error> {
52 Self::init_all(DEV_TESTNET_PRIVATE_KEY_HEX, NetworkChoice::Devnet).await
53 }
54
55 pub async fn init_alphanet(private_key_hex: &str) -> Result<Self, Error> {
56 Self::init_all(private_key_hex, NetworkChoice::Alphanet).await
57 }
58
59 pub async fn init_public_alphanet() -> Result<Self, Error> {
60 Self::init_all(DEV_TESTNET_PRIVATE_KEY_HEX, NetworkChoice::Alphanet).await
61 }
62
63 pub async fn put(
64 &self,
65 user_key: &str,
66 data_bytes: Arc<Vec<u8>>,
67 mode: StorageMode,
68 public: bool,
69 no_verify: bool,
70 put_callback: Option<PutCallback>,
71 ) -> Result<ScratchpadAddress, Error> {
72 self.data
73 .read()
74 .await
75 .put(user_key, data_bytes, mode, public, no_verify, put_callback)
76 .await
77 }
78
79 pub async fn get(
80 &self,
81 user_key: &str,
82 get_callback: Option<GetCallback>,
83 ) -> Result<Vec<u8>, Error> {
84 self.data.read().await.get(user_key, get_callback).await
85 }
86
87 pub async fn get_public(
88 &self,
89 address: &ScratchpadAddress,
90 get_callback: Option<GetCallback>,
91 ) -> Result<Vec<u8>, Error> {
92 self.data
93 .read()
94 .await
95 .get_public(address, get_callback)
96 .await
97 }
98
99 pub async fn rm(&self, user_key: &str) -> Result<(), Error> {
100 self.index.write().await.remove_key(user_key)?;
101 Ok(())
102 }
103
104 pub async fn list(&self) -> Result<BTreeMap<String, IndexEntry>, Error> {
105 let keys = self.index.read().await.list();
106 Ok(keys)
107 }
108
109 pub async fn contains_key(&self, user_key: &str) -> bool {
110 self.index.read().await.contains_key(user_key)
111 }
112
113 pub async fn get_public_index_address(&self, user_key: &str) -> Result<String, Error> {
117 let index_guard = self.index.read().await;
118
119 if !index_guard.is_public(user_key) {
121 return Err(Error::Internal(format!("Key '{}' is not a public key", user_key)));
122 }
123
124 if let Some(index_pad) = index_guard.extract_public_index_pad(user_key) {
126 Ok(index_pad.address.to_hex())
127 } else {
128 Err(Error::Internal(format!("Failed to get index pad for key '{}'", user_key)))
129 }
130 }
131
132 pub async fn export_raw_pads_private_key(&self) -> Result<Vec<PadInfo>, Error> {
133 let pads_hex = self.index.read().await.export_raw_pads_private_key()?;
134 Ok(pads_hex)
135 }
136
137 pub async fn import_raw_pads_private_key(&self, pads_hex: Vec<PadInfo>) -> Result<(), Error> {
138 self.index
139 .write()
140 .await
141 .import_raw_pads_private_key(pads_hex)?;
142
143 Ok(())
144 }
145
146 pub async fn purge(
147 &self,
148 aggressive: bool,
149 purge_callback: Option<PurgeCallback>,
150 ) -> Result<PurgeResult, Error> {
151 self.data
152 .read()
153 .await
154 .purge(aggressive, purge_callback)
155 .await
156 }
157
158 pub async fn get_storage_stats(&self) -> StorageStats {
159 self.index.read().await.get_storage_stats()
160 }
161
162 pub async fn health_check(
163 &self,
164 key_name: &str,
165 recycle: bool,
166 health_check_callback: Option<HealthCheckCallback>,
167 ) -> Result<HealthCheckResult, Error> {
168 self.data
169 .read()
170 .await
171 .health_check(key_name, recycle, health_check_callback)
172 .await
173 }
174
175 pub async fn sync(
176 &self,
177 force: bool,
178 sync_callback: Option<SyncCallback>,
179 ) -> Result<SyncResult, Error> {
180 self.data.read().await.sync(force, sync_callback).await
181 }
182}
183
184#[cfg(test)]
185mod tests {
186 use super::*;
187 use rand::{distributions::Alphanumeric, Rng};
188
189 fn generate_random_string(len: usize) -> String {
190 rand::thread_rng()
191 .sample_iter(&Alphanumeric)
192 .take(len)
193 .map(char::from)
194 .collect()
195 }
196
197 fn generate_random_bytes(len: usize) -> Vec<u8> {
198 let mut vec = vec![0u8; len];
199 rand::thread_rng().fill(&mut vec[..]);
200 vec
201 }
202
203 async fn setup_mutant() -> MutAnt {
204 MutAnt::init_local()
205 .await
206 .expect("Failed to initialize MutAnt for test")
207 }
208
209 #[tokio::test]
210 async fn test_store_basic() {
211 let mutant = setup_mutant().await;
212 let user_key = generate_random_string(10);
213 let data_bytes = generate_random_bytes(128);
214
215 let result = mutant
216 .put(
217 &user_key,
218 Arc::new(data_bytes.clone()),
219 StorageMode::Medium,
220 false,
221 false,
222 None,
223 )
224 .await;
225
226 assert!(result.is_ok(), "Store operation failed: {:?}", result.err());
227 let keys = mutant.list().await.unwrap();
233 assert!(keys.contains_key(&user_key));
234
235 let data = mutant.get(&user_key, None).await.unwrap();
237 assert_eq!(data, data_bytes);
238 }
239
240 #[tokio::test]
241 async fn test_store_update() {
242 let mutant = setup_mutant().await;
243 let user_key = generate_random_string(10);
244 let data_bytes_initial = generate_random_bytes(128);
245
246 let result = mutant
247 .put(
248 &user_key,
249 Arc::new(data_bytes_initial),
250 StorageMode::Medium,
251 false,
252 false,
253 None,
254 )
255 .await;
256
257 assert!(result.is_ok(), "Store operation failed: {:?}", result.err());
258
259 let data_bytes_updated = generate_random_bytes(128);
260
261 let result = mutant
262 .put(
263 &user_key,
264 Arc::new(data_bytes_updated.clone()),
265 StorageMode::Medium,
266 false,
267 false,
268 None,
269 )
270 .await;
271
272 assert!(result.is_ok(), "Store operation failed: {:?}", result.err());
273
274 let data = mutant.get(&user_key, None).await.unwrap();
275 assert_eq!(data, data_bytes_updated);
276 }
277
278 #[tokio::test]
279 async fn test_store_resume() {
280 let mutant = setup_mutant().await;
281 let user_key = generate_random_string(10);
282 let data_bytes = generate_random_bytes(128);
283 let data_bytes_clone_for_put = data_bytes.clone();
284
285 let first_put = mutant.put(
287 &user_key,
288 Arc::new(data_bytes_clone_for_put),
289 StorageMode::Medium,
290 false,
291 false,
292 None,
293 );
294
295 drop(first_put);
297
298 let result = mutant
300 .put(
301 &user_key,
302 Arc::new(data_bytes.clone()),
303 StorageMode::Medium,
304 false,
305 false,
306 None,
307 )
308 .await;
309
310 assert!(result.is_ok(), "Store operation failed: {:?}", result.err());
311
312 let data = mutant.get(&user_key, None).await.unwrap();
314 assert_eq!(data, data_bytes);
315 }
316}