Enforce maximum image size

diff --git a/src/main.rs b/src/main.rs
index c9c48b8338d13255985ec58eaba7af688ffb1f96..96cb1b30d025f21f1161d000d766ec2acc04c72d 100644 (file)
--- 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;
     }