feat: enhance image download and validation

- Add max image size limit from config - Improve image download with timeout and redirect handling - Validate content type and file size during image download - Update test cases to use httptest for mocking HTTP servers
aldinokemal · Feb 4, 2025 · 9c67c09 · 9c67c09
1 parent 37796b9
commit 9c67c09
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 20 deletions.
diff --git a/src/config/settings.go b/src/config/settings.go
@@ -24,6 +24,7 @@ var (
 	WhatsappWebhook                []string
 	WhatsappWebhookSecret                = "secret"
 	WhatsappLogLevel                     = "ERROR"
+	WhatsappSettingMaxImageSize    int64 = 20000000  // 20MB
 	WhatsappSettingMaxFileSize     int64 = 50000000  // 50MB
 	WhatsappSettingMaxVideoSize    int64 = 100000000 // 100MB
 	WhatsappSettingMaxDownloadSize int64 = 500000000 // 500MB

diff --git a/src/internal/rest/send.go b/src/internal/rest/send.go
@@ -52,6 +52,14 @@ func (controller *Send) SendImage(c *fiber.Ctx) error {
 	utils.PanicIfNeeded(err)
 
 	file, err := c.FormFile("image")
+	if err != nil && err != fiber.ErrNotFound {
+		return c.Status(fiber.StatusBadRequest).JSON(utils.ResponseData{
+			Status:  fiber.StatusBadRequest,
+			Code:    "ERROR",
+			Message: "Failed to process image file",
+			Results: err.Error(),
+		})
+	}
 	if err == nil {
 		request.Image = file
 	}

diff --git a/src/pkg/utils/general.go b/src/pkg/utils/general.go
@@ -13,6 +13,7 @@ import (
 	"time"
 
 	"github.com/PuerkitoBio/goquery"
+	"github.com/aldinokemal/go-whatsapp-web-multidevice/config"
 )
 
 // RemoveFile is removing file with delay
@@ -133,17 +134,34 @@ func ContainsMention(message string) []string {
 }
 
 func DownloadImageFromURL(url string) ([]byte, string, error) {
-	response, err := http.Get(url)
+	client := &http.Client{
+		Timeout: 30 * time.Second,
+		CheckRedirect: func(req *http.Request, via []*http.Request) error {
+			if len(via) >= 10 {
+				return fmt.Errorf("too many redirects")
+			}
+			return nil
+		},
+	}
+	response, err := client.Get(url)
 	if err != nil {
 		return nil, "", err
 	}
 	defer response.Body.Close()
-
+	contentType := response.Header.Get("Content-Type")
+	if !strings.HasPrefix(contentType, "image/") {
+		return nil, "", fmt.Errorf("invalid content type: %s", contentType)
+	}
+	// Check content length if available
+	if contentLength := response.ContentLength; contentLength > int64(config.WhatsappSettingMaxImageSize) {
+		return nil, "", fmt.Errorf("image size %d exceeds maximum allowed size %d", contentLength, config.WhatsappSettingMaxImageSize)
+	}
+	// Limit the size from config
+	reader := io.LimitReader(response.Body, int64(config.WhatsappSettingMaxImageSize))
 	// Extract the file name from the URL and remove query parameters if present
 	segments := strings.Split(url, "/")
 	fileName := segments[len(segments)-1]
 	fileName = strings.Split(fileName, "?")[0]
-
 	// Check if the file extension is supported
 	allowedExtensions := map[string]bool{
 		".jpg":  true,
@@ -155,11 +173,9 @@ func DownloadImageFromURL(url string) ([]byte, string, error) {
 	if !allowedExtensions[extension] {
 		return nil, "", fmt.Errorf("unsupported file type: %s", extension)
 	}
-
-	imageData, err := io.ReadAll(response.Body)
+	imageData, err := io.ReadAll(reader)
 	if err != nil {
 		return nil, "", err
 	}
-
 	return imageData, fileName, nil
 }
diff --git a/src/pkg/utils/general_test.go b/src/pkg/utils/general_test.go
@@ -2,9 +2,9 @@ package utils_test
 
 import (
 	"net/http"
+	"net/http/httptest"
 	"os"
 	"testing"
-	"time"
 
 	"github.com/aldinokemal/go-whatsapp-web-multidevice/pkg/utils"
 	"github.com/stretchr/testify/assert"
@@ -87,28 +87,31 @@ func (suite *UtilsTestSuite) TestStrToFloat64() {
 }
 
 func (suite *UtilsTestSuite) TestGetMetaDataFromURL() {
-	// Mock HTTP server
-	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+	// Use httptest.NewServer to mock HTTP server
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Write([]byte(`<!DOCTYPE html><html><head><title>Test Title</title><meta name='description' content='Test Description'><meta property='og:image' content='http://example.com/image.jpg'></head><body></body></html>`))
-	})
-	go http.ListenAndServe(":8080", nil)
-	time.Sleep(1 * time.Second) // Allow server to start
+	}))
+	defer server.Close() // Ensure the server is closed when the test ends
 
-	meta := utils.GetMetaDataFromURL("http://localhost:8080")
+	meta := utils.GetMetaDataFromURL(server.URL)
 	assert.Equal(suite.T(), "Test Title", meta.Title)
 	assert.Equal(suite.T(), "Test Description", meta.Description)
 	assert.Equal(suite.T(), "http://example.com/image.jpg", meta.Image)
 }
 
 func (suite *UtilsTestSuite) TestDownloadImageFromURL() {
-	// Mock HTTP server
-	http.HandleFunc("/image.jpg", func(w http.ResponseWriter, r *http.Request) {
-		w.Write([]byte("image data"))
-	})
-	go http.ListenAndServe(":8081", nil)
-	time.Sleep(1 * time.Second) // Allow server to start
+	// Use httptest.NewServer to mock HTTP server
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/image.jpg" {
+			w.Header().Set("Content-Type", "image/jpeg") // Set content type to image
+			w.Write([]byte("image data"))
+		} else {
+			http.NotFound(w, r)
+		}
+	}))
+	defer server.Close() // Ensure the server is closed when the test ends
 
-	imageData, fileName, err := utils.DownloadImageFromURL("http://localhost:8081/image.jpg")
+	imageData, fileName, err := utils.DownloadImageFromURL(server.URL + "/image.jpg")
 	assert.NoError(suite.T(), err)
 	assert.Equal(suite.T(), []byte("image data"), imageData)
 	assert.Equal(suite.T(), "image.jpg", fileName)