leech/thumbnail/thumbnail.go

package thumbnail

import (
	"bytes"
	"image"
	"image/color"
	"image/png"
	"leech/config"
	"os"
	"path/filepath"
	"slices"
	"sync"

	"github.com/disintegration/imaging"
	"github.com/gofiber/fiber/v2"
)

var SupportedFileTypes = []string{".png", ".PNG", ".jpg", ".JPG", ".jpeg", ".JPEG"}
var FileTypesMap = map[string]imaging.Format{
	".png":  imaging.PNG,
	".PNG":  imaging.PNG,
	".jpg":  imaging.JPEG,
	".JPG":  imaging.JPEG,
	".jpeg": imaging.JPEG,
	".JPEG": imaging.JPEG,
}
var thumbnailSize = 48

var thumbnailCache = map[string][]byte{}
var thumbnailCacheMutex = &sync.RWMutex{}
var memLimiterMutex = &sync.RWMutex{}
var jobCounter = 0
var checkAgain = make(chan bool, 5)

func IsSupportedFileType(completePath string) bool {
	fileExt := filepath.Ext(completePath)
	return slices.Contains(SupportedFileTypes, fileExt)
}

func WaitForAvailable() {
	memLimiterMutex.RLock()
	for config.Config.ThumbnailJobLimit == jobCounter {
		memLimiterMutex.RUnlock()
		<-checkAgain
		memLimiterMutex.RLock()
	}
	memLimiterMutex.RUnlock()
	memLimiterMutex.Lock()
	jobCounter += 1
	memLimiterMutex.Unlock()
}

func Free() {
	memLimiterMutex.Lock()
	jobCounter -= 1
	memLimiterMutex.Unlock()
	select {
	case checkAgain <- true:
		return
	default:
		return
	}
}

func GetThumbnail(c *fiber.Ctx, completePath string) {
	c.Set(fiber.HeaderContentType, "image")
	thumbnailCacheMutex.RLock()
	bytesThumb, ok := thumbnailCache[completePath]
	thumbnailCacheMutex.RUnlock()
	if ok {
		c.Write(bytesThumb)
		return
	}
	WaitForAvailable()

	fileExt := filepath.Ext(completePath)
	if !slices.Contains(SupportedFileTypes, fileExt) {
		Free()
		return
	}
	f, err := os.Open(completePath)
	if err != nil {
		Free()
		return
	}
	img, _, err := image.Decode(f)
	if err != nil {
		Free()
		return
	}
	f.Close()
	// load images and make 64x64 thumbnails of them
	thumbnail := imaging.Thumbnail(img, thumbnailSize, thumbnailSize, imaging.CatmullRom)
	// create a new blank image
	dst := imaging.New(thumbnailSize, thumbnailSize, color.NRGBA{0, 0, 0, 0})
	// paste thumbnails into the new image side by side
	dst = imaging.Paste(dst, thumbnail, image.Pt(0, 0))
	// save the combined image to buffer
	var buf bytes.Buffer
	if FileTypesMap[fileExt] == imaging.PNG {
		err = png.Encode(&buf, dst)
	} else {
		err = imaging.Encode(&buf, dst, FileTypesMap[fileExt])
	}
	if err != nil {
		Free()
		return
	}
	thumbnailCacheMutex.Lock()
	thumbnailCache[completePath] = buf.Bytes()
	thumbnailCacheMutex.Unlock()
	c.Write(buf.Bytes())
	Free()
}