zt-closet-backend/pkg/upload/upload.go

// Package upload 文件上传、校验、落盘、公开 token 生成的工具函数。
package upload

import (
	"context"
	"crypto/md5"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"mime/multipart"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"git.lzcwy.cn/zhengchiliu/zt-closet-backend/global"
)

func FileGetExt(name string) string {
	return strings.ToLower(path.Ext(name))
}

// 校验扩展名是否在白名单内。
func FileCheckAllowExt(ext string, allowed []string) bool {
	ext = strings.TrimPrefix(strings.ToLower(ext), ".")
	for _, e := range allowed {
		if strings.ToLower(e) == ext {
			return true
		}
	}
	return false
}

// 校验文件大小。maxMB <= 0 表示不限制。
func FileCheckMaxSize(size, maxMB int64) bool {
	if maxMB <= 0 {
		return true
	}
	return size <= maxMB*1024*1024
}

// FileCheckImageMagic 校验文件头 magic bytes，判断内容是否为真实图片（jpeg/png/gif/webp）。
// 用于防止伪造扩展名上传非图片文件（如脚本、可执行文件、HTML）。
func FileCheckImageMagic(data []byte) bool {
	if len(data) < 4 {
		return false
	}

	// JPEG: FF D8 FF
	if data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF {
		return true
	}

	// PNG: 89 50 4E 47 0D 0A 1A 0A
	if len(data) >= 8 &&
		data[0] == 0x89 && data[1] == 0x50 && data[2] == 0x4E && data[3] == 0x47 &&
		data[4] == 0x0D && data[5] == 0x0A && data[6] == 0x1A && data[7] == 0x0A {
		return true
	}

	// GIF: "GIF8"
	if data[0] == 'G' && data[1] == 'I' && data[2] == 'F' && data[3] == '8' {
		return true
	}

	// WebP: "RIFF" .... "WEBP"
	if len(data) >= 12 &&
		data[0] == 'R' && data[1] == 'I' && data[2] == 'F' && data[3] == 'F' &&
		data[8] == 'W' && data[9] == 'E' && data[10] == 'B' && data[11] == 'P' {
		return true
	}

	return false
}

// GenFileName 生成防冲突随机相对文件名。
func FileGenName(originName string) (string, error) {
	ext := FileGetExt(originName)
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return hex.EncodeToString(b) + ext, nil
}

// 临时保存从前端获取的图片
func FileSaveTmp(ctx context.Context, file *multipart.FileHeader, fileName string, userID int64) (string, error) {
	ctx, span := global.Tracer.Start(ctx, "upload.FileSaveTmp")
	defer span.End()

	dir := filepath.Join(global.UploadSetting.TmpSavePath, fmt.Sprint(userID))
	if err := os.MkdirAll(dir, 0o755); err != nil {
		span.RecordError(err)
		return "", err
	}

	tmpFile := filepath.Join(dir, fileName)

	dst, err := os.Create(tmpFile)
	if err != nil {
		return "", err
	}
	defer dst.Close()

	src, err := file.Open()
	if err != nil {
		return "", err
	}
	defer src.Close()

	if _, err := io.Copy(dst, src); err != nil {
		return "", err
	}

	return tmpFile, nil

}

// 从后端抠图返回的图片保存到临时文件（subDir 为相对子目录，如 "用户UUID/用户ID"）。
func FileSaveTmpBytes(ctx context.Context, data []byte, fileName string, subDir string) (string, error) {
	_, span := global.Tracer.Start(ctx, "upload.FileSaveTmpBytes")
	defer span.End()

	dir := filepath.Join(global.UploadSetting.TmpSavePath, subDir)
	if err := os.MkdirAll(dir, 0o755); err != nil {
		span.RecordError(err)
		return "", err
	}

	tmpFile := filepath.Join(dir, fileName)
	if err := os.WriteFile(tmpFile, data, 0o644); err != nil {
		span.RecordError(err)
		return "", err
	}
	return tmpFile, nil
}

// 将抠图后的文件上传到cos
func FileUploadCos(ctx context.Context, fileName string) error {
	ctx, span := global.Tracer.Start(ctx, "upload.FileUploadCos")
	defer span.End()

	// COS 客户端初始化失败时降级为 nil，此处做保护避免 nil 指针 panic。
	if global.CosClient == nil {
		err := fmt.Errorf("cos client not initialized")
		span.RecordError(err)
		return err
	}

	_, _, err := global.CosClient.Object.Upload(ctx, fileName, fmt.Sprintf("%s/%s", global.UploadSetting.TmpSavePath, fileName), nil)
	if err != nil {
		span.RecordError(err)
		return err
	}

	return nil

}

// 删除 COS 对象。key 为对象键（即 processed_image 存储值）。
// key 为空直接返回 nil；client 未初始化返回 error。
// 删除为尽力而为的清理操作，内部使用短超时避免网络异常长时间阻塞调用方。
func FileDeleteCos(ctx context.Context, key string) error {
	ctx, span := global.Tracer.Start(ctx, "upload.FileDeleteCos")
	defer span.End()

	if key == "" {
		return nil
	}

	// COS 客户端初始化失败时降级为 nil，此处做保护避免 nil 指针 panic。
	if global.CosClient == nil {
		err := fmt.Errorf("cos client not initialized")
		span.RecordError(err)
		return err
	}

	delCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()

	_, err := global.CosClient.Object.Delete(delCtx, key)
	if err != nil {
		span.RecordError(err)
		return err
	}

	return nil
}

// 获取cos访问url
func GetCosUrl(fileName string) string {
	ts := time.Now().Unix()
	raw := global.CosSetting.CdnDownloadSecretKey + "/" + fileName + strconv.FormatInt(ts, 10)
	sign := fmt.Sprintf("%x", md5.Sum([]byte(raw)))

	return fmt.Sprintf("https://%s/%s?sign=%s&t=%d",
		global.CosSetting.CosDownloadDomain,
		fileName,
		sign,
		ts,
	)
}

// AvatarURL 将头像存储值转为可访问 URL：
// 历史通过修改资料接口传入的 http(s) URL 直接透传；COS key 统一转为 CDN 签名 URL。
func AvatarURL(raw string) string {
	if raw == "" {
		return ""
	}
	if strings.HasPrefix(raw, "http://") || strings.HasPrefix(raw, "https://") {
		return raw
	}
	return GetCosUrl(raw)
}