ryg_rans_rs_cli/container/
reader.rs

1//! # Container reader — parse and validate RYGRANS containers
2
3use crate::container::block::{Block, BlockHeaderInfo};
4use crate::container::footer::FileFooter;
5use crate::container::header::FileHeader;
6use crate::container::{BLOCK_HEADER_SIZE, FOOTER_SIZE, HEADER_SIZE};
7use crate::error::{AppError, FormatError};
8use crate::limits::Limits;
9use sha2::{Digest, Sha256};
10use std::io::Read;
11
12/// Container reader that validates every field, bound, and hash.
13pub struct ContainerReader<R: Read> {
14    reader: R,
15    limits: Limits,
16    header: Option<FileHeader>,
17    bytes_read: u64,
18    blocks_seen: u64,
19    total_uncompressed: u64,
20    total_payload: u64,
21    hasher: Sha256,
22    decoded_hasher: Sha256,
23}
24
25impl<R: Read> ContainerReader<R> {
26    /// Create a new container reader.
27    pub fn new(reader: R, limits: Limits) -> Self {
28        Self {
29            reader,
30            limits,
31            header: None,
32            bytes_read: 0,
33            blocks_seen: 0,
34            total_uncompressed: 0,
35            total_payload: 0,
36            hasher: Sha256::new(),
37            decoded_hasher: Sha256::new(),
38        }
39    }
40
41    /// Read and validate the file header.
42    pub fn read_header(&mut self) -> Result<FileHeader, AppError> {
43        let mut header_buf = [0u8; HEADER_SIZE];
44        self.read_exact(&mut header_buf)?;
45        let header = FileHeader::from_bytes(&header_buf)?;
46
47        // Check declared block size against limits
48        self.limits.check_block_size(header.declared_block_size)?;
49
50        // Feed header to container hasher
51        self.hasher.update(&header_buf);
52
53        self.header = Some(header.clone());
54        Ok(header)
55    }
56
57    /// Read and process the next block.  Returns `None` when no more blocks
58    /// (footer will be read separately).
59    pub fn read_block<F>(&mut self, mut decode_fn: F) -> Result<Option<Block>, AppError>
60    where
61        F: FnMut(&BlockHeaderInfo, &[u8], &[u8]) -> Result<Vec<u8>, AppError>,
62    {
63        // Peek ahead to see if we've reached the footer
64        let mut peek_tag = [0u8; 4];
65        match self.peek_exact(&mut peek_tag) {
66            Ok(()) => {
67                if &peek_tag == b"END1" {
68                    return Ok(None); // footer follows
69                }
70                if &peek_tag != b"BLK1" {
71                    return Err(AppError::Format(FormatError {
72                        detail: format!("expected block or footer, got {:02x?}", peek_tag),
73                        block_index: Some(self.blocks_seen),
74                        offset: Some(self.bytes_read),
75                    }));
76                }
77            }
78            Err(AppError::Format(_)) => {
79                // Reached end before footer
80                return Err(AppError::Format(FormatError {
81                    detail: "truncated container: missing footer".into(),
82                    block_index: Some(self.blocks_seen),
83                    offset: Some(self.bytes_read),
84                }));
85            }
86            Err(e) => return Err(e),
87        }
88
89        // Read block header
90        let mut header_buf = [0u8; BLOCK_HEADER_SIZE];
91        self.read_exact(&mut header_buf)?;
92        let (info, _) = Block::parse_header(&header_buf, self.blocks_seen)?;
93
94        // Check limits
95        self.limits.check_block_size(info.uncompressed_length)?;
96        self.limits.check_payload_size(info.payload_length)?;
97        self.limits.check_model_size(info.model_length)?;
98        self.limits.check_block_count(self.blocks_seen + 1)?;
99        self.limits
100            .check_output_total(self.total_uncompressed, info.uncompressed_length as u64)?;
101
102        // Feed header to container hasher
103        self.hasher.update(&header_buf);
104
105        // Read model data
106        let model_data = if info.model_length > 0 {
107            let mut buf = vec![0u8; info.model_length as usize];
108            self.read_exact(&mut buf)?;
109            self.hasher.update(&buf);
110            buf
111        } else {
112            Vec::new()
113        };
114
115        // Read payload
116        let mut payload = vec![0u8; info.payload_length as usize];
117        self.read_exact(&mut payload)?;
118        self.hasher.update(&payload);
119
120        // Verify payload SHA-256
121        let mut pay_hasher = Sha256::new();
122        pay_hasher.update(&payload);
123        let pay_hash: [u8; 32] = pay_hasher.finalize().into();
124        if pay_hash != info.payload_sha256 {
125            return Err(AppError::Integrity(crate::error::IntegrityError {
126                detail: format!("payload hash mismatch at block {}", self.blocks_seen),
127                block_index: Some(self.blocks_seen),
128            }));
129        }
130
131        // Decode
132        let decoded = decode_fn(&info, &model_data, &payload)?;
133
134        // Verify decoded SHA-256
135        let mut dec_hasher = Sha256::new();
136        dec_hasher.update(&decoded);
137        let dec_hash: [u8; 32] = dec_hasher.finalize().into();
138        if info.block_kind != crate::container::BLOCK_KIND_RAW && dec_hash != info.decoded_sha256 {
139            return Err(AppError::Integrity(crate::error::IntegrityError {
140                detail: format!("decoded hash mismatch at block {}", self.blocks_seen),
141                block_index: Some(self.blocks_seen),
142            }));
143        }
144
145        // Accumulate
146        self.blocks_seen += 1;
147        self.total_uncompressed = self
148            .total_uncompressed
149            .checked_add(info.uncompressed_length as u64)
150            .ok_or_else(|| {
151                AppError::Format(FormatError {
152                    detail: "total uncompressed overflow".into(),
153                    block_index: Some(self.blocks_seen - 1),
154                    offset: None,
155                })
156            })?;
157        self.total_payload = self
158            .total_payload
159            .checked_add(info.payload_length as u64)
160            .ok_or_else(|| {
161                AppError::Format(FormatError {
162                    detail: "total payload overflow".into(),
163                    block_index: Some(self.blocks_seen - 1),
164                    offset: None,
165                })
166            })?;
167
168        self.decoded_hasher.update(&decoded);
169
170        Ok(Some(Block {
171            block_index: info.block_index,
172            block_kind: info.block_kind,
173            codec_id: info.codec_id,
174            scale_bits: info.scale_bits,
175            state_count: info.state_count,
176            uncompressed_length: info.uncompressed_length,
177            payload,
178            model_data,
179            payload_sha256: info.payload_sha256,
180            decoded_sha256: info.decoded_sha256,
181        }))
182    }
183
184    /// Read and verify the footer.
185    pub fn read_footer(&mut self) -> Result<FileFooter, AppError> {
186        let mut footer_buf = [0u8; FOOTER_SIZE];
187        self.read_exact(&mut footer_buf)?;
188        let footer = FileFooter::from_bytes(&footer_buf)?;
189
190        // Verify footer totals
191        footer.verify_totals(
192            self.blocks_seen,
193            self.total_uncompressed,
194            self.total_payload,
195        )?;
196
197        // Verify container hash
198        use sha2::Digest;
199        let hasher_clone = std::mem::replace(&mut self.hasher, Sha256::new());
200        let container_hash: [u8; 32] = hasher_clone.finalize().into();
201        if container_hash != footer.container_sha256 {
202            return Err(AppError::Integrity(crate::error::IntegrityError {
203                detail: "container hash mismatch".into(),
204                block_index: None,
205            }));
206        }
207
208        // Verify decoded stream hash
209        let dec_hasher_clone = std::mem::replace(&mut self.decoded_hasher, Sha256::new());
210        let decoded_hash: [u8; 32] = dec_hasher_clone.finalize().into();
211        if decoded_hash != footer.decoded_stream_sha256 {
212            return Err(AppError::Integrity(crate::error::IntegrityError {
213                detail: "decoded stream hash mismatch".into(),
214                block_index: None,
215            }));
216        }
217
218        Ok(footer)
219    }
220
221    /// Check that there is no trailing data after the footer.
222    pub fn check_trailing_data(&mut self) -> Result<(), AppError> {
223        let mut buf = [0u8; 1];
224        match self.reader.read_exact(&mut buf) {
225            Ok(()) => Err(AppError::Format(FormatError {
226                detail: "trailing data after footer".into(),
227                block_index: Some(self.blocks_seen),
228                offset: Some(self.bytes_read),
229            })),
230            Err(ref e) if e.kind() == std::io::ErrorKind::UnexpectedEof => Ok(()),
231            Err(e) => Err(AppError::Io(crate::error::IoError {
232                path: None,
233                detail: e.to_string(),
234            })),
235        }
236    }
237
238    // ---- Private helpers ----
239
240    fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), AppError> {
241        self.reader.read_exact(buf).map_err(|e| {
242            AppError::Format(FormatError {
243                detail: format!("read error: {}", e),
244                block_index: Some(self.blocks_seen),
245                offset: Some(self.bytes_read),
246            })
247        })?;
248        self.bytes_read += buf.len() as u64;
249        Ok(())
250    }
251
252    fn peek_exact(&mut self, buf: &mut [u8]) -> Result<(), AppError> {
253        // Use std::io::Read::read to peek without consuming
254        // This requires buffered I/O.  For simplicity, we read and
255        // the caller is expected to use a BufReader.
256        self.read_exact(buf)
257    }
258}
ryg_rans_rs_cli/container/reader.rs

ryg_rans_rs_cli/container/
reader.rs
