tavern/asset/agent.go

package asset

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"image"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"

	"github.com/buckket/go-blurhash"

	"github.com/ngerakines/tavern/common"
	"github.com/ngerakines/tavern/config"
	"github.com/ngerakines/tavern/storage"
)

type Agent struct {
	AssetStorage       Storage
	DataStorage        storage.Storage
	HTTPClient         common.HTTPClient
	AssetStorageConfig config.AssetStorageConfig
	FedConfig          config.FedConfig
}

func (a Agent) HandleImage(ctx context.Context, location string) (storage.ImageAsset, error) {
	// A whole lot of wrong can happen here. Need to log errors and handle cleanup gracefully.

	existingImages, err := a.DataStorage.GetImagesByAlias(ctx, []string{location})
	if err != nil {
		return storage.ImageAsset{}, err
	}
	if len(existingImages) == 1 {
		return existingImages[0], nil
	}

	if !a.allow(location) {
		return storage.ImageAsset{}, nil
	}

	tmpfile, err := ioutil.TempFile("", "image")
	if err != nil {
		return storage.ImageAsset{}, err
	}

	hasher := sha256.New()

	w := io.MultiWriter(tmpfile, hasher)

	sourceReader, err := a.download(ctx, location)
	if err != nil {
		return storage.ImageAsset{}, err
	}

	lsr := io.LimitReader(sourceReader, a.AssetStorageConfig.MaxFileSize+1) // 5 megs

	written, err := io.Copy(w, lsr)
	if err != nil {
		return storage.ImageAsset{}, err
	}
	if written > a.AssetStorageConfig.MaxFileSize {
		return storage.ImageAsset{}, fmt.Errorf("file size limit of 5 megabytes reached")
	}

	tmpFileName := tmpfile.Name()

	checksum := hex.EncodeToString(hasher.Sum(nil))

	err = tmpfile.Close()
	if err != nil {
		return storage.ImageAsset{}, err
	}

	bounds, ff, blurHash, err := a.imageBounds(tmpFileName)
	if err != nil {
		return storage.ImageAsset{}, err
	}

	contentType := storage.ContentTypeUnknown
	switch ff {
	case "png":
		contentType = storage.ContentTypePNG
	case "jpg", "jpeg":
		contentType = storage.ContentTypeJPG
	case "svg":
		// This is put in here as a placeholder to do some basic parsing and detection of SVG images.
		contentType = storage.ContentTypeSVG
	}

	fullLocation, err := a.AssetStorage.Upload(ctx, checksum, tmpFileName)
	if err != nil {
		return storage.ImageAsset{}, err
	}

	return a.DataStorage.CreateImage(ctx, fullLocation, checksum, blurHash, int(written), contentType, bounds.Max.Y, bounds.Max.X, []string{location})
}

func (a Agent) download(ctx context.Context, location string) (io.ReadCloser, error) {
	req, err := http.NewRequestWithContext(ctx, "GET", location, nil)
	if err != nil {
		return nil, err
	}

	resp, err := a.HTTPClient.Do(req)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode >= http.StatusOK && resp.StatusCode < http.StatusMultipleChoices {
		return resp.Body, nil
	}

	return nil, fmt.Errorf("unexpected response code: %d", resp.StatusCode)
}

func (a Agent) imageBounds(tempFileLocation string) (image.Rectangle, string, string, error) {
	file, err := os.Open(tempFileLocation)
	if err != nil {
		return image.Rectangle{}, "", "", err
	}
	defer file.Close()
	img, ff, err := image.Decode(file)
	if err != nil {
		return image.Rectangle{}, "", "", err
	}

	blurHash, err := blurhash.Encode(4, 3, &img)
	if err != nil {
		return image.Rectangle{}, "", "", err
	}

	return img.Bounds(), ff, blurHash, nil
}

func (a Agent) allow(location string) bool {
	u, err := url.Parse(location)
	if err != nil {
		return false
	}
	domain := u.Hostname()
	for _, matcher := range a.AssetStorageConfig.AllowDomains {
		if matcher.Match(domain) {
			return true
		}
	}
	for _, matcher := range a.AssetStorageConfig.DenyDomains {
		if matcher.Match(domain) {
			return false
		}
	}
	return true
}