yumechi-no-kuni-proxy-worker/src/post_process/image_processing.rs

use std::io::{BufReader, Cursor};

use image::{
    codecs::{png::PngDecoder, webp::WebPDecoder},
    AnimationDecoder, DynamicImage, GenericImageView, ImageDecoder, ImageReader, ImageResult,
    Limits,
};

use crate::ImageOptions;

const fn clamp_width(input: (u32, u32), max_width: u32) -> (u32, u32) {
    if input.0 > max_width {
        (max_width, input.1 * max_width / input.0)
    } else {
        input
    }
}

const fn clamp_height(input: (u32, u32), max_height: u32) -> (u32, u32) {
    if input.1 > max_height {
        (input.0 * max_height / input.1, max_height)
    } else {
        input
    }
}

const fn clamp_dimensions(input: (u32, u32), max_width: u32, max_height: u32) -> (u32, u32) {
    clamp_height(clamp_width(input, max_width), max_height)
}

// All constants are following https://github.com/misskey-dev/media-proxy/blob/master/SPECIFICATION.md

/// Postprocesses an WebP image using the given options
pub fn postprocess_webp_image(
    data: &[u8],
    opt: &ImageOptions,
) -> ImageResult<Option<DynamicImage>> {
    let mut dec = WebPDecoder::new(Cursor::new(data))?;
    dec.set_limits(default_limits())
        .expect("Failed to set limits");

    if !dec.has_animation() && opt.requested_resize() {
        return Ok(Some(postprocess_static_image(
            data,
            opt,
            Some(image::ImageFormat::WebP),
        )?));
    }

    if opt.static_ == Some(true) {
        let first_frame = match dec.into_frames().next() {
            Some(Ok(frame)) => frame,
            _ => return Ok(None),
        };
        return Ok(Some(process_static_image_impl(
            first_frame.into_buffer().into(),
            opt,
        )));
    }

    Ok(None)
}

/// Postprocesses an PNG image using the given options
pub fn postprocess_png_image(data: &[u8], opt: &ImageOptions) -> ImageResult<Option<DynamicImage>> {
    let span = tracing::span!(tracing::Level::INFO, "postprocess_png_image");
    let _guard = span.enter();

    let mut dec = PngDecoder::new(Cursor::new(data))?;
    dec.set_limits(default_limits())
        .expect("Failed to set limits");

    let apng = dec.is_apng()?;

    if !apng && opt.requested_resize() {
        return Ok(Some(postprocess_static_image(
            data,
            opt,
            Some(image::ImageFormat::Png),
        )?));
    }

    if apng && opt.static_ == Some(true) {
        tracing::info!("Processing APNG image into static image");
        let first_frame = match dec.apng()?.into_frames().next() {
            Some(Ok(frame)) => frame,
            _ => return Ok(None),
        };
        return Ok(Some(process_static_image_impl(
            first_frame.into_buffer().into(),
            opt,
        )));
    }

    Ok(None)
}

