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
8pub 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 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 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 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}