tatara_engine/p2p/
chunk.rs

1use serde::{Deserialize, Serialize};
2
3/// Default chunk size: 256 KB.
4pub const DEFAULT_CHUNK_SIZE: usize = 256 * 1024;
5
6/// A content-addressed chunk of data.
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct Chunk {
9    /// BLAKE3 hash of the content.
10    pub hash: String,
11    /// Raw bytes.
12    pub data: Vec<u8>,
13    /// Size in bytes.
14    pub size: usize,
15}
16
17/// Metadata describing a complete object split into chunks.
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct ChunkManifest {
20    /// BLAKE3 hash of the complete object.
21    pub root_hash: String,
22    /// Total size in bytes.
23    pub total_size: usize,
24    /// Ordered list of chunk hashes.
25    pub chunks: Vec<String>,
26    /// Content type tag (e.g., "job_spec", "node_report", "log_archive", "nix_closure").
27    pub content_type: String,
28    /// Human label.
29    pub label: String,
30}
31
32impl Chunk {
33    pub fn new(data: Vec<u8>) -> Self {
34        let hash = blake3_hash(&data);
35        let size = data.len();
36        Self { hash, data, size }
37    }
38
39    pub fn verify(&self) -> bool {
40        blake3_hash(&self.data) == self.hash
41    }
42}
43
44impl ChunkManifest {
45    /// Split data into content-addressed chunks and produce a manifest.
46    pub fn from_data(data: &[u8], content_type: &str, label: &str) -> (Self, Vec<Chunk>) {
47        let root_hash = blake3_hash(data);
48        let mut chunks = Vec::new();
49        let mut chunk_hashes = Vec::new();
50
51        for chunk_data in data.chunks(DEFAULT_CHUNK_SIZE) {
52            let chunk = Chunk::new(chunk_data.to_vec());
53            chunk_hashes.push(chunk.hash.clone());
54            chunks.push(chunk);
55        }
56
57        let manifest = Self {
58            root_hash,
59            total_size: data.len(),
60            chunks: chunk_hashes,
61            content_type: content_type.to_string(),
62            label: label.to_string(),
63        };
64
65        (manifest, chunks)
66    }
67
68    /// Reassemble data from ordered chunks. Verifies each chunk hash.
69    pub fn reassemble(&self, chunks: &[Chunk]) -> Result<Vec<u8>, String> {
70        if chunks.len() != self.chunks.len() {
71            return Err(format!(
72                "Expected {} chunks, got {}",
73                self.chunks.len(),
74                chunks.len()
75            ));
76        }
77
78        let mut data = Vec::with_capacity(self.total_size);
79
80        for (i, chunk) in chunks.iter().enumerate() {
81            if chunk.hash != self.chunks[i] {
82                return Err(format!(
83                    "Chunk {} hash mismatch: expected {}, got {}",
84                    i, self.chunks[i], chunk.hash
85                ));
86            }
87            if !chunk.verify() {
88                return Err(format!("Chunk {} content verification failed", i));
89            }
90            data.extend_from_slice(&chunk.data);
91        }
92
93        // Verify reassembled data
94        let actual_hash = blake3_hash(&data);
95        if actual_hash != self.root_hash {
96            return Err(format!(
97                "Root hash mismatch: expected {}, got {}",
98                self.root_hash, actual_hash
99            ));
100        }
101
102        Ok(data)
103    }
104}
105
106fn blake3_hash(data: &[u8]) -> String {
107    let hash = blake3::hash(data);
108    hash.to_hex().to_string()
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    #[test]
116    fn test_chunk_roundtrip() {
117        let data = b"Hello, tatara! This is some test data for chunking.";
118        let (manifest, chunks) = ChunkManifest::from_data(data, "test", "test-data");
119
120        assert_eq!(manifest.total_size, data.len());
121        assert_eq!(manifest.chunks.len(), 1); // Small data = 1 chunk
122
123        let reassembled = manifest.reassemble(&chunks).unwrap();
124        assert_eq!(reassembled, data);
125    }
126
127    #[test]
128    fn test_multi_chunk_roundtrip() {
129        // Generate data larger than chunk size
130        let data: Vec<u8> = (0..DEFAULT_CHUNK_SIZE * 3 + 100)
131            .map(|i| (i % 256) as u8)
132            .collect();
133
134        let (manifest, chunks) = ChunkManifest::from_data(&data, "test", "big-data");
135
136        assert_eq!(manifest.chunks.len(), 4); // 3 full + 1 partial
137        assert_eq!(manifest.total_size, data.len());
138
139        let reassembled = manifest.reassemble(&chunks).unwrap();
140        assert_eq!(reassembled, data);
141    }
142
143    #[test]
144    fn test_chunk_verification() {
145        let chunk = Chunk::new(b"test data".to_vec());
146        assert!(chunk.verify());
147
148        let mut bad_chunk = chunk.clone();
149        bad_chunk.data = b"tampered".to_vec();
150        assert!(!bad_chunk.verify());
151    }
152}
tatara_engine/p2p/chunk.rs

tatara_engine/p2p/
chunk.rs
