j2k_jpeg/
segment.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! Public JPEG marker-level utilities.
4
5use crate::allocation::{
6    checked_add_allocation_bytes, checked_allocation_bytes, try_vec_with_capacity,
7};
8use crate::error::{JpegError, MarkerKind, UnsupportedReason};
9use crate::info::{SamplingFactors, SofKind};
10use crate::parse::markers::next_marker_after_entropy;
11use alloc::vec::Vec;
12use core::ops::Range;
13use memchr::memchr;
14
15mod table_normalization;
16
17use table_normalization::for_each_normalized_segment;
18
19/// One marker segment in a JPEG byte stream.
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub struct JpegSegment<'a> {
22    /// Raw marker byte after the `0xff` prefix.
23    pub marker: u8,
24    /// Offset of the marker prefix byte.
25    pub marker_offset: usize,
26    /// Offset of the segment payload. Standalone markers use the byte after the marker.
27    pub payload_offset: usize,
28    /// Segment payload excluding marker and length bytes.
29    pub payload: &'a [u8],
30}
31
32/// Iterator over marker segments in a JPEG byte stream.
33#[derive(Debug)]
34pub struct JpegSegmentIter<'a> {
35    input: &'a [u8],
36    pos: usize,
37    started: bool,
38    finished: bool,
39    scan_entropy: bool,
40}
41
42/// Parsed Start-of-Frame facts.
43#[derive(Debug, Clone, PartialEq, Eq)]
44pub struct JpegSofInfo {
45    /// Raw SOF marker byte.
46    pub marker: u8,
47    /// Supported SOF kind classification.
48    pub sof_kind: SofKind,
49    /// Component sample precision.
50    pub bit_depth: u8,
51    /// Width and height from the SOF payload.
52    pub dimensions: (u16, u16),
53    /// Component identifiers in declaration order.
54    pub component_ids: Vec<u8>,
55    /// Component sampling factors in declaration order.
56    pub sampling: SamplingFactors,
57    /// Quantization-table selectors in declaration order.
58    pub quant_table_ids: Vec<u8>,
59}
60
61/// Byte ranges around the first Start-of-Scan marker.
62#[derive(Debug, Clone, PartialEq, Eq)]
63pub struct JpegScanRanges {
64    /// Offset of the SOS marker prefix byte.
65    pub sos_marker_offset: usize,
66    /// Payload range of the SOS segment, excluding marker and length bytes.
67    pub sos_payload_range: Range<usize>,
68    /// Entropy-coded scan data range after SOS and before EOI or the next marker.
69    pub entropy_range: Range<usize>,
70    /// Offset of EOI when present.
71    pub eoi_marker_offset: Option<usize>,
72}
73
74/// Options for preparing TIFF/WSI JPEG tile payloads.
75#[derive(Debug, Clone, Copy, PartialEq, Eq)]
76pub struct JpegTilePrepareOptions {
77    /// Container-derived expected tile dimensions.
78    pub expected_dimensions: Option<(u16, u16)>,
79    /// Duplicate DQT/DHT handling policy within supplied `JPEGTables` bytes.
80    pub duplicate_table_policy: DuplicateTablePolicy,
81    /// Repair zero SOF dimensions using `expected_dimensions`.
82    pub repair_zero_sof_dimensions: bool,
83    /// Validate restart marker order in scan data.
84    pub validate_restart_markers: bool,
85}
86
87impl Default for JpegTilePrepareOptions {
88    fn default() -> Self {
89        Self {
90            expected_dimensions: None,
91            duplicate_table_policy: DuplicateTablePolicy::RejectConflicting,
92            repair_zero_sof_dimensions: false,
93            validate_restart_markers: false,
94        }
95    }
96}
97
98/// Duplicate JPEG table handling policy.
99#[derive(Debug, Clone, Copy, PartialEq, Eq)]
100pub enum DuplicateTablePolicy {
101    /// Coalesce byte-identical duplicate definitions while rejecting conflicts.
102    AllowIdentical,
103    /// Preserve byte-identical redefinitions but reject conflicting definitions.
104    ///
105    /// This mode retains the source marker bytes when no conflict is present.
106    RejectConflicting,
107}
108
109/// Prepared JPEG bytes, borrowed when unchanged and owned when normalized.
110#[derive(Debug, PartialEq, Eq)]
111pub enum PreparedJpeg<'a> {
112    /// Original tile bytes can be decoded directly.
113    Borrowed(&'a [u8]),
114    /// Preparation changed the byte stream.
115    Owned(Vec<u8>),
116}
117
118impl PreparedJpeg<'_> {
119    /// Return decode-ready JPEG bytes.
120    #[must_use]
121    pub fn as_bytes(&self) -> &[u8] {
122        match self {
123            Self::Borrowed(bytes) => bytes,
124            Self::Owned(bytes) => bytes,
125        }
126    }
127
128    /// Fallibly duplicate this prepared JPEG.
129    ///
130    /// Borrowed inputs remain borrowed without allocation. Owned normalized
131    /// payloads are copied through the same typed, capped allocation contract
132    /// used by dimension rewrite and TIFF assembly.
133    ///
134    /// # Errors
135    ///
136    /// Returns [`JpegError::MemoryCapExceeded`] when the owned payload exceeds
137    /// the JPEG host cap, or [`JpegError::HostAllocationFailed`] when a
138    /// cap-valid copy cannot be reserved.
139    pub fn try_clone(&self) -> Result<Self, JpegError> {
140        match self {
141            Self::Borrowed(bytes) => Ok(Self::Borrowed(bytes)),
142            Self::Owned(bytes) => Ok(Self::Owned(try_copy_bytes(bytes)?)),
143        }
144    }
145}
146
147impl AsRef<[u8]> for PreparedJpeg<'_> {
148    fn as_ref(&self) -> &[u8] {
149        self.as_bytes()
150    }
151}
152
153/// Iterate over JPEG marker segments.
154#[must_use]
155pub fn iter_segments(input: &[u8]) -> JpegSegmentIter<'_> {
156    JpegSegmentIter {
157        input,
158        pos: 0,
159        started: false,
160        finished: false,
161        scan_entropy: false,
162    }
163}
164
165/// Return true for JPEG SOF marker classes.
166#[must_use]
167pub const fn is_sof_marker(marker: u8) -> bool {
168    matches!(
169        marker,
170        0xc0..=0xc3 | 0xc5..=0xc7 | 0xc9..=0xcb | 0xcd..=0xcf
171    )
172}
173
174/// Parse a Start-of-Frame payload.
175///
176/// # Errors
177///
178/// Returns an error when the marker or SOF payload is malformed.
179pub fn parse_sof_info(marker: u8, payload: &[u8]) -> Result<JpegSofInfo, JpegError> {
180    parse_sof_info_at(marker, payload, 0, false)
181}
182
183pub(crate) fn parse_sof_info_allowing_zero_dimensions(
184    marker: u8,
185    payload: &[u8],
186    payload_offset: usize,
187) -> Result<JpegSofInfo, JpegError> {
188    parse_sof_info_at(marker, payload, payload_offset, true)
189}
190
191/// Parse a Define Restart Interval payload.
192///
193/// # Errors
194///
195/// Returns an error when the DRI payload length is not exactly two bytes.
196pub fn parse_dri(payload: &[u8]) -> Result<Option<u16>, JpegError> {
197    if payload.len() != 2 {
198        return Err(JpegError::InvalidSegmentLength {
199            offset: 0,
200            marker: 0xdd,
201            length: u16::try_from(payload.len() + 2).unwrap_or(u16::MAX),
202        });
203    }
204    let interval = u16::from_be_bytes([payload[0], payload[1]]);
205    Ok((interval > 0).then_some(interval))
206}
207
208/// Find the first SOS header and scan data ranges.
209///
210/// # Errors
211///
212/// Returns an error when segment syntax is malformed or no SOS marker exists.
213pub fn find_scan_ranges(input: &[u8]) -> Result<JpegScanRanges, JpegError> {
214    for segment in iter_segments(input) {
215        let segment = segment?;
216        if segment.marker == 0xda {
217            let entropy_start = segment.payload_offset + segment.payload.len();
218            let next_marker = next_marker_after_entropy(input, entropy_start)?;
219            let (entropy_end, eoi_marker_offset) = match next_marker {
220                Some((marker_offset, 0xd9)) => (marker_offset, Some(marker_offset)),
221                Some((marker_offset, _)) => (marker_offset, None),
222                None => (input.len(), None),
223            };
224            return Ok(JpegScanRanges {
225                sos_marker_offset: segment.marker_offset,
226                sos_payload_range: segment.payload_offset..entropy_start,
227                entropy_range: entropy_start..entropy_end,
228                eoi_marker_offset,
229            });
230        }
231    }
232    Err(JpegError::MissingMarker {
233        marker: MarkerKind::Sos,
234    })
235}
236
237/// Return a copy of `input` with the first SOF dimensions rewritten.
238///
239/// # Errors
240///
241/// Returns an error for zero dimensions, malformed segment syntax, or a
242/// missing SOF marker.
243pub fn rewrite_sof_dimensions(input: &[u8], dimensions: (u16, u16)) -> Result<Vec<u8>, JpegError> {
244    if dimensions.0 == 0 || dimensions.1 == 0 {
245        return Err(JpegError::ZeroDimension {
246            width: dimensions.0,
247            height: dimensions.1,
248        });
249    }
250    for segment in iter_segments(input) {
251        let segment = segment?;
252        if is_sof_marker(segment.marker) {
253            if segment.payload.len() < 5 {
254                return Err(JpegError::Truncated {
255                    offset: segment.payload_offset + segment.payload.len(),
256                    expected: 5 - segment.payload.len(),
257                });
258            }
259            let mut out = try_copy_bytes(input)?;
260            let width = dimensions.0.to_be_bytes();
261            let height = dimensions.1.to_be_bytes();
262            out[segment.payload_offset + 1] = height[0];
263            out[segment.payload_offset + 2] = height[1];
264            out[segment.payload_offset + 3] = width[0];
265            out[segment.payload_offset + 4] = width[1];
266            return Ok(out);
267        }
268    }
269    Err(JpegError::MissingMarker {
270        marker: MarkerKind::Sof,
271    })
272}
273
274fn try_copy_bytes(input: &[u8]) -> Result<Vec<u8>, JpegError> {
275    checked_allocation_bytes::<u8>(input.len())?;
276    let mut out = try_vec_with_capacity(input.len())?;
277    out.extend_from_slice(input);
278    Ok(out)
279}
280
281/// Prepare a TIFF/WSI JPEG tile for decode.
282///
283/// # Errors
284///
285/// Returns an error when a complete tile is invalid or an abbreviated tile
286/// cannot be assembled safely with the supplied tables and options.
287pub fn prepare_tiff_jpeg_tile<'a>(
288    tile: &'a [u8],
289    tables: Option<&'a [u8]>,
290    opts: JpegTilePrepareOptions,
291) -> Result<PreparedJpeg<'a>, JpegError> {
292    if is_complete_jpeg(tile) {
293        validate_complete_tile(tile, opts)
294    } else {
295        assemble_abbreviated_tile(tile, tables, opts)
296    }
297}
298
299fn is_complete_jpeg(input: &[u8]) -> bool {
300    input.len() >= 4
301        && input[0] == 0xff
302        && input[1] == 0xd8
303        && input[input.len() - 2] == 0xff
304        && input[input.len() - 1] == 0xd9
305}
306
307fn validate_complete_tile(
308    tile: &[u8],
309    opts: JpegTilePrepareOptions,
310) -> Result<PreparedJpeg<'_>, JpegError> {
311    if let Some(repaired) = finalize_prepared_bytes(tile, opts)? {
312        return Ok(PreparedJpeg::Owned(repaired));
313    }
314    Ok(PreparedJpeg::Borrowed(tile))
315}
316
317fn finalize_prepared_bytes(
318    bytes: &[u8],
319    opts: JpegTilePrepareOptions,
320) -> Result<Option<Vec<u8>>, JpegError> {
321    let repaired = repair_or_validate_dimensions(bytes, opts)?;
322    let validation_input = repaired.as_deref().unwrap_or(bytes);
323    let _ = find_scan_ranges(validation_input)?;
324    if opts.validate_restart_markers {
325        validate_restart_markers(validation_input)?;
326    }
327    Ok(repaired)
328}
329
330fn finalize_owned_prepared_bytes(
331    bytes: &mut [u8],
332    opts: JpegTilePrepareOptions,
333) -> Result<(), JpegError> {
334    if let Some(repair) = planned_dimension_repair(bytes, opts)? {
335        apply_dimension_repair(bytes, repair);
336        validate_nonzero_sof_dimensions(bytes, opts)?;
337    }
338    let _ = find_scan_ranges(bytes)?;
339    if opts.validate_restart_markers {
340        validate_restart_markers(bytes)?;
341    }
342    Ok(())
343}
344
345fn repair_or_validate_dimensions(
346    bytes: &[u8],
347    opts: JpegTilePrepareOptions,
348) -> Result<Option<Vec<u8>>, JpegError> {
349    let Some(repair) = planned_dimension_repair(bytes, opts)? else {
350        return Ok(None);
351    };
352    let mut repaired = try_copy_bytes(bytes)?;
353    apply_dimension_repair(&mut repaired, repair);
354    validate_nonzero_sof_dimensions(&repaired, opts)?;
355    Ok(Some(repaired))
356}
357
358#[derive(Clone, Copy)]
359struct SofDimensionRepair {
360    payload_offset: usize,
361    dimensions: (u16, u16),
362}
363
364fn planned_dimension_repair(
365    bytes: &[u8],
366    opts: JpegTilePrepareOptions,
367) -> Result<Option<SofDimensionRepair>, JpegError> {
368    let mut saw_sof = false;
369    for segment in iter_segments(bytes) {
370        let segment = segment?;
371        if segment.marker == 0xda {
372            break;
373        }
374        if is_sof_marker(segment.marker) {
375            saw_sof = true;
376            let sof = parse_sof_info_allowing_zero_dimensions(
377                segment.marker,
378                segment.payload,
379                segment.payload_offset,
380            )?;
381            if sof.dimensions.0 == 0 || sof.dimensions.1 == 0 {
382                let Some(expected) = opts.expected_dimensions else {
383                    return Err(JpegError::ExpectedDimensionsRequired {
384                        offset: segment.marker_offset,
385                    });
386                };
387                if !opts.repair_zero_sof_dimensions {
388                    return Err(JpegError::ZeroDimension {
389                        width: sof.dimensions.0,
390                        height: sof.dimensions.1,
391                    });
392                }
393                if expected.0 == 0 || expected.1 == 0 {
394                    return Err(JpegError::ZeroDimension {
395                        width: expected.0,
396                        height: expected.1,
397                    });
398                }
399                return Ok(Some(SofDimensionRepair {
400                    payload_offset: segment.payload_offset,
401                    dimensions: expected,
402                }));
403            }
404            if let Some(expected) = opts.expected_dimensions {
405                if expected != sof.dimensions {
406                    return Err(JpegError::ConflictingExpectedDimensions {
407                        offset: segment.marker_offset,
408                        expected,
409                        actual: sof.dimensions,
410                    });
411                }
412            }
413        }
414    }
415    if !saw_sof {
416        return Err(JpegError::MissingMarker {
417            marker: MarkerKind::Sof,
418        });
419    }
420    Ok(None)
421}
422
423fn apply_dimension_repair(bytes: &mut [u8], repair: SofDimensionRepair) {
424    let width = repair.dimensions.0.to_be_bytes();
425    let height = repair.dimensions.1.to_be_bytes();
426    bytes[repair.payload_offset + 1] = height[0];
427    bytes[repair.payload_offset + 2] = height[1];
428    bytes[repair.payload_offset + 3] = width[0];
429    bytes[repair.payload_offset + 4] = width[1];
430}
431
432fn validate_nonzero_sof_dimensions(
433    bytes: &[u8],
434    opts: JpegTilePrepareOptions,
435) -> Result<(), JpegError> {
436    let mut saw_sof = false;
437    for segment in iter_segments(bytes) {
438        let segment = segment?;
439        if is_sof_marker(segment.marker) {
440            saw_sof = true;
441            let sof = parse_sof_info(segment.marker, segment.payload)?;
442            if let Some(expected) = opts.expected_dimensions {
443                if expected != sof.dimensions {
444                    return Err(JpegError::ConflictingExpectedDimensions {
445                        offset: segment.marker_offset,
446                        expected,
447                        actual: sof.dimensions,
448                    });
449                }
450            }
451        }
452    }
453    if !saw_sof {
454        return Err(JpegError::MissingMarker {
455            marker: MarkerKind::Sof,
456        });
457    }
458    Ok(())
459}
460
461fn validate_restart_markers(bytes: &[u8]) -> Result<(), JpegError> {
462    let ranges = find_scan_ranges(bytes)?;
463    let mut expected = 0u8;
464    let mut pos = ranges.entropy_range.start;
465    while pos < ranges.entropy_range.end {
466        let Some(relative) = memchr(0xff, &bytes[pos..ranges.entropy_range.end]) else {
467            break;
468        };
469        let prefix = pos + relative;
470        let mut marker_pos = prefix + 1;
471        while marker_pos < ranges.entropy_range.end && bytes[marker_pos] == 0xff {
472            marker_pos += 1;
473        }
474        if marker_pos >= ranges.entropy_range.end {
475            return Err(JpegError::Truncated {
476                offset: prefix,
477                expected: 1,
478            });
479        }
480        let marker = bytes[marker_pos];
481        match marker {
482            0x00 => pos = marker_pos + 1,
483            0xd0..=0xd7 => {
484                let found = marker & 0x07;
485                if found != expected {
486                    return Err(JpegError::RestartMismatch {
487                        offset: marker_pos - 1,
488                        expected,
489                        found: marker,
490                    });
491                }
492                expected = (expected + 1) & 0x07;
493                pos = marker_pos + 1;
494            }
495            0xd9 => break,
496            _ => {
497                return Err(JpegError::UnexpectedMarker {
498                    offset: marker_pos - 1,
499                    expected: MarkerKind::Eoi,
500                    found: marker,
501                });
502            }
503        }
504    }
505    Ok(())
506}
507
508fn assemble_abbreviated_tile<'a>(
509    tile: &'a [u8],
510    tables: Option<&'a [u8]>,
511    opts: JpegTilePrepareOptions,
512) -> Result<PreparedJpeg<'a>, JpegError> {
513    let Some(tables) = tables else {
514        return Err(JpegError::InvalidJpegAssembly {
515            offset: 0,
516            reason: "abbreviated JPEG tile requires JPEGTables",
517        });
518    };
519    let tile_body = normalized_abbreviated_tile_body(tile)?;
520    let output_len = abbreviated_output_len(tables, tile_body.len(), opts.duplicate_table_policy)?;
521    let mut out = try_vec_with_capacity(output_len)?;
522    out.extend_from_slice(&[0xff, 0xd8]);
523    for_each_normalized_segment(tables, opts.duplicate_table_policy, |segment| {
524        segment.append_to(&mut out)
525    })?;
526    out.extend_from_slice(tile_body);
527    out.extend_from_slice(&[0xff, 0xd9]);
528    if out.len() != output_len {
529        return Err(JpegError::InternalInvariant {
530            reason: "normalized TIFF JPEG output length diverged from its preflight plan",
531        });
532    }
533    finalize_owned_prepared_bytes(&mut out, opts)?;
534    Ok(PreparedJpeg::Owned(out))
535}
536
537fn abbreviated_output_len(
538    tables: &[u8],
539    tile_body_len: usize,
540    policy: DuplicateTablePolicy,
541) -> Result<usize, JpegError> {
542    let mut table_bytes = 0usize;
543    for_each_normalized_segment(tables, policy, |segment| {
544        table_bytes = checked_add_allocation_bytes(table_bytes, segment.byte_len())?;
545        Ok(())
546    })?;
547    checked_abbreviated_output_len(table_bytes, tile_body_len)
548}
549
550fn checked_abbreviated_output_len(
551    normalized_table_bytes: usize,
552    tile_body_len: usize,
553) -> Result<usize, JpegError> {
554    let mut output_len = checked_allocation_bytes::<u8>(2)?;
555    output_len = checked_add_allocation_bytes(output_len, normalized_table_bytes)?;
556    output_len = checked_add_allocation_bytes(output_len, tile_body_len)?;
557    checked_add_allocation_bytes(output_len, 2)
558}
559
560fn normalized_abbreviated_tile_body(tile: &[u8]) -> Result<&[u8], JpegError> {
561    let start = if tile.starts_with(&[0xff, 0xd8]) {
562        2
563    } else {
564        0
565    };
566    let end = if tile.len() >= start + 2 && tile[tile.len() - 2..] == [0xff, 0xd9] {
567        tile.len() - 2
568    } else {
569        tile.len()
570    };
571    if start >= end {
572        return Err(JpegError::InvalidJpegAssembly {
573            offset: 0,
574            reason: "abbreviated JPEG tile is empty",
575        });
576    }
577    Ok(&tile[start..end])
578}
579
580impl<'a> Iterator for JpegSegmentIter<'a> {
581    type Item = Result<JpegSegment<'a>, JpegError>;
582
583    fn next(&mut self) -> Option<Self::Item> {
584        if self.finished {
585            return None;
586        }
587        if !self.started {
588            self.started = true;
589            return Some(self.read_soi());
590        }
591        if self.scan_entropy {
592            match next_marker_after_entropy(self.input, self.pos) {
593                Ok(Some((marker_offset, marker))) => {
594                    self.pos = marker_offset;
595                    self.scan_entropy = false;
596                    if marker == 0xd9 {
597                        return Some(self.read_standalone_marker());
598                    }
599                }
600                Ok(None) => {
601                    self.finished = true;
602                    return None;
603                }
604                Err(error) => {
605                    self.finished = true;
606                    return Some(Err(error));
607                }
608            }
609        }
610        Some(self.read_segment())
611    }
612}
613
614impl<'a> JpegSegmentIter<'a> {
615    fn read_soi(&mut self) -> Result<JpegSegment<'a>, JpegError> {
616        if self.input.len() < 2 {
617            return Err(JpegError::Truncated {
618                offset: 0,
619                expected: 2 - self.input.len(),
620            });
621        }
622        if self.input[0] != 0xff || self.input[1] != 0xd8 {
623            return Err(JpegError::UnexpectedMarker {
624                offset: 0,
625                expected: MarkerKind::Soi,
626                found: self.input.get(1).copied().unwrap_or(0),
627            });
628        }
629        self.pos = 2;
630        Ok(JpegSegment {
631            marker: 0xd8,
632            marker_offset: 0,
633            payload_offset: 2,
634            payload: &[],
635        })
636    }
637
638    fn read_segment(&mut self) -> Result<JpegSegment<'a>, JpegError> {
639        if self.pos >= self.input.len() {
640            return Err(JpegError::Truncated {
641                offset: self.pos,
642                expected: 2,
643            });
644        }
645        if self.input[self.pos] != 0xff {
646            return Err(JpegError::InvalidMarker {
647                offset: self.pos,
648                marker: self.input[self.pos],
649            });
650        }
651        while self.pos < self.input.len() && self.input[self.pos] == 0xff {
652            self.pos += 1;
653        }
654        if self.pos >= self.input.len() {
655            return Err(JpegError::Truncated {
656                offset: self.pos,
657                expected: 1,
658            });
659        }
660        let marker = self.input[self.pos];
661        let marker_offset = self.pos - 1;
662        self.pos += 1;
663
664        match marker {
665            0x01 | 0xd0..=0xd9 => {
666                if marker == 0xd9 {
667                    self.finished = true;
668                }
669                Ok(JpegSegment {
670                    marker,
671                    marker_offset,
672                    payload_offset: self.pos,
673                    payload: &[],
674                })
675            }
676            0x00 => Err(JpegError::InvalidMarker {
677                offset: marker_offset,
678                marker,
679            }),
680            _ => {
681                if self.pos + 2 > self.input.len() {
682                    return Err(JpegError::Truncated {
683                        offset: self.pos,
684                        expected: self.pos + 2 - self.input.len(),
685                    });
686                }
687                let length = u16::from_be_bytes([self.input[self.pos], self.input[self.pos + 1]]);
688                if length < 2 {
689                    return Err(JpegError::InvalidSegmentLength {
690                        offset: self.pos,
691                        marker,
692                        length,
693                    });
694                }
695                let payload_offset = self.pos + 2;
696                let payload_end = self.pos.checked_add(usize::from(length)).ok_or(
697                    JpegError::InvalidSegmentLength {
698                        offset: self.pos,
699                        marker,
700                        length,
701                    },
702                )?;
703                if payload_end > self.input.len() {
704                    return Err(JpegError::Truncated {
705                        offset: payload_offset,
706                        expected: payload_end - self.input.len(),
707                    });
708                }
709                self.pos = payload_end;
710                if marker == 0xda {
711                    self.scan_entropy = true;
712                }
713                Ok(JpegSegment {
714                    marker,
715                    marker_offset,
716                    payload_offset,
717                    payload: &self.input[payload_offset..payload_end],
718                })
719            }
720        }
721    }
722
723    fn read_standalone_marker(&mut self) -> Result<JpegSegment<'a>, JpegError> {
724        let marker_offset = self.pos;
725        if self.pos + 1 >= self.input.len() {
726            return Err(JpegError::Truncated {
727                offset: self.pos,
728                expected: self.pos + 2 - self.input.len(),
729            });
730        }
731        let marker = self.input[self.pos + 1];
732        self.pos += 2;
733        if marker == 0xd9 {
734            self.finished = true;
735        }
736        Ok(JpegSegment {
737            marker,
738            marker_offset,
739            payload_offset: self.pos,
740            payload: &[],
741        })
742    }
743}
744
745fn parse_sof_info_at(
746    marker: u8,
747    payload: &[u8],
748    payload_offset: usize,
749    allow_zero_dimensions: bool,
750) -> Result<JpegSofInfo, JpegError> {
751    if payload.len() < 8 {
752        return Err(JpegError::Truncated {
753            offset: payload_offset + payload.len(),
754            expected: 8 - payload.len(),
755        });
756    }
757
758    let bit_depth = payload[0];
759    let height = u16::from_be_bytes([payload[1], payload[2]]);
760    let width = u16::from_be_bytes([payload[3], payload[4]]);
761    let component_count = payload[5];
762    let expected_len = 6 + usize::from(component_count) * 3;
763    if payload.len() < expected_len {
764        return Err(JpegError::Truncated {
765            offset: payload_offset + payload.len(),
766            expected: expected_len - payload.len(),
767        });
768    }
769
770    let sof_kind = match (marker, bit_depth) {
771        (0xc0, 8) => SofKind::Baseline8,
772        (0xc1, 8) => SofKind::Extended8,
773        (0xc1, 12) => SofKind::Extended12,
774        (0xc2, 8) => SofKind::Progressive8,
775        (0xc2, 12) => SofKind::Progressive12,
776        (0xc3, 2..=16) => SofKind::Lossless,
777        (0xc5, _) => {
778            return Err(JpegError::UnsupportedSof {
779                marker,
780                reason: UnsupportedReason::DifferentialBaseline,
781            });
782        }
783        (0xc6 | 0xc7, _) => {
784            return Err(JpegError::UnsupportedSof {
785                marker,
786                reason: UnsupportedReason::Hierarchical,
787            });
788        }
789        (0xc9..=0xcb, _) => {
790            return Err(JpegError::UnsupportedSof {
791                marker,
792                reason: UnsupportedReason::ArithmeticCoding,
793            });
794        }
795        (0xcd..=0xcf, _) => {
796            return Err(JpegError::UnsupportedSof {
797                marker,
798                reason: UnsupportedReason::ArithmeticAndHierarchical,
799            });
800        }
801        (_, bad_precision) => {
802            return Err(JpegError::UnsupportedBitDepth {
803                depth: bad_precision,
804            })
805        }
806    };
807
808    if !allow_zero_dimensions && (width == 0 || height == 0) {
809        return Err(JpegError::ZeroDimension { width, height });
810    }
811    if width > 65_500 || height > 65_500 {
812        return Err(JpegError::DimensionOverflow {
813            width: u32::from(width),
814            height: u32::from(height),
815        });
816    }
817    if !matches!(component_count, 1 | 3 | 4) {
818        return Err(JpegError::UnsupportedComponentCount {
819            count: component_count,
820        });
821    }
822
823    let components = parse_sof_components(payload, component_count)?;
824
825    Ok(JpegSofInfo {
826        marker,
827        sof_kind,
828        bit_depth,
829        dimensions: (width, height),
830        component_ids: components.ids,
831        sampling: SamplingFactors::from_validated_components(
832            &components.sampling[..usize::from(component_count)],
833        ),
834        quant_table_ids: components.quant_table_ids,
835    })
836}
837
838struct ParsedSofComponents {
839    ids: Vec<u8>,
840    sampling: [(u8, u8); 4],
841    quant_table_ids: Vec<u8>,
842}
843
844#[expect(
845    clippy::cast_possible_truncation,
846    reason = "validated SOF component indices fit the JPEG component-count byte"
847)]
848fn parse_sof_components(
849    payload: &[u8],
850    component_count: u8,
851) -> Result<ParsedSofComponents, JpegError> {
852    let count = usize::from(component_count);
853    let mut sampling = [(0u8, 0u8); 4];
854    let mut ids = try_vec_with_capacity(count)?;
855    let mut quant_table_ids = try_vec_with_capacity(count)?;
856    for (i, sampling_slot) in sampling.iter_mut().take(count).enumerate() {
857        let base = 6 + i * 3;
858        let sampling_byte = payload[base + 1];
859        let (h, v) = (sampling_byte >> 4, sampling_byte & 0x0f);
860        if !(1..=4).contains(&h) || !(1..=4).contains(&v) {
861            return Err(JpegError::InvalidSampling {
862                component: i as u8,
863                h,
864                v,
865            });
866        }
867        ids.push(payload[base]);
868        *sampling_slot = (h, v);
869        quant_table_ids.push(payload[base + 2]);
870    }
871    Ok(ParsedSofComponents {
872        ids,
873        sampling,
874        quant_table_ids,
875    })
876}
877
878#[cfg(test)]
879#[path = "segment/allocation_tests.rs"]
880mod allocation_tests;
j2k_jpeg/segment.rs

j2k_jpeg/
segment.rs
