rusthound_ce/storage/
iter.rs

1use std::error::Error;
2use std::fs::File;
3use std::io::{BufReader, Read};
4use std::marker::PhantomData;
5
6pub type DiskStorageReader<T> = BincodeIterator<T, BufReader<File>>;
7
8/// Lazy iterator for bincode-encoded, length-prefixed data
9pub struct BincodeIterator<T, R: Read> {
10    reader: R,
11    _phantom: PhantomData<T>,
12}
13
14impl<T> BincodeIterator<T, BufReader<File>>
15where
16    T: bincode::Decode<()>,
17{
18    /// Create a new iterator from a file path
19    pub fn from_path(file_path: impl AsRef<std::path::Path>) -> std::io::Result<Self> {
20        let file = File::open(file_path)?;
21        let reader = BufReader::new(file);
22        Ok(Self {
23            reader,
24            _phantom: PhantomData,
25        })
26    }
27
28    pub fn from_file(file: std::fs::File) -> Self {
29        Self {
30            reader: BufReader::new(file),
31            _phantom: PhantomData,
32        }
33    }
34}
35
36impl<T, R: Read> BincodeIterator<T, R>
37where
38    T: bincode::Decode<()>,
39{
40    /// Create a new iterator from any reader
41    pub fn new(reader: R) -> Self {
42        Self {
43            reader,
44            _phantom: PhantomData,
45        }
46    }
47}
48
49impl<T, R: Read> BincodeIterator<T, R> {
50    /// Read the next length-prefixed record as raw bytes, WITHOUT decoding it.
51    ///
52    /// Decoding (bincode) is CPU-heavy; reading the raw blob here (cheap I/O)
53    /// lets the caller push the decode work onto worker threads instead of
54    /// paying for it single-threaded on the reader thread.
55    pub fn next_raw(&mut self) -> Option<std::io::Result<Vec<u8>>> {
56        let mut len_bytes = [0u8; 4];
57        match self.reader.read_exact(&mut len_bytes) {
58            Ok(()) => {}
59            Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return None,
60            Err(e) => return Some(Err(e)),
61        }
62
63        let len = u32::from_le_bytes(len_bytes) as usize;
64
65        let mut data = vec![0u8; len];
66        if let Err(e) = self.reader.read_exact(&mut data) {
67            return Some(Err(e));
68        }
69
70        Some(Ok(data))
71    }
72}
73
74impl<T, R: Read> Iterator for BincodeIterator<T, R>
75where
76    T: bincode::Decode<()>,
77{
78    type Item = Result<T, Box<dyn Error>>;
79
80    fn next(&mut self) -> Option<Self::Item> {
81        // Try to read length prefix
82        let mut len_bytes = [0u8; 4];
83        match self.reader.read_exact(&mut len_bytes) {
84            Ok(()) => {}
85            Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
86                // Clean EOF - no more records
87                return None;
88            }
89            Err(e) => return Some(Err(e.into())),
90        }
91
92        let len = u32::from_le_bytes(len_bytes) as usize;
93
94        // Validate length to prevent excessive allocation
95        // if len > 100_000_000 {
96        //     // 100MB limit, adjust as needed
97        //     return Some(Err(format!(
98        //         "Item length {len} exceeds maximum allowed size"
99        //     )
100        //     .into()));
101        // }
102
103        // Read the exact amount of data for this item
104        let mut data = vec![0u8; len];
105        if let Err(e) = self.reader.read_exact(&mut data) {
106            return Some(Err(format!("Failed to read {len} bytes: {e}").into()));
107        }
108
109        // Decode the item
110        match bincode::decode_from_slice::<T, _>(&data, bincode::config::standard()) {
111            Ok((item, _)) => Some(Ok(item)),
112            Err(e) => Some(Err(format!("Failed to decode item: {e:?}").into())),
113        }
114    }
115}
rusthound_ce/storage/iter.rs

rusthound_ce/storage/
iter.rs
