Refactored image downloading and watermark preparation

download.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"bytes"
 	"context"
 	"crypto/tls"
 	"fmt"
@@ -21,7 +20,6 @@ import (
 
 var (
 	downloadClient  *http.Client
-	imageTypeCtxKey = ctxKey("imageType")
 	imageDataCtxKey = ctxKey("imageData")
 
 	errSourceDimensionsTooBig      = newError(422, "Source image dimensions are too big", "Invalid source image")
@@ -34,8 +32,21 @@ const msgSourceImageIsUnreachable = "Source image is unreachable"
 
 var downloadBufPool *bufPool
 
+type imageData struct {
+	Data []byte
+	Type imageType
+
+	cancel context.CancelFunc
+}
+
+func (d *imageData) Close() {
+	if d.cancel != nil {
+		d.cancel()
+	}
+}
+
 type limitReader struct {
-	r    io.ReadCloser
+	r    io.Reader
 	left int
 }
 
@@ -50,10 +61,6 @@ func (lr *limitReader) Read(p []byte) (n int, err error) {
 	return
 }
 
-func (lr *limitReader) Close() error {
-	return lr.r.Close()
-}
-
 func initDownloading() {
 	transport := &http.Transport{
 		Proxy:               http.ProxyFromEnvironment,
@@ -120,45 +127,56 @@ func checkTypeAndDimensions(r io.Reader) (imageType, error) {
 	return imgtype, nil
 }
 
-func readAndCheckImage(ctx context.Context, res *http.Response) (context.Context, context.CancelFunc, error) {
-	var contentLength int
-
-	if res.ContentLength > 0 {
-		contentLength = int(res.ContentLength)
-
-		if conf.MaxSrcFileSize > 0 && contentLength > conf.MaxSrcFileSize {
-			return ctx, func() {}, errSourceFileTooBig
-		}
+func readAndCheckImage(r io.Reader, contentLength int) (*imageData, error) {
+	if conf.MaxSrcFileSize > 0 && contentLength > conf.MaxSrcFileSize {
+		return nil, errSourceFileTooBig
 	}
 
 	buf := downloadBufPool.Get(contentLength)
-	cancel := func() {
-		downloadBufPool.Put(buf)
+	cancel := func() { downloadBufPool.Put(buf) }
+
+	if conf.MaxSrcFileSize > 0 {
+		r = &limitReader{r: r, left: conf.MaxSrcFileSize}
 	}
 
-	body := res.Body
+	imgtype, err := checkTypeAndDimensions(io.TeeReader(r, buf))
+	if err != nil {
+		cancel()
+		return nil, err
+	}
 
-	if conf.MaxSrcFileSize > 0 {
-		body = &limitReader{r: body, left: conf.MaxSrcFileSize}
+	if _, err = buf.ReadFrom(r); err != nil {
+		cancel()
+		return nil, newError(404, err.Error(), msgSourceImageIsUnreachable)
 	}
 
-	imgtype, err := checkTypeAndDimensions(io.TeeReader(body, buf))
+	return &imageData{buf.Bytes(), imgtype, cancel}, nil
+}
+
+func requestImage(imageURL string) (*http.Response, error) {
+	req, err := http.NewRequest("GET", imageURL, nil)
 	if err != nil {
-		return ctx, cancel, err
+		return nil, newError(404, err.Error(), msgSourceImageIsUnreachable).MarkAsUnexpected()
 	}
 
-	if _, err = buf.ReadFrom(body); err != nil {
-		return ctx, cancel, newError(404, err.Error(), msgSourceImageIsUnreachable)
+	req.Header.Set("User-Agent", conf.UserAgent)
+
+	res, err := downloadClient.Do(req)
+	if err != nil {
+		return res, newError(404, err.Error(), msgSourceImageIsUnreachable).MarkAsUnexpected()
 	}
 
-	ctx = context.WithValue(ctx, imageTypeCtxKey, imgtype)
-	ctx = context.WithValue(ctx, imageDataCtxKey, buf)
+	if res.StatusCode != 200 {
+		body, _ := ioutil.ReadAll(res.Body)
+		msg := fmt.Sprintf("Can't download image; Status: %d; %s", res.StatusCode, string(body))
+		return res, newError(404, msg, msgSourceImageIsUnreachable).MarkAsUnexpected()
+	}
 
-	return ctx, cancel, nil
+	return res, nil
 }
 
 func downloadImage(ctx context.Context) (context.Context, context.CancelFunc, error) {
-	url := getImageURL(ctx)
+	imageURL := getImageURL(ctx)
 
 	if newRelicEnabled {
 		newRelicCancel := startNewRelicSegment(ctx, "Downloading image")
@@ -169,34 +187,24 @@ func downloadImage(ctx context.Context) (context.Context, context.CancelFunc, er
 		defer startPrometheusDuration(prometheusDownloadDuration)()
 	}
 
-	req, err := http.NewRequest("GET", url, nil)
-	if err != nil {
-		return ctx, func() {}, newError(404, err.Error(), msgSourceImageIsUnreachable).MarkAsUnexpected()
-	}
-
-	req.Header.Set("User-Agent", conf.UserAgent)
-
-	res, err := downloadClient.Do(req)
+	res, err := requestImage(imageURL)
 	if res != nil {
 		defer res.Body.Close()
 	}
 	if err != nil {
-		return ctx, func() {}, newError(404, err.Error(), msgSourceImageIsUnreachable).MarkAsUnexpected()
+		return ctx, func() {}, err
 	}
 
-	if res.StatusCode != 200 {
-		body, _ := ioutil.ReadAll(res.Body)
-		msg := fmt.Sprintf("Can't download image; Status: %d; %s", res.StatusCode, string(body))
-		return ctx, func() {}, newError(404, msg, msgSourceImageIsUnreachable).MarkAsUnexpected()
+	imgdata, err := readAndCheckImage(res.Body, int(res.ContentLength))
+	if err != nil {
+		return ctx, func() {}, err
 	}
 
-	return readAndCheckImage(ctx, res)
-}
+	ctx = context.WithValue(ctx, imageDataCtxKey, imgdata)
 
-func getImageType(ctx context.Context) imageType {
-	return ctx.Value(imageTypeCtxKey).(imageType)
+	return ctx, imgdata.Close, err
 }
 
-func getImageData(ctx context.Context) *bytes.Buffer {
-	return ctx.Value(imageDataCtxKey).(*bytes.Buffer)
+func getImageData(ctx context.Context) *imageData {
+	return ctx.Value(imageDataCtxKey).(*imageData)
 }

etag.go

@@ -31,7 +31,7 @@ func calcETag(ctx context.Context) string {
 	defer eTagCalcPool.Put(c)
 
 	c.hash.Reset()
-	c.hash.Write(getImageData(ctx).Bytes())
+	c.hash.Write(getImageData(ctx).Data)
 	footprint := c.hash.Sum(nil)
 
 	c.hash.Reset()

process.go

@@ -195,8 +195,8 @@ func cropImage(img *vipsImage, cropWidth, cropHeight int, gravity *gravityOption
 	return img.Crop(left, top, cropWidth, cropHeight)
 }
 
-func prepareWatermark(wm *vipsImage, wmData *watermarkData, opts *watermarkOptions, imgWidth, imgHeight int) error {
-	if err := wm.Load(wmData.data, wmData.imgtype, 1, 1.0, 1); err != nil {
+func prepareWatermark(wm *vipsImage, wmData *imageData, opts *watermarkOptions, imgWidth, imgHeight int) error {
+	if err := wm.Load(wmData.Data, wmData.Type, 1, 1.0, 1); err != nil {
 		return err
 	}
 
@@ -204,14 +204,14 @@ func prepareWatermark(wm *vipsImage, wmData *watermarkData, opts *watermarkOptio
 	po.Resize = resizeFit
 	po.Dpr = 1
 	po.Enlarge = true
-	po.Format = wmData.imgtype
+	po.Format = wmData.Type
 
 	if opts.Scale > 0 {
 		po.Width = maxInt(scaleInt(imgWidth, opts.Scale), 1)
 		po.Height = maxInt(scaleInt(imgHeight, opts.Scale), 1)
 	}
 
-	if err := transformImage(context.Background(), wm, wmData.data, po, wmData.imgtype); err != nil {
+	if err := transformImage(context.Background(), wm, wmData.Data, po, wmData.Type); err != nil {
 		return err
 	}
 
@@ -226,7 +226,7 @@ func prepareWatermark(wm *vipsImage, wmData *watermarkData, opts *watermarkOptio
 	return wm.Embed(opts.Gravity, imgWidth, imgHeight, opts.OffsetX, opts.OffsetY, rgbColor{0, 0, 0})
 }
 
-func applyWatermark(img *vipsImage, wmData *watermarkData, opts *watermarkOptions, framesCount int) error {
+func applyWatermark(img *vipsImage, wmData *imageData, opts *watermarkOptions, framesCount int) error {
 	wm := new(vipsImage)
 	defer wm.Clear()
 
@@ -541,15 +541,14 @@ func processImage(ctx context.Context) ([]byte, context.CancelFunc, error) {
 	defer vipsCleanup()
 
 	po := getProcessingOptions(ctx)
-	data := getImageData(ctx).Bytes()
-	imgtype := getImageType(ctx)
+	imgdata := getImageData(ctx)
 
 	if po.Format == imageTypeUnknown {
 		switch {
 		case po.PreferWebP && vipsTypeSupportSave[imageTypeWEBP]:
 			po.Format = imageTypeWEBP
-		case vipsTypeSupportSave[imgtype] && imgtype != imageTypeHEIC:
-			po.Format = imgtype
+		case vipsTypeSupportSave[imgdata.Type] && imgdata.Type != imageTypeHEIC:
+			po.Format = imgdata.Type
 		default:
 			po.Format = imageTypeJPEG
 		}
@@ -577,7 +576,7 @@ func processImage(ctx context.Context) ([]byte, context.CancelFunc, error) {
 		po.Width, po.Height = 0, 0
 	}
 
-	animationSupport := conf.MaxAnimationFrames > 1 && vipsSupportAnimation(imgtype) && vipsSupportAnimation(po.Format)
+	animationSupport := conf.MaxAnimationFrames > 1 && vipsSupportAnimation(imgdata.Type) && vipsSupportAnimation(po.Format)
 
 	pages := 1
 	if animationSupport {
@@ -587,16 +586,16 @@ func processImage(ctx context.Context) ([]byte, context.CancelFunc, error) {
 	img := new(vipsImage)
 	defer img.Clear()
 
-	if err := img.Load(data, imgtype, 1, 1.0, pages); err != nil {
+	if err := img.Load(imgdata.Data, imgdata.Type, 1, 1.0, pages); err != nil {
 		return nil, func() {}, err
 	}
 
 	if animationSupport && img.IsAnimated() {
-		if err := transformAnimated(ctx, img, data, po, imgtype); err != nil {
+		if err := transformAnimated(ctx, img, imgdata.Data, po, imgdata.Type); err != nil {
 			return nil, func() {}, err
 		}
 	} else {
-		if err := transformImage(ctx, img, data, po, imgtype); err != nil {
+		if err := transformImage(ctx, img, imgdata.Data, po, imgdata.Type); err != nil {
 			return nil, func() {}, err
 		}
 	}

vips.go

@@ -25,7 +25,7 @@ var (
 	vipsTypeSupportLoad  = make(map[imageType]bool)
 	vipsTypeSupportSave  = make(map[imageType]bool)
 
-	watermark *watermarkData
+	watermark *imageData
 )
 
 var vipsConf struct {

watermark_data.go

@@ -2,99 +2,73 @@ package main
 
 import (
 	"bytes"
-	"context"
 	"encoding/base64"
 	"fmt"
-	"io/ioutil"
 	"os"
 )
 
-type watermarkData struct {
-	data    []byte
-	imgtype imageType
-}
-
-func getWatermarkData() (*watermarkData, error) {
+func getWatermarkData() (*imageData, error) {
 	if len(conf.WatermarkData) > 0 {
-		data, imgtype, err := base64WatermarkData()
-
-		if err != nil {
-			return nil, err
-		}
-
-		return &watermarkData{data, imgtype}, err
+		return base64WatermarkData(conf.WatermarkData)
 	}
 
 	if len(conf.WatermarkPath) > 0 {
-		data, imgtype, err := fileWatermarkData()
-
-		if err != nil {
-			return nil, err
-		}
-
-		return &watermarkData{data, imgtype}, err
+		return fileWatermarkData(conf.WatermarkPath)
 	}
 
 	if len(conf.WatermarkURL) > 0 {
-		b, imgtype, cancel, err := remoteWatermarkData()
-		defer cancel()
-
-		if err != nil {
-			return nil, err
-		}
-
-		data := make([]byte, len(b))
-		copy(data, b)
-
-		return &watermarkData{data, imgtype}, err
+		return remoteWatermarkData(conf.WatermarkURL)
 	}
 
 	return nil, nil
 }
 
-func base64WatermarkData() ([]byte, imageType, error) {
-	data, err := base64.StdEncoding.DecodeString(conf.WatermarkData)
+func base64WatermarkData(encoded string) (*imageData, error) {
+	data, err := base64.StdEncoding.DecodeString(encoded)
 	if err != nil {
-		return nil, imageTypeUnknown, fmt.Errorf("Can't decode watermark data: %s", err)
+		return nil, fmt.Errorf("Can't decode watermark data: %s", err)
 	}
 
 	imgtype, err := checkTypeAndDimensions(bytes.NewReader(data))
 	if err != nil {
-		return nil, imageTypeUnknown, fmt.Errorf("Can't decode watermark: %s", err)
+		return nil, fmt.Errorf("Can't decode watermark: %s", err)
 	}
 
-	return data, imgtype, nil
+	return &imageData{Data: data, Type: imgtype}, nil
 }
 
-func fileWatermarkData() ([]byte, imageType, error) {
-	f, err := os.Open(conf.WatermarkPath)
+func fileWatermarkData(path string) (*imageData, error) {
+	f, err := os.Open(path)
 	if err != nil {
-		return nil, imageTypeUnknown, fmt.Errorf("Can't read watermark: %s", err)
+		return nil, fmt.Errorf("Can't read watermark: %s", err)
 	}
 
-	imgtype, err := checkTypeAndDimensions(f)
+	fi, err := f.Stat()
 	if err != nil {
-		return nil, imageTypeUnknown, fmt.Errorf("Can't decode watermark: %s", err)
+		return nil, fmt.Errorf("Can't read watermark: %s", err)
 	}
 
-	// Return to the beginning of the file
-	f.Seek(0, 0)
-
-	data, err := ioutil.ReadAll(f)
+	imgdata, err := readAndCheckImage(f, int(fi.Size()))
 	if err != nil {
-		return nil, imageTypeUnknown, fmt.Errorf("Can't read watermark: %s", err)
+		return nil, fmt.Errorf("Can't read watermark: %s", err)
 	}
 
-	return data, imgtype, nil
+	return imgdata, err
 }
 
-func remoteWatermarkData() ([]byte, imageType, context.CancelFunc, error) {
-	ctx := context.WithValue(context.Background(), imageURLCtxKey, conf.WatermarkURL)
-	ctx, cancel, err := downloadImage(ctx)
+func remoteWatermarkData(imageURL string) (*imageData, error) {
+	res, err := requestImage(imageURL)
+	if res != nil {
+		defer res.Body.Close()
+	}
+	if err != nil {
+		return nil, fmt.Errorf("Can't download watermark: %s", err)
+	}
 
+	imgdata, err := readAndCheckImage(res.Body, int(res.ContentLength))
 	if err != nil {
-		return nil, imageTypeUnknown, cancel, fmt.Errorf("Can't download watermark: %s", err)
+		return nil, fmt.Errorf("Can't download watermark: %s", err)
 	}
 
-	return getImageData(ctx).Bytes(), getImageType(ctx), cancel, err
+	return imgdata, err
 }