From 2aa0d40c6197c6d23b47f8160e5d6befbe7a822c Mon Sep 17 00:00:00 2001
From: Lazy Nina <81658138+lazynina@users.noreply.github.com>
Date: Wed, 31 Jan 2024 12:44:39 -0500
Subject: [PATCH] Validate image dimensions (#534)

---
 routes/media.go | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/routes/media.go b/routes/media.go
index dcaffeca..dc502821 100644
--- a/routes/media.go
+++ b/routes/media.go
@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"github.com/gorilla/mux"
+	"image"
 	"io"
 	"net/http"
 	"strconv"
@@ -84,10 +85,31 @@ func resizeAndConvertToWebp(encodedImageString string, maxDim uint) (_image []by
 
 }
 
+// Prevent pixel flood attack. This function checks the image size before processing it.
+// Reference: https://github.com/h2non/bimg/issues/394#issuecomment-1015932411
+func validateImageSize(encodedImageContentBytes []byte) error {
+	byteReader := bytes.NewReader(encodedImageContentBytes)
+	imageConfig, _, err := image.DecodeConfig(byteReader)
+	if err != nil {
+		return err
+	}
+	if imageConfig.Width > bimg.MaxSize || imageConfig.Height > bimg.MaxSize {
+		return fmt.Errorf("image too large. Max dimensions are %v x %v. ImageConfig dimensions are %v x %v",
+			bimg.MaxSize, bimg.MaxSize, imageConfig.Width, imageConfig.Height)
+	}
+	return nil
+}
+
 func resizeAndConvertFromEncodedImageContent(encodedImageContent string, maxDim uint) (_image []byte, _err error) {
 	// always strip metadata
 	processOptions := bimg.Options{StripMetadata: true}
 	decodedBytes, err := base64.StdEncoding.DecodeString(encodedImageContent)
+	if err != nil {
+		return nil, err
+	}
+	if err = validateImageSize(decodedBytes); err != nil {
+		return nil, err
+	}
 	imgBytes, err := bimg.NewImage(decodedBytes).Process(processOptions)
 	if err != nil {
 		return nil, err