From f5f4fb64879692b179f124ee9b14ed1cf07f95f6 Mon Sep 17 00:00:00 2001
From: DarthSim <darthsim@gmail.com>
Date: Sat, 15 Oct 2022 19:31:04 +0600
Subject: [PATCH] Borrow buffer from imagedata pool when sanitizing SVG

---
 imagedata/read.go          |  9 +++++++++
 processing_handler.go      |  7 ++-----
 processing_handler_test.go |  4 ++--
 svg/svg.go                 | 19 ++++++++++++++-----
 4 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/imagedata/read.go b/imagedata/read.go
index 24bbcd99..363707e4 100644
--- a/imagedata/read.go
+++ b/imagedata/read.go
@@ -1,6 +1,8 @@
 package imagedata
 
 import (
+	"bytes"
+	"context"
 	"io"
 
 	"github.com/imgproxy/imgproxy/v3/bufpool"
@@ -86,3 +88,10 @@ func readAndCheckImage(r io.Reader, contentLength int) (*ImageData, error) {
 		cancel: cancel,
 	}, nil
 }
+
+func BorrowBuffer() (*bytes.Buffer, context.CancelFunc) {
+	buf := downloadBufPool.Get(0, false)
+	cancel := func() { downloadBufPool.Put(buf) }
+
+	return buf, cancel
+}
diff --git a/processing_handler.go b/processing_handler.go
index d55b8092..fa0c5221 100644
--- a/processing_handler.go
+++ b/processing_handler.go
@@ -344,17 +344,14 @@ func handleProcessing(reqID string, rw http.ResponseWriter, r *http.Request) {
 		// Don't process SVG
 		if originData.Type == imagetype.SVG {
 			if config.SanitizeSvg {
-				sanitized, svgErr := svg.Satitize(originData.Data)
+				sanitized, svgErr := svg.Satitize(originData)
 				checkErr(ctx, "svg_processing", svgErr)
 
 				// Since we'll replace origin data, it's better to close it to return
 				// it's buffer to the pool
 				originData.Close()
 
-				originData = &imagedata.ImageData{
-					Data: sanitized,
-					Type: imagetype.SVG,
-				}
+				originData = sanitized
 			}
 
 			respondWithImage(reqID, r, rw, statusCode, originData, po, imageURL, originData)
diff --git a/processing_handler_test.go b/processing_handler_test.go
index b241cb85..68e827d0 100644
--- a/processing_handler_test.go
+++ b/processing_handler_test.go
@@ -291,11 +291,11 @@ func (s *ProcessingHandlerTestSuite) TestSkipProcessingSVG() {
 	require.Equal(s.T(), 200, res.StatusCode)
 
 	actual := s.readBody(res)
-	expected, err := svg.Satitize(s.readTestFile("test1.svg"))
+	expected, err := svg.Satitize(&imagedata.ImageData{Data: s.readTestFile("test1.svg")})
 
 	require.Nil(s.T(), err)
 
-	require.True(s.T(), bytes.Equal(expected, actual))
+	require.True(s.T(), bytes.Equal(expected.Data, actual))
 }
 
 func (s *ProcessingHandlerTestSuite) TestNotSkipProcessingSVGToJPG() {
diff --git a/svg/svg.go b/svg/svg.go
index d3178a06..a9a02aa3 100644
--- a/svg/svg.go
+++ b/svg/svg.go
@@ -7,14 +7,15 @@ import (
 
 	"github.com/tdewolff/parse/v2"
 	"github.com/tdewolff/parse/v2/xml"
+
+	"github.com/imgproxy/imgproxy/v3/imagedata"
 )
 
-func Satitize(data []byte) ([]byte, error) {
-	r := bytes.NewReader(data)
+func Satitize(data *imagedata.ImageData) (*imagedata.ImageData, error) {
+	r := bytes.NewReader(data.Data)
 	l := xml.NewLexer(parse.NewInput(r))
 
-	buf := new(bytes.Buffer)
-	buf.Grow(len(data))
+	buf, cancel := imagedata.BorrowBuffer()
 
 	ignoreTag := 0
 
@@ -35,9 +36,17 @@ func Satitize(data []byte) ([]byte, error) {
 		switch tt {
 		case xml.ErrorToken:
 			if l.Err() != io.EOF {
+				cancel()
 				return nil, l.Err()
 			}
-			return buf.Bytes(), nil
+
+			newData := imagedata.ImageData{
+				Data: buf.Bytes(),
+				Type: data.Type,
+			}
+			newData.SetCancel(cancel)
+
+			return &newData, nil
 		case xml.StartTagToken:
 			if strings.ToLower(string(l.Text())) == "script" {
 				ignoreTag++