digital-psychology/apps/api/internal/avatar/store.go

// Package avatar stores and serves user profile photos.
package avatar

import (
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"os"
	"path/filepath"
	"strings"

	"github.com/google/uuid"
)

const MaxBytes = 2 << 20 // 2 MiB

var (
	ErrTooLarge   = errors.New("图片过大，请选择 2MB 以内")
	ErrBadType    = errors.New("仅支持 JPG / PNG / WebP")
	ErrEmpty      = errors.New("请选择图片")
	ErrBadName    = errors.New("无效头像")
)

// Store writes multipart image under dir as {userID}{ext}; returns public API path.
func Store(dir string, userID uuid.UUID, fh *multipart.FileHeader) (publicPath string, err error) {
	if fh == nil || fh.Size <= 0 {
		return "", ErrEmpty
	}
	if fh.Size > MaxBytes {
		return "", ErrTooLarge
	}
	ext, err := extFor(fh.Filename, fh.Header.Get("Content-Type"))
	if err != nil {
		return "", err
	}
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return "", err
	}
	src, err := fh.Open()
	if err != nil {
		return "", err
	}
	defer src.Close()

	name := userID.String() + ext
	dstPath := filepath.Join(dir, name)
	tmp := dstPath + ".tmp"
	out, err := os.Create(tmp)
	if err != nil {
		return "", err
	}
	n, copyErr := io.Copy(out, io.LimitReader(src, MaxBytes+1))
	_ = out.Close()
	if copyErr != nil {
		_ = os.Remove(tmp)
		return "", copyErr
	}
	if n > MaxBytes {
		_ = os.Remove(tmp)
		return "", ErrTooLarge
	}
	// remove previous extensions for same user
	for _, old := range []string{".jpg", ".jpeg", ".png", ".webp"} {
		p := filepath.Join(dir, userID.String()+old)
		if p != dstPath {
			_ = os.Remove(p)
		}
	}
	if err := os.Rename(tmp, dstPath); err != nil {
		_ = os.Remove(tmp)
		return "", err
	}
	return "/api/v1/media/avatars/" + name, nil
}

// ResolveAbs validates file name and returns absolute path under dir.
func ResolveAbs(dir, name string) (string, error) {
	base := filepath.Base(name)
	if base != name || strings.Contains(base, "..") {
		return "", ErrBadName
	}
	ext := strings.ToLower(filepath.Ext(base))
	if ext != ".jpg" && ext != ".jpeg" && ext != ".png" && ext != ".webp" {
		return "", ErrBadName
	}
	idPart := strings.TrimSuffix(base, ext)
	if _, err := uuid.Parse(idPart); err != nil {
		return "", ErrBadName
	}
	full := filepath.Join(dir, base)
	if _, err := os.Stat(full); err != nil {
		return "", ErrBadName
	}
	return full, nil
}

func extFor(filename, contentType string) (string, error) {
	ct := strings.ToLower(contentType)
	switch {
	case strings.Contains(ct, "jpeg"), strings.Contains(ct, "jpg"):
		return ".jpg", nil
	case strings.Contains(ct, "png"):
		return ".png", nil
	case strings.Contains(ct, "webp"):
		return ".webp", nil
	}
	ext := strings.ToLower(filepath.Ext(filename))
	switch ext {
	case ".jpg", ".jpeg":
		return ".jpg", nil
	case ".png":
		return ".png", nil
	case ".webp":
		return ".webp", nil
	default:
		return "", fmt.Errorf("%w", ErrBadType)
	}
}