thread-master/apps/backend/internal/logic/media/upload_logic.go

package media

import (
	"context"
	"encoding/base64"
	"fmt"
	"path"
	"strings"

	"apps/backend/internal/middleware"
	"apps/backend/internal/response"
	"apps/backend/internal/svc"
	"apps/backend/internal/types"

	"github.com/google/uuid"
	"github.com/zeromicro/go-zero/core/logx"
)

const maxUploadBytes = 5 * 1024 * 1024

var allowedImageTypes = map[string]string{
	"image/jpeg": ".jpg",
	"image/jpg":  ".jpg",
	"image/png":  ".png",
	"image/webp": ".webp",
	"image/gif":  ".gif",
}

type UploadLogic struct {
	logx.Logger
	ctx    context.Context
	svcCtx *svc.ServiceContext
}

func NewUploadLogic(ctx context.Context, svcCtx *svc.ServiceContext) *UploadLogic {
	return &UploadLogic{Logger: logx.WithContext(ctx), ctx: ctx, svcCtx: svcCtx}
}

// Upload 解 data URL → object storage → 回公開 URL（Mongo 只存此 URL）。
func (l *UploadLogic) Upload(req *types.MediaUploadReq) (*types.MediaUploadData, error) {
	if l.svcCtx.Storage == nil || !l.svcCtx.Storage.Enabled() {
		return nil, response.Biz(503, 503001, "object storage not configured")
	}
	uid, ok := middleware.UIDFrom(l.ctx)
	if !ok || uid <= 0 {
		return nil, response.Biz(401, 401001, "missing authorization")
	}

	contentType, raw, err := decodeImageContent(req.Content)
	if err != nil {
		return nil, response.Biz(400, 400001, err.Error())
	}
	if len(raw) > maxUploadBytes {
		return nil, response.Biz(400, 400001, "file size exceeds 5MB limit")
	}
	ext, ok := allowedImageTypes[contentType]
	if !ok {
		return nil, response.Biz(400, 400001, "unsupported image type (jpeg/png/webp/gif only)")
	}

	kind := strings.ToLower(strings.TrimSpace(req.Kind))
	if kind == "" {
		kind = "avatar"
	}
	if kind != "avatar" && kind != "other" {
		kind = "other"
	}

	id := uuid.NewString()
	key := path.Join(kind, fmt.Sprintf("%d", uid), id+ext)

	url, err := l.svcCtx.Storage.Upload(l.ctx, key, raw, contentType)
	if err != nil {
		l.Errorf("storage upload: %v", err)
		return nil, response.Biz(500, 500000, "file upload failed")
	}
	return &types.MediaUploadData{Url: url, Key: key}, nil
}

func decodeImageContent(content string) (contentType string, data []byte, err error) {
	content = strings.TrimSpace(content)
	if content == "" {
		return "", nil, fmt.Errorf("empty content")
	}
	// data:image/jpeg;base64,....
	if strings.HasPrefix(content, "data:") {
		parts := strings.SplitN(content, ",", 2)
		if len(parts) != 2 {
			return "", nil, fmt.Errorf("invalid data url")
		}
		meta := parts[0] // data:image/jpeg;base64
		meta = strings.TrimPrefix(meta, "data:")
		semi := strings.Index(meta, ";")
		if semi < 0 {
			return "", nil, fmt.Errorf("invalid data url meta")
		}
		contentType = strings.TrimSpace(meta[:semi])
		payload := parts[1]
		data, err = base64.StdEncoding.DecodeString(payload)
		if err != nil {
			// try raw std without padding issues
			data, err = base64.RawStdEncoding.DecodeString(payload)
		}
		if err != nil {
			return "", nil, fmt.Errorf("invalid base64")
		}
		return contentType, data, nil
	}
	// pure base64 — sniff
	data, err = base64.StdEncoding.DecodeString(content)
	if err != nil {
		return "", nil, fmt.Errorf("invalid base64")
	}
	contentType = sniffImage(data)
	if contentType == "" {
		return "", nil, fmt.Errorf("cannot detect image type")
	}
	return contentType, data, nil
}

func sniffImage(b []byte) string {
	if len(b) >= 3 && b[0] == 0xff && b[1] == 0xd8 && b[2] == 0xff {
		return "image/jpeg"
	}
	if len(b) >= 8 && b[0] == 0x89 && b[1] == 0x50 && b[2] == 0x4e && b[3] == 0x47 {
		return "image/png"
	}
	if len(b) >= 6 && (string(b[:6]) == "GIF87a" || string(b[:6]) == "GIF89a") {
		return "image/gif"
	}
	if len(b) >= 12 && string(b[0:4]) == "RIFF" && string(b[8:12]) == "WEBP" {
		return "image/webp"
	}
	return ""
}