Skip to content
Clone
b453d5d
462 lines 15.8 KB Rust
//! One-shot image preview generation shared by every image upload path:
//! avatars, space images and chat attachments.
//!
//! The originals are always kept; this crate only produces the smaller
//! derivative that list/grid views load. Anything it cannot handle — animated
//! GIF/APNG/WebP, unsupported formats, oversized sources, decode failures —
//! reports `Skipped`/`UseOriginal` instead of failing the caller's upload, so
//! the feature degrades to "serve the original" rather than a broken request.
mod error;
pub mod layout;
pub mod preset;
pub use error::PreviewSkip;
pub use preset::{JPEG_QUALITY, MAX_DECODE_PIXELS, MAX_SOURCE_BYTES, PreviewPreset};
use std::io::Cursor;
use image::codecs::gif::GifDecoder;
use image::codecs::jpeg::JpegEncoder;
use image::codecs::png::PngDecoder;
use image::codecs::webp::{WebPDecoder, WebPEncoder};
use image::imageops::FilterType;
use image::metadata::Orientation;
use image::{AnimationDecoder, DynamicImage, ImageDecoder, ImageFormat, ImageReader, Limits};
/// Decoder allocation ceiling, kept well under the process budget: a decompression
/// bomb otherwise turns one request into gigabytes of resident memory.
const DECODER_MAX_ALLOC: u64 = 256 * 1024 * 1024;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PreviewImage {
pub bytes: Vec<u8>,
pub content_type: &'static str,
pub extension: &'static str,
pub width: u32,
pub height: u32,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PreviewOutcome {
Preview(PreviewImage),
/// Nothing about the original is worth changing (animated or already
/// smaller than the preset): keep serving the original as the preview.
UseOriginal,
Skipped(PreviewSkip),
}
impl PreviewOutcome {
#[must_use]
pub fn preview(self) -> Option<PreviewImage> {
match self {
Self::Preview(image) => Some(image),
Self::UseOriginal | Self::Skipped(_) => None,
}
}
#[must_use]
pub fn is_skipped(&self) -> bool {
matches!(self, Self::Skipped(_))
}
}
/// Content types this crate is willing to decode. Everything else is skipped
/// untouched by the callers.
#[must_use]
pub fn is_previewable_content_type(content_type: &str) -> bool {
matches!(
canonical_mime(content_type),
Some(ImageFormat::Jpeg | ImageFormat::Png | ImageFormat::WebP | ImageFormat::Gif)
)
}
fn canonical_mime(content_type: &str) -> Option<ImageFormat> {
let main = content_type.split(';').next()?.trim();
let normalized = match main.to_ascii_lowercase().as_str() {
"image/jpg" => "image/jpeg".to_string(),
other => other.to_string(),
};
ImageFormat::from_mime_type(normalized)
}
/// Build a preview for `bytes`.
///
/// This is the synchronous core; call it through [`render_preview`] from an
/// async context so the decode/encode never blocks the runtime.
#[must_use]
pub fn render_preview_sync(
bytes: &[u8],
content_type: &str,
preset: PreviewPreset,
) -> PreviewOutcome {
let Some(format) = canonical_mime(content_type) else {
return PreviewOutcome::Skipped(PreviewSkip::UnsupportedFormat);
};
if !matches!(
format,
ImageFormat::Jpeg | ImageFormat::Png | ImageFormat::WebP | ImageFormat::Gif
) {
return PreviewOutcome::Skipped(PreviewSkip::UnsupportedFormat);
}
if bytes.is_empty() {
return PreviewOutcome::Skipped(PreviewSkip::Empty);
}
let source_len = bytes.len();
if bytes.len() as u64 > preset.max_source_bytes() {
return PreviewOutcome::Skipped(PreviewSkip::SourceTooLarge);
}
if is_animated(bytes, format) {
return PreviewOutcome::UseOriginal;
}
let mut limits = Limits::no_limits();
limits.max_image_width = Some(20_000);
limits.max_image_height = Some(20_000);
limits.max_alloc = Some(DECODER_MAX_ALLOC);
let mut reader = ImageReader::new(Cursor::new(bytes));
reader.set_format(format);
reader.limits(limits);
let mut decoder = match reader.into_decoder() {
Ok(decoder) => decoder,
Err(error) => {
return PreviewOutcome::Skipped(PreviewSkip::Decode(format!("{error}")));
}
};
let (raw_width, raw_height) = decoder.dimensions();
if u64::from(raw_width) * u64::from(raw_height) > preset.max_decode_pixels() {
return PreviewOutcome::Skipped(PreviewSkip::SourceTooLarge);
}
let orientation = decoder.orientation().unwrap_or(Orientation::NoTransforms);
let mut image = match DynamicImage::from_decoder(decoder) {
Ok(image) => image,
Err(error) => {
return PreviewOutcome::Skipped(PreviewSkip::Decode(format!("{error}")));
}
};
image.apply_orientation(orientation);
let (target_width, target_height) =
fit_within(image.width(), image.height(), preset.max_edge());
if target_width != image.width() || target_height != image.height() {
image = image.resize(target_width, target_height, FilterType::Triangle);
}
let has_alpha = image.color().has_alpha() && preset.keep_alpha();
let encoded = if has_alpha {
encode_webp_lossless(&image)
} else {
encode_jpeg(&image, preset.jpeg_quality())
};
let (bytes, content_type, extension) = match encoded {
Ok(value) => value,
Err(error) => {
return PreviewOutcome::Skipped(PreviewSkip::Encode(error.to_string()));
}
};
if bytes.is_empty() {
return PreviewOutcome::Skipped(PreviewSkip::Encode("empty encoder output".into()));
}
// A derivative that is not smaller than the original saves nothing on the
// wire; hand the original back instead of storing a useless twin.
if !is_smaller(bytes.len(), source_len) {
return PreviewOutcome::UseOriginal;
}
PreviewOutcome::Preview(PreviewImage {
bytes,
content_type,
extension,
width: image.width(),
height: image.height(),
})
}
/// Async wrapper: the decode/encode runs on the blocking pool so a burst of
/// image traffic cannot stall request handling.
pub async fn render_preview(
bytes: Vec<u8>,
content_type: String,
preset: PreviewPreset,
) -> PreviewOutcome {
if !is_previewable_content_type(&content_type) {
return PreviewOutcome::Skipped(PreviewSkip::UnsupportedFormat);
}
let task =
tokio::task::spawn_blocking(move || render_preview_sync(&bytes, &content_type, preset));
match task.await {
Ok(outcome) => outcome,
Err(error) => {
PreviewOutcome::Skipped(PreviewSkip::Decode(format!("preview task failed: {error}")))
}
}
}
fn encode_jpeg(
image: &DynamicImage,
quality: u8,
) -> Result<(Vec<u8>, &'static str, &'static str), String> {
// JPEG has no alpha channel; flattening to RGB is what browsers do for a
// transparent JPEG anyway.
let rgb = image.to_rgb8();
let mut out = Vec::new();
let mut encoder = JpegEncoder::new_with_quality(&mut out, quality);
encoder
.encode(
rgb.as_raw(),
rgb.width(),
rgb.height(),
image::ExtendedColorType::Rgb8,
)
.map_err(|error| error.to_string())?;
Ok((out, "image/jpeg", "jpg"))
}
fn encode_webp_lossless(
image: &DynamicImage,
) -> Result<(Vec<u8>, &'static str, &'static str), String> {
let rgba = image.to_rgba8();
let mut out = Vec::new();
WebPEncoder::new_lossless(&mut out)
.encode(
rgba.as_raw(),
rgba.width(),
rgba.height(),
image::ExtendedColorType::Rgba8,
)
.map_err(|error| error.to_string())?;
Ok((out, "image/webp", "webp"))
}
/// Animated containers are handed back untouched: re-encoding them would either
/// flatten the animation or explode the output size.
fn is_animated(bytes: &[u8], format: ImageFormat) -> bool {
let cursor = Cursor::new(bytes);
match format {
ImageFormat::Gif => GifDecoder::new(cursor)
.map(|decoder| decoder.into_frames().take(2).count() > 1)
.unwrap_or(false),
ImageFormat::WebP => WebPDecoder::new(cursor)
.map(|decoder| decoder.has_animation())
.unwrap_or(false),
ImageFormat::Png => PngDecoder::new(cursor)
.ok()
.and_then(|decoder| decoder.is_apng().ok())
.unwrap_or(false),
_ => false,
}
}
/// Whether a re-encoded derivative is actually worth storing. Requiring a
/// meaningful margin avoids keeping a "preview" that is byte-for-byte as heavy
/// as the original.
fn is_smaller(preview_len: usize, source_len: usize) -> bool {
preview_len < source_len
}
/// Longest edge at most `max_edge`, preserving the aspect ratio; images already
/// inside the box are returned unchanged.
fn fit_within(width: u32, height: u32, max_edge: u32) -> (u32, u32) {
let longest = width.max(height);
if longest <= max_edge || longest == 0 {
return (width, height);
}
let scale = f64::from(max_edge) / f64::from(longest);
let scaled_width = (f64::from(width) * scale).round().max(1.0) as u32;
let scaled_height = (f64::from(height) * scale).round().max(1.0) as u32;
(scaled_width, scaled_height)
}
#[cfg(test)]
mod tests {
use super::*;
use image::{Rgb, RgbImage, Rgba};
fn jpeg_bytes(width: u32, height: u32) -> Vec<u8> {
let image = RgbImage::from_fn(width, height, |x, y| {
Rgb([(x % 251) as u8, (y % 241) as u8, ((x + y) % 229) as u8])
});
let mut out = Vec::new();
let mut encoder = JpegEncoder::new_with_quality(&mut out, 92);
encoder
.encode(
image.as_raw(),
width,
height,
image::ExtendedColorType::Rgb8,
)
.expect("encode jpeg fixture");
out
}
fn png_with_alpha(width: u32, height: u32) -> Vec<u8> {
let image = image::RgbaImage::from_fn(width, height, |x, y| {
Rgba([(x % 200) as u8, (y % 200) as u8, 40, ((x + y) % 255) as u8])
});
let mut out = Vec::new();
image
.write_to(&mut Cursor::new(&mut out), ImageFormat::Png)
.expect("encode png fixture");
out
}
fn gif_two_frames(width: u32, height: u32) -> Vec<u8> {
let frame = image::RgbaImage::from_pixel(width, height, Rgba([10, 20, 30, 255]));
let mut out = Vec::new();
{
use image::codecs::gif::GifEncoder;
let mut encoder = GifEncoder::new(&mut out);
encoder
.encode_frame(image::Frame::new(frame.clone()))
.expect("gif frame 1");
encoder
.encode_frame(image::Frame::new(frame))
.expect("gif frame 2");
}
out
}
#[test]
fn downscales_to_preset_edge() {
let bytes = jpeg_bytes(3000, 1500);
let outcome = render_preview_sync(&bytes, "image/jpeg", PreviewPreset::Content);
let PreviewOutcome::Preview(preview) = outcome else {
panic!("expected a preview");
};
assert_eq!(preview.width, 1920);
assert_eq!(preview.height, 960);
assert_eq!(preview.content_type, "image/jpeg");
assert!(preview.bytes.len() < bytes.len());
}
#[test]
fn avatar_preset_uses_512_edge() {
let bytes = jpeg_bytes(2000, 2000);
let PreviewOutcome::Preview(preview) =
render_preview_sync(&bytes, "image/jpeg", PreviewPreset::Avatar)
else {
panic!("expected a preview");
};
assert_eq!(preview.width, 512);
assert_eq!(preview.height, 512);
}
#[test]
fn alpha_keeps_webp_alpha() {
let bytes = png_with_alpha(400, 300);
let PreviewOutcome::Preview(preview) =
render_preview_sync(&bytes, "image/png", PreviewPreset::Avatar)
else {
panic!("expected a preview");
};
assert_eq!(preview.content_type, "image/webp");
assert_eq!(preview.extension, "webp");
let decoded = image::load_from_memory(&preview.bytes).expect("decode preview");
assert!(decoded.color().has_alpha());
}
#[test]
fn animated_gif_is_left_alone() {
let bytes = gif_two_frames(64, 64);
let outcome = render_preview_sync(&bytes, "image/gif", PreviewPreset::Content);
assert_eq!(outcome, PreviewOutcome::UseOriginal);
}
#[test]
fn unsupported_type_is_skipped() {
let outcome = render_preview_sync(b"not an image", "image/avif", PreviewPreset::Content);
assert!(outcome.is_skipped());
}
#[test]
fn undecodable_payload_is_skipped() {
let outcome = render_preview_sync(
b"definitely not a jpeg",
"image/jpeg",
PreviewPreset::Content,
);
assert!(outcome.is_skipped());
}
#[test]
fn preview_smaller_rule_rejects_equal_or_larger_output() {
assert!(is_smaller(10, 11));
assert!(!is_smaller(11, 11));
assert!(!is_smaller(12, 11));
}
/// EXIF orientation 6 means "rotate 90° clockwise", so a wide camera frame
/// must come out of the preview upright (portrait) and still within the box.
#[test]
fn exif_rotated_jpeg_is_uprighted() {
let source = jpeg_bytes(1200, 600);
let rotated = embed_exif_orientation(&source, 6);
let PreviewOutcome::Preview(preview) =
render_preview_sync(&rotated, "image/jpeg", PreviewPreset::Content)
else {
panic!("expected a preview");
};
assert_eq!((preview.width, preview.height), (600, 1200));
}
/// Longest-edge capping must apply after the rotation, not before.
#[test]
fn exif_rotated_oversized_jpeg_is_capped_after_rotation() {
let source = jpeg_bytes(3000, 1200);
let rotated = embed_exif_orientation(&source, 6);
let PreviewOutcome::Preview(preview) =
render_preview_sync(&rotated, "image/jpeg", PreviewPreset::Content)
else {
panic!("expected a preview");
};
assert_eq!((preview.width, preview.height), (768, 1920));
}
/// Splice a minimal APP1/EXIF segment carrying `orientation` (tag 0x0112)
/// directly behind the JPEG SOI marker.
fn embed_exif_orientation(jpeg: &[u8], orientation: u16) -> Vec<u8> {
assert_eq!(&jpeg[..2], &[0xff, 0xd8], "fixture must start with SOI");
let mut tiff = Vec::new();
tiff.extend_from_slice(b"II");
tiff.extend_from_slice(&42u16.to_le_bytes());
tiff.extend_from_slice(&8u32.to_le_bytes());
tiff.extend_from_slice(&1u16.to_le_bytes());
tiff.extend_from_slice(&0x0112u16.to_le_bytes());
tiff.extend_from_slice(&3u16.to_le_bytes());
tiff.extend_from_slice(&1u32.to_le_bytes());
tiff.extend_from_slice(&orientation.to_le_bytes());
tiff.extend_from_slice(&[0, 0]);
tiff.extend_from_slice(&0u32.to_le_bytes());
let mut payload = Vec::new();
payload.extend_from_slice(b"Exif\0\0");
payload.extend_from_slice(&tiff);
let mut out = Vec::new();
out.extend_from_slice(&jpeg[..2]);
out.extend_from_slice(&[0xff, 0xe1]);
let segment_len = (payload.len() + 2) as u16;
out.extend_from_slice(&segment_len.to_be_bytes());
out.extend_from_slice(&payload);
out.extend_from_slice(&jpeg[2..]);
out
}
#[test]
fn small_image_is_not_upscaled() {
let bytes = jpeg_bytes(320, 200);
let PreviewOutcome::Preview(preview) =
render_preview_sync(&bytes, "image/jpeg", PreviewPreset::Content)
else {
panic!("expected a preview");
};
assert_eq!((preview.width, preview.height), (320, 200));
}
#[test]
fn extra_large_source_is_skipped() {
let bypassed = PreviewPreset::Content;
assert!(bypassed.max_source_bytes() > 0);
assert_eq!(fit_within(10, 10, 1920), (10, 10));
assert_eq!(fit_within(0, 0, 1920), (0, 0));
}
}