Skip to main content

sdkt_storage/
analyzer.rs

1use crate::error::StorageError;
2use crate::snapshot::EXPIRING_SOON_LEDGERS;
3use crate::types::{StorageClass, StorageEntry, StorageReport, TtlInfoSummary};
4use base64::Engine;
5use sdkt_rpc::SorobanRpcClient;
6use stellar_xdr::{ContractDataDurability, LedgerKey, ReadXdr, ScVal};
7
8/// Classify a storage entry from its base64 XDR `LedgerKey`.
9///
10/// - `LedgerKey::ContractData` whose `key` is `ScVal::LedgerKeyContractInstance`
11///   is the contract **instance** singleton.
12/// - Other `LedgerKey::ContractData` entries are categorized by their
13///   `durability` (`Persistent` / `Temporary`).
14/// - Anything else (account, trustline, contract code, etc.) is `Other`.
15///
16/// Returns `StorageClass::Other` if the key cannot be decoded — never errors,
17/// so a single malformed entry does not abort the whole analysis.
18pub fn classify_key(base64_key: &str) -> StorageClass {
19    let Ok(bytes) = base64::engine::general_purpose::STANDARD.decode(base64_key.trim()) else {
20        return StorageClass::Other;
21    };
22    let mut cursor = std::io::Cursor::new(bytes);
23    let mut limited = stellar_xdr::Limited::new(&mut cursor, stellar_xdr::Limits::none());
24    let Ok(key) = LedgerKey::read_xdr(&mut limited) else {
25        return StorageClass::Other;
26    };
27
28    match key {
29        LedgerKey::ContractData(cd) => {
30            if matches!(cd.key, ScVal::LedgerKeyContractInstance) {
31                StorageClass::Instance
32            } else {
33                match cd.durability {
34                    ContractDataDurability::Temporary => StorageClass::Temporary,
35                    ContractDataDurability::Persistent => StorageClass::Persistent,
36                }
37            }
38        }
39        _ => StorageClass::Other,
40    }
41}
42
43pub struct StorageAnalyzer {
44    client: SorobanRpcClient,
45}
46
47impl StorageAnalyzer {
48    pub fn new(client: SorobanRpcClient) -> Self {
49        Self { client }
50    }
51
52    pub async fn inspect_contract_storage(
53        &self,
54        contract_id: &str,
55    ) -> Result<StorageReport, StorageError> {
56        self.inspect_contract_storage_keys(contract_id, &[]).await
57    }
58
59    pub async fn inspect_contract_storage_keys(
60        &self,
61        contract_id: &str,
62        extra_keys: &[String],
63    ) -> Result<StorageReport, StorageError> {
64        if contract_id.is_empty() {
65            return Err(StorageError::InvalidContractId(
66                "Contract ID cannot be empty".to_string(),
67            ));
68        }
69
70        let ttl_info =
71            sdkt_rpc::get_ttl_info_for_keys(&self.client, contract_id, extra_keys).await?;
72
73        if ttl_info.entries.is_empty() {
74            return Ok(StorageReport {
75                contract_id: contract_id.to_string(),
76                ..Default::default()
77            });
78        }
79
80        let mut min_ttl = u32::MAX;
81        let mut max_ttl = 0u32;
82        let mut total_ttl: u64 = 0;
83        let mut expiring_soon = 0;
84        let mut total_cost: u64 = 0;
85        let mut ttl_bearing_entries = 0usize;
86
87        let mut instance_entries = 0;
88        let mut persistent_entries = 0;
89        let mut temporary_entries = 0;
90        let mut other_entries = 0;
91        let mut detailed: Vec<StorageEntry> = Vec::with_capacity(ttl_info.entries.len());
92
93        for entry in &ttl_info.entries {
94            let class = classify_key(&entry.key);
95            match class {
96                StorageClass::Instance => instance_entries += 1,
97                StorageClass::Persistent => persistent_entries += 1,
98                StorageClass::Temporary => temporary_entries += 1,
99                StorageClass::Other => other_entries += 1,
100            }
101
102            if class != StorageClass::Other {
103                let ttl = entry.current_ttl;
104                if ttl < min_ttl {
105                    min_ttl = ttl;
106                }
107                if ttl > max_ttl {
108                    max_ttl = ttl;
109                }
110                total_ttl += ttl as u64;
111
112                if ttl < EXPIRING_SOON_LEDGERS {
113                    expiring_soon += 1;
114                }
115                total_cost += entry.extension_cost_stroops;
116                ttl_bearing_entries += 1;
117            }
118
119            detailed.push(StorageEntry {
120                key: entry.key.clone(),
121                label: None,
122                class,
123                current_ttl: entry.current_ttl,
124                days_remaining: entry.days_remaining,
125                extension_cost_stroops: entry.extension_cost_stroops,
126            });
127        }
128
129        let count = ttl_info.entries.len();
130        let ttl_summary = if ttl_bearing_entries > 0 {
131            let average_ttl = (total_ttl / ttl_bearing_entries as u64) as u32;
132            let minimum_ttl = if min_ttl == u32::MAX { 0 } else { min_ttl };
133            Some(TtlInfoSummary {
134                minimum_ttl,
135                maximum_ttl: max_ttl,
136                average_ttl,
137                expiring_entries_count: expiring_soon,
138                estimated_rent_cost: Some(total_cost),
139            })
140        } else {
141            None
142        };
143
144        Ok(StorageReport {
145            contract_id: contract_id.to_string(),
146            total_entries: count,
147            instance_entries,
148            persistent_entries,
149            temporary_entries,
150            other_entries,
151            total_size_bytes: None,
152            ttl_summary,
153            entries: detailed,
154        })
155    }
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161    use sdkt_rpc::SorobanRpcClient;
162    use std::io::{Read, Write};
163    use std::net::TcpListener;
164    use std::thread;
165    use stellar_xdr::{
166        AccountId, ContractDataDurability, LedgerKey, LedgerKeyAccount, LedgerKeyContractData,
167        PublicKey, ScAddress, ScVal, Uint256, WriteXdr,
168    };
169
170    const TEST_CONTRACT: &str = "CAE3U7JKESRWZHPEQ72DVNGOQ6WPA7HSPQZL5YV46NPCE4TMUPAGYMEC";
171
172    fn encode_ledger_key(key: &LedgerKey) -> String {
173        let mut buf = Vec::new();
174        let mut limited = stellar_xdr::Limited::new(&mut buf, stellar_xdr::Limits::none());
175        key.write_xdr(&mut limited).unwrap();
176        base64::engine::general_purpose::STANDARD.encode(&buf)
177    }
178
179    fn contract_address() -> ScAddress {
180        // All-zero contract address (valid XDR shape, value irrelevant for classification).
181        ScAddress::Contract(stellar_xdr::ContractId(stellar_xdr::Hash([0u8; 32])))
182    }
183
184    #[test]
185    fn test_classify_instance() {
186        let key = LedgerKey::ContractData(LedgerKeyContractData {
187            contract: contract_address(),
188            key: ScVal::LedgerKeyContractInstance,
189            durability: ContractDataDurability::Persistent,
190        });
191        assert_eq!(
192            classify_key(&encode_ledger_key(&key)),
193            StorageClass::Instance
194        );
195    }
196
197    #[test]
198    fn test_classify_persistent() {
199        let key = LedgerKey::ContractData(LedgerKeyContractData {
200            contract: contract_address(),
201            key: ScVal::U32(1),
202            durability: ContractDataDurability::Persistent,
203        });
204        assert_eq!(
205            classify_key(&encode_ledger_key(&key)),
206            StorageClass::Persistent
207        );
208    }
209
210    #[test]
211    fn test_classify_temporary() {
212        let key = LedgerKey::ContractData(LedgerKeyContractData {
213            contract: contract_address(),
214            key: ScVal::U32(2),
215            durability: ContractDataDurability::Temporary,
216        });
217        assert_eq!(
218            classify_key(&encode_ledger_key(&key)),
219            StorageClass::Temporary
220        );
221    }
222
223    #[test]
224    fn test_classify_invalid_base64_is_other() {
225        assert_eq!(classify_key("not-valid-base64!!!"), StorageClass::Other);
226    }
227
228    #[test]
229    fn test_storage_class_label() {
230        assert_eq!(StorageClass::Instance.label(), "instance");
231        assert_eq!(StorageClass::Persistent.label(), "persistent");
232        assert_eq!(StorageClass::Temporary.label(), "temporary");
233        assert_eq!(StorageClass::Other.label(), "other");
234    }
235
236    #[tokio::test]
237    async fn test_inspect_empty_contract_id() {
238        let client = SorobanRpcClient::new("http://localhost");
239        let analyzer = StorageAnalyzer::new(client);
240        let err = analyzer.inspect_contract_storage("").await.unwrap_err();
241        assert!(matches!(err, StorageError::InvalidContractId(_)));
242    }
243
244    #[tokio::test]
245    async fn test_inspect_contract_storage_keys_rejects_invalid_key_offline() {
246        let client = SorobanRpcClient::new("http://127.0.0.1:1");
247        let analyzer = StorageAnalyzer::new(client);
248        let err = analyzer
249            .inspect_contract_storage_keys(TEST_CONTRACT, &["invalid-key".to_string()])
250            .await
251            .unwrap_err();
252        assert!(matches!(err, StorageError::Rpc(_)));
253    }
254
255    #[tokio::test]
256    async fn test_merged_classification_multiple_classes() {
257        let instance_key = encode_ledger_key(&LedgerKey::ContractData(LedgerKeyContractData {
258            contract: contract_address(),
259            key: ScVal::LedgerKeyContractInstance,
260            durability: ContractDataDurability::Persistent,
261        }));
262        let persistent_key = encode_ledger_key(&LedgerKey::ContractData(LedgerKeyContractData {
263            contract: contract_address(),
264            key: ScVal::U32(1),
265            durability: ContractDataDurability::Persistent,
266        }));
267        let temporary_key = encode_ledger_key(&LedgerKey::ContractData(LedgerKeyContractData {
268            contract: contract_address(),
269            key: ScVal::U32(2),
270            durability: ContractDataDurability::Temporary,
271        }));
272
273        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
274        let addr = listener.local_addr().unwrap();
275        let url = format!("http://{addr}");
276
277        let ik = instance_key.clone();
278        let pk = persistent_key.clone();
279        let tk = temporary_key.clone();
280
281        thread::spawn(move || {
282            for conn in listener.incoming() {
283                let Ok(mut sock) = conn else { break };
284                let mut buf = [0u8; 16384];
285                let n = sock.read(&mut buf).unwrap_or(0);
286                let req = String::from_utf8_lossy(&buf[..n]).to_string();
287                let body = if req.contains("getLatestLedger") {
288                    r#"{"jsonrpc":"2.0","id":1,"result":{"id":"mock","sequence":100}}"#.to_string()
289                } else if req.contains("getLedgerEntries") {
290                    format!(
291                        r#"{{"jsonrpc":"2.0","id":1,"result":{{"entries":[
292                            {{"key":"{ik}","xdr":"AAAAAQAAAABpc25nAAAA","lastModifiedLedgerSeq":100,"liveUntilLedgerSeq":200}},
293                            {{"key":"{pk}","xdr":"AAAAAQAAAABpc25nAAAA","lastModifiedLedgerSeq":100,"liveUntilLedgerSeq":300}},
294                            {{"key":"{tk}","xdr":"AAAAAQAAAABpc25nAAAA","lastModifiedLedgerSeq":100,"liveUntilLedgerSeq":400}}
295                        ],"latestLedger":100}}}}"#
296                    )
297                } else {
298                    r#"{"jsonrpc":"2.0","id":1,"error":{"code":-32601,"message":"not found"}}"#
299                        .to_string()
300                };
301                let resp = format!(
302                    "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
303                    body.len(),
304                    body
305                );
306                let _ = sock.write_all(resp.as_bytes());
307            }
308        });
309
310        let client = SorobanRpcClient::new(&url);
311        let analyzer = StorageAnalyzer::new(client);
312
313        let report = analyzer
314            .inspect_contract_storage_keys(
315                TEST_CONTRACT,
316                &[persistent_key.clone(), temporary_key.clone()],
317            )
318            .await
319            .unwrap();
320
321        assert_eq!(report.total_entries, 3);
322        assert_eq!(report.instance_entries, 1);
323        assert_eq!(report.persistent_entries, 1);
324        assert_eq!(report.temporary_entries, 1);
325        assert_eq!(report.other_entries, 0);
326
327        assert_eq!(report.entries.len(), 3);
328        assert_eq!(report.entries[0].class, StorageClass::Instance);
329        assert_eq!(report.entries[0].current_ttl, 100);
330        assert_eq!(report.entries[1].class, StorageClass::Persistent);
331        assert_eq!(report.entries[1].current_ttl, 200);
332        assert_eq!(report.entries[2].class, StorageClass::Temporary);
333        assert_eq!(report.entries[2].current_ttl, 300);
334
335        let summary = report.ttl_summary.unwrap();
336        assert_eq!(summary.minimum_ttl, 100);
337        assert_eq!(summary.maximum_ttl, 300);
338        assert_eq!(summary.average_ttl, 200);
339    }
340
341    #[tokio::test]
342    async fn test_inspect_contract_storage_no_keys_instance_only() {
343        let instance_key = encode_ledger_key(&LedgerKey::ContractData(LedgerKeyContractData {
344            contract: contract_address(),
345            key: ScVal::LedgerKeyContractInstance,
346            durability: ContractDataDurability::Persistent,
347        }));
348
349        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
350        let addr = listener.local_addr().unwrap();
351        let url = format!("http://{addr}");
352
353        let ik = instance_key.clone();
354
355        thread::spawn(move || {
356            for conn in listener.incoming() {
357                let Ok(mut sock) = conn else { break };
358                let mut buf = [0u8; 16384];
359                let n = sock.read(&mut buf).unwrap_or(0);
360                let req = String::from_utf8_lossy(&buf[..n]).to_string();
361                let body = if req.contains("getLatestLedger") {
362                    r#"{"jsonrpc":"2.0","id":1,"result":{"id":"mock","sequence":100}}"#.to_string()
363                } else if req.contains("getLedgerEntries") {
364                    format!(
365                        r#"{{"jsonrpc":"2.0","id":1,"result":{{"entries":[
366                            {{"key":"{ik}","xdr":"AAAAAQAAAABpc25nAAAA","lastModifiedLedgerSeq":100,"liveUntilLedgerSeq":200}}
367                        ],"latestLedger":100}}}}"#
368                    )
369                } else {
370                    r#"{"jsonrpc":"2.0","id":1,"error":{"code":-32601,"message":"not found"}}"#
371                        .to_string()
372                };
373                let resp = format!(
374                    "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
375                    body.len(),
376                    body
377                );
378                let _ = sock.write_all(resp.as_bytes());
379            }
380        });
381
382        let client = SorobanRpcClient::new(&url);
383        let analyzer = StorageAnalyzer::new(client);
384
385        let report = analyzer
386            .inspect_contract_storage(TEST_CONTRACT)
387            .await
388            .unwrap();
389
390        assert_eq!(report.total_entries, 1);
391        assert_eq!(report.instance_entries, 1);
392        assert_eq!(report.persistent_entries, 0);
393        assert_eq!(report.temporary_entries, 0);
394        assert_eq!(report.entries.len(), 1);
395        assert_eq!(report.entries[0].class, StorageClass::Instance);
396    }
397
398    #[tokio::test]
399    async fn test_other_entries_excluded_from_ttl_aggregates() {
400        let instance_key = encode_ledger_key(&LedgerKey::ContractData(LedgerKeyContractData {
401            contract: contract_address(),
402            key: ScVal::LedgerKeyContractInstance,
403            durability: ContractDataDurability::Persistent,
404        }));
405        let account_key = encode_ledger_key(&LedgerKey::Account(LedgerKeyAccount {
406            account_id: AccountId(PublicKey::PublicKeyTypeEd25519(Uint256([1u8; 32]))),
407        }));
408
409        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
410        let addr = listener.local_addr().unwrap();
411        let url = format!("http://{addr}");
412
413        let ik = instance_key.clone();
414        let ak = account_key.clone();
415
416        thread::spawn(move || {
417            for conn in listener.incoming() {
418                let Ok(mut sock) = conn else { break };
419                let mut buf = [0u8; 16384];
420                let n = sock.read(&mut buf).unwrap_or(0);
421                let req = String::from_utf8_lossy(&buf[..n]).to_string();
422                let body = if req.contains("getLatestLedger") {
423                    r#"{"jsonrpc":"2.0","id":1,"result":{"id":"mock","sequence":100}}"#.to_string()
424                } else if req.contains("getLedgerEntries") {
425                    format!(
426                        r#"{{"jsonrpc":"2.0","id":1,"result":{{"entries":[
427                            {{"key":"{ik}","xdr":"AAAAAQAAAABpc25nAAAA","lastModifiedLedgerSeq":100,"liveUntilLedgerSeq":20100}},
428                            {{"key":"{ak}","xdr":"AAAAAQAAAABpc25nAAAA","lastModifiedLedgerSeq":100}}
429                        ],"latestLedger":100}}}}"#
430                    )
431                } else {
432                    r#"{"jsonrpc":"2.0","id":1,"error":{"code":-32601,"message":"not found"}}"#
433                        .to_string()
434                };
435                let resp = format!(
436                    "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
437                    body.len(),
438                    body
439                );
440                let _ = sock.write_all(resp.as_bytes());
441            }
442        });
443
444        let client = SorobanRpcClient::new(&url);
445        let analyzer = StorageAnalyzer::new(client);
446
447        let report = analyzer
448            .inspect_contract_storage_keys(TEST_CONTRACT, &[account_key])
449            .await
450            .unwrap();
451
452        assert_eq!(report.total_entries, 2);
453        assert_eq!(report.instance_entries, 1);
454        assert_eq!(report.other_entries, 1);
455
456        // Account entry without liveUntilLedgerSeq has current_ttl = 0 and class = Other.
457        // It must NOT drag down minimum_ttl to 0, nor be counted in average_ttl or expiring_soon.
458        let summary = report
459            .ttl_summary
460            .expect("ttl_summary present for instance entry");
461        assert_eq!(summary.minimum_ttl, 20000);
462        assert_eq!(summary.maximum_ttl, 20000);
463        assert_eq!(summary.average_ttl, 20000);
464        assert_eq!(summary.expiring_entries_count, 0);
465    }
466
467    #[tokio::test]
468    async fn test_all_other_entries_omits_ttl_summary() {
469        let account_key = encode_ledger_key(&LedgerKey::Account(LedgerKeyAccount {
470            account_id: AccountId(PublicKey::PublicKeyTypeEd25519(Uint256([2u8; 32]))),
471        }));
472
473        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
474        let addr = listener.local_addr().unwrap();
475        let url = format!("http://{addr}");
476
477        let ak = account_key.clone();
478
479        thread::spawn(move || {
480            for conn in listener.incoming() {
481                let Ok(mut sock) = conn else { break };
482                let mut buf = [0u8; 16384];
483                let n = sock.read(&mut buf).unwrap_or(0);
484                let req = String::from_utf8_lossy(&buf[..n]).to_string();
485                let body = if req.contains("getLatestLedger") {
486                    r#"{"jsonrpc":"2.0","id":1,"result":{"id":"mock","sequence":100}}"#.to_string()
487                } else if req.contains("getLedgerEntries") {
488                    // Only the account entry is returned by the RPC
489                    format!(
490                        r#"{{"jsonrpc":"2.0","id":1,"result":{{"entries":[
491                            {{"key":"{ak}","xdr":"AAAAAQAAAABpc25nAAAA","lastModifiedLedgerSeq":100}}
492                        ],"latestLedger":100}}}}"#
493                    )
494                } else {
495                    r#"{"jsonrpc":"2.0","id":1,"error":{"code":-32601,"message":"not found"}}"#
496                        .to_string()
497                };
498                let resp = format!(
499                    "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
500                    body.len(),
501                    body
502                );
503                let _ = sock.write_all(resp.as_bytes());
504            }
505        });
506
507        let client = SorobanRpcClient::new(&url);
508        let analyzer = StorageAnalyzer::new(client);
509
510        let report = analyzer
511            .inspect_contract_storage_keys(TEST_CONTRACT, &[account_key])
512            .await
513            .unwrap();
514
515        assert_eq!(report.total_entries, 1);
516        assert_eq!(report.other_entries, 1);
517        assert_eq!(report.instance_entries, 0);
518        assert!(report.ttl_summary.is_none());
519    }
520}