#[derive(Debug, thiserror::Error)]
/// An error that occurred during SVG postprocessing
pub enum SvgPostprocessError {
    /// The SVG image is unsupported
    #[error("SVG unsupported")]
    Unsupported,
    /// The complexity of the SVG image exceeded the limit
    #[error("SVG complexity limit exceeded")]
    LimitExceeded,
    /// Disallowed content was found in the SVG image
    #[error("Disallowed content found in SVG")]
    DisallowedContent,
    /// An error that occurred during SVG parsing
    #[error("XML read error: {0}")]
    XmlRead(#[from] xml::reader::Error),
    /// An error that occurred during SVG writing
    #[error("XML write error: {0}")]
    XmlWrite(#[from] xml::writer::Error),
}

/// Preprocesses an SVG image using the given options. By specs SVG must be rasterized as opposed to passing the SVG data directly for security reasons.
pub fn postprocess_svg_image(
    data: &[u8],
    opt: &ImageOptions,
) -> Result<Vec<u8>, SvgPostprocessError> {
    let span = tracing::span!(tracing::Level::INFO, "postprocess_svg_image");
    let _guard = span.enter();

    let static_ = opt.static_ == Some(true);
    let rdr = BufReader::new(Cursor::new(data));
    let tree = xml::EventReader::new(rdr);
    let mut out = Vec::new();
    let mut wtr = xml::EventWriter::new(&mut out);

    let mut complexity = 0;
    let mut ignore_stack = None;
    for r in tree {
        match r {
            Ok(o) => match o {
                xml::reader::XmlEvent::StartDocument { .. } => {
                    complexity += 1;
                    wtr.write(
                        o.as_writer_event()
                            .ok_or(SvgPostprocessError::Unsupported)?,
                    )?;
                }
                xml::reader::XmlEvent::EndDocument => {
                    break;
                }
                xml::reader::XmlEvent::ProcessingInstruction { ref name, .. } => {
                    complexity += 1 + name.len();
                    if ignore_stack.is_some() {
                        continue;
                    }
                    if name != "xml" {
                        return Err(SvgPostprocessError::DisallowedContent);
                    }

                    wtr.write(
                        o.as_writer_event()
                            .ok_or(SvgPostprocessError::Unsupported)?,
                    )?;
                }
                xml::reader::XmlEvent::StartElement {
                    name,
                    mut attributes,
                    namespace,
                } => {
                    complexity += 1 + attributes.len() + namespace.0.len();

                    if ignore_stack.is_some() {
                        ignore_stack = ignore_stack.map(|v| v + 1);
                        continue;
                    }

                    let name_str = name.local_name.to_ascii_lowercase();
                    if name_str.contains("script") {
                        return Err(SvgPostprocessError::DisallowedContent);
                    }

                    if static_ {
                        if name_str == "set" || name_str.starts_with("animate") {
                            ignore_stack = Some(0);
                            continue;
                        }
                    }
                    attributes.retain(|attr| {
                        let lname = attr.name.local_name.to_ascii_lowercase();
                        if lname.starts_with("on") {
                            return false;
                        }

                        if lname.ends_with("src") || lname.ends_with("href") {
                            if !attr.value.starts_with("data:image/")
                                && !attr.value.contains("svg")
                                && !attr.value.contains("xml")
                            {
                                return false;
                            }
                        }

                        if lname.eq_ignore_ascii_case("style") {
                            let style = attr.value.as_str();
                            if style.to_ascii_lowercase().contains("url") {
                                return false;
                            }
                        }
                        return true;
                    });

                    wtr.write(
                        xml::reader::XmlEvent::StartElement {
                            name,
                            attributes,
                            namespace,
                        }
                        .as_writer_event()
                        .ok_or(SvgPostprocessError::Unsupported)?,
                    )?;
                }
                xml::reader::XmlEvent::EndElement { .. } => {
                    complexity += 1;
                    if ignore_stack.is_some() {
                        ignore_stack = ignore_stack.map(|v| v - 1);
                        if ignore_stack == Some(0) {
                            ignore_stack = None;
                        }
                        continue;
                    }
                    wtr.write(
                        o.as_writer_event()
                            .ok_or(SvgPostprocessError::Unsupported)?,
                    )?;
                }
                xml::reader::XmlEvent::CData(ref s) => {
                    complexity += 1 + s.len() / 32;
                    if ignore_stack.is_some() {
                        continue;
                    }
                    wtr.write(
                        o.as_writer_event()
                            .ok_or(SvgPostprocessError::Unsupported)?,
                    )?;
                }
                xml::reader::XmlEvent::Comment(s) => complexity += 1 + s.len() / 64,
                xml::reader::XmlEvent::Characters(ref s) => {
                    complexity += 1 + s.len() / 32;
                    if ignore_stack.is_some() {
                        continue;
                    }
                    wtr.write(
                        o.as_writer_event()
                            .ok_or(SvgPostprocessError::Unsupported)?,
                    )?;
                }
                xml::reader::XmlEvent::Whitespace(s) => complexity += 1 + s.len() / 64,
            },

            Err(e) => return Err(SvgPostprocessError::XmlRead(e)),
        }

        tracing::debug!("SVG complexity: {}", complexity);

        if complexity > 65536 {
            return Err(SvgPostprocessError::LimitExceeded);
        }
    }

    Ok(out)
}

fn default_limits() -> Limits {
    let mut l = Limits::default();
    l.max_alloc = Some(64 << 20);
    l.max_image_height = Some(16384);
    l.max_image_width = Some(16384);

    l
}

/// Preprocesses a static image using the given options
pub fn postprocess_static_image(
    data: &[u8],
    opt: &ImageOptions,
    format: Option<image::ImageFormat>,
) -> ImageResult<DynamicImage> {
    let span = tracing::span!(tracing::Level::INFO, "postprocess_static_image");
    let _guard = span.enter();

    let mut r = match format {
        Some(f) => ImageReader::with_format(Cursor::new(data), f),
        None => ImageReader::new(Cursor::new(data))
            .with_guessed_format()
            .inspect_err(|e| {
                tracing::error!("Failed to guess image format: {}", e);
            })?,
    };

    r.limits(default_limits());

    Ok(process_static_image_impl(
        r.decode().inspect_err(|e| {
            tracing::error!("Failed to decode image: {}", e);
        })?,
        opt,
    ))
}

fn process_static_image_impl(mut image: DynamicImage, opt: &ImageOptions) -> DynamicImage {
    let mut out_dim = image.dimensions();

    if opt.badge == Some(true) {
        return image.resize_exact(96, 96, image::imageops::FilterType::Nearest);
    }

    let mut use_bilinear = true;

    if opt.emoji == Some(true) {
        use_bilinear = false;
        out_dim = clamp_height(out_dim, 128);
    } else if opt.avatar == Some(true) {
        use_bilinear = false;
        out_dim = clamp_height(out_dim, 320);
    }

    if opt.preview == Some(true) {
        use_bilinear = false;
        out_dim = clamp_dimensions(out_dim, 200, 200);
    }

    if opt.static_ == Some(true) {
        out_dim = clamp_dimensions(out_dim, 498, 422);
    }

    if out_dim != image.dimensions() {
        image = image.resize_exact(
            out_dim.0,
            out_dim.1,
            if use_bilinear {
                image::imageops::FilterType::Triangle
            } else {
                image::imageops::FilterType::Nearest
            },
        );
    }

    image
}