1use super::{DeltaDecoder, DeltaLimits, DeltaStatus};
3use crate::{Error, ObjectFormat, ObjectId, ObjectInfo, Result};
4
5#[derive(Clone, Copy, Debug, Eq, PartialEq)]
7pub enum ObjectKey {
8 Id(ObjectId),
9 Staged(u64),
10}
11
12#[derive(Clone, Copy, Debug, Eq, PartialEq)]
13pub enum ObjectRange {
14 All,
15 Bytes { offset: u64, length: u64 },
16}
17
18#[derive(Clone, Copy, Debug, Eq, PartialEq)]
21pub enum RangeRepresentation {
22 Plain,
23 Delta {
24 base: ObjectKey,
25 base_info: ObjectInfo,
26 base_depth: u32,
27 instructions_size: u64,
28 },
29}
30
31#[derive(Clone, Copy, Debug, Eq, PartialEq)]
32pub struct RangeObject {
33 pub info: ObjectInfo,
34 pub representation: RangeRepresentation,
35}
36
37#[derive(Clone, Copy, Debug, Eq, PartialEq)]
38pub struct RangeLimits {
39 pub max_output: u64,
40 pub max_work_bytes: u64,
43 pub max_instructions: u64,
44 pub max_depth: u32,
45 pub max_read_bytes: usize,
46}
47
48#[derive(Clone, Copy, Debug, Eq, PartialEq)]
51pub enum RangeInput {
52 Plain,
53 Delta,
54}
55
56#[derive(Clone, Copy, Debug, Eq, PartialEq)]
57pub enum RangeRequest {
58 Describe(ObjectKey),
59 Read {
60 object: ObjectKey,
61 kind: RangeInput,
62 offset: u64,
63 length: usize,
64 },
65 NeedOutput,
66 Finished,
67}
68
69#[derive(Clone, Copy, Debug, Eq, PartialEq)]
70pub struct RangeProgress {
71 pub consumed: usize,
72 pub written: usize,
73 pub request: RangeRequest,
74}
75
76pub struct RangeFrame {
79 key: ObjectKey,
80 range: ObjectRange,
81 expected: Option<(ObjectInfo, u32)>,
82 object: Option<RangeObject>,
83 delta: Option<DeltaDecoder>,
84 cursor: u64,
85 start: u64,
86 end: u64,
87 remaining: u64,
88 complete: bool,
89}
90
91impl RangeFrame {
92 pub const fn empty() -> Self {
93 Self {
94 key: ObjectKey::Staged(0),
95 range: ObjectRange::All,
96 expected: None,
97 object: None,
98 delta: None,
99 cursor: 0,
100 start: 0,
101 end: 0,
102 remaining: 0,
103 complete: false,
104 }
105 }
106
107 fn pending(key: ObjectKey, range: ObjectRange, expected: Option<(ObjectInfo, u32)>) -> Self {
108 Self {
109 key,
110 range,
111 expected,
112 ..Self::empty()
113 }
114 }
115
116 fn describe(
117 &mut self,
118 object: RangeObject,
119 format: ObjectFormat,
120 limits: RangeLimits,
121 root: bool,
122 ) -> Result<u64> {
123 let depth = object_depth(object)?;
124 if depth > limits.max_depth {
125 return Err(Error::LimitExceeded);
126 }
127 if self
128 .expected
129 .is_some_and(|expected| expected != (object.info, depth))
130 {
131 return Err(Error::InvalidObject);
132 }
133 let (start, length) = match self.range {
134 ObjectRange::All => (0, object.info.size),
135 ObjectRange::Bytes { offset, length } => (offset, length),
136 };
137 let end = start
138 .checked_add(length)
139 .filter(|end| *end <= object.info.size)
140 .ok_or(Error::InvalidObject)?;
141 if root && length > limits.max_output {
142 return Err(Error::LimitExceeded);
143 }
144 self.start = start;
145 self.end = end;
146 self.remaining = length;
147 match object.representation {
148 RangeRepresentation::Plain => self.cursor = start,
149 RangeRepresentation::Delta {
150 base,
151 base_info,
152 base_depth,
153 ..
154 } => {
155 check_key(format, base)?;
156 if base_info.kind != object.info.kind {
157 return Err(Error::InvalidObject);
158 }
159 self.delta = Some(DeltaDecoder::new(
160 base_info.size,
161 base_depth,
162 DeltaLimits {
163 max_output: object.info.size,
164 max_instructions: limits.max_instructions,
165 max_depth: limits.max_depth,
166 },
167 )?);
168 }
169 }
170 self.object = Some(object);
171 Ok(length)
172 }
173
174 fn next_read(
175 &mut self,
176 max_read_bytes: usize,
177 remaining_work: u64,
178 ) -> Result<Option<RangeRequest>> {
179 let Some(object) = self.object else {
180 return Ok(Some(RangeRequest::Describe(self.key)));
181 };
182 let (kind, remaining) = match object.representation {
183 RangeRepresentation::Plain => (RangeInput::Plain, self.remaining),
184 RangeRepresentation::Delta {
185 instructions_size, ..
186 } => {
187 if self.cursor == instructions_size {
188 let delta = self.delta.as_mut().ok_or(Error::InvalidObject)?;
189 let step = delta.update(&[], true)?;
190 if step.status != DeltaStatus::Finished || self.remaining != 0 {
191 return Err(Error::CorruptPack);
192 }
193 delta.finish()?;
194 self.complete = true;
195 return Ok(None);
196 }
197 (RangeInput::Delta, instructions_size - self.cursor)
198 }
199 };
200 if remaining_work == 0 {
201 return Err(Error::LimitExceeded);
202 }
203 let length = remaining.min(max_read_bytes as u64).min(remaining_work) as usize;
204 Ok(Some(RangeRequest::Read {
205 object: self.key,
206 kind,
207 offset: self.cursor,
208 length,
209 }))
210 }
211
212 fn done(&self) -> bool {
213 self.complete
214 || self.object.is_some_and(|object| {
215 self.remaining == 0
216 && (matches!(object.representation, RangeRepresentation::Plain)
217 || self.end < object.info.size)
218 })
219 }
220}
221
222pub struct RangeResolver<'frames> {
237 frames: &'frames mut [RangeFrame],
238 active: usize,
239 format: ObjectFormat,
240 limits: RangeLimits,
241 written: u64,
242 work: u64,
243 expected: Option<u64>,
244 finished: bool,
245 failure: Option<Error>,
246}
247
248impl<'frames> RangeResolver<'frames> {
249 pub fn new(
250 format: ObjectFormat,
251 key: ObjectKey,
252 range: ObjectRange,
253 frames: &'frames mut [RangeFrame],
254 limits: RangeLimits,
255 ) -> Result<Self> {
256 check_key(format, key)?;
257 if frames.is_empty() {
258 return Err(Error::BufferTooSmall);
259 }
260 if limits.max_read_bytes == 0 {
261 return Err(Error::InvalidRequest);
262 }
263 if let ObjectRange::Bytes { offset, length } = range {
264 offset.checked_add(length).ok_or(Error::InvalidObject)?;
265 if length > limits.max_output {
266 return Err(Error::LimitExceeded);
267 }
268 }
269 frames[0] = RangeFrame::pending(key, range, None);
270 Ok(Self {
271 frames,
272 active: 0,
273 format,
274 limits,
275 written: 0,
276 work: 0,
277 expected: None,
278 finished: false,
279 failure: None,
280 })
281 }
282
283 pub fn position(&self) -> u64 {
284 self.written
285 }
286 pub fn work_bytes(&self) -> u64 {
287 self.work
288 }
289
290 pub fn request(&mut self) -> Result<RangeRequest> {
293 let result = self.next();
294 self.record(result)
295 }
296
297 pub fn describe(&mut self, object: RangeObject) -> Result<()> {
298 let result = self.accept_object(object);
299 self.record(result)
300 }
301
302 pub fn update(&mut self, input: &[u8], output: &mut [u8]) -> Result<RangeProgress> {
306 let result = self.read(input, output);
307 self.record(result)
308 }
309
310 pub fn finish(&self) -> Result<()> {
311 if let Some(error) = self.failure {
312 return Err(error);
313 }
314 if self.finished && self.expected == Some(self.written) {
315 Ok(())
316 } else {
317 Err(Error::CorruptPack)
318 }
319 }
320
321 fn record<T>(&mut self, result: Result<T>) -> Result<T> {
322 if let Err(error) = result {
323 self.failure = Some(error);
324 }
325 result
326 }
327
328 fn next(&mut self) -> Result<RangeRequest> {
329 if let Some(error) = self.failure {
330 return Err(error);
331 }
332 if self.finished {
333 return Ok(RangeRequest::Finished);
334 }
335 loop {
336 let frame = &mut self.frames[self.active];
337 if frame.done() {
338 if self.active == 0 {
339 if self.expected != Some(self.written) {
340 return Err(Error::CorruptPack);
341 }
342 self.finished = true;
343 return Ok(RangeRequest::Finished);
344 }
345 self.active -= 1;
346 continue;
347 }
348 if let Some(request) = frame.next_read(
349 self.limits.max_read_bytes,
350 self.limits.max_work_bytes - self.work,
351 )? {
352 return Ok(request);
353 }
354 }
355 }
356
357 fn accept_object(&mut self, object: RangeObject) -> Result<()> {
358 if !matches!(self.next()?, RangeRequest::Describe(_)) {
359 return Err(Error::InvalidRequest);
360 }
361 let length = self.frames[self.active].describe(
362 object,
363 self.format,
364 self.limits,
365 self.active == 0,
366 )?;
367 if self.active == 0 {
368 self.expected = Some(length);
369 }
370 Ok(())
371 }
372
373 fn read(&mut self, input: &[u8], output: &mut [u8]) -> Result<RangeProgress> {
374 let request = self.next()?;
375 let RangeRequest::Read { kind, length, .. } = request else {
376 if request == RangeRequest::Finished && input.is_empty() {
377 return Ok(RangeProgress {
378 consumed: 0,
379 written: 0,
380 request,
381 });
382 }
383 return Err(Error::InvalidRequest);
384 };
385 if output.is_empty() {
386 return Ok(RangeProgress {
387 consumed: 0,
388 written: 0,
389 request: RangeRequest::NeedOutput,
390 });
391 }
392 if input.is_empty() {
393 return Err(Error::CorruptPack);
394 }
395 let take = input.len().min(output.len()).min(length);
396 let (consumed, written) = match kind {
397 RangeInput::Plain => {
398 output[..take].copy_from_slice(&input[..take]);
399 let frame = &mut self.frames[self.active];
400 frame.cursor += take as u64;
401 frame.remaining -= take as u64;
402 (take, take)
403 }
404 RangeInput::Delta => self.instructions(&input[..take], output)?,
405 };
406 self.work += consumed as u64;
407 self.written += written as u64;
408 Ok(RangeProgress {
409 consumed,
410 written,
411 request: self.next()?,
412 })
413 }
414
415 fn instructions(&mut self, input: &[u8], output: &mut [u8]) -> Result<(usize, usize)> {
416 let frame = &mut self.frames[self.active];
417 let object = frame.object.ok_or(Error::InvalidObject)?;
418 let RangeRepresentation::Delta {
419 base,
420 base_info,
421 base_depth,
422 instructions_size,
423 } = object.representation
424 else {
425 return Err(Error::InvalidObject);
426 };
427 let delta = frame.delta.as_mut().ok_or(Error::InvalidObject)?;
428 let before = delta.position();
429 let step = delta.update(
430 input,
431 frame.cursor + input.len() as u64 == instructions_size,
432 )?;
433 frame.cursor += step.consumed as u64;
434 let mut written = 0;
435 match step.status {
436 DeltaStatus::Header { result_size, .. } => {
437 if result_size != object.info.size {
438 return Err(Error::CorruptPack);
439 }
440 }
441 DeltaStatus::Insert(bytes) => {
442 if let Some((start, length)) =
443 overlap(before, bytes.len() as u64, frame.start, frame.end)
444 {
445 written = length as usize;
446 let start = start as usize;
447 output[..written].copy_from_slice(&bytes[start..start + written]);
448 frame.remaining -= length;
449 }
450 }
451 DeltaStatus::Copy { offset, length } => {
452 if let Some((start, length)) = overlap(before, length, frame.start, frame.end) {
453 frame.remaining -= length;
454 if self.active + 1 == self.frames.len() {
455 return Err(Error::BufferTooSmall);
456 }
457 self.active += 1;
458 self.frames[self.active] = RangeFrame::pending(
459 base,
460 ObjectRange::Bytes {
461 offset: offset + start,
462 length,
463 },
464 Some((base_info, base_depth)),
465 );
466 }
467 }
468 DeltaStatus::Finished => {
469 if frame.remaining != 0 {
470 return Err(Error::CorruptPack);
471 }
472 frame.complete = true;
473 }
474 DeltaStatus::NeedInput => {}
475 }
476 Ok((step.consumed, written))
477 }
478}
479
480fn overlap(position: u64, length: u64, start: u64, end: u64) -> Option<(u64, u64)> {
481 let first = position.max(start);
482 let last = (position + length).min(end);
483 (first < last).then_some((first - position, last.saturating_sub(first)))
484}
485
486fn check_key(format: ObjectFormat, key: ObjectKey) -> Result<()> {
487 if matches!(key, ObjectKey::Id(id) if id.format() != format) {
488 Err(Error::InvalidObjectId)
489 } else {
490 Ok(())
491 }
492}
493
494fn object_depth(object: RangeObject) -> Result<u32> {
495 match object.representation {
496 RangeRepresentation::Plain => Ok(0),
497 RangeRepresentation::Delta { base_depth, .. } => {
498 base_depth.checked_add(1).ok_or(Error::LimitExceeded)
499 }
500 }
501}