avatars/main.go

package main

import (
	"bytes"
	"embed"
	"hash/fnv"
	"image"
	"image/color"
	"image/jpeg"
	_ "image/png"
	"io/fs"
	"log"
	"net/http"
	"sync"

	"strconv"
	"strings"
	"time"

	"golang.org/x/image/draw"
)

type cacheItem struct {
	value      []byte
	lastAccess int64
}
type AvatarCache struct {
	cache map[string]*cacheItem
	mu    *sync.RWMutex
}

func (ac *AvatarCache) Clear() {
	ac.mu.Lock()
	for k, v := range ac.cache {
		if time.Now().Unix()-v.lastAccess > int64(100) {
			delete(ac.cache, k)
		}
	}
	ac.mu.Unlock()
}

// insert to cache
func (ac *AvatarCache) Insert(id string, size int, value []byte) {
	ac.mu.Lock()
	defer ac.mu.Unlock()
	ac.cache[id+strconv.Itoa(size)] = &cacheItem{value: value, lastAccess: time.Now().Unix()}
	// remove oldest item if cache is full
	if len(ac.cache) > 1000 {
		for k := range ac.cache {
			delete(ac.cache, k)
			break
		}
	}
}

// get from cache
func (ac *AvatarCache) Get(id string, size int) []byte {
	ac.mu.Lock()
	defer ac.mu.Unlock()
	item, ok := ac.cache[id+strconv.Itoa(size)]
	if !ok {
		return nil
	}
	item.lastAccess = time.Now().Unix()
	return item.value
}

// create new cache
var avatarCache = AvatarCache{cache: make(map[string]*cacheItem), mu: &sync.RWMutex{}}

//go:embed assets/*
var files embed.FS

var BaseRects [11]draw.Image
var EYES = make([]image.Image, 0)
var NOSE = make([]image.Image, 0)
var MOUTH = make([]image.Image, 0)
var emptyRect = image.Rect(0, 0, 400, 400)

func getPartsById(id string) (draw.Image, image.Image, image.Image, image.Image) {
	idBytes := []byte(id)
	colorHash := fnv.New32a()
	colorHash.Write(idBytes)
	eyesHash := fnv.New32a()
	eyesHash.Write(idBytes)
	noseHash := fnv.New32a()
	noseHash.Write(idBytes)
	mouthHash := fnv.New32a()
	mouthHash.Write(idBytes)
	return BaseRects[colorHash.Sum32()%11],
		EYES[eyesHash.Sum32()%10],
		NOSE[noseHash.Sum32()%9],
		MOUTH[mouthHash.Sum32()%10]
}

func serveAvatar(w http.ResponseWriter, r *http.Request) {
	startTime := time.Now()
	// trollin' attackers
	w.Header().Set("x-powered-by", "Undertow/1")
	w.Header().Set("server", "JBoss-EAP/7")

	// get the id and size from the request
	id := r.PathValue("id")
	sizeStr := r.PathValue("size")
	if sizeStr == "" {
		sizeStr = r.URL.Query().Get("s")
	}
	size, err := strconv.Atoi(sizeStr)
	if err != nil {
		size = 400
	}
	// min size is 32, max size is 2048
	if size < 32 {
		size = 32
	} else if size > 2048 {
		size = 2048
	}
	if avatarCache.Get(id, size) != nil {
		// get from cache
		w.Header().Set("content-type", "image/jpeg")
		w.WriteHeader(http.StatusFound)
		byteSize, err := w.Write(avatarCache.Get(id, size))
		if err != nil {
			log.Println(err.Error())
		}
		log.Printf("[id:%s][size:%d] [%s] [%s] [%d bytes] [%v] [cache hit]\n",
			id, size, r.RemoteAddr, r.Proto, byteSize, time.Since(startTime))
		return
	}

	img := image.NewRGBA(emptyRect)

	// draw parts over base(color) image
	base, eyes, nose, mouth := getPartsById(id)
	draw.Draw(img, emptyRect, base, emptyRect.Min, draw.Src)
	draw.Draw(img, emptyRect, eyes, emptyRect.Min, draw.Over)
	draw.Draw(img, emptyRect, nose, emptyRect.Min, draw.Over)
	draw.Draw(img, emptyRect, mouth, emptyRect.Min, draw.Over)
	if size != 400 {
		// resize the image
		dst := image.NewRGBA(image.Rect(0, 0, size, size))
		draw.CatmullRom.Scale(dst, dst.Bounds(), img, emptyRect, draw.Over, nil)
		img = dst
	}

	// write the image to the response
	w.Header().Set("content-type", "image/jpeg")
	w.WriteHeader(http.StatusCreated)
	buf := new(bytes.Buffer)
	err = jpeg.Encode(buf, img, nil)
	if err != nil {
		log.Println(err.Error())
	}
	byteSize, err := w.Write(buf.Bytes())
	if err != nil {
		log.Println(err.Error())
	}
	avatarCache.Insert(id, size, buf.Bytes())
	log.Printf("[id:%s][size:%d] [%s] [%s] [%d bytes] [%v]\n",
		id, size, r.RemoteAddr, r.Proto, byteSize, time.Since(startTime))
}

func main() {
	go func() {
		for range time.Tick(time.Minute) {
			avatarCache.Clear()
		}
	}()
	mux := http.NewServeMux()
	mux.HandleFunc("GET /{id}", serveAvatar)
	mux.HandleFunc("GET /{size}/{id}", serveAvatar)
	http.ListenAndServe("0.0.0.0:3382", mux)
	log.Println("Server started at 0.0.0.0:3382")
}

func init() {
	_ = fs.WalkDir(files, "assets", func(path string, d fs.DirEntry, _ error) error {
		if !d.IsDir() {
			file, _ := files.ReadFile(path)
			img, _, _ := image.Decode(bytes.NewReader(file))
			if strings.Contains(path, "/eyes/eyes") {
				EYES = append(EYES, img)
			} else if strings.Contains(path, "/nose/nose") {
				NOSE = append(NOSE, img)
			} else if strings.Contains(path, "/mouth/mouth") {
				MOUTH = append(MOUTH, img)
			}
		}
		return nil
	})
	colors := [11]color.RGBA{
		{129, 190, 241, 255},
		{173, 139, 242, 255},
		{191, 242, 136, 255},
		{222, 120, 120, 255},
		{165, 170, 197, 255},
		{111, 242, 197, 255},
		{240, 218, 94, 255},
		{235, 89, 114, 255},
		{246, 190, 93, 255},
		{0, 185, 189, 255},
		{251, 105, 0, 255},
	}
	//init Base Rectangles
	for i := 0; i < 11; i++ {
		BaseRects[i] = image.NewRGBA(image.Rect(0, 0, 400, 400))
		for t := 0; t < 200; t++ {
			// draw horizontal lines
			for x := 0; x <= 399; x++ {
				BaseRects[i].Set(x, 0+t, colors[i])
				BaseRects[i].Set(x, 399-t, colors[i])
			}
			// draw vertical lines
			for y := 0; y <= 399; y++ {
				BaseRects[i].Set(0+t, y, colors[i])
				BaseRects[i].Set(399-t, y, colors[i])
			}
		}
	}
	log.Printf("%d different combinations loaded\n", len(BaseRects)*len(EYES)*len(NOSE)*len(MOUTH))
}