terdut-server/internal/api/avatars.go

package api

import (
	"bytes"
	"database/sql"
	"errors"
	"image"
	"image/color"
	"image/draw"
	_ "image/jpeg" // registers the JPEG decoder with image.Decode
	"image/png"
	"io"
	"net/http"
	"strconv"
	"time"

	"github.com/go-chi/chi/v5"
)

const (
	// avatarSize is the side of the stored square picture. The web UI crops to
	// it before uploading; the server enforces it anyway.
	avatarSize = 256

	// maxAvatarBytes caps the upload. The web UI sends a ~256x256 PNG, a few
	// tens of KiB; this leaves room for a client that sends a photo as is.
	maxAvatarBytes = 1 << 20

	// maxAvatarPixels bounds the decoded size, checked from the header before
	// decoding: a small file can declare an enormous canvas.
	maxAvatarPixels = 16 << 20
)

// normalizeAvatar decodes a PNG or JPEG, crops it to its centred square, scales
// it to avatarSize and encodes it as PNG. Decoding and re-encoding is what
// makes the stored bytes trustworthy: EXIF, trailing data and anything that is
// not pixels do not survive it.
func normalizeAvatar(raw []byte) ([]byte, error) {
	cfg, _, err := image.DecodeConfig(bytes.NewReader(raw))
	if err != nil {
		return nil, errors.New("the picture must be a PNG or JPEG image")
	}
	if cfg.Width <= 0 || cfg.Height <= 0 || cfg.Width*cfg.Height > maxAvatarPixels {
		return nil, errors.New("the picture is too large")
	}
	img, _, err := image.Decode(bytes.NewReader(raw))
	if err != nil {
		return nil, errors.New("the picture must be a PNG or JPEG image")
	}

	b := img.Bounds()
	side := min(b.Dx(), b.Dy())
	src := image.Rect(
		b.Min.X+(b.Dx()-side)/2, b.Min.Y+(b.Dy()-side)/2,
		b.Min.X+(b.Dx()-side)/2+side, b.Min.Y+(b.Dy()-side)/2+side)

	var out *image.NRGBA
	if side <= avatarSize {
		out = image.NewNRGBA(image.Rect(0, 0, side, side))
		draw.Draw(out, out.Bounds(), img, src.Min, draw.Src)
	} else {
		out = boxScale(img, src, avatarSize)
	}

	var buf bytes.Buffer
	if err := png.Encode(&buf, out); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// boxScale shrinks the square region src of img to size x size by averaging the
// source pixels each output pixel covers. Plain area averaging is enough for a
// downscale and keeps the standard library the only dependency.
func boxScale(img image.Image, src image.Rectangle, size int) *image.NRGBA {
	out := image.NewNRGBA(image.Rect(0, 0, size, size))
	n := src.Dx()
	for y := 0; y < size; y++ {
		y0, y1 := src.Min.Y+y*n/size, src.Min.Y+(y+1)*n/size
		for x := 0; x < size; x++ {
			x0, x1 := src.Min.X+x*n/size, src.Min.X+(x+1)*n/size
			var r, g, b, a, cnt uint64
			for sy := y0; sy < y1; sy++ {
				for sx := x0; sx < x1; sx++ {
					pr, pg, pb, pa := img.At(sx, sy).RGBA()
					r, g, b, a, cnt = r+uint64(pr), g+uint64(pg), b+uint64(pb), a+uint64(pa), cnt+1
				}
			}
			if cnt == 0 {
				continue
			}
			// RGBA() is alpha-premultiplied 16-bit; undo that for NRGBA.
			a /= cnt
			if a == 0 {
				continue
			}
			r, g, b = r/cnt*0xffff/a, g/cnt*0xffff/a, b/cnt*0xffff/a
			out.SetNRGBA(x, y, color.NRGBA{uint8(r >> 8), uint8(g >> 8), uint8(b >> 8), uint8(a >> 8)})
		}
	}
	return out
}

func handleGetAvatar(db *sql.DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		id, err := strconv.ParseInt(chi.URLParam(r, "id"), 10, 64)
		if err != nil {
			respond(w, http.StatusBadRequest, errResp("invalid user id"))
			return
		}
		var data []byte
		var version int64
		err = db.QueryRowContext(r.Context(),
			"SELECT data, updated_at FROM user_avatars WHERE user_id = $1", id).Scan(&data, &version)
		if errors.Is(err, sql.ErrNoRows) {
			respond(w, http.StatusNotFound, errResp("no avatar"))
			return
		}
		if err != nil {
			serverError(w, r, err)
			return
		}
		etag := `"` + strconv.FormatInt(version, 10) + `"`
		w.Header().Set("ETag", etag)
		// private: it is behind a login. The URL carries ?v=, so the browser can
		// keep it for a long time without ever showing a stale picture.
		w.Header().Set("Cache-Control", "private, max-age=86400")
		if r.Header.Get("If-None-Match") == etag {
			w.WriteHeader(http.StatusNotModified)
			return
		}
		w.Header().Set("Content-Type", "image/png")
		w.Header().Set("X-Content-Type-Options", "nosniff")
		_, _ = w.Write(data)
	}
}

// handleSetAvatar stores the request body, a PNG or JPEG, as the user's picture.
func handleSetAvatar(db *sql.DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		id, err := strconv.ParseInt(chi.URLParam(r, "id"), 10, 64)
		if err != nil {
			respond(w, http.StatusBadRequest, errResp("invalid user id"))
			return
		}
		if !requireSelfOrAdmin(w, r, id) {
			return
		}
		raw, err := io.ReadAll(http.MaxBytesReader(w, r.Body, maxAvatarBytes))
		if err != nil {
			var tooBig *http.MaxBytesError
			if errors.As(err, &tooBig) {
				respond(w, http.StatusRequestEntityTooLarge, errResp("the picture is too large"))
				return
			}
			respond(w, http.StatusBadRequest, errResp("could not read the picture"))
			return
		}
		data, err := normalizeAvatar(raw)
		if err != nil {
			respond(w, http.StatusBadRequest, errResp(err.Error()))
			return
		}
		version := time.Now().UnixMilli()
		res, err := db.ExecContext(r.Context(),
			`INSERT INTO user_avatars (user_id, data, updated_at)
			 SELECT id, $2, $3 FROM users WHERE id = $1
			 ON CONFLICT (user_id) DO UPDATE SET data = EXCLUDED.data, updated_at = EXCLUDED.updated_at`,
			id, data, version)
		if err != nil {
			serverError(w, r, err)
			return
		}
		if n, _ := res.RowsAffected(); n == 0 {
			respond(w, http.StatusNotFound, errResp("user not found"))
			return
		}
		respond(w, http.StatusOK, map[string]int64{"avatar_version": version})
	}
}

func handleDeleteAvatar(db *sql.DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		id, err := strconv.ParseInt(chi.URLParam(r, "id"), 10, 64)
		if err != nil {
			respond(w, http.StatusBadRequest, errResp("invalid user id"))
			return
		}
		if !requireSelfOrAdmin(w, r, id) {
			return
		}
		if _, err := db.ExecContext(r.Context(),
			"DELETE FROM user_avatars WHERE user_id = $1", id); err != nil {
			serverError(w, r, err)
			return
		}
		w.WriteHeader(http.StatusNoContent)
	}
}