From 4d0d760f9278021549ba3c6655b3d101378b382f Mon Sep 17 00:00:00 2001
From: Ivan Smirnov <i.s.smirnov@gmail.com>
Date: Wed, 29 Dec 2021 16:00:16 +0300
Subject: [PATCH] Enforce maximum size of 400Mp (QOI_PIXELS_MAX)

---
 src/consts.rs | 2 ++
 src/encode.rs | 4 +++-
 src/error.rs  | 7 ++++++-
 src/header.rs | 4 +++-
 4 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/consts.rs b/src/consts.rs
index ee13ee0..204903b 100644
--- a/src/consts.rs
+++ b/src/consts.rs
@@ -24,3 +24,5 @@ pub const QOI_PADDING: usize = 4;
 
 pub const QOI_MAGIC: u32 =
     (b'q' as u32) << 24 | (b'o' as u32) << 16 | (b'i' as u32) << 8 | (b'f' as u32);
+
+pub const QOI_PIXELS_MAX: usize = 400_000_000;
diff --git a/src/encode.rs b/src/encode.rs
index 24cc598..93cff5f 100644
--- a/src/encode.rs
+++ b/src/encode.rs
@@ -3,7 +3,7 @@ use std::slice;
 use crate::colorspace::ColorSpace;
 use crate::consts::{
     QOI_COLOR, QOI_DIFF_16, QOI_DIFF_24, QOI_DIFF_8, QOI_HEADER_SIZE, QOI_INDEX, QOI_PADDING,
-    QOI_RUN_16, QOI_RUN_8,
+    QOI_PIXELS_MAX, QOI_RUN_16, QOI_RUN_8,
 };
 use crate::error::{Error, Result};
 use crate::header::Header;
@@ -137,6 +137,8 @@ where
     let n_pixels = (width as usize) * (height as usize);
     if unlikely(data.is_empty()) {
         return Err(Error::EmptyImage { width, height });
+    } else if unlikely(n_pixels > QOI_PIXELS_MAX) {
+        return Err(Error::ImageTooLarge { width, height });
     } else if unlikely(n_pixels * CHANNELS != data.len()) {
         return Err(Error::BadEncodingDataSize { size: data.len(), expected: n_pixels * CHANNELS });
     }
diff --git a/src/error.rs b/src/error.rs
index b1f8e5b..049c02d 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -2,12 +2,13 @@ use std::error::Error as StdError;
 use std::fmt::{self, Display};
 use std::result::Result as StdResult;
 
-use crate::consts::QOI_MAGIC;
+use crate::consts::{QOI_MAGIC, QOI_PIXELS_MAX};
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub enum Error {
     InvalidChannels { channels: u8 },
     EmptyImage { width: u32, height: u32 },
+    ImageTooLarge { width: u32, height: u32 },
     BadEncodingDataSize { size: usize, expected: usize },
     InputBufferTooSmall { size: usize, required: usize },
     OutputBufferTooSmall { size: usize, required: usize },
@@ -27,6 +28,10 @@ impl Display for Error {
             Self::EmptyImage { width, height } => {
                 write!(f, "image contains no pixels: {}x{}", width, height)
             }
+            Self::ImageTooLarge { width, height } => {
+                let mp = QOI_PIXELS_MAX / 1_000_000;
+                write!(f, "image is too large: {}x{} (max={}Mp)", width, height, mp)
+            }
             Self::BadEncodingDataSize { size, expected } => {
                 write!(f, "bad data size when encoding: {} (expected: {})", size, expected)
             }
diff --git a/src/header.rs b/src/header.rs
index 443634b..3e6071a 100644
--- a/src/header.rs
+++ b/src/header.rs
@@ -1,5 +1,5 @@
 use crate::colorspace::ColorSpace;
-use crate::consts::{QOI_HEADER_SIZE, QOI_MAGIC};
+use crate::consts::{QOI_HEADER_SIZE, QOI_MAGIC, QOI_PIXELS_MAX};
 use crate::error::{Error, Result};
 use crate::utils::unlikely;
 
@@ -90,6 +90,8 @@ impl Header {
             return Err(Error::InvalidMagic { magic: self.magic });
         } else if unlikely(self.height == 0 || self.width == 0) {
             return Err(Error::EmptyImage { width: self.width, height: self.height });
+        } else if unlikely((self.height as usize) * (self.width as usize) > QOI_PIXELS_MAX) {
+            return Err(Error::ImageTooLarge { width: self.width, height: self.height });
         } else if unlikely(self.channels < 3 || self.channels > 4) {
             return Err(Error::InvalidChannels { channels: self.channels });
         }