Skip to main content

hadris_iso/read/
chunk.rs

1use alloc::{vec, vec::Vec};
2
3use super::{DirEntry, Extent, IsoImage, Read, Seek, io};
4
5/// Iterator over file chunks.
6///
7/// Yields chunks of up to `N` bytes from a file, handling multi-extent files.
8///
9pub struct FileChunkIterator<'a, DATA: Read + Seek, const N: usize> {
10    pub(crate) image: &'a IsoImage<DATA>,
11    pub(crate) extents: alloc::vec::Vec<Extent>,
12    pub(crate) current_extent: usize,
13    pub(crate) offset_in_extent: usize,
14    pub(crate) bytes_remaining: u64,
15    pub(crate) total_size: u64,
16}
17
18io_transform! {
19impl<'a, DATA: Read + Seek, const N: usize> FileChunkIterator<'a, DATA, N> {
20    /// Creates a new chunked iterator for a file entry.
21    pub(crate) fn new(image: &'a IsoImage<DATA>, entry: &DirEntry) -> Self {
22        let extents: Vec<Extent> = entry.extents().collect();
23        let bytes_remaining = entry.total_size();
24        let total_size = entry.total_size();
25
26        Self {
27            image,
28            extents,
29            current_extent: 0,
30            offset_in_extent: 0,
31            bytes_remaining,
32            total_size
33        }
34    }
35
36    /// Returns the next chunk of data.
37    ///
38    /// Returns `Ok(Some(Vec<u8>))` with up to N bytes, `Ok(None)` at EOF.
39    pub async fn next_chunk(&mut self) -> io::Result<Option<Vec<u8>>> {
40        if self.bytes_remaining == 0 {
41            return Ok(None);
42        }
43        if N == 0 {
44            return Err(io::Error::new(
45                io::ErrorKind::InvalidInput,
46                "chunk size must be greater than zero",
47            ));
48        }
49
50        let chunk_size = N.min(usize::try_from(self.bytes_remaining).unwrap_or(usize::MAX));
51        let mut buffer = vec![0; chunk_size];
52        let mut read = 0;
53
54        while read < chunk_size {
55            let extent = match self.extents.get(self.current_extent) {
56                Some(e) => e,
57                None => {
58                    return Err(io::Error::new(
59                        io::ErrorKind::InvalidData,
60                        "file extent chain ended before its declared size",
61                    ));
62                }
63            };
64
65            let remaining_in_extent = (extent.length as usize).saturating_sub(self.offset_in_extent);
66            let to_read = (chunk_size - read).min(remaining_in_extent);
67
68            let byte_offset = (extent.sector.0 as u64 * 2048) + self.offset_in_extent as u64;
69            self.image
70                .read_bytes_at(byte_offset, &mut buffer[read..read + to_read])
71                .await?;
72
73            read += to_read;
74            self.offset_in_extent += to_read;
75            self.bytes_remaining -= to_read as u64;
76
77            if self.offset_in_extent >= extent.length as usize {
78                self.current_extent += 1;
79                self.offset_in_extent = 0;
80            }
81        }
82
83        buffer.truncate(read);
84        Ok(Some(buffer))
85    }
86
87    /// Returns the total size of the file in bytes.
88    pub fn total_size(&self) -> u64 {
89        self.total_size
90    }
91
92    /// Returns the current position in the file (bytes read so far).
93    pub fn position(&self) -> u64 {
94        self.total_size - self.bytes_remaining
95    }
96}
97} // io_transform!
98
99sync_only! {
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use alloc::{string::ToString, sync::Arc, vec};
104    use crate::{
105        IsoImage,
106        directory::FileFlags,
107        read::PathSeparator,
108        write::{options::*, *},
109    };
110    use std::io::Cursor;
111
112    fn write_both_endian(bytes: &mut [u8], offset: usize, value: u32) {
113        bytes[offset..offset + 4].copy_from_slice(&value.to_le_bytes());
114        bytes[offset + 4..offset + 8].copy_from_slice(&value.to_be_bytes());
115    }
116
117    #[test]
118    fn should_read_multi_extent_file_in_chunks() {
119        const CHUNK_SIZE: usize = 3000;
120        const EXTENT_SIZE: usize = 4096;
121        const SIZE: usize = EXTENT_SIZE * 2;
122
123        let input = InputFiles {
124            path_separator: PathSeparator::ForwardSlash,
125            files: vec![File::File {
126                name: Arc::new("TESTFILE".into()),
127                contents: vec![0xAA; SIZE],
128            }],
129        };
130        let options = IsoFormatOptions {
131            volume_name: "EMPTY".to_string(),
132            system_id: Some("SYSTEM".to_string()),
133            volume_set_id: Some("VOL_SET_ID".to_string()),
134            publisher_id: Some("PUBLISHER_ID".to_string()),
135            preparer_id: Some("PREPARER_ID".to_string()),
136            application_id: Some("APP_ID".to_string()),
137            sector_size: 2048,
138            path_separator: PathSeparator::ForwardSlash,
139            features: CreationFeatures {
140                ..Default::default()
141            },
142            strict_charset: false,
143        };
144
145        let cursor = Cursor::new(vec![0u8; 1024 * 1024]);
146        let output = IsoImageWriter::create(cursor, input, options).unwrap();
147        let mut bytes = output.into_inner();
148        let root_record = 16 * 2048 + 156;
149        let root_sector = u32::from_le_bytes(
150            bytes[root_record + 2..root_record + 6]
151                .try_into()
152                .unwrap(),
153        );
154        let mut file_record = root_sector as usize * 2048;
155        file_record += bytes[file_record] as usize;
156        file_record += bytes[file_record] as usize;
157        let record_len = bytes[file_record] as usize;
158        let original_extent = u32::from_le_bytes(
159            bytes[file_record + 2..file_record + 6]
160                .try_into()
161                .unwrap(),
162        );
163        let mut continuation = bytes[file_record..file_record + record_len].to_vec();
164
165        bytes[file_record + 25] |= FileFlags::NOT_FINAL.bits();
166        write_both_endian(&mut bytes, file_record + 10, EXTENT_SIZE as u32);
167        write_both_endian(&mut continuation, 2, original_extent + 2);
168        write_both_endian(&mut continuation, 10, EXTENT_SIZE as u32);
169        continuation[25] &= !FileFlags::NOT_FINAL.bits();
170        let continuation_start = file_record + record_len;
171        bytes[continuation_start..continuation_start + record_len]
172            .copy_from_slice(&continuation);
173
174        let image = IsoImage::open(Cursor::new(bytes)).expect("Failed to parse ISO image");
175        let root_dir = image.root_dir();
176        let iso_dir = root_dir.iter(&image);
177
178        let mut entries = iso_dir.entries();
179        entries.next().unwrap().expect("Failed to parse iso dir");
180        entries.next().unwrap().expect("Failed to parse iso dir");
181        let file = entries.next().unwrap().expect("Failed to parse iso file");
182        assert_eq!(file.additional_extents.len(), 1);
183
184        let mut zero_sized = image.read_file_chunked::<0>(&file).unwrap();
185        assert_eq!(
186            zero_sized.next_chunk().unwrap_err().kind(),
187            io::ErrorKind::InvalidInput
188        );
189
190        let mut iter = image.read_file_chunked::<CHUNK_SIZE>(&file).unwrap();
191        let mut total_read = 0;
192        let mut chunk_count = 0;
193
194        while let Some(chunk) = iter.next_chunk().unwrap() {
195            chunk_count += 1;
196            total_read += chunk.len();
197
198            assert!(chunk.iter().all(|&b| b == 0xAA));
199        }
200
201        assert_eq!(total_read, SIZE);
202        let expected_chunks = SIZE.div_ceil(CHUNK_SIZE);
203        assert_eq!(chunk_count, expected_chunks);
204        assert_eq!(iter.position(), SIZE as u64);
205        assert_eq!(iter.total_size(), SIZE as u64);
206    }
207}
208}