1use std::cell::RefCell;
30use std::fs::File;
31use std::io::{Read, Seek, SeekFrom};
32use std::path::Path;
33
34use crate::container::directory::RecordSite;
35use crate::container::header::{HEADER_LEN, Header};
36use crate::container::observation::{ObservationIndex, SECTION_OP_TABLE, primary_dependency};
37use crate::container::record::{RECORD_HEADER_LEN, RECORD_TRAILER_LEN, RecordTag, read_record_at};
38use crate::dra::Program;
39use crate::entropy::codec::EntropyChannelDescriptor;
40use crate::entropy::model::EntropyModel;
41use crate::error::{Error, Result};
42use crate::limits::Limits;
43use crate::materialize::observation::{
44 select_ops_from_lengths, selection_references, serve_selection,
45};
46use crate::materialize::seek::CountingReader;
47
48use super::dag::SourceServer;
49
50pub struct PartialDescriptor {
56 reader: RefCell<CountingReader<File>>,
57 source_len: u64,
58 source_sha256: [u8; 32],
59 object_sites: Vec<RecordSite>,
60 channel_sites: Vec<RecordSite>,
61 model_sites: Vec<RecordSite>,
62 graph_site: Option<RecordSite>,
67 index: ObservationIndex,
69 program: RefCell<Option<Program>>,
71 per_op: Vec<u64>,
73}
74
75impl std::fmt::Debug for PartialDescriptor {
76 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77 f.debug_struct("PartialDescriptor")
78 .field("source_len", &self.source_len)
79 .field("objects", &self.object_sites.len())
80 .field("channels", &self.channel_sites.len())
81 .field("models", &self.model_sites.len())
82 .field("graph_ops", &self.index.ops.len())
83 .finish_non_exhaustive()
84 }
85}
86
87pub enum PartialLoad {
89 Ready(Box<PartialDescriptor>),
91 Ineligible { bytes_read: u64 },
94}
95
96impl PartialDescriptor {
97 pub fn open(path: &Path, limits: Limits) -> Result<PartialLoad> {
102 let file = File::open(path).map_err(|e| Error::io(format!("opening descriptor: {e}")))?;
103 let mut reader = CountingReader::new(file);
104
105 let mut hdr = [0u8; HEADER_LEN];
106 reader
107 .seek(SeekFrom::Start(0))
108 .map_err(|e| Error::io(format!("seek to header failed: {e}")))?;
109 reader
110 .read_exact(&mut hdr)
111 .map_err(|_| Error::invalid_container("truncated header"))?;
112 let header = Header::decode(&hdr)?;
113
114 let mut universe: Option<String> = None;
115 let mut graph_site: Option<RecordSite> = None;
116 let mut index: Option<ObservationIndex> = None;
117 let mut source_sha256: Option<[u8; 32]> = None;
118 let mut source_len: Option<u64> = None;
119 let mut object_sites: Vec<RecordSite> = Vec::new();
120 let mut channel_sites: Vec<RecordSite> = Vec::new();
121 let mut model_sites: Vec<RecordSite> = Vec::new();
122
123 let file_len = reader
124 .seek(SeekFrom::End(0))
125 .map_err(|e| Error::io(format!("seek to end failed: {e}")))?;
126 let mut pos: u64 = HEADER_LEN as u64;
127 while pos < file_len {
128 reader
129 .seek(SeekFrom::Start(pos))
130 .map_err(|e| Error::io(format!("seek to record failed: {e}")))?;
131 let mut rec_hdr = [0u8; RECORD_HEADER_LEN];
132 reader
133 .read_exact(&mut rec_hdr)
134 .map_err(|_| Error::invalid_container("truncated record header"))?;
135 let tag = rec_hdr[0];
136 let _flags = rec_hdr[1];
137 let reserved = u16::from_le_bytes([rec_hdr[2], rec_hdr[3]]);
138 if reserved != 0 {
139 return Err(Error::invalid_container(
140 "record reserved field must be zero",
141 ));
142 }
143 let len = u32::from_le_bytes([rec_hdr[4], rec_hdr[5], rec_hdr[6], rec_hdr[7]]);
144 if len > limits.max_record_len {
145 return Err(Error::resource_limit(format!(
146 "record payload length {len} exceeds limit {}",
147 limits.max_record_len
148 )));
149 }
150 let total = RECORD_HEADER_LEN as u64 + u64::from(len) + RECORD_TRAILER_LEN as u64;
151 let end = pos
152 .checked_add(total)
153 .ok_or_else(|| Error::invalid_container("record length overflow"))?;
154 if end > file_len {
155 return Err(Error::invalid_container("truncated record"));
156 }
157
158 let site = RecordSite {
159 tag,
160 offset: pos,
161 payload_len: len,
162 };
163 match RecordTag::from_u8(tag) {
164 Some(RecordTag::Universe) => {
165 let rec = read_record_at(&mut reader, pos, limits)?;
166 universe = Some(
167 String::from_utf8(rec.payload)
168 .map_err(|_| Error::invalid_container("UNIVERSE is not UTF-8"))?,
169 );
170 }
171 Some(RecordTag::Graph) => graph_site = Some(site),
172 Some(RecordTag::ObservationIndex) => {
173 let rec = read_record_at(&mut reader, pos, limits)?;
174 index = Some(ObservationIndex::decode(&rec.payload, limits)?);
175 }
176 Some(RecordTag::Integrity) => {
177 let rec = read_record_at(&mut reader, pos, limits)?;
178 if rec.payload.len() != 40 {
179 return Err(Error::invalid_container(
180 "INTEGRITY payload must be 40 bytes",
181 ));
182 }
183 let mut sha = [0u8; 32];
184 sha.copy_from_slice(&rec.payload[0..32]);
185 source_sha256 = Some(sha);
186 source_len = Some(u64::from_le_bytes([
187 rec.payload[32],
188 rec.payload[33],
189 rec.payload[34],
190 rec.payload[35],
191 rec.payload[36],
192 rec.payload[37],
193 rec.payload[38],
194 rec.payload[39],
195 ]));
196 }
197 Some(RecordTag::Object) => object_sites.push(site),
198 Some(RecordTag::EntropyChannel) => channel_sites.push(site),
199 Some(RecordTag::Model) => model_sites.push(site),
200 Some(RecordTag::ExternalRef) => {
201 return Ok(PartialLoad::Ineligible {
204 bytes_read: reader.bytes_read(),
205 });
206 }
207 Some(RecordTag::Trailer) => {
208 let _ = read_record_at(&mut reader, pos, limits)?;
210 break;
211 }
212 _ => {}
214 }
215 pos = end;
216 }
217
218 let (Some(universe), Some(index), Some(source_sha256), Some(source_len)) =
222 (universe, index, source_sha256, source_len)
223 else {
224 return Ok(PartialLoad::Ineligible {
225 bytes_read: reader.bytes_read(),
226 });
227 };
228 if header.universe_id != crate::container::universe_id_from_str(&universe) {
229 return Err(Error::invalid_container(
230 "UNIVERSE record does not match the header universe id",
231 ));
232 }
233 if index.section_flags & SECTION_OP_TABLE == 0 {
234 return Ok(PartialLoad::Ineligible {
236 bytes_read: reader.bytes_read(),
237 });
238 }
239 let per_op = per_op_from_index(&index, source_len)?;
240
241 Ok(PartialLoad::Ready(Box::new(PartialDescriptor {
242 reader: RefCell::new(reader),
243 source_len,
244 source_sha256,
245 object_sites,
246 channel_sites,
247 model_sites,
248 graph_site,
249 index,
250 program: RefCell::new(None),
251 per_op,
252 })))
253 }
254
255 pub fn bytes_read(&self) -> u64 {
257 self.reader.borrow().bytes_read()
258 }
259
260 pub fn source_len(&self) -> u64 {
262 self.source_len
263 }
264
265 pub fn source_sha256(&self) -> [u8; 32] {
267 self.source_sha256
268 }
269
270 pub fn object_count(&self) -> usize {
272 self.object_sites.len()
273 }
274
275 pub fn graph_ops(&self) -> usize {
277 self.index.ops.len()
278 }
279
280 fn ensure_program(&self, limits: Limits) -> Result<std::cell::Ref<'_, Program>> {
287 if self.program.borrow().is_none() {
288 let site = self.graph_site.ok_or_else(|| {
289 Error::unsupported_feature("partial descriptor has no reconstruction graph")
290 })?;
291 let payload = self.read_site(&site, limits)?;
292 let program = Program::decode(&payload, limits)?;
293 let object_lens: Vec<u64> = self
294 .object_sites
295 .iter()
296 .map(|s| u64::from(s.payload_len))
297 .collect();
298 let checked = validate_op_table(&program, &self.index, &object_lens, self.source_len)?;
299 if checked != self.per_op {
300 return Err(Error::invalid_container(
301 "observation index op table disagrees with the program",
302 ));
303 }
304 *self.program.borrow_mut() = Some(program);
305 }
306 Ok(std::cell::Ref::map(self.program.borrow(), |p| {
307 p.as_ref().expect("program set above")
308 }))
309 }
310
311 fn read_site(&self, site: &RecordSite, limits: Limits) -> Result<Vec<u8>> {
312 let mut reader = self.reader.borrow_mut();
313 let rec = read_record_at(&mut *reader, site.offset, limits)?;
314 if rec.tag != site.tag || rec.payload.len() as u64 != u64::from(site.payload_len) {
315 return Err(Error::invalid_container(format!(
316 "record framing at offset {} disagrees with the walk (tag {:#04x}, len {})",
317 site.offset, site.tag, site.payload_len
318 )));
319 }
320 Ok(rec.payload)
321 }
322}
323
324impl SourceServer for PartialDescriptor {
325 fn serve_range(&self, offset: u64, len: u64, limits: Limits) -> Result<Vec<u8>> {
326 let end = offset
327 .checked_add(len)
328 .ok_or_else(|| Error::usage("source slice end overflows"))?;
329 if end > self.source_len {
330 return Err(Error::usage(format!(
331 "source slice {offset}..{end} exceeds source length {}",
332 self.source_len
333 )));
334 }
335 let program = self.ensure_program(limits)?;
336 let window = select_ops_from_lengths(&program, &self.per_op, offset, end)?;
337 let (objects_used, channels_used) = selection_references(
338 &window.ops,
339 self.object_sites.len(),
340 self.channel_sites.len(),
341 );
342
343 let mut objects: Vec<Vec<u8>> = vec![Vec::new(); self.object_sites.len()];
344 for (id, used) in objects_used.iter().enumerate() {
345 if *used {
346 objects[id] = self.read_site(&self.object_sites[id], limits)?;
347 }
348 }
349
350 let mut channels: Vec<EntropyChannelDescriptor> =
351 vec![placeholder_channel(); self.channel_sites.len()];
352 let mut models_needed = vec![false; self.model_sites.len()];
353 for (id, used) in channels_used.iter().enumerate() {
354 if *used {
355 let payload = self.read_site(&self.channel_sites[id], limits)?;
356 let channel = EntropyChannelDescriptor::decode(&payload, limits)?;
357 if channel.model_id as usize >= self.model_sites.len() {
358 return Err(Error::invalid_model(format!(
359 "entropy channel {id} references missing model {}",
360 channel.model_id
361 )));
362 }
363 models_needed[channel.model_id as usize] = true;
364 channels[id] = channel;
365 }
366 }
367 let mut models: Vec<EntropyModel> = vec![placeholder_model(); self.model_sites.len()];
368 for (id, used) in models_needed.iter().enumerate() {
369 if *used {
370 let payload = self.read_site(&self.model_sites[id], limits)?;
371 models[id] = EntropyModel::decode(&payload)?;
372 }
373 }
374
375 let served = serve_selection(
376 &objects,
377 &channels,
378 &models,
379 window,
380 &objects_used,
381 &channels_used,
382 offset,
383 end,
384 limits,
385 )?;
386 Ok(served.bytes)
387 }
388
389 fn serve_document(&self, _limits: Limits) -> Result<Vec<u8>> {
390 Err(Error::unsupported_feature(
394 "partial descriptor cannot serve the whole document",
395 ))
396 }
397}
398
399fn per_op_from_index(index: &ObservationIndex, source_len: u64) -> Result<Vec<u64>> {
403 let mut per_op: Vec<u64> = Vec::with_capacity(index.ops.len());
404 let mut total: u64 = 0;
405 for entry in &index.ops {
406 let claimed = u64::from(entry.out_len);
407 total = total
408 .checked_add(claimed)
409 .ok_or_else(|| Error::resource_limit("observation op length overflow"))?;
410 per_op.push(claimed);
411 }
412 if total != source_len {
413 return Err(Error::invalid_container(format!(
414 "observation index op table totals {total} bytes but the source declares {source_len}"
415 )));
416 }
417 Ok(per_op)
418}
419
420fn validate_op_table(
428 program: &Program,
429 index: &ObservationIndex,
430 object_lens: &[u64],
431 source_len: u64,
432) -> Result<Vec<u64>> {
433 use crate::dra::op::{DEFLATE_SOURCE_CHANNEL, Op};
434
435 let mut per_op: Vec<u64> = Vec::with_capacity(index.ops.len());
436 let mut total: u64 = 0;
437 for (i, entry) in index.ops.iter().enumerate() {
438 let op = &program.ops[i];
439 let claimed = u64::from(entry.out_len);
440 let (dep_kind, dep_id) = primary_dependency(op);
441 if entry.dep_kind != dep_kind
442 || (dep_kind != crate::container::observation::DEP_NONE && entry.dep_id != dep_id)
443 {
444 return Err(Error::invalid_container(format!(
445 "observation index op {i} dependency disagrees with the program"
446 )));
447 }
448 let derived: Option<u64> = match op {
449 Op::EmitObject { object_id } => {
450 Some(*object_lens.get(*object_id as usize).ok_or_else(|| {
451 Error::invalid_graph(format!("graph references missing object {object_id}"))
452 })?)
453 }
454 Op::Inline { bytes } => Some(bytes.len() as u64),
455 Op::MarkOffset { .. } => Some(0),
456 Op::EmitOffset { width, .. } => Some(u64::from(*width)),
457 Op::DeflateReplay {
458 source_kind,
459 declared_output_len,
460 ..
461 } => {
462 if *source_kind == DEFLATE_SOURCE_CHANNEL {
463 None
464 } else {
465 Some(u64::from(*declared_output_len))
466 }
467 }
468 _ => None,
469 };
470 if let Some(expected) = derived
471 && expected != claimed
472 {
473 return Err(Error::invalid_container(format!(
474 "observation index op {i} out_len {claimed} disagrees with the program's {expected}"
475 )));
476 }
477 total = total
478 .checked_add(claimed)
479 .ok_or_else(|| Error::resource_limit("observation op length overflow"))?;
480 per_op.push(claimed);
481 }
482 if total != source_len {
483 return Err(Error::invalid_container(format!(
484 "observation index op table totals {total} bytes but the source declares {source_len}"
485 )));
486 }
487 Ok(per_op)
488}
489
490fn placeholder_channel() -> EntropyChannelDescriptor {
492 use crate::entropy::{CODER_ORDER0_BYTE_RANS, CODER_VERSION_1};
493 EntropyChannelDescriptor {
494 coder: CODER_ORDER0_BYTE_RANS,
495 coder_version: CODER_VERSION_1,
496 scale_bits: 0,
497 lane_count: 1,
498 model_id: 0,
499 symbol_count: 0,
500 decoded_length: 0,
501 initial_state: 0,
502 payload: Vec::new(),
503 }
504}
505
506fn placeholder_model() -> EntropyModel {
508 EntropyModel {
509 scale_bits: 0,
510 frequencies: Vec::new(),
511 }
512}