mutant_lib/api/
mod.rs

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/// The main entry point for interacting with the MutAnt distributed storage system.
22///
23/// This struct encapsulates the different managers (data, index, pad lifecycle) and the network adapter.
24/// Instances are typically created using the `init` or `init_with_progress` associated functions.
25#[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    /// Get the public index address for a key
114    ///
115    /// This is used to get the address that can be used to fetch the key publicly
116    pub async fn get_public_index_address(&self, user_key: &str) -> Result<String, Error> {
117        let index_guard = self.index.read().await;
118
119        // Check if the key exists and is public
120        if !index_guard.is_public(user_key) {
121            return Err(Error::Internal(format!("Key '{}' is not a public key", user_key)));
122        }
123
124        // Extract the index pad
125        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        // Ideally, we'd also check if the data is retrievable here,
228        // but that requires a `get` method which is not yet implemented.
229        // For now, we just check if the store operation completed without error.
230
231        // check of the index
232        let keys = mutant.list().await.unwrap();
233        assert!(keys.contains_key(&user_key));
234
235        // check of the data
236        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        // Start the first put operation but do not await its completion
286        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        // Simulate an interruption by dropping the future before it completes
296        drop(first_put);
297
298        // Now attempt to resume the operation with the same data
299        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        // Verify that the data is correctly stored
313        let data = mutant.get(&user_key, None).await.unwrap();
314        assert_eq!(data, data_bytes);
315    }
316}