fakejpeg 0.2.0

// SPDX-FileCopyrightText: Alun Jones
// SPDX-FileCopyrightText: Gergely Nagy
// SPDX-FileContributor: Gergely Nagy
//
// SPDX-License-Identifier: MIT

use winnow::{
    ModalResult, Parser,
    binary::{be_u16, length_take, u8},
    combinator::{alt, cut_err, dispatch, empty, fail, peek, preceded, repeat, terminated, trace},
    error::StrContext,
    stream::{Offset, Stream},
    token::{rest, take, take_till},
};

use crate::template::chunk::{Chunk, Marker};

fn no_remaining(input: &mut &[u8]) -> ModalResult<usize> {
    trace("no_remaining", empty.value(0)).parse_next(input)
}

fn no_length(input: &mut &[u8]) -> ModalResult<&'static [u8]> {
    trace("no_length", empty.value::<&[u8]>(&[])).parse_next(input)
}

fn marker(input: &mut &[u8]) -> ModalResult<u16> {
    trace(
        "jpeg_marker",
        be_u16
            .verify(Marker::is_valid)
            .context(StrContext::Label("marker")),
    )
    .parse_next(input)
}

fn chunk_length(input: &mut &[u8]) -> ModalResult<u16> {
    trace("jpeg_chunk_length", be_u16.map(|n| n - 2)).parse_next(input)
}

fn marker_only_chunk(input: &mut &[u8]) -> ModalResult<Chunk> {
    trace(
        "jpeg_marker_only_chunk",
        ((marker, no_length), no_remaining).map(Chunk::build),
    )
    .parse_next(input)
}

fn header_with_length<'i>(input: &mut &'i [u8]) -> ModalResult<(u16, &'i [u8])> {
    trace(
        "jpeg_header_with_length",
        (marker, length_take(chunk_length)),
    )
    .parse_next(input)
}

fn chunk_with_length(input: &mut &[u8]) -> ModalResult<Chunk> {
    trace(
        "jpeg_chunk_with_length",
        (header_with_length, no_remaining).map(Chunk::build),
    )
    .parse_next(input)
}

fn discard_chunk(input: &mut &[u8]) -> ModalResult<Chunk> {
    trace(
        "jpeg_discard_chunk",
        header_with_length.void().map(|()| Chunk::default()),
    )
    .parse_next(input)
}

fn sos_escaped_eos(input: &mut &[u8]) -> ModalResult<(u8, u8)> {
    trace("jpeg_sos_escaped_eos", (0xff, 0x00)).parse_next(input)
}

fn sos_rst_marker(input: &mut &[u8]) -> ModalResult<(u8, u8)> {
    trace(
        "jpeg_sos_rst_marker",
        (0xff, u8.verify(|b| *b >= 0xd0 && *b <= 0xd7)),
    )
    .parse_next(input)
}

fn sos_discarded_variable_part(input: &mut &[u8]) -> ModalResult<()> {
    trace(
        "jpeg_sos_discarded_variable_part",
        terminated(take_till(0.., 0xff), alt((sos_escaped_eos, sos_rst_marker))).void(),
    )
    .parse_next(input)
}

fn sos_discarded_final_part(input: &mut &[u8]) -> ModalResult<()> {
    trace("jpeg_sos_discarded_final_part", take_till(0.., 0xff).void()).parse_next(input)
}

fn sos_discarded_variable_sequence(input: &mut &[u8]) -> ModalResult<()> {
    trace(
        "jpeg_sos_discarded_variable_sequence",
        repeat(0.., sos_discarded_variable_part).map(|()| ()).void(),
    )
    .parse_next(input)
}

fn sos_data_length(input: &mut &[u8]) -> ModalResult<usize> {
    trace("sos_data_length", |i: &mut &[u8]| {
        let start = i.checkpoint();
        let _ = (sos_discarded_variable_sequence, sos_discarded_final_part).parse_next(i)?;
        Ok(i.offset_from(&start))
    })
    .parse_next(input)
}

