Warp/server/internal/mediaapi/thumbnail.go

package mediaapi

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"image"
	_ "image/gif" // register the GIF decoder with image.Decode
	"image/jpeg"
	"image/png"
	"io"
	"math"
	"mime"
	"time"

	"golang.org/x/image/draw"
	_ "golang.org/x/image/webp" // register the WebP decoder with image.Decode

	"github.com/dariuszjaniak/warp/server/internal/config"
	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// The two thumbnail methods.
const (
	MethodCrop  = "crop"
	MethodScale = "scale"
)

// ThumbnailRequest is the thumbnail a client asked for.
type ThumbnailRequest struct {
	Width  int
	Height int

	// Method is MethodCrop or MethodScale.
	Method string

	// Animated asks for an animated thumbnail of animated media. Warp
	// makes still thumbnails only, which the spec allows: a server that
	// cannot animate one returns a still image.
	Animated bool
}

type thumbnailSize = config.ThumbnailSize

// Thumbnail returns a thumbnail of media.
func (a *API) Thumbnail(ctx context.Context, origin, mediaID string, req ThumbnailRequest, opts FetchOptions) (Content, error) {
	if !validServerName(origin) || !ValidMediaID(mediaID) {
		return Content{}, ErrNotFound
	}
	size, ok := a.pickSize(req)
	if !ok {
		return Content{}, ErrCannotThumbnail
	}

	if origin == a.serverName {
		m, err := a.localMedia(ctx, mediaID, opts.Wait)
		if err != nil {
			return Content{}, err
		}
		return a.thumbnailOf(ctx, m, size)
	}
	if opts.LocalOnly {
		return Content{}, ErrNotFound
	}

	if t, err := a.store.Thumbnail(ctx, origin, mediaID, size.Width, size.Height, size.Method); err == nil {
		return a.open(ctx, t.ContentType, "", t.Size, t.SHA256)
	} else if !errors.Is(err, storage.ErrNoSuchMedia) {
		return Content{}, err
	}
	// The original is already here: make the thumbnail rather than
	// fetch one.
	if m, err := a.store.Media(ctx, origin, mediaID); err == nil && !m.Pending {
		return a.thumbnailOf(ctx, m, size)
	} else if err != nil && !errors.Is(err, storage.ErrNoSuchMedia) {
		return Content{}, err
	}
	remote := a.remoteClient()
	if remote == nil {
		return Content{}, ErrNotFound
	}

	wait := min(max(opts.Wait, 0), MaxWait)
	key := fmt.Sprintf("thumbnail\x00%s\x00%dx%d\x00%s", mediaName(origin, mediaID), size.Width, size.Height, size.Method)
	v, err := a.shared(ctx, key, func(ctx context.Context) (any, error) {
		return a.fetchRemoteThumbnail(ctx, remote, origin, mediaID, size, wait)
	})
	if err != nil {
		return Content{}, err
	}
	t := v.(storage.Thumbnail)
	return a.open(ctx, t.ContentType, "", t.Size, t.SHA256)
}

// pickSize chooses the configured size that answers a request: the
// smallest of the requested method that is at least as large in both
// dimensions, or failing that the largest. A method with no configured
// sizes is answered from the other.
func (a *API) pickSize(req ThumbnailRequest) (thumbnailSize, bool) {
	var candidates []thumbnailSize
	for _, s := range a.cfg.ThumbnailSizes {
		if s.Method == req.Method {
			candidates = append(candidates, s)
		}
	}
	if len(candidates) == 0 {
		candidates = a.cfg.ThumbnailSizes
	}
	if len(candidates) == 0 {
		return thumbnailSize{}, false
	}

	area := func(s thumbnailSize) int { return s.Width * s.Height }
	var best, largest thumbnailSize
	found := false
	for _, s := range candidates {
		if area(s) > area(largest) {
			largest = s
		}
		if s.Width >= req.Width && s.Height >= req.Height && (!found || area(s) < area(best)) {
			best, found = s, true
		}
	}
	if !found {
		return largest, true
	}
	return best, true
}

