imgproxy/download.go

package main

import (
	"bytes"
	"context"
	"crypto/tls"
	"fmt"
	"image"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"strconv"
	"time"

	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"

	_ "github.com/mat/besticon/ico"
)

var (
	downloadClient  *http.Client
	imageTypeCtxKey = ctxKey("imageType")
	imageDataCtxKey = ctxKey("imageData")

	errSourceDimensionsTooBig      = newError(422, "Source image dimensions are too big", "Invalid source image")
	errSourceResolutionTooBig      = newError(422, "Source image resolution is too big", "Invalid source image")
	errSourceFileTooBig            = newError(422, "Source image file is too big", "Invalid source image")
	errSourceImageTypeNotSupported = newError(422, "Source image type not supported", "Invalid source image")
)

const msgSourceImageIsUnreachable = "Source image is unreachable"

var downloadBufPool *bufPool

type limitReader struct {
	r    io.ReadCloser
	left int
}

func (lr *limitReader) Read(p []byte) (n int, err error) {
	n, err = lr.r.Read(p)
	lr.left = lr.left - n

	if err == nil && lr.left < 0 {
		err = errSourceFileTooBig
	}

	return
}

func (lr *limitReader) Close() error {
	return lr.r.Close()
}

func initDownloading() {
	transport := &http.Transport{
		Proxy:               http.ProxyFromEnvironment,
		MaxIdleConns:        conf.Concurrency,
		MaxIdleConnsPerHost: conf.Concurrency,
		DisableCompression:  true,
		Dial:                (&net.Dialer{KeepAlive: 600 * time.Second}).Dial,
	}

	if conf.IgnoreSslVerification {
		transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
	}

	if conf.LocalFileSystemRoot != "" {
		transport.RegisterProtocol("local", http.NewFileTransport(http.Dir(conf.LocalFileSystemRoot)))
	}

	if conf.S3Enabled {
		transport.RegisterProtocol("s3", newS3Transport())
	}

	if len(conf.GCSKey) > 0 {
		transport.RegisterProtocol("gs", newGCSTransport())
	}

	downloadClient = &http.Client{
		Timeout:   time.Duration(conf.DownloadTimeout) * time.Second,
		Transport: transport,
	}

	downloadBufPool = newBufPool("download", conf.Concurrency, conf.DownloadBufferSize)
}

func checkDimensions(width, height int) error {
	if conf.MaxSrcDimension > 0 && (width > conf.MaxSrcDimension || height > conf.MaxSrcDimension) {
		return errSourceDimensionsTooBig
	}

	if width*height > conf.MaxSrcResolution {
		return errSourceResolutionTooBig
	}

	return nil
}

func checkTypeAndDimensions(r io.Reader) (imageType, error) {
	imgconf, imgtypeStr, err := image.DecodeConfig(r)
	if err == image.ErrFormat {
		return imageTypeUnknown, errSourceImageTypeNotSupported
	}
	if err != nil {
		return imageTypeUnknown, err
	}

	imgtype, imgtypeOk := imageTypes[imgtypeStr]
	if !imgtypeOk || !vipsTypeSupportLoad[imgtype] {
		return imageTypeUnknown, errSourceImageTypeNotSupported
	}

	if err = checkDimensions(imgconf.Width, imgconf.Height); err != nil {
		return imageTypeUnknown, err
	}

	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)
	} else {
		// ContentLength wasn't set properly, trying to parse the header
		contentLength, _ = strconv.Atoi(res.Header.Get("Content-Length"))
	}

	buf := downloadBufPool.Get(contentLength)
	cancel := func() {
		downloadBufPool.Put(buf)
	}

	if contentLength > buf.Cap() {
		buf.Grow(contentLength - buf.Len())
	}

	body := res.Body

	if conf.MaxSrcFileSize > 0 {
		body = &limitReader{r: body, left: conf.MaxSrcFileSize}
	}

	imgtype, err := checkTypeAndDimensions(io.TeeReader(body, buf))
	if err != nil {
		return ctx, cancel, err
	}

	if _, err = buf.ReadFrom(body); err != nil {
		return ctx, cancel, newError(404, err.Error(), msgSourceImageIsUnreachable)
	}

	ctx = context.WithValue(ctx, imageTypeCtxKey, imgtype)
	ctx = context.WithValue(ctx, imageDataCtxKey, buf)

	return ctx, cancel, nil
}

func downloadImage(ctx context.Context) (context.Context, context.CancelFunc, error) {
	url := getImageURL(ctx)

	if newRelicEnabled {
		newRelicCancel := startNewRelicSegment(ctx, "Downloading image")
		defer newRelicCancel()
	}

	if prometheusEnabled {
		defer startPrometheusDuration(prometheusDownloadDuration)()
	}

	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return ctx, func() {}, newError(404, err.Error(), msgSourceImageIsUnreachable)
	}

	req.Header.Set("User-Agent", conf.UserAgent)

	res, err := downloadClient.Do(req)
	if res != nil {
		defer res.Body.Close()
	}
	if err != nil {
		return ctx, func() {}, newError(404, err.Error(), msgSourceImageIsUnreachable)
	}

	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)
	}

	return readAndCheckImage(ctx, res)
}

func getImageType(ctx context.Context) imageType {
	return ctx.Value(imageTypeCtxKey).(imageType)
}

func getImageData(ctx context.Context) *bytes.Buffer {
	return ctx.Value(imageDataCtxKey).(*bytes.Buffer)
}