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> Iterator for BincodeIterator<T, R>
50where
51    T: bincode::Decode<()>,
52{
53    type Item = Result<T, Box<dyn Error>>;
54
55    fn next(&mut self) -> Option<Self::Item> {
56        // Try to read length prefix
57        let mut len_bytes = [0u8; 4];
58        match self.reader.read_exact(&mut len_bytes) {
59            Ok(()) => {}
60            Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
61                // Clean EOF - no more records
62                return None;
63            }
64            Err(e) => return Some(Err(e.into())),
65        }
66
67        let len = u32::from_le_bytes(len_bytes) as usize;
68
69        // Validate length to prevent excessive allocation
70        // if len > 100_000_000 {
71        //     // 100MB limit, adjust as needed
72        //     return Some(Err(format!(
73        //         "Item length {len} exceeds maximum allowed size"
74        //     )
75        //     .into()));
76        // }
77
78        // Read the exact amount of data for this item
79        let mut data = vec![0u8; len];
80        if let Err(e) = self.reader.read_exact(&mut data) {
81            return Some(Err(format!("Failed to read {len} bytes: {e}").into()));
82        }
83
84        // Decode the item
85        match bincode::decode_from_slice::<T, _>(&data, bincode::config::standard()) {
86            Ok((item, _)) => Some(Ok(item)),
87            Err(e) => Some(Err(format!("Failed to decode item: {e:?}").into())),
88        }
89    }
90}
rusthound_ce/storage/iter.rs

rusthound_ce/storage/
iter.rs