fn chunk_sos(input: &mut &[u8]) -> ModalResult<Chunk> {
    trace(
        "jpeg_chunk_sos",
        (header_with_length, sos_data_length).map(Chunk::build),
    )
    .parse_next(input)
}

fn chunk_app0_identifier<'a>(input: &mut &'a [u8]) -> ModalResult<&'a [u8]> {
    trace(
        "jpeg_chunk_app0_identifier",
        terminated(alt(("JFIF", "JFXX")), 0),
    )
    .parse_next(input)
}

fn jfif_app0_info<'a>(input: &mut &'a [u8]) -> ModalResult<&'a [u8]> {
    trace("jfif_app0_info", take(7u8)).parse_next(input)
}

fn jfif_discarded_thumbnail(input: &mut &[u8]) -> ModalResult<()> {
    trace(
        "jfif_discarded_thumbnail",
        length_take((u8, u8).map(|(x, y)| 3 * x * y)).void(),
    )
    .parse_next(input)
}

fn chunk_app0_jfif<'a>(input: &mut &'a [u8]) -> ModalResult<&'a [u8]> {
    trace(
        "jpeg_chunk_app0_jfif",
        terminated(jfif_app0_info, jfif_discarded_thumbnail),
    )
    .parse_next(input)
}

fn chunk_app0(input: &mut &[u8]) -> ModalResult<Chunk> {
    trace("jpeg_chunk_app0", |i: &mut &[u8]| {
        let (marker, length) = (marker, chunk_length).parse_next(i)?;

        dispatch!(chunk_app0_identifier;
                  // If the identifier is JFIF, grab the useful headers
                  // (anything before the thumbnail info & data), then discard
                  // the thumbnail, and reconstruct the payload.
                  b"JFIF" => chunk_app0_jfif.map(|data| {
                      let mut bytes = Vec::with_capacity(14);
                      bytes.extend(b"JFIF\0");
                      bytes.extend(data);
                      bytes.extend(&[0, 0]);

                      Chunk::build(((marker, &bytes), 0))
                  }),
                  // If the identifier is JFXX, take the specified length (minus
                  // 5, for the null-terminated identifier).
                  b"JFXX" => take(length - 5).void().map(|()| Chunk::default()),
                  // If all else fails - which shouldn't happen, because I used
                  // an alt(): fail.
                  _ => fail
        )
        .parse_next(i)
    })
    .parse_next(input)
}

fn chunk_soi(input: &mut &[u8]) -> ModalResult<Chunk> {
    trace(
        "jpeg_chunk_soi",
        preceded(
            peek(
                marker
                    .verify(|m| *m == Marker::SOI)
                    .context(StrContext::Label("SOI marker")),
            ),
            marker_only_chunk,
        ),
    )
    .parse_next(input)
}

fn chunk_eoi(input: &mut &[u8]) -> ModalResult<Chunk> {
    trace(
        "jpeg_chunk_eoi",
        preceded(
            peek(cut_err(
                marker
                    .verify(|m| *m == Marker::EOI)
                    .context(StrContext::Label("EOI marker")),
            )),
            marker_only_chunk,
        ),
    )
    .parse_next(input)
}

fn chunk(input: &mut &[u8]) -> ModalResult<Chunk> {
    trace(
        "jpeg_chunk",
        dispatch!(peek(marker);
                  Marker::EOI => fail,
                  Marker::RST0..=Marker::RST7 => marker_only_chunk,
                  Marker::SOS => chunk_sos,
                  Marker::APP0 => chunk_app0,
                  Marker::APP1..=Marker::COM | Marker::DRI => discard_chunk,
                 _ => chunk_with_length,
        ),
    )
    .parse_next(input)
}

pub fn parse_jpeg(input: &mut &[u8]) -> ModalResult<(Chunk, Vec<Chunk>, Chunk)> {
    trace(
        "jpeg",
        terminated((chunk_soi, repeat(1.., chunk), chunk_eoi), rest),
    )
    .parse_next(input)
}