dcp/stream/
dcp_stream.rs

1//! DCP Stream with Blake3 checksum for integrity verification.
2//!
3//! Provides streaming support with chunk sequencing, integrity verification,
4//! and retransmission request handling.
5
6use blake3::Hasher;
7use std::sync::atomic::{AtomicU32, Ordering};
8use std::sync::Mutex;
9
10use crate::binary::{ChunkFlags, StreamChunk};
11use crate::stream::ring_buffer::StreamRingBuffer;
12use crate::DCPError;
13
14/// Retransmission request for lost chunks
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub struct RetransmitRequest {
17    /// Starting sequence number
18    pub start_seq: u32,
19    /// Number of chunks to retransmit
20    pub count: u32,
21}
22
23/// DCP Stream with integrity verification
24pub struct DcpStream {
25    /// Underlying ring buffer
26    buffer: StreamRingBuffer,
27    /// Running Blake3 hasher for integrity
28    hasher: Mutex<Hasher>,
29    /// Serializes producer-side state transitions.
30    write_lock: Mutex<()>,
31    /// Next expected sequence number (consumer side)
32    next_seq: AtomicU32,
33    /// Next sequence to send (producer side)
34    send_seq: AtomicU32,
35    /// Stream ID
36    stream_id: u32,
37    /// Whether the stream is complete
38    complete: std::sync::atomic::AtomicBool,
39}
40
41impl DcpStream {
42    /// Create a new DCP stream with the given buffer capacity
43    pub fn new(stream_id: u32, capacity: usize) -> Self {
44        Self {
45            buffer: StreamRingBuffer::new(capacity),
46            hasher: Mutex::new(Hasher::new()),
47            write_lock: Mutex::new(()),
48            next_seq: AtomicU32::new(0),
49            send_seq: AtomicU32::new(0),
50            stream_id,
51            complete: std::sync::atomic::AtomicBool::new(false),
52        }
53    }
54
55    /// Get the stream ID
56    pub fn stream_id(&self) -> u32 {
57        self.stream_id
58    }
59
60    /// Get the underlying buffer
61    pub fn buffer(&self) -> &StreamRingBuffer {
62        &self.buffer
63    }
64
65    /// Check if the stream is complete
66    pub fn is_complete(&self) -> bool {
67        self.complete.load(Ordering::Acquire)
68    }
69
70    /// Write a chunk to the stream (producer side)
71    /// Returns the chunk header that should be sent
72    pub fn write_chunk(&self, data: &[u8], is_last: bool) -> Result<StreamChunk, DCPError> {
73        let _write_guard = self.write_lock.lock().unwrap();
74        if self.is_complete() {
75            return Err(DCPError::ValidationFailed);
76        }
77
78        if data.len() > u16::MAX as usize {
79            return Err(DCPError::OutOfBounds);
80        }
81        let total_len = StreamChunk::SIZE
82            .checked_add(data.len())
83            .ok_or(DCPError::OutOfBounds)?;
84        if self.buffer.available_space() < total_len {
85            return Err(DCPError::Backpressure);
86        }
87
88        let seq = self.send_seq.fetch_add(1, Ordering::AcqRel);
89        let is_first = seq == 0;
90
91        let flags = if is_first && is_last {
92            ChunkFlags::FIRST | ChunkFlags::LAST
93        } else if is_first {
94            ChunkFlags::FIRST
95        } else if is_last {
96            ChunkFlags::LAST
97        } else {
98            ChunkFlags::CONTINUE
99        };
100
101        // Write chunk header + data to buffer
102        let chunk = StreamChunk::new(seq, flags, data.len() as u16);
103
104        // Write header
105        self.buffer.push(chunk.as_bytes())?;
106
107        // Write payload
108        if !data.is_empty() {
109            self.buffer.push(data)?;
110        }
111
112        // Update hasher
113        {
114            let mut hasher = self.hasher.lock().unwrap();
115            hasher.update(data);
116        }
117
118        if is_last {
119            self.complete.store(true, Ordering::Release);
120        }
121
122        Ok(chunk)
123    }
124
125    /// Read a chunk from the stream (consumer side)
126    /// Returns the chunk header and data, or None if no data available
127    pub fn read_chunk(&self) -> Result<Option<(StreamChunk, Vec<u8>)>, DCPError> {
128        // First peek at the header
129        let mut header_buf = [0u8; StreamChunk::SIZE];
130        let peeked = self.buffer.peek(&mut header_buf);
131
132        if peeked < StreamChunk::SIZE {
133            return Ok(None);
134        }
135
136        let chunk = StreamChunk::from_bytes(&header_buf)?;
137        let chunk_len = chunk.len as usize;
138        let total_len = StreamChunk::SIZE + chunk_len;
139
140        // Check if we have the full chunk
141        if self.buffer.len() < total_len {
142            return Ok(None);
143        }
144
145        // Verify sequence number
146        let expected_seq = self.next_seq.load(Ordering::Acquire);
147        if chunk.sequence != expected_seq {
148            // Out of order - request retransmission
149            return Err(DCPError::ChecksumMismatch);
150        }
151
152        // Read the full chunk
153        let mut full_buf = vec![0u8; total_len];
154        self.buffer.pop(&mut full_buf);
155
156        // Extract payload
157        let payload = full_buf[StreamChunk::SIZE..].to_vec();
158
159        // Update sequence
160        self.next_seq.store(expected_seq + 1, Ordering::Release);
161
162        // Copy chunk values before returning
163        let result_chunk = StreamChunk::new(chunk.sequence, chunk.flags, chunk.len);
164
165        Ok(Some((result_chunk, payload)))
166    }
167
168    /// Get the current checksum of all data written so far
169    pub fn checksum(&self) -> [u8; 32] {
170        let hasher = self.hasher.lock().unwrap();
171        *hasher.finalize().as_bytes()
172    }
173
174    /// Verify the stream checksum against an expected value
175    pub fn verify_checksum(&self, expected: &[u8; 32]) -> bool {
176        &self.checksum() == expected
177    }
178
179    /// Create a retransmission request for missing chunks
180    pub fn request_retransmit(&self, missing_seq: u32) -> RetransmitRequest {
181        let expected = self.next_seq.load(Ordering::Acquire);
182        RetransmitRequest {
183            start_seq: expected,
184            count: missing_seq.saturating_sub(expected) + 1,
185        }
186    }
187
188    /// Get the next expected sequence number
189    pub fn next_expected_seq(&self) -> u32 {
190        self.next_seq.load(Ordering::Acquire)
191    }
192
193    /// Get the next sequence number to send
194    pub fn next_send_seq(&self) -> u32 {
195        self.send_seq.load(Ordering::Acquire)
196    }
197
198    /// Check if backpressure is active
199    pub fn is_backpressure(&self) -> bool {
200        self.buffer.backpressure().is_full()
201    }
202
203    /// Get available buffer space
204    pub fn available_space(&self) -> usize {
205        self.buffer.available_space()
206    }
207
208    /// Reset the stream for reuse
209    pub fn reset(&self) {
210        self.buffer.clear();
211        self.next_seq.store(0, Ordering::Release);
212        self.send_seq.store(0, Ordering::Release);
213        self.complete.store(false, Ordering::Release);
214        *self.hasher.lock().unwrap() = Hasher::new();
215    }
216}
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221
222    #[test]
223    fn test_stream_basic() {
224        let stream = DcpStream::new(1, 1024);
225
226        // Write first chunk
227        let chunk1 = stream.write_chunk(b"hello", false).unwrap();
228        assert!(chunk1.is_first());
229        assert!(!chunk1.is_last());
230        let seq1 = chunk1.sequence;
231        assert_eq!(seq1, 0);
232
233        // Write last chunk
234        let chunk2 = stream.write_chunk(b"world", true).unwrap();
235        assert!(!chunk2.is_first());
236        assert!(chunk2.is_last());
237        let seq2 = chunk2.sequence;
238        assert_eq!(seq2, 1);
239
240        assert!(stream.is_complete());
241    }
242
243    #[test]
244    fn test_stream_read_write() {
245        let stream = DcpStream::new(1, 1024);
246
247        stream.write_chunk(b"test", false).unwrap();
248        stream.write_chunk(b"data", true).unwrap();
249
250        // Read first chunk
251        let (chunk1, data1) = stream.read_chunk().unwrap().unwrap();
252        let seq1 = chunk1.sequence;
253        assert_eq!(seq1, 0);
254        assert_eq!(data1, b"test");
255
256        // Read second chunk
257        let (chunk2, data2) = stream.read_chunk().unwrap().unwrap();
258        let seq2 = chunk2.sequence;
259        assert_eq!(seq2, 1);
260        assert_eq!(data2, b"data");
261
262        // No more data
263        assert!(stream.read_chunk().unwrap().is_none());
264    }
265
266    #[test]
267    fn test_stream_checksum() {
268        let stream = DcpStream::new(1, 1024);
269
270        stream.write_chunk(b"hello", false).unwrap();
271        let checksum1 = stream.checksum();
272
273        stream.write_chunk(b"world", true).unwrap();
274        let checksum2 = stream.checksum();
275
276        // Checksums should be different
277        assert_ne!(checksum1, checksum2);
278
279        // Verify checksum
280        assert!(stream.verify_checksum(&checksum2));
281        assert!(!stream.verify_checksum(&checksum1));
282    }
283
284    #[test]
285    fn test_stream_single_chunk() {
286        let stream = DcpStream::new(1, 1024);
287
288        // Single chunk that is both first and last
289        let chunk = stream.write_chunk(b"single", true).unwrap();
290        assert!(chunk.is_first());
291        assert!(chunk.is_last());
292        let seq = chunk.sequence;
293        assert_eq!(seq, 0);
294    }
295
296    #[test]
297    fn test_retransmit_request() {
298        let stream = DcpStream::new(1, 1024);
299
300        let req = stream.request_retransmit(5);
301        assert_eq!(req.start_seq, 0);
302        assert_eq!(req.count, 6);
303    }
304
305    #[test]
306    fn test_stream_reset() {
307        let stream = DcpStream::new(1, 1024);
308
309        stream.write_chunk(b"data", true).unwrap();
310        assert!(stream.is_complete());
311        assert_eq!(stream.next_send_seq(), 1);
312
313        stream.reset();
314        assert!(!stream.is_complete());
315        assert_eq!(stream.next_send_seq(), 0);
316        assert_eq!(stream.next_expected_seq(), 0);
317    }
318}
dcp/stream/dcp_stream.rs

dcp/stream/
dcp_stream.rs
