From 9c67c09d47f74450fbca22f2c8c39c8454b920ef Mon Sep 17 00:00:00 2001 From: Aldino Kemal Date: Wed, 5 Feb 2025 06:56:52 +0700 Subject: [PATCH] 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 --- src/config/settings.go | 1 + src/internal/rest/send.go | 8 ++++++++ src/pkg/utils/general.go | 28 ++++++++++++++++++++++------ src/pkg/utils/general_test.go | 31 +++++++++++++++++-------------- 4 files changed, 48 insertions(+), 20 deletions(-) diff --git a/src/config/settings.go b/src/config/settings.go index 5d45e7d..f35b027 100644 --- 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 index e30e9e4..9dc8bce 100644 --- 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 index aa64d55..90c8b80 100644 --- 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 index 9eb5b1a..b16c2bf 100644 --- 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(`Test Title`)) - }) - 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)