// thumbnailable lists the content types that are decoded for thumbnails.
var thumbnailable = map[string]bool{
	"image/jpeg": true,
	"image/png":  true,
	"image/gif":  true,
	"image/webp": true,
}

// thumbnailOf returns the thumbnail of media held here, making it on first
// use.
func (a *API) thumbnailOf(ctx context.Context, m storage.Media, size thumbnailSize) (Content, error) {
	t, err := a.store.Thumbnail(ctx, m.Origin, m.MediaID, size.Width, size.Height, size.Method)
	if err == nil {
		return a.open(ctx, t.ContentType, "", t.Size, t.SHA256)
	}
	if !errors.Is(err, storage.ErrNoSuchMedia) {
		return Content{}, err
	}
	mediaType, _, _ := mime.ParseMediaType(m.ContentType)
	if !thumbnailable[mediaType] {
		return Content{}, ErrCannotThumbnail
	}

	key := fmt.Sprintf("resize\x00%s\x00%dx%d\x00%s", mediaName(m.Origin, m.MediaID), size.Width, size.Height, size.Method)
	v, err := a.shared(ctx, key, func(ctx context.Context) (any, error) {
		return a.makeThumbnail(ctx, m, size)
	})
	if err != nil {
		return Content{}, err
	}
	t = v.(storage.Thumbnail)
	return a.open(ctx, t.ContentType, "", t.Size, t.SHA256)
}

// makeThumbnail decodes media, resizes it and stores the result.
func (a *API) makeThumbnail(ctx context.Context, m storage.Media, size thumbnailSize) (storage.Thumbnail, error) {
	select {
	case a.resizing <- struct{}{}:
		defer func() { <-a.resizing }()
	case <-ctx.Done():
		return storage.Thumbnail{}, ctx.Err()
	}

	src, err := a.blobs.Open(ctx, m.SHA256)
	if err != nil {
		return storage.Thumbnail{}, err
	}
	defer func() { _ = src.Close() }()

	started := time.Now()
	thumb, err := a.resize(src, m, size)
	if err != nil {
		return storage.Thumbnail{}, err
	}

	t := storage.Thumbnail{
		Origin: m.Origin, MediaID: m.MediaID,
		Width: size.Width, Height: size.Height, Method: size.Method,
		CreatedAtMS: a.nowMS(),
	}
	if thumb == nil {
		// The original is already small enough, and is its own thumbnail.
		// The row points at the original's blob, so it is neither decoded
		// again nor stored twice.
		t.ContentType, t.Size, t.SHA256 = m.ContentType, m.Size, m.SHA256
	} else {
		blob, err := a.blobs.Put(ctx, bytes.NewReader(thumb.data), int64(len(thumb.data)))
		if err != nil {
			return storage.Thumbnail{}, err
		}
		t.ContentType, t.Size, t.SHA256 = thumb.contentType, blob.Size, blob.SHA256
	}
	if err := a.store.StoreThumbnail(ctx, t); err != nil {
		return storage.Thumbnail{}, err
	}
	a.log.Debug("made a thumbnail", "media", mediaName(m.Origin, m.MediaID),
		"size", fmt.Sprintf("%dx%d", size.Width, size.Height), "method", size.Method,
		"bytes", t.Size, "duration", time.Since(started))
	return t, nil
}

// encoded is a finished thumbnail.
type encoded struct {
	contentType string
	data        []byte
}

// resize makes the thumbnail of an image. It returns nil, and no error,
// when the original already fits and can be served as it is.
func (a *API) resize(src io.ReadSeeker, m storage.Media, size thumbnailSize) (*encoded, error) {
	// The header first: the pixel count is checked before a single pixel
	// is decoded.
	header, format, err := image.DecodeConfig(src)
	if err != nil {
		return nil, ErrCannotThumbnail
	}
	if int64(header.Width)*int64(header.Height) > a.cfg.MaxThumbnailPixels {
		return nil, fmt.Errorf("%w: the image is %d × %d", ErrCannotThumbnail, header.Width, header.Height)
	}

	orientation := 1
	if format == "jpeg" {
		if _, err := src.Seek(0, io.SeekStart); err != nil {
			return nil, err
		}
		orientation = jpegOrientation(io.LimitReader(src, 256<<10))
	}

	// Everything below works on the image as stored. An EXIF rotation by
	// a quarter turn swaps the box, the crop is centred either way, and
	// the small result is turned upright at the end.
	boxW, boxH := size.Width, size.Height
	if orientation >= 5 {
		boxW, boxH = boxH, boxW
	}
	srcW, srcH := header.Width, header.Height
	outW, outH, crop := fit(srcW, srcH, boxW, boxH, size.Method)

	// A GIF is always re-encoded: it may be animated, and a still
	// thumbnail was asked for.
	if outW == srcW && outH == srcH && crop == image.Rect(0, 0, srcW, srcH) && orientation == 1 && format != "gif" {
		return nil, nil
	}

	if _, err := src.Seek(0, io.SeekStart); err != nil {
		return nil, err
	}
	img, _, err := image.Decode(src)
	if err != nil {
		return nil, ErrCannotThumbnail
	}
	// The decoded bounds need not start at the origin.
	crop = crop.Add(img.Bounds().Min)

	dst := image.NewRGBA(image.Rect(0, 0, outW, outH))
	draw.CatmullRom.Scale(dst, dst.Bounds(), img, crop, draw.Src, nil)
	out := orient(dst, orientation)

	var buf bytes.Buffer
	if format == "jpeg" {
		if err := jpeg.Encode(&buf, out, &jpeg.Options{Quality: 85}); err != nil {
			return nil, err
		}
		return &encoded{contentType: "image/jpeg", data: buf.Bytes()}, nil
	}
	// PNG for everything else: the source may have transparency, and a
	// thumbnail is small enough that PNG's size does not matter.
	if err := png.Encode(&buf, out); err != nil {
		return nil, err
	}
	return &encoded{contentType: "image/png", data: buf.Bytes()}, nil
}

// fit works out a thumbnail's size and the part of the source it is made
// from. Neither method enlarges: a source smaller than the box gives a
// thumbnail no larger than itself.
//
// "scale" keeps all of the source and its aspect ratio, fitting inside the
// box. "crop" keeps the box's aspect ratio, cutting the source's longer
// side down equally on both ends.
func fit(srcW, srcH, boxW, boxH int, method string) (outW, outH int, crop image.Rectangle) {
	full := image.Rect(0, 0, srcW, srcH)
	if method == MethodCrop {
		ratio := math.Min(1, math.Min(float64(srcW)/float64(boxW), float64(srcH)/float64(boxH)))
		outW = max(1, int(math.Round(float64(boxW)*ratio)))
		outH = max(1, int(math.Round(float64(boxH)*ratio)))

		// The largest centred rectangle of the box's aspect ratio.
		cropW, cropH := srcW, srcH
		if srcW*boxH > srcH*boxW {
			cropW = max(1, int(math.Round(float64(srcH)*float64(boxW)/float64(boxH))))
		} else {
			cropH = max(1, int(math.Round(float64(srcW)*float64(boxH)/float64(boxW))))
		}
		x0, y0 := (srcW-cropW)/2, (srcH-cropH)/2
		return outW, outH, image.Rect(x0, y0, x0+cropW, y0+cropH)
	}

	ratio := math.Min(1, math.Min(float64(boxW)/float64(srcW), float64(boxH)/float64(srcH)))
	outW = max(1, int(math.Round(float64(srcW)*ratio)))
	outH = max(1, int(math.Round(float64(srcH)*ratio)))
	return outW, outH, full
}

// orient turns an image upright according to its EXIF orientation:
//
//	1 as stored          5 transposed
//	2 mirrored           6 a quarter turn clockwise
//	3 half a turn        7 transversed
//	4 flipped            8 a quarter turn anticlockwise
//
// It runs on the finished thumbnail, which is small.
func orient(img *image.RGBA, orientation int) image.Image {
	if orientation < 2 || orientation > 8 {
		return img
	}
	b := img.Bounds()
	w, h := b.Dx(), b.Dy()
	outW, outH := w, h
	if orientation >= 5 {
		outW, outH = h, w
	}
	out := image.NewRGBA(image.Rect(0, 0, outW, outH))
	for y := range h {
		for x := range w {
			var dx, dy int
			switch orientation {
			case 2:
				dx, dy = w-1-x, y
			case 3:
				dx, dy = w-1-x, h-1-y
			case 4:
				dx, dy = x, h-1-y
			case 5:
				dx, dy = y, x
			case 6:
				dx, dy = h-1-y, x
			case 7:
				dx, dy = h-1-y, w-1-x
			case 8:
				dx, dy = y, w-1-x
			}
			out.SetRGBA(dx, dy, img.RGBAAt(b.Min.X+x, b.Min.Y+y))
		}
	}
	return out
}

// jpegOrientation reads the EXIF orientation tag of a JPEG, or 1 when it
// has none. It walks the markers up to the start of the image data and
// gives up quietly on anything it does not understand: a wrong guess here
// is a sideways thumbnail, never a failure.
func jpegOrientation(r io.Reader) int {
	data, err := io.ReadAll(r)
	if err != nil || len(data) < 4 || data[0] != 0xFF || data[1] != 0xD8 {
		return 1
	}
	pos := 2
	for pos+4 <= len(data) {
		if data[pos] != 0xFF {
			return 1
		}
		marker := data[pos+1]
		if marker == 0xD8 || (marker >= 0xD0 && marker <= 0xD7) || marker == 0x01 || marker == 0xFF {
			pos++
			continue
		}
		if marker == 0xDA || marker == 0xD9 { // start of scan, end of image
			return 1
		}
		length := int(data[pos+2])<<8 | int(data[pos+3])
		if length < 2 || pos+2+length > len(data) {
			return 1
		}
		segment := data[pos+4 : pos+2+length]
		if marker == 0xE1 && len(segment) > 6 && string(segment[:6]) == "Exif\x00\x00" {
			return exifOrientation(segment[6:])
		}
		pos += 2 + length
	}
	return 1
}

// exifOrientation finds tag 0x0112 in the first IFD of a TIFF structure.
func exifOrientation(tiff []byte) int {
	if len(tiff) < 8 {
		return 1
	}
	var u16 func([]byte) int
	var u32 func([]byte) int
	switch string(tiff[:2]) {
	case "II":
		u16 = func(b []byte) int { return int(b[0]) | int(b[1])<<8 }
		u32 = func(b []byte) int { return int(b[0]) | int(b[1])<<8 | int(b[2])<<16 | int(b[3])<<24 }
	case "MM":
		u16 = func(b []byte) int { return int(b[0])<<8 | int(b[1]) }
		u32 = func(b []byte) int { return int(b[0])<<24 | int(b[1])<<16 | int(b[2])<<8 | int(b[3]) }
	default:
		return 1
	}
	if u16(tiff[2:4]) != 42 {
		return 1
	}
	ifd := u32(tiff[4:8])
	if ifd < 8 || ifd+2 > len(tiff) {
		return 1
	}
	entries := u16(tiff[ifd : ifd+2])
	for i := range entries {
		entry := ifd + 2 + i*12
		if entry+12 > len(tiff) {
			return 1
		}
		if u16(tiff[entry:entry+2]) == 0x0112 {
			if o := u16(tiff[entry+8 : entry+10]); o >= 1 && o <= 8 {
				return o
			}
			return 1
		}
	}
	return 1
}