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

gufo_jpeg/
lib.rs
