gufo_jpeg/
lib.rs

1#![doc = include_str!("../README.md")]
2
3#[cfg(feature = "encoder")]
4mod encoder;
5mod segments;
6
7use std::io::{Cursor, Read};
8use std::ops::Range;
9
10use gufo_common::error::ErrorWithData;
11use gufo_common::math::*;
12use gufo_common::physical_dimension::PixelDensity;
13use gufo_common::prelude::*;
14use indexmap::IndexMap;
15pub use segments::*;
16use zerocopy::FromBytes;
17
18use crate::Marker::APP14;
19
20pub const EXIF_IDENTIFIER_STRING: &[u8] = b"Exif\0\0";
21pub const XMP_IDENTIFIER_STRING: &[u8] = b"http://ns.adobe.com/xap/1.0/\0";
22
23pub const MAGIC_BYTES: &[u8] = &[0xFF, 0xD8, 0xFF];
24
25pub const MARKER_START: u8 = 0xFF;
26
27#[derive(Debug)]
28pub struct Jpeg {
29    segments: Vec<RawSegment>,
30    data: Vec<u8>,
31}
32
33impl ImageFormat for Jpeg {
34    fn is_filetype(data: &[u8]) -> bool {
35        data.starts_with(MAGIC_BYTES)
36    }
37}
38
39impl ImageMetadata for Jpeg {
40    fn exif(&self) -> Vec<Vec<u8>> {
41        self.exif_data().map(|x| x.to_vec()).collect()
42    }
43
44    fn xmp(&self) -> Vec<Vec<u8>> {
45        self.xmp_data().map(|x| x.to_vec()).collect()
46    }
47
48    fn pixel_density(&self) -> Option<PixelDensity> {
49        self.jfif().ok().and_then(|(jfif, _)| jfif.pixel_density())
50    }
51}
52
53impl Jpeg {
54    pub fn new(data: Vec<u8>) -> Result<Self, ErrorWithData<Error>> {
55        match Self::find_segments(&data) {
56            Ok(segments) => Ok(Self { segments, data }),
57            Err(err) => Err(ErrorWithData::new(err, data)),
58        }
59    }
60
61    pub fn into_inner(self) -> Vec<u8> {
62        self.data
63    }
64
65    /// List all segments in their order of appearance
66    pub fn segments(&self) -> Vec<Segment<'_>> {
67        self.segments.iter().map(|x| x.segment(self)).collect()
68    }
69
70    /// List all segments with the given marker
71    pub fn segments_marker(&self, marker: Marker) -> impl Iterator<Item = Segment<'_>> {
72        self.segments
73            .iter()
74            .filter(move |x| x.marker == Some(marker))
75            .map(|x| x.segment(self))
76    }
77
78    /// Quantization tables with `Tq` value as key
79    pub fn dqts(&self) -> Result<IndexMap<u8, Dqt>, Error> {
80        let segments = self.segments();
81
82        let mut dqts = Vec::new();
83        for i in segments
84            .into_iter()
85            .filter(|x| x.marker == Some(Marker::DQT))
86        {
87            let data = i.data();
88            dqts.push(Dqt::from_data(data)?);
89        }
90
91        let mut map = IndexMap::new();
92        for dqt in dqts.into_iter().flatten() {
93            map.insert(dqt.tq(), dqt);
94        }
95
96        Ok(map)
97    }
98
99    pub fn sof(&self) -> Result<Sof, Error> {
100        let segment = self
101            .segments()
102            .into_iter()
103            .find(|x| x.marker.is_some_and(|x| x.is_sof()))
104            .ok_or(Error::NoSofSegmentFound)?;
105
106        Sof::from_data(segment.data())
107    }
108
109    pub fn is_progressive(&self) -> Result<bool, Error> {
110        let sof_marker = self
111            .segments()
112            .into_iter()
113            .flat_map(|x| x.marker())
114            .find(|x| x.is_sof())
115            .ok_or(Error::NoSofSegmentFound)?;
116
117        sof_marker.is_progressive_sof()
118    }
119
120    /// Number of SOS segments
121    ///
122    /// For `is_progressive()` being true, this is the number of scans.
123    pub fn n_sos(&self) -> usize {
124        self.segments()
125            .into_iter()
126            .filter(|x| matches!(x.marker, Some(Marker::SOS)))
127            .count()
128    }
129
130    pub fn sos(&self) -> Result<Sos, Error> {
131        let segment = self
132            .segment_by_marker(Marker::SOS)
133            .ok_or(Error::NoSosSegmentFound)?;
134
135        Sos::from_data(segment.data())
136    }
137
138    pub fn components_specification_parameters(
139        &self,
140        component: usize,
141    ) -> Result<ComponentSpecificationParameters, Error> {
142        let cs = self
143            .sos()?
144            .components_specifications
145            .get(component)
146            .ok_or(Error::MissingComponentSpecification)?
147            .cs;
148        self.sof()?
149            .parameters
150            .iter()
151            .find(|x| x.c == cs)
152            .ok_or(Error::MissingComponentSpecificationParameters)
153            .cloned()
154    }
155
156    pub fn color_model(&self) -> Result<ColorModel, Error> {
157        let sof = self.sof()?;
158        let n_components = sof.parameters.len();
159
160        if let Some(app14) = self.segment_by_marker(Marker::APP14) {
161            if app14.data().starts_with(b"Adobe\0") {
162                if let Some(color_model) = app14.data().get(11) {
163                    return match *color_model {
164                        0 if n_components == 4 => Ok(ColorModel::Cmyk),
165                        0 if n_components == 3 => Ok(ColorModel::Rgb),
166                        1 => Ok(ColorModel::YCbCr),
167                        2 => Ok(ColorModel::Ycck),
168                        _ => Err(Error::UnknownColorModel),
169                    };
170                }
171            }
172        }
173
174        match n_components {
175            1 => Ok(ColorModel::Grayscale),
176            3 => Ok(ColorModel::YCbCr),
177            _ => Err(Error::UnknownColorModel),
178        }
179    }
180
181    pub fn segment_by_marker(&self, marker: Marker) -> Option<Segment<'_>> {
182        self.segments
183            .iter()
184            .find(|x| x.marker == Some(marker))
185            .map(|x| x.segment(self))
186    }
187
188    pub fn exif_segments(&self) -> impl Iterator<Item = Segment<'_>> {
189        self.segments_marker(Marker::APP1)
190            .filter(|x| x.data().starts_with(EXIF_IDENTIFIER_STRING))
191    }
192
193    pub fn exif_data(&self) -> impl Iterator<Item = &[u8]> {
194        self.exif_segments()
195            .filter_map(|x| x.data().get(EXIF_IDENTIFIER_STRING.len()..))
196    }
197
198    pub fn xmp_segments(&self) -> impl Iterator<Item = Segment<'_>> {
199        self.segments_marker(Marker::APP1)
200            .filter(|x| x.data().starts_with(XMP_IDENTIFIER_STRING))
201    }
202
203    pub fn xmp_data(&self) -> impl Iterator<Item = &[u8]> {
204        self.xmp_segments()
205            .filter_map(|x| x.data().get(XMP_IDENTIFIER_STRING.len()..))
206    }
207
208    pub fn jfif(&self) -> Result<(&Jfif, &[u8]), Error> {
209        if let Some(jfif) = self.segments().get(1) {
210            if jfif.data().get(0..5) == Some(b"JFIF\0") {
211                if let Some(data) = jfif.data().get(5..) {
212                    return Jfif::ref_from_prefix(data)
213                        .map_err(|err| Error::JfifUnavailable(format!("{err:?}")));
214                }
215            }
216        }
217
218        Err(Error::JfifUnavailable("Not found".to_string()))
219    }
220
221    fn find_segments(data: &[u8]) -> Result<Vec<RawSegment>, Error> {
222        let mut cur = Cursor::new(data);
223
224        let buf = &mut [0; 2];
225        cur.read_exact(buf).map_err(|_| Error::UnexpectedEof)?;
226
227        if data.get(..MAGIC_BYTES.len()) != Some(MAGIC_BYTES) {
228            return Err(Error::InvalidMagicBytes(*buf));
229        }
230
231        let mut segments = Vec::new();
232        segments.push(RawSegment {
233            marker: Some(Marker::SOI),
234            data: 2..2,
235        });
236
237        let mut entropy_coded_segment = false;
238        let byte = &mut [0; 1];
239        loop {
240            if entropy_coded_segment {
241                let data_start = cur.position().usize()?;
242                loop {
243                    cur.read_exact(byte).map_err(|_| Error::UnexpectedEof)?;
244                    if byte == &[MARKER_START] {
245                        cur.read_exact(byte).map_err(|_| Error::UnexpectedEof)?;
246
247                        if byte == &[0] {
248                            continue;
249                        } else {
250                            let data_end = cur.position().safe_sub(2)?.usize()?;
251                            segments.push(RawSegment {
252                                marker: None,
253                                data: data_start..data_end,
254                            });
255                            break;
256                        }
257                    }
258                }
259            } else {
260                // Read tag
261                cur.read_exact(byte).map_err(|_| Error::UnexpectedEof)?;
262
263                if byte != &[MARKER_START] {
264                    return Err(Error::ExpectedMarkerStart(buf[0]));
265                }
266
267                cur.read_exact(byte).map_err(|_| Error::UnexpectedEof)?;
268
269                tracing::debug!("Found tag {byte:0>2X?}");
270            }
271
272            let marker = Marker::from(byte[0]);
273            let len_start = cur.position();
274
275            let (data_start, len) = if marker.is_standalone() {
276                (len_start.usize()?, 0)
277            } else {
278                // Read length. The length includes the two length bytes, but not the marker.
279                cur.read_exact(buf).map_err(|_| Error::UnexpectedEof)?;
280                (len_start.usize()?.safe_add(2)?, u16::from_be_bytes(*buf))
281            };
282
283            let data_end = len_start.usize()?.safe_add(len.into())?;
284
285            let segment = RawSegment {
286                marker: Some(marker),
287                data: data_start..data_end,
288            };
289
290            tracing::debug!("Found segment {segment:?}");
291
292            segments.push(segment);
293
294            if marker == Marker::EOI {
295                break;
296            } else if marker == Marker::SOS {
297                entropy_coded_segment = true;
298            }
299
300            cur.set_position(len_start.safe_add(len.into())?);
301        }
302
303        Ok(segments)
304    }
305
306    pub fn replace_segment(
307        &mut self,
308        old_segment: RawSegment,
309        new_segment: NewSegment,
310    ) -> Result<(), Error> {
311        let old_range = old_segment.complete_data();
312
313        let mut new = Vec::new();
314        new.extend_from_slice(&self.data[..old_range.start]);
315        new_segment.write_to(&mut new);
316        new.extend_from_slice(&self.data[old_range.end..]);
317
318        self.data = new;
319        self.segments = Self::find_segments(&self.data)?;
320        Ok(())
321    }
322
323    /// Replaces this PNG's image data with those from another
324    ///
325    /// Keeps all the metadata from this image but replaces the `IHDR` and
326    /// `IDAT` chunks with the ones from `other`.
327    pub fn replace_image_data(&mut self, other: &Self) -> Result<(), Error> {
328        let mut buf = Vec::with_capacity(other.data.len());
329        buf.extend_from_slice(&MAGIC_BYTES[0..2]);
330
331        for segment in &self.segments {
332            let data = &self.data[segment.complete_data()];
333            // APP14 Adobe defines what color format an image is using which can clash with
334            // the `other` data.
335            if segment.marker.is_some_and(|x| x.is_metadata())
336                && (segment.marker != Some(APP14) || data.get(4..10) != Some(b"Adobe\0"))
337            {
338                buf.extend_from_slice(data);
339            }
340        }
341
342        for segment in &other.segments {
343            if !matches!(segment.marker, Some(Marker::SOI)) {
344                buf.extend_from_slice(&other.data[segment.complete_data()]);
345            }
346        }
347
348        self.segments = Self::find_segments(&buf).unwrap();
349        self.data = buf;
350
351        Ok(())
352    }
353}
354
355#[derive(Debug)]
356pub struct NewSegment<'a> {
357    marker: Marker,
358    data: &'a [u8],
359    total_len: u16,
360}
361
362impl<'a> NewSegment<'a> {
363    pub fn new(marker: Marker, data: &'a [u8]) -> Result<Self, Error> {
364        let total_len = data.len().u16()?.safe_add(2)?;
365
366        Ok(Self {
367            marker,
368            data,
369            total_len,
370        })
371    }
372
373    pub fn write_to(&self, vec: &mut Vec<u8>) {
374        vec.push(MARKER_START);
375        vec.push(self.marker.into());
376        vec.extend_from_slice(&self.total_len.to_be_bytes());
377        vec.extend_from_slice(self.data);
378    }
379}
380
381#[derive(Debug)]
382pub struct RawSegment {
383    marker: Option<Marker>,
384    data: Range<usize>,
385}
386
387impl RawSegment {
388    pub fn segment<'a>(&self, jpeg: &'a Jpeg) -> Segment<'a> {
389        Segment {
390            marker: self.marker,
391            data: self.data.clone(),
392            jpeg,
393        }
394    }
395
396    /// Complete segment including marker and length
397    pub fn complete_data(&self) -> Range<usize> {
398        let sub = if let Some(marker) = self.marker {
399            if marker.is_standalone() { 2 } else { 4 }
400        } else {
401            0
402        };
403
404        self.data.start.checked_sub(sub).expect(&format!(
405            "Unreachable: Marker and length fields always exist: {self:?}"
406        ))..self.data.end
407    }
408}
409
410#[derive(Clone, Debug)]
411pub struct Segment<'a> {
412    marker: Option<Marker>,
413    data: Range<usize>,
414    jpeg: &'a Jpeg,
415}
416
417impl<'a> Segment<'a> {
418    pub fn marker(&self) -> Option<Marker> {
419        self.marker
420    }
421
422    pub fn data_pos(&self) -> usize {
423        self.data.start
424    }
425
426    pub fn data(&self) -> &'a [u8] {
427        self.jpeg
428            .data
429            .get(self.data.clone())
430            .expect("Unreachable: This data must exist after successful loading")
431    }
432
433    pub fn unsafe_raw_segment(self) -> RawSegment {
434        RawSegment {
435            data: self.data,
436            marker: self.marker,
437        }
438    }
439}
440
441#[derive(Debug, Clone, thiserror::Error)]
442pub enum Error {
443    #[error("Invalid magic bytes: {0:x?}")]
444    InvalidMagicBytes([u8; 2]),
445    #[error("Unexpected end of file")]
446    UnexpectedEof,
447    #[error("Expected marker start: {0:x}")]
448    ExpectedMarkerStart(u8),
449    #[error("Math error: {0}")]
450    Math(#[from] MathError),
451    #[error("Unknown uantization table element precision {0}")]
452    UnknownPq(u8),
453    #[error("No SOS segment found")]
454    NoSosSegmentFound,
455    #[error("No SOF segment found")]
456    NoSofSegmentFound,
457    #[error("Couldn't detemine a color model")]
458    UnknownColorModel,
459    #[error("Missing component specification")]
460    MissingComponentSpecification,
461    #[error("Missing component specification parameters")]
462    MissingComponentSpecificationParameters,
463    #[error("Missing quantization table")]
464    MissingDqt,
465    #[error("JFIF unavailable: {0}")]
466    JfifUnavailable(String),
467}
468
469gufo_common::utils::convertible_enum!(
470    #[repr(u8)]
471    #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
472    #[non_exhaustive]
473    /// Segment marker
474    pub enum Marker {
475        TEM = 0x01,
476
477        SOF0 = 0xC0,
478        SOF1 = 0xC1,
479        SOF2 = 0xC2,
480        /// Define Huffman table
481        DHT = 0xC4,
482        RST0 = 0xD0,
483        RST1 = 0xD1,
484        RST2 = 0xD2,
485        RST3 = 0xD3,
486        RST4 = 0xD4,
487        RST5 = 0xD5,
488        RST6 = 0xD6,
489        RST7 = 0xD7,
490        /// Start of image
491        SOI = 0xD8,
492        /// End of image
493        EOI = 0xD9,
494        /// Start of scan
495        SOS = 0xDA,
496        /// Define quantization table(s)
497        DQT = 0xDB,
498
499        /// JFIF (pixel density, aspect ratio)
500        APP0 = 0xE0,
501        /// Exif, XMP
502        APP1 = 0xE1,
503        /// ICC color profile
504        APP2 = 0xE2,
505        APP3 = 0xE3,
506        APP4 = 0xE4,
507        APP5 = 0xE5,
508        APP6 = 0xE6,
509        APP7 = 0xE7,
510        APP8 = 0xE8,
511        APP9 = 0xE9,
512        APP10 = 0xEA,
513        APP11 = 0xEB,
514        APP12 = 0xEC,
515        APP13 = 0xED,
516        APP14 = 0xEE,
517        APP15 = 0xEF,
518        /// Define Restart Interval
519        DRI = 0xDD,
520
521        JPG0 = 0xF0,
522        JPG1 = 0xF1,
523        JPG2 = 0xF2,
524        JPG3 = 0xF3,
525        JPG4 = 0xF4,
526        JPG5 = 0xF5,
527        JPG6 = 0xF6,
528        JPG7 = 0xF7,
529        JPG8 = 0xF8,
530        JPG9 = 0xF9,
531        JPG10 = 0xFA,
532        JPG11 = 0xFB,
533        JPG12 = 0xFC,
534        JPG13 = 0xFD,
535        /// Comment
536        COM = 0xFE,
537    }
538);
539
540impl Marker {
541    pub fn is_standalone(&self) -> bool {
542        matches!(
543            self,
544            Self::RST0
545                | Self::RST1
546                | Self::RST2
547                | Self::RST3
548                | Self::RST4
549                | Self::RST5
550                | Self::RST6
551                | Self::RST7
552                | Self::SOI
553                | Self::EOI
554        )
555    }
556
557    pub fn is_sof(&self) -> bool {
558        matches!(self, Self::SOF0 | Self::SOF1 | Self::SOF2)
559    }
560
561    pub fn is_progressive_sof(&self) -> Result<bool, Error> {
562        match self {
563            Self::SOF0 | Self::SOF1 => Ok(false),
564            Self::SOF2 => Ok(true),
565            _ => Err(Error::NoSofSegmentFound),
566        }
567    }
568
569    pub fn is_metadata(&self) -> bool {
570        matches!(
571            self,
572            Self::COM
573                | Self::APP0
574                | Self::APP1
575                | Self::APP2
576                | Self::APP3
577                | Self::APP4
578                | Self::APP5
579                | Self::APP6
580                | Self::APP7
581                | Self::APP8
582                | Self::APP9
583                | Self::APP10
584                | Self::APP11
585                | Self::APP12
586                | Self::APP13
587                | Self::APP14
588                | Self::APP15
589        )
590    }
591}
592
593#[derive(Debug)]
594pub enum ColorModel {
595    Grayscale,
596    YCbCr,
597    Cmyk,
598    Rgb,
599    Ycck,
600}
gufo_jpeg/lib.rs

gufo_jpeg/
lib.rs
