Convert thumbnails over to the new middle layer

Part of #58

src/github.com/turt2live/matrix-media-repo/api/r0/thumbnail.go

@@ -9,7 +9,7 @@ import (
 	"github.com/turt2live/matrix-media-repo/api"
 	"github.com/turt2live/matrix-media-repo/common"
 	"github.com/turt2live/matrix-media-repo/common/config"
-	"github.com/turt2live/matrix-media-repo/old_middle_layer/media_cache"
+	"github.com/turt2live/matrix-media-repo/controllers/thumbnail_controller"
 )
 
 func ThumbnailMedia(r *http.Request, log *logrus.Entry, user api.UserInfo) interface{} {
@@ -75,9 +75,7 @@ func ThumbnailMedia(r *http.Request, log *logrus.Entry, user api.UserInfo) inter
 		"requestedAnimated": animated,
 	})
 
-	mediaCache := media_cache.Create(r.Context(), log)
-
-	streamedThumbnail, err := mediaCache.GetThumbnail(server, mediaId, width, height, method, animated, downloadRemote)
+	streamedThumbnail, err := thumbnail_controller.GetThumbnail(server, mediaId, width, height, animated, method, downloadRemote, r.Context(), log)
 	if err != nil {
 		if err == common.ErrMediaNotFound {
 			return api.NotFoundError()

src/github.com/turt2live/matrix-media-repo/controllers/thumbnail_controller/thumbnail_controller.go

+package thumbnail_controller
+
+import (
+	"bytes"
+	"context"
+	"database/sql"
+	"fmt"
+	"image"
+	"image/color"
+	"math"
+	"os"
+	"time"
+
+	"github.com/disintegration/imaging"
+	"github.com/fogleman/gg"
+	"github.com/golang/freetype/truetype"
+	"github.com/patrickmn/go-cache"
+	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
+	"github.com/turt2live/matrix-media-repo/common"
+	"github.com/turt2live/matrix-media-repo/common/config"
+	"github.com/turt2live/matrix-media-repo/controllers/download_controller"
+	"github.com/turt2live/matrix-media-repo/internal_cache"
+	"github.com/turt2live/matrix-media-repo/storage"
+	"github.com/turt2live/matrix-media-repo/types"
+	"github.com/turt2live/matrix-media-repo/util"
+	"golang.org/x/image/font/gofont/gosmallcaps"
+)
+
+// These are the content types that we can actually thumbnail
+var supportedThumbnailTypes = []string{"image/jpeg", "image/jpg", "image/png", "image/gif", "image/svg+xml"}
+
+// Of the SupportedThumbnailTypes, these are the 'animated' types
+var animatedTypes = []string{"image/gif"}
+
+var localCache = cache.New(30*time.Second, 60*time.Second)
+
+func GetThumbnail(origin string, mediaId string, desiredWidth int, desiredHeight int, animated bool, method string, downloadRemote bool, ctx context.Context, log *logrus.Entry) (*types.StreamedThumbnail, error) {
+	media, err := download_controller.FindMediaRecord(origin, mediaId, downloadRemote, ctx, log)
+	if err != nil {
+		return nil, err
+	}
+
+	if !util.ArrayContains(supportedThumbnailTypes, media.ContentType) {
+		log.Warn("Cannot generate thumbnail for " + media.ContentType + " because it is not supported")
+		return nil, errors.New("cannot generate thumbnail for this media's content type")
+	}
+
+	if !util.ArrayContains(config.Get().Thumbnails.Types, media.ContentType) {
+		log.Warn("Cannot generate thumbnail for " + media.ContentType + " because it is not listed in the config")
+		return nil, errors.New("cannot generate thumbnail for this media's content type")
+	}
+
+	if media.Quarantined {
+		log.Warn("Quarantined media accessed")
+
+		if config.Get().Quarantine.ReplaceThumbnails {
+			log.Info("Replacing thumbnail with a quarantined one")
+
+			img, err := GenerateQuarantineThumbnail(desiredWidth, desiredHeight)
+			if err != nil {
+				return nil, err
+			}
+
+			data := &bytes.Buffer{}
+			imaging.Encode(data, img, imaging.PNG)
+			return &types.StreamedThumbnail{
+				Stream: util.BufferToStream(data),
+				Thumbnail: &types.Thumbnail{
+					// We lie about the details to ensure we keep our contract
+					Width:       img.Bounds().Max.X,
+					Height:      img.Bounds().Max.Y,
+					MediaId:     media.MediaId,
+					Origin:      media.Origin,
+					Location:    "",
+					ContentType: "image/png",
+					Animated:    false,
+					Method:      method,
+					CreationTs:  util.NowMillis(),
+					SizeBytes:   int64(data.Len()),
+				},
+			}, nil
+		}
+
+		return nil, common.ErrMediaQuarantined
+	}
+
+	if animated && config.Get().Thumbnails.MaxAnimateSizeBytes > 0 && config.Get().Thumbnails.MaxAnimateSizeBytes < media.SizeBytes {
+		log.Warn("Attempted to animate a media record that is too large. Assuming animated=false")
+		animated = false
+	}
+
+	if animated && !util.ArrayContains(animatedTypes, media.ContentType) {
+		log.Warn("Attempted to animate a media record that isn't an animated type. Assuming animated=false")
+		animated = false
+	}
+
+	if media.SizeBytes > config.Get().Thumbnails.MaxSourceBytes {
+		log.Warn("Media too large to thumbnail")
+		return nil, common.ErrMediaTooLarge
+	}
+
+	db := storage.GetDatabase().GetThumbnailStore(ctx, log)
+
+	width, height, method, err := pickThumbnailDimensions(desiredWidth, desiredHeight, method)
+	if err != nil {
+		return nil, err
+	}
+
+	cacheKey := fmt.Sprintf("%s/%s?w=%d&h=%d&m=%s&a=%t", media.Origin, media.MediaId, width, height, method, animated)
+
+	var thumbnail *types.Thumbnail
+	item, found := localCache.Get(cacheKey)
+	if found {
+		thumbnail = item.(*types.Thumbnail)
+	} else {
+		log.Info("Getting thumbnail record from database")
+		dbThumb, err := db.Get(media.Origin, media.MediaId, width, height, method, animated)
+		if err != nil {
+			if err == sql.ErrNoRows {
+				log.Info("Thumbnail does not exist, attempting to generate it")
+				genThumb, err2 := GetOrGenerateThumbnail(media, width, height, animated, method, ctx, log)
+				if err2 != nil {
+					return nil, err2
+				}
+
+				thumbnail = genThumb
+			} else {
+				return nil, err
+			}
+		} else {
+			thumbnail = dbThumb
+		}
+	}
+
+	if thumbnail == nil {
+		log.Warn("Despite all efforts, a thumbnail record could not be found or generated")
+		return nil, common.ErrMediaNotFound
+	}
+
+	localCache.Set(cacheKey, thumbnail, cache.DefaultExpiration)
+	internal_cache.Get().IncrementDownloads(*thumbnail.Sha256Hash)
+
+	cached, err := internal_cache.Get().GetThumbnail(thumbnail, log)
+	if err != nil {
+		return nil, err
+	}
+	if cached != nil && cached.Contents != nil {
+		return &types.StreamedThumbnail{
+			Thumbnail: thumbnail,
+			Stream:    util.BufferToStream(cached.Contents),
+		}, nil
+	}
+
+	log.Info("Reading thumbnail from disk")
+	stream, err := os.Open(thumbnail.Location)
+	if err != nil {
+		return nil, err
+	}
+
+	return &types.StreamedThumbnail{Thumbnail: thumbnail, Stream: stream}, nil
+}
+
+func GetOrGenerateThumbnail(media *types.Media, width int, height int, animated bool, method string, ctx context.Context, log *logrus.Entry) (*types.Thumbnail, error) {
+	db := storage.GetDatabase().GetThumbnailStore(ctx, log)
+	thumbnail, err := db.Get(media.Origin, media.MediaId, width, height, method, animated)
+	if err != nil && err != sql.ErrNoRows {
+		return nil, err
+	}
+	if err != sql.ErrNoRows {
+		log.Info("Using thumbnail from database")
+		return thumbnail, nil
+	}
+
+	log.Info("Generating thumbnail")
+
+	result := <-getResourceHandler().GenerateThumbnail(media, width, height, method, animated)
+	return result.thumbnail, result.err
+}
+
+func GenerateQuarantineThumbnail(width int, height int) (image.Image, error) {
+	var centerImage image.Image
+	var err error
+	if config.Get().Quarantine.ThumbnailPath != "" {
+		centerImage, err = imaging.Open(config.Get().Quarantine.ThumbnailPath)
+	} else {
+		centerImage, err = generateDefaultQuarantineThumbnail()
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	c := gg.NewContext(width, height)
+
+	centerImage = imaging.Fit(centerImage, width, height, imaging.Lanczos)
+
+	c.DrawImageAnchored(centerImage, width/2, height/2, 0.5, 0.5)
+
+	buf := &bytes.Buffer{}
+	c.EncodePNG(buf)
+
+	return imaging.Decode(buf)
+}
+
+func generateDefaultQuarantineThumbnail() (image.Image, error) {
+	c := gg.NewContext(700, 700)
+	c.Clear()
+
+	red := color.RGBA{R: 190, G: 26, B: 25, A: 255}
+	x := 350.0
+	y := 300.0
+	r := 256.0
+	w := 55.0
+	p := 64.0
+	m := "media not allowed"
+
+	c.SetColor(red)
+	c.DrawCircle(x, y, r)
+	c.Fill()
+
+	c.SetColor(color.White)
+	c.DrawCircle(x, y, r-w)
+	c.Fill()
+
+	lr := r - (w / 2)
+	sx := x + (lr * math.Cos(gg.Radians(225.0)))
+	sy := y + (lr * math.Sin(gg.Radians(225.0)))
+	ex := x + (lr * math.Cos(gg.Radians(45.0)))
+	ey := y + (lr * math.Sin(gg.Radians(45.0)))
+	c.SetLineCap(gg.LineCapButt)
+	c.SetLineWidth(w)
+	c.SetColor(red)
+	c.DrawLine(sx, sy, ex, ey)
+	c.Stroke()
+
+	f, err := truetype.Parse(gosmallcaps.TTF)
+	if err != nil {
+		panic(err)
+	}
+
+	c.SetColor(color.Black)
+	c.SetFontFace(truetype.NewFace(f, &truetype.Options{Size: 64}))
+	c.DrawStringAnchored(m, x, y+r+p, 0.5, 0.5)
+
+	buf := &bytes.Buffer{}
+	c.EncodePNG(buf)
+
+	return imaging.Decode(buf)
+}
+
+func pickThumbnailDimensions(desiredWidth int, desiredHeight int, desiredMethod string) (int, int, string, error) {
+	if desiredWidth <= 0 {
+		return 0, 0, "", errors.New("width must be positive")
+	}
+	if desiredHeight <= 0 {
+		return 0, 0, "", errors.New("height must be positive")
+	}
+	if desiredMethod != "crop" && desiredMethod != "scale" {
+		return 0, 0, "", errors.New("method must be crop or scale")
+	}
+
+	foundSize := false
+	targetWidth := 0
+	targetHeight := 0
+	largestWidth := 0
+	largestHeight := 0
+
+	for _, size := range config.Get().Thumbnails.Sizes {
+		largestWidth = util.MaxInt(largestWidth, size.Width)
+		largestHeight = util.MaxInt(largestHeight, size.Height)
+
+		// Unlikely, but if we get the exact dimensions then just use that
+		if desiredWidth == size.Width && desiredHeight == size.Height {
+			return size.Width, size.Height, desiredMethod, nil
+		}
+
+		// If we come across a size that's smaller, try and use that
+		if desiredWidth < size.Width && desiredHeight < size.Height {
+			// Only use our new found size if it's smaller than one we've already picked
+			if !foundSize || (targetWidth > size.Width && targetHeight > size.Height) {
+				targetWidth = size.Width
+				targetHeight = size.Height
+				foundSize = true
+			}
+		}
+	}
+
+	// Use the largest dimensions available if we didn't find anything
+	if !foundSize {
+		targetWidth = largestWidth
+		targetHeight = largestHeight
+	}
+
+	return targetWidth, targetHeight, desiredMethod, nil
+}

src/github.com/turt2live/matrix-media-repo/controllers/thumbnail_controller/thumbnail_resource_handler.go

+package thumbnail_controller
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"image"
+	"image/draw"
+	"image/gif"
+	"io/ioutil"
+	"os"
+	"os/exec"
+	"path"
+	"sync"
+
+	"github.com/disintegration/imaging"
+	"github.com/sirupsen/logrus"
+	"github.com/turt2live/matrix-media-repo/common/config"
+	"github.com/turt2live/matrix-media-repo/storage"
+	"github.com/turt2live/matrix-media-repo/types"
+	"github.com/turt2live/matrix-media-repo/util"
+	"github.com/turt2live/matrix-media-repo/util/resource_handler"
+)
+
+type thumbnailResourceHandler struct {
+	resourceHandler *resource_handler.ResourceHandler
+}
+
+type thumbnailRequest struct {
+	media    *types.Media
+	width    int
+	height   int
+	method   string
+	animated bool
+}
+
+type thumbnailResponse struct {
+	thumbnail *types.Thumbnail
+	err       error
+}
+
+type GeneratedThumbnail struct {
+	ContentType  string
+	DiskLocation string
+	SizeBytes    int64
+	Animated     bool
+	Sha256Hash   *string
+}
+
+var resHandlerInstance *thumbnailResourceHandler
+var resHandlerSingletonLock = &sync.Once{}
+
+func getResourceHandler() (*thumbnailResourceHandler) {
+	if resHandlerInstance == nil {
+		resHandlerSingletonLock.Do(func() {
+			handler, err := resource_handler.New(config.Get().Thumbnails.NumWorkers, thumbnailWorkFn)
+			if err != nil {
+				panic(err)
+			}
+
+			resHandlerInstance = &thumbnailResourceHandler{handler}
+		})
+	}
+
+	return resHandlerInstance
+}
+
+func thumbnailWorkFn(request *resource_handler.WorkRequest) interface{} {
+	info := request.Metadata.(*thumbnailRequest)
+	log := logrus.WithFields(logrus.Fields{
+		"worker_requestId": request.Id,
+		"worker_media":     info.media.Origin + "/" + info.media.MediaId,
+		"worker_width":     info.width,
+		"worker_height":    info.height,
+		"worker_method":    info.method,
+		"worker_animated":  info.animated,
+	})
+	log.Info("Processing thumbnail request")
+
+	ctx := context.TODO() // TODO: Should we use a real context?
+
+	generated, err := GenerateThumbnail(info.media, info.width, info.height, info.method, info.animated, ctx, log)
+	if err != nil {
+		return &thumbnailResponse{err: err}
+	}
+
+	newThumb := &types.Thumbnail{
+		Origin:      info.media.Origin,
+		MediaId:     info.media.MediaId,
+		Width:       info.width,
+		Height:      info.height,
+		Method:      info.method,
+		Animated:    generated.Animated,
+		CreationTs:  util.NowMillis(),
+		ContentType: generated.ContentType,
+		Location:    generated.DiskLocation,
+		SizeBytes:   generated.SizeBytes,
+		Sha256Hash:  generated.Sha256Hash,
+	}
+
+	db := storage.GetDatabase().GetThumbnailStore(ctx, log)
+	err = db.Insert(newThumb)
+	if err != nil {
+		log.Error("Unexpected error caching thumbnail: " + err.Error())
+		return &thumbnailResponse{err: err}
+	}
+
+	return &thumbnailResponse{thumbnail: newThumb}
+}
+
+func (h *thumbnailResourceHandler) GenerateThumbnail(media *types.Media, width int, height int, method string, animated bool) chan *thumbnailResponse {
+	resultChan := make(chan *thumbnailResponse)
+	go func() {
+		reqId := fmt.Sprintf("thumbnail_%s_%s_%d_%d_%s_%t", media.Origin, media.MediaId, width, height, method, animated)
+		result := <-h.resourceHandler.GetResource(reqId, &thumbnailRequest{
+			media:    media,
+			width:    width,
+			height:   height,
+			method:   method,
+			animated: animated,
+		})
+		resultChan <- result.(*thumbnailResponse)
+	}()
+	return resultChan
+}
+
+func GenerateThumbnail(media *types.Media, width int, height int, method string, animated bool, ctx context.Context, log *logrus.Entry) (*GeneratedThumbnail, error) {
+	var src image.Image
+	var err error
+
+	if media.ContentType == "image/svg+xml" {
+		src, err = svgToImage(media)
+	} else {
+		src, err = imaging.Open(media.Location)
+	}
+
+	if err != nil {
+		return nil, err
+	}
+
+	srcWidth := src.Bounds().Max.X
+	srcHeight := src.Bounds().Max.Y
+
+	aspectRatio := float32(srcHeight) / float32(srcWidth)
+	targetAspectRatio := float32(width) / float32(height)
+	if aspectRatio == targetAspectRatio {
+		// Highly unlikely, but if the aspect ratios match then just resize
+		method = "scale"
+		log.Info("Aspect ratio is the same, converting method to 'scale'")
+	}
+
+	thumb := &GeneratedThumbnail{
+		Animated: animated,
+	}
+
+	if srcWidth <= width && srcHeight <= height {
+		if animated {
+			log.Warn("Image is too small but the image should be animated. Adjusting dimensions to fit image exactly.")
+			width = srcWidth
+			height = srcHeight
+		} else {
+			// Image is too small - don't upscale
+			thumb.ContentType = media.ContentType
+			thumb.DiskLocation = media.Location
+			thumb.SizeBytes = media.SizeBytes
+			thumb.Sha256Hash = &media.Sha256Hash
+			log.Warn("Image too small, returning raw image")
+			return thumb, nil
+		}
+	}
+
+	var orientation *util.ExifOrientation = nil
+	if media.ContentType == "image/jpeg" || media.ContentType == "image/jpg" {
+		orientation, err = util.GetExifOrientation(media)
+		if err != nil {
+			log.Warn("Non-fatal error getting EXIF orientation: " + err.Error())
+			orientation = nil // just in case
+		}
+	}
+
+	contentType := "image/png"
+	imgData := &bytes.Buffer{}
+	if config.Get().Thumbnails.AllowAnimated && animated {
+		log.Info("Generating animated thumbnail")
+		contentType = "image/gif"
+
+		// Animated GIFs are a bit more special because we need to do it frame by frame.
+		// This is fairly resource intensive. The calling code is responsible for limiting this case.
+
+		inputFile, err := os.Open(media.Location)
+		if err != nil {
+			log.Error("Error generating animated thumbnail: " + err.Error())
+			return nil, err
+		}
+		defer inputFile.Close()
+
+		g, err := gif.DecodeAll(inputFile)
+		if err != nil {
+			log.Error("Error generating animated thumbnail: " + err.Error())
+			return nil, err
+		}
+
+		// Prepare a blank frame to use as swap space
+		frameImg := image.NewRGBA(g.Image[0].Bounds())
+
+		for i := range g.Image {
+			img := g.Image[i]
+
+			// Clear the transparency of the previous frame
+			draw.Draw(frameImg, frameImg.Bounds(), image.Transparent, image.ZP, draw.Src)
+
+			// Copy the frame to a new image and use that
+			draw.Draw(frameImg, frameImg.Bounds(), img, image.ZP, draw.Over)
+
+			// Do the thumbnailing on the copied frame
+			frameThumb, err := thumbnailFrame(frameImg, method, width, height, imaging.Linear, nil)
+			if err != nil {
+				log.Error("Error generating animated thumbnail frame: " + err.Error())
+				return nil, err
+			}
+
+			//t.log.Info(fmt.Sprintf("Width = %d    Height = %d    FW=%d    FH=%d", width, height, frameThumb.Bounds().Max.X, frameThumb.Bounds().Max.Y))
+
+			targetImg := image.NewPaletted(frameThumb.Bounds(), img.Palette)
+			draw.FloydSteinberg.Draw(targetImg, frameThumb.Bounds(), frameThumb, image.ZP)
+			g.Image[i] = targetImg
+		}
+
+		// Set the image size to the first frame's size
+		g.Config.Width = g.Image[0].Bounds().Max.X
+		g.Config.Height = g.Image[0].Bounds().Max.Y
+
+		err = gif.EncodeAll(imgData, g)
+		if err != nil {
+			log.Error("Error generating animated thumbnail: " + err.Error())
+			return nil, err
+		}
+	} else {
+		src, err = thumbnailFrame(src, method, width, height, imaging.Lanczos, orientation)
+		if err != nil {
+			log.Error("Error generating thumbnail: " + err.Error())
+			return nil, err
+		}
+
+		// Put the image bytes into a memory buffer
+		err = imaging.Encode(imgData, src, imaging.PNG)
+		if err != nil {
+			log.Error("Unexpected error encoding thumbnail: " + err.Error())
+			return nil, err
+		}
+	}
+
+	// Reset the buffer pointer and store the file
+	location, err := storage.PersistFile(imgData, ctx, log)
+	if err != nil {
+		log.Error("Unexpected error saving thumbnail: " + err.Error())
+		return nil, err
+	}
+
+	fileSize, err := util.FileSize(location)
+	if err != nil {
+		log.Error("Unexpected error getting the size of the thumbnail: " + err.Error())
+		return nil, err
+	}
+
+	hash, err := storage.GetFileHash(location)
+	if err != nil {
+		log.Error("Unexpected error getting the hash for the thumbnail: ", err.Error())
+		return nil, err
+	}
+
+	thumb.DiskLocation = location
+	thumb.ContentType = contentType
+	thumb.SizeBytes = fileSize
+	thumb.Sha256Hash = &hash
+
+	return thumb, nil
+}
+
+func thumbnailFrame(src image.Image, method string, width int, height int, filter imaging.ResampleFilter, orientation *util.ExifOrientation) (image.Image, error) {
+	var result image.Image
+	if method == "scale" {
+		result = imaging.Fit(src, width, height, filter)
+	} else if method == "crop" {
+		result = imaging.Fill(src, width, height, imaging.Center, filter)
+	} else {
+		return nil, errors.New("unrecognized method: " + method)
+	}
+
+	if orientation != nil {
+		// Rotate first
+		if orientation.RotateDegrees == 90 {
+			result = imaging.Rotate90(result)
+		} else if orientation.RotateDegrees == 180 {
+			result = imaging.Rotate180(result)
+		} else if orientation.RotateDegrees == 270 {
+			result = imaging.Rotate270(result)
+		} // else we don't care to rotate
+
+		// Flip second
+		if orientation.FlipHorizontal {
+			result = imaging.FlipH(result)
+		}
+		if orientation.FlipVertical {
+			result = imaging.FlipV(result)
+		}
+	}
+
+	return result, nil
+}
+
+func svgToImage(media *types.Media) (image.Image, error) {
+	tempFile := path.Join(os.TempDir(), "media_repo."+media.Origin+"."+media.MediaId+".png")
+	defer os.Remove(tempFile)
+
+	// requires imagemagick
+	err := exec.Command("convert", media.Location, tempFile).Run()
+	if err != nil {
+		return nil, err
+	}
+
+	b, err := ioutil.ReadFile(tempFile)
+	if err != nil {
+		return nil, err
+	}
+
+	imgData := bytes.NewBuffer(b)
+	return imaging.Decode(imgData)
+}

src/github.com/turt2live/matrix-media-repo/internal_cache/media_cache.go

@@ -103,6 +103,23 @@ func (c *MediaCache) GetMedia(media *types.Media, log *logrus.Entry) (*cachedFil
 	return c.updateItemInCache(media.Sha256Hash, media.SizeBytes, cacheFn, log)
 }
 
+func (c *MediaCache) GetThumbnail(thumbnail *types.Thumbnail, log *logrus.Entry) (*cachedFile, error) {
+	if !c.enabled {
+		return nil, nil
+	}
+
+	cacheFn := func() (*cachedFile, error) {
+		data, err := ioutil.ReadFile(thumbnail.Location)
+		if err != nil {
+			return nil, err
+		}
+
+		return &cachedFile{thumbnail: thumbnail, Contents: bytes.NewBuffer(data)}, nil
+	}
+
+	return c.updateItemInCache(*thumbnail.Sha256Hash, thumbnail.SizeBytes, cacheFn, log)
+}
+
 func (c *MediaCache) updateItemInCache(recordId string, mediaSize int64, cacheFn func() (*cachedFile, error), log *logrus.Entry) (*cachedFile, error) {
 	downloads := c.tracker.NumDownloads(recordId)
 	enoughDownloads := downloads >= config.Get().Downloads.Cache.MinDownloads