kra-new/internal/biz/media_upload.go

package biz

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/google/uuid"
)

var ErrUploadSessionNotFound = errors.New("upload session not found")

func (uc *MediaUsecase) chunkPrefix(uploadID uint) string {
	directory := "uploads/chunks"
	if uc.settings != nil {
		if configured := strings.Trim(uc.settings.MediaSettings().ChunkDir, "/\\ "); configured != "" {
			directory = configured
		}
	}
	return path.Join(directory, fmt.Sprintf("%d", uploadID))
}

func (uc *MediaUsecase) chunkKey(uploadID uint, index int) string {
	return path.Join(uc.chunkPrefix(uploadID), fmt.Sprintf("%08d", index))
}

func (uc *MediaUsecase) InitUpload(ctx context.Context, userID uint, name, hash string, size, chunkSize int64, total int) (*UploadSession, *MediaFile, []int, error) {
	if err := validateMediaName(name); err != nil {
		return nil, nil, nil, err
	}
	if completed, err := uc.FindCompletedSession(ctx, userID, hash); err == nil && completed.MediaID != 0 {
		if media, findErr := uc.FindMedia(ctx, completed.MediaID); findErr == nil {
			copy := &MediaFile{Name: name, URL: media.URL, Tag: media.Tag, Key: media.Key}
			if createErr := uc.CreateMedia(ctx, copy); createErr == nil {
				return nil, copy, nil, nil
			}
		}
	}
	session, err := uc.FindUploadingSession(ctx, userID, hash)
	if errors.Is(err, ErrUploadSessionNotFound) {
		session = &UploadSession{UserID: userID, FileName: name, FileHash: hash, FileSize: size, ChunkSize: chunkSize, ChunkTotal: total, Status: "uploading"}
		if err = uc.CreateUploadSession(ctx, session); err != nil {
			return nil, nil, nil, err
		}
	} else if err != nil {
		return nil, nil, nil, err
	}
	chunks, err := uc.ListChunks(ctx, session.ID)
	if err != nil {
		return nil, nil, nil, err
	}
	indices := make([]int, 0, len(chunks))
	for _, v := range chunks {
		indices = append(indices, v.Index)
	}
	sort.Ints(indices)
	return session, nil, indices, nil
}
func (uc *MediaUsecase) SaveChunk(ctx context.Context, userID, uploadID uint, index int, expected string, reader io.Reader) error {
	session, err := uc.FindUploadSession(ctx, uploadID)
	if err != nil {
		return errors.New("上传会话不存在")
	}
	if session.UserID != userID {
		return errors.New("无权操作该上传")
	}
	if session.Status != "uploading" {
		return errors.New("上传会话状态不允许收片")
	}
	hash := md5.New()
	key := uc.chunkKey(uploadID, index)
	temporary, err := os.CreateTemp("", "kra-upload-chunk-*")
	if err != nil {
		return err
	}
	temporaryName := temporary.Name()
	defer os.Remove(temporaryName)
	defer temporary.Close()
	if _, err = io.Copy(io.MultiWriter(temporary, hash), reader); err != nil {
		return err
	}
	actual := hex.EncodeToString(hash.Sum(nil))
	if actual != expected {
		return fmt.Errorf("分片 %d 校验失败", index)
	}
	if _, err = temporary.Seek(0, io.SeekStart); err != nil {
		return err
	}
	stored, err := uc.files.Put(ctx, key, temporary)
	if err != nil {
		return err
	}
	return uc.UpsertChunk(ctx, uploadID, &UploadChunk{Index: index, Hash: actual, Size: stored.Size})
}
func (uc *MediaUsecase) CompleteUpload(ctx context.Context, userID, uploadID uint, mime string) (*MediaFile, error) {
	session, err := uc.FindUploadSession(ctx, uploadID)
	if err != nil {
		return nil, errors.New("上传会话不存在")
	}
	if session.UserID != userID {
		return nil, errors.New("无权操作该上传")
	}
	if err = validateMediaName(session.FileName); err != nil {
		return nil, err
	}
	claimed, err := uc.ClaimUploadSession(ctx, uploadID)
	if err != nil {
		return nil, err
	}
	if !claimed {
		return nil, errors.New("上传不在可合并状态(可能已在合并或已完成)")
	}
	fail := func(value error) (*MediaFile, error) {
		_ = uc.FailUploadSession(ctx, uploadID)
		return nil, value
	}
	chunks, err := uc.ListChunks(ctx, uploadID)
	if err != nil || len(chunks) != session.ChunkTotal {
		return fail(fmt.Errorf("分片不全: %d/%d", len(chunks), session.ChunkTotal))
	}
	sort.Slice(chunks, func(i, j int) bool { return chunks[i].Index < chunks[j].Index })
	names := make([]string, 0, len(chunks))
	for index, chunk := range chunks {
		if chunk.Index != index {
			return fail(errors.New("分片序号不连续"))
		}
		names = append(names, uc.chunkKey(uploadID, index))
	}
	ext := strings.ToLower(filepath.Ext(session.FileName))
	key := time.Now().Format("20060102") + "/" + uuid.NewString() + ext
	stored, hash, err := uc.files.Compose(ctx, names, key)
	if err != nil {
		return fail(err)
	}
	if hash != session.FileHash {
		_ = uc.files.Delete(ctx, key)
		return fail(errors.New("整文件校验失败"))
	}
	// Keep the size declared when the upload session was created. The full-file
	// MD5 above remains the integrity check even if storage reports another
	// byte count.
	media := &MediaFile{Name: session.FileName, URL: stored.URL, Tag: strings.TrimPrefix(filepath.Ext(session.FileName), "."), Key: key, Size: session.FileSize, Mime: mime, MD5: hash, UserID: userID}
	if err = uc.CreateMedia(ctx, media); err != nil {
		return fail(err)
	}
	_ = uc.CompleteUploadSession(ctx, uploadID, key, media.ID)
	_ = uc.DeleteChunks(ctx, uploadID)
	_ = uc.files.DeletePrefix(ctx, uc.chunkPrefix(uploadID))
	return media, nil
}
func (uc *MediaUsecase) CancelUpload(ctx context.Context, userID, uploadID uint) error {
	session, err := uc.FindUploadSession(ctx, uploadID)
	if err != nil {
		return errors.New("上传会话不存在")
	}
	if session.UserID != userID {
		return errors.New("无权操作该上传")
	}
	_ = uc.DeleteChunks(ctx, uploadID)
	_ = uc.files.DeletePrefix(ctx, uc.chunkPrefix(uploadID))
	return uc.DeleteUploadSession(ctx, uploadID)
}
func (uc *MediaUsecase) CleanupStale(ctx context.Context, ttlHours int) error {
	if ttlHours <= 0 {
		ttlHours = 24
	}
	ids, err := uc.StaleUploadSessionIDs(ctx, time.Now().Add(-time.Duration(ttlHours)*time.Hour))
	if err != nil {
		return err
	}
	for _, id := range ids {
		if err = uc.DeleteUploadData(ctx, id); err != nil {
			return err
		}
		if err = uc.files.DeletePrefix(ctx, uc.chunkPrefix(id)); err != nil {
			return err
		}
	}
	return nil
}