1use crate::dht::content_addressing::ContentAddress;
6use anyhow::{Result, anyhow};
7use chrono::{DateTime, Utc};
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10use std::sync::Arc;
11use tokio::sync::RwLock;
12
13#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
15pub struct OperationId(String);
16
17impl Default for OperationId {
18 fn default() -> Self {
19 Self::new()
20 }
21}
22
23impl OperationId {
24 pub fn new() -> Self {
25 Self(uuid::Uuid::new_v4().to_string())
26 }
27}
28
29#[derive(Debug, Clone, Serialize, Deserialize)]
31pub enum OperationType {
32 Store,
33 Retrieve,
34 Verify,
35 Delete,
36 Batch(Vec<OperationType>),
37}
38
39#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
41pub struct NodeId(String);
42
43impl NodeId {
44 pub fn new(id: &str) -> Self {
45 Self(id.to_string())
46 }
47}
48
49#[derive(Debug, Clone, Serialize, Deserialize)]
51pub struct OperationMetadata {
52 pub size_bytes: usize,
53 pub chunk_count: Option<usize>,
54 pub redundancy_level: Option<f64>,
55 pub custom: HashMap<String, String>,
56}
57
58#[derive(Debug, Clone, Serialize, Deserialize)]
60pub struct MlKemSignature {
61 data: Vec<u8>,
62}
63
64impl MlKemSignature {
65 pub fn placeholder() -> Self {
66 Self {
67 data: vec![0u8; 64],
68 }
69 }
70
71 fn _verify(&self, _data: &[u8], _public_key: &MlKemPublicKey) -> bool {
72 true
74 }
75}
76
77#[derive(Debug, Clone, Serialize, Deserialize)]
79pub struct MlKemPublicKey {
80 data: Vec<u8>,
81}
82
83#[derive(Debug, Clone)]
85pub struct MlKemPrivateKey {
86 _data: Vec<u8>,
87}
88
89impl MlKemPrivateKey {
90 fn sign(&self, _data: &[u8]) -> MlKemSignature {
91 MlKemSignature::placeholder()
92 }
93}
94
95#[derive(Debug, Clone, Serialize, Deserialize)]
97pub struct StorageProof {
98 pub chunk_hashes: Vec<ContentAddress>,
99 pub storage_commitment: ContentAddress,
100 pub merkle_proof: Vec<ContentAddress>,
101}
102
103#[derive(Debug, Clone, Serialize, Deserialize)]
105pub struct RetrievalProof {
106 pub retrieved_hash: ContentAddress,
107 pub retrieval_time: DateTime<Utc>,
108 pub bandwidth_used: usize,
109}
110
111#[derive(Debug, Clone, Serialize, Deserialize)]
113pub struct WitnessProof {
114 pub node_id: NodeId,
115 pub operation_hash: ContentAddress,
116 pub node_signature: MlKemSignature,
117 pub storage_proof: Option<StorageProof>,
118 pub retrieval_proof: Option<RetrievalProof>,
119}
120
121#[derive(Debug, Clone, Serialize, Deserialize)]
123pub struct WitnessReceipt {
124 pub operation_id: OperationId,
125 pub operation_type: OperationType,
126 pub content_hash: ContentAddress,
127 pub timestamp: DateTime<Utc>,
128 pub participating_nodes: Vec<NodeId>,
129 pub operation_metadata: OperationMetadata,
130 pub signature: MlKemSignature,
131 pub witness_proofs: Vec<WitnessProof>,
132}
133
134impl WitnessReceipt {
135 pub fn hash(&self) -> Result<ContentAddress> {
137 let json = serde_json::to_vec(self)?;
138 Ok(ContentAddress::from_bytes(blake3::hash(&json).as_bytes()))
139 }
140
141 pub fn verify_consistency(&self) -> bool {
143 for proof in &self.witness_proofs {
145 if !self.participating_nodes.contains(&proof.node_id) {
146 return false;
147 }
148 }
149 true
150 }
151}
152
153pub struct DhtOperation {
155 pub operation_type: OperationType,
156 pub content_hash: ContentAddress,
157 pub nodes: Vec<NodeId>,
158 pub metadata: OperationMetadata,
159}
160
161struct ReceiptStorage {
163 receipts: HashMap<OperationId, WitnessReceipt>,
164 by_content: HashMap<ContentAddress, Vec<OperationId>>,
165}
166
167impl ReceiptStorage {
168 fn new() -> Self {
169 Self {
170 receipts: HashMap::new(),
171 by_content: HashMap::new(),
172 }
173 }
174
175 fn store(&mut self, receipt: WitnessReceipt) {
176 let content_hash = receipt.content_hash.clone();
177 let operation_id = receipt.operation_id.clone();
178
179 self.by_content
180 .entry(content_hash)
181 .or_default()
182 .push(operation_id.clone());
183
184 self.receipts.insert(operation_id, receipt);
185 }
186
187 fn _get(&self, operation_id: &OperationId) -> Option<&WitnessReceipt> {
188 self.receipts.get(operation_id)
189 }
190
191 fn get_by_content(&self, content_hash: &ContentAddress) -> Vec<&WitnessReceipt> {
192 self.by_content
193 .get(content_hash)
194 .map(|ids| ids.iter().filter_map(|id| self.receipts.get(id)).collect())
195 .unwrap_or_default()
196 }
197}
198
199pub struct WitnessReceiptSystem {
201 signing_key: MlKemPrivateKey,
202 _verification_keys: Arc<RwLock<HashMap<NodeId, MlKemPublicKey>>>,
203 receipt_store: Arc<RwLock<ReceiptStorage>>,
204}
205
206impl WitnessReceiptSystem {
207 pub fn new() -> Self {
209 Self {
210 signing_key: MlKemPrivateKey {
211 _data: vec![0u8; 32],
212 },
213 _verification_keys: Arc::new(RwLock::new(HashMap::new())),
214 receipt_store: Arc::new(RwLock::new(ReceiptStorage::new())),
215 }
216 }
217
218 pub async fn create_receipt(&self, operation: &DhtOperation) -> Result<WitnessReceipt> {
220 let operation_id = OperationId::new();
221
222 let witness_proofs = operation
224 .nodes
225 .iter()
226 .map(|node_id| WitnessProof {
227 node_id: node_id.clone(),
228 operation_hash: operation.content_hash.clone(),
229 node_signature: MlKemSignature::placeholder(),
230 storage_proof: match operation.operation_type {
231 OperationType::Store => Some(StorageProof {
232 chunk_hashes: vec![operation.content_hash.clone()],
233 storage_commitment: operation.content_hash.clone(),
234 merkle_proof: vec![],
235 }),
236 _ => None,
237 },
238 retrieval_proof: match operation.operation_type {
239 OperationType::Retrieve => Some(RetrievalProof {
240 retrieved_hash: operation.content_hash.clone(),
241 retrieval_time: Utc::now(),
242 bandwidth_used: operation.metadata.size_bytes,
243 }),
244 _ => None,
245 },
246 })
247 .collect();
248
249 let receipt = WitnessReceipt {
250 operation_id: operation_id.clone(),
251 operation_type: operation.operation_type.clone(),
252 content_hash: operation.content_hash.clone(),
253 timestamp: Utc::now(),
254 participating_nodes: operation.nodes.clone(),
255 operation_metadata: operation.metadata.clone(),
256 signature: self.signing_key.sign(b"receipt_data"),
257 witness_proofs,
258 };
259
260 let mut store = self.receipt_store.write().await;
262 store.store(receipt.clone());
263
264 Ok(receipt)
265 }
266
267 pub async fn verify_receipt(&self, receipt: &WitnessReceipt) -> Result<bool> {
269 if !receipt.verify_consistency() {
271 return Ok(false);
272 }
273
274 for proof in &receipt.witness_proofs {
279 if proof.node_signature.data.is_empty() {
281 return Ok(false);
282 }
283 }
284
285 Ok(true)
286 }
287
288 pub async fn create_batch_receipt(
290 &self,
291 operations: &[DhtOperation],
292 ) -> Result<WitnessReceipt> {
293 if operations.is_empty() {
294 return Err(anyhow!("Cannot create batch receipt for empty operations"));
295 }
296
297 let batch_type = OperationType::Batch(
299 operations
300 .iter()
301 .map(|op| op.operation_type.clone())
302 .collect(),
303 );
304
305 let batch_operation = DhtOperation {
307 operation_type: batch_type,
308 content_hash: operations[0].content_hash.clone(),
309 nodes: operations[0].nodes.clone(),
310 metadata: OperationMetadata {
311 size_bytes: operations.iter().map(|op| op.metadata.size_bytes).sum(),
312 chunk_count: Some(operations.len()),
313 redundancy_level: operations[0].metadata.redundancy_level,
314 custom: HashMap::new(),
315 },
316 };
317
318 self.create_receipt(&batch_operation).await
319 }
320
321 pub async fn get_audit_trail(
323 &self,
324 content_hash: &ContentAddress,
325 ) -> Result<Vec<WitnessReceipt>> {
326 let store = self.receipt_store.read().await;
327 let receipts = store
328 .get_by_content(content_hash)
329 .into_iter()
330 .cloned()
331 .collect();
332 Ok(receipts)
333 }
334
335 pub async fn verify_chain_integrity(&self, receipts: &[WitnessReceipt]) -> Result<bool> {
337 if receipts.is_empty() {
338 return Ok(true);
339 }
340
341 for receipt in receipts {
343 if !self.verify_receipt(receipt).await? {
344 return Ok(false);
345 }
346 }
347
348 for window in receipts.windows(2) {
350 if window[0].timestamp > window[1].timestamp {
351 return Ok(false);
352 }
353 }
354
355 Ok(true)
356 }
357}
358
359impl Default for WitnessReceiptSystem {
360 fn default() -> Self {
361 Self::new()
362 }
363}
364
365#[cfg(test)]
366mod tests {
367 use super::*;
368
369 #[tokio::test]
370 async fn test_receipt_creation_and_verification() -> Result<()> {
371 let receipt_system = WitnessReceiptSystem::new();
372 let operation = DhtOperation {
373 operation_type: OperationType::Store,
374 content_hash: ContentAddress::from_bytes(&[1u8; 32]),
375 nodes: vec![NodeId::new("node1"), NodeId::new("node2")],
376 metadata: OperationMetadata {
377 size_bytes: 1024,
378 chunk_count: Some(1),
379 redundancy_level: Some(0.5),
380 custom: HashMap::new(),
381 },
382 };
383
384 let receipt = receipt_system.create_receipt(&operation).await?;
385 assert!(receipt_system.verify_receipt(&receipt).await?);
386
387 Ok(())
388 }
389
390 #[tokio::test]
391 async fn test_audit_trail_reconstruction() -> Result<()> {
392 let receipt_system = WitnessReceiptSystem::new();
393 let content_hash = ContentAddress::from_bytes(&[42u8; 32]);
394
395 for op_type in [
397 OperationType::Store,
398 OperationType::Retrieve,
399 OperationType::Verify,
400 ] {
401 let operation = DhtOperation {
402 operation_type: op_type,
403 content_hash: content_hash.clone(),
404 nodes: vec![NodeId::new("node1")],
405 metadata: OperationMetadata {
406 size_bytes: 1024,
407 chunk_count: None,
408 redundancy_level: None,
409 custom: HashMap::new(),
410 },
411 };
412 receipt_system.create_receipt(&operation).await?;
413 }
414
415 let audit_trail = receipt_system.get_audit_trail(&content_hash).await?;
416 assert_eq!(audit_trail.len(), 3);
417 assert!(receipt_system.verify_chain_integrity(&audit_trail).await?);
418
419 Ok(())
420 }
421}