From: MegaBrutal <code+git@megabrutal.com>
Date: Sun, 26 Jun 2022 00:20:00 +0000 (+0200)
Subject: Enforce maximum image size
X-Git-Url: http://git.megabrutal.com/?p=litoprism.git;a=commitdiff_plain;h=7e084bcf82e10b7836e96dcc383ae6d81eae0942

Enforce maximum image size
---

diff --git a/src/main.rs b/src/main.rs
index c9c48b8..96cb1b3 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -17,6 +17,12 @@ use image::error::{LimitError, LimitErrorKind};
 #[derive(Debug)]
 struct ImageError(image::ImageError);
 
+impl ImageError {
+    fn dimension_error() -> ImageError {
+        ImageError(image::ImageError::Limits(LimitError::from_kind(LimitErrorKind::DimensionError)))
+    }
+}
+
 impl Display for ImageError {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "{}", &self.0)
@@ -71,11 +77,16 @@ async fn greet(name: web::Path<String>, req: HttpRequest) -> impl Responder {
 }
 
 fn make_png(dim_x: u32, dim_y: u32, scale: u32, data: &mut dyn Iterator<Item = u8>) -> Result<Cursor<Vec<u8>>, ImageError> {
+    // Image must not be larger than 1 megapixel
+    if dim_x * dim_y * scale > 1048576 {
+        return Err(ImageError::dimension_error())
+    }
+
     let mut img: RgbImage = ImageBuffer::new(dim_x, dim_y);
     let mut pixels = img.pixels_mut();
 
     for sp in ToRgbIter::new(data) {
-        let mut dp = pixels.next().ok_or(ImageError(image::ImageError::Limits(LimitError::from_kind(LimitErrorKind::DimensionError))))?;
+        let mut dp = pixels.next().ok_or(ImageError::dimension_error())?;
         println!("{:?}", sp);
         dp.0 = sp.0;
     }