zogo/fileupload/service.go

package fileupload

import (
	"crypto/md5"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"os"
	"path/filepath"
	"strconv"
	"time"
)

// service 核心业务，持 Options + Repository。
type service struct {
	opts Options
	repo Repository
}

func newService(opts Options, repo Repository) *service {
	return &service{opts: opts, repo: repo}
}

func (s *service) chunkDir(uploadID string) string {
	return filepath.Join(s.opts.ChunkRoot, uploadID)
}

func (s *service) fileURL(storagePath string) string {
	return s.opts.URLPrefix + "/" + storagePath
}

// ── init ────────────────────────────────────────────────────────────

type InitReq struct {
	FileName   string `json:"fileName" binding:"required,max=255"`
	FileSize   int64  `json:"fileSize" binding:"min=0"`
	FileMD5    string `json:"fileMD5" binding:"required,len=32"`
	MimeType   string `json:"mimeType" binding:"max=128"`
	ChunkSize  int64  `json:"chunkSize" binding:"required,min=1"`
	TotalChunk int    `json:"totalChunk" binding:"required,min=1"`
}

func (s *service) init(req *InitReq, operator string, tenantID int64) (any, error) {
	if req.ChunkSize <= 0 || req.ChunkSize > s.opts.maxChunkBytes() {
		return nil, fmt.Errorf("chunkSize 超出范围 (0, %d]", s.opts.maxChunkBytes())
	}
	if req.FileSize < 0 {
		return nil, errors.New("fileSize 非法")
	}
	if len(req.FileMD5) != 32 {
		return nil, errors.New("fileMD5 必须为 32 位十六进制")
	}
	if expect := expectChunks(req.FileSize, req.ChunkSize); req.TotalChunk != expect {
		return nil, fmt.Errorf("totalChunk(%d) 与 fileSize(%d)/chunkSize(%d) 不匹配，应为 %d",
			req.TotalChunk, req.FileSize, req.ChunkSize, expect)
	}
	if s.opts.MaxTotalSizeMB > 0 && req.FileSize > s.opts.MaxTotalSizeMB*1024*1024 {
		return nil, fmt.Errorf("文件超过总大小上限 %dMB", s.opts.MaxTotalSizeMB)
	}

	// 1. 秒传：同 MD5 已完成任务，复用其文件
	done, err := s.repo.FindDoneByMD5(req.FileMD5)
	if err != nil {
		return nil, err
	}
	if done != nil && done.StoragePath != "" {
		if s.opts.OnInstantHit != nil {
			if err := s.opts.OnInstantHit(*done); err != nil {
				return nil, err
			}
		}
		return map[string]any{
			"uploadId":   "",
			"skipUpload": true,
			"fileUrl":    s.fileURL(done.StoragePath),
		}, nil
	}

	// 2. 断点续传：同 MD5 进行中任务
	resuming, err := s.repo.FindUploadingByMD5(req.FileMD5)
	if err != nil {
		return nil, err
	}
	if resuming != nil {
		indexes, err := s.repo.ListChunkIndexes(resuming.UploadId)
		if err != nil {
			return nil, err
		}
		return map[string]any{
			"uploadId":       resuming.UploadId,
			"skipUpload":     false,
			"uploadedChunks": indexes,
		}, nil
	}

	// 3. 新建任务
	name := sanitizeFileName(req.FileName)
	if name == "" {
		return nil, errors.New("文件名非法")
	}
	task := TaskInfo{
		UploadId:   newUploadID(),
		FileName:   name,
		FileSize:   req.FileSize,
		FileMd5:    req.FileMD5,
		MimeType:   req.MimeType,
		ChunkSize:  req.ChunkSize,
		TotalChunk: int32(req.TotalChunk),
		Status:     StatusUploading,
		Operator:   operator,
		TenantId:   tenantID,
	}
	if err := s.repo.CreateTask(&task); err != nil {
		return nil, err
	}
	return map[string]any{
		"uploadId":       task.UploadId,
		"skipUpload":     false,
		"uploadedChunks": []int32{},
	}, nil
}

// ── chunk ───────────────────────────────────────────────────────────

func (s *service) chunk(uploadID string, chunkIndex int, fh *multipart.FileHeader) (any, error) {
	if uploadID == "" {
		return nil, errors.New("缺少 uploadId")
	}
	task, err := s.repo.GetTaskByUploadID(uploadID)
	if err != nil {
		return nil, err
	}
	if task == nil {
		return nil, errors.New("上传任务不存在")
	}
	if task.Status != StatusUploading {
		return nil, fmt.Errorf("任务状态异常: %s", task.Status)
	}
	if chunkIndex < 0 || int32(chunkIndex) >= task.TotalChunk {
		return nil, fmt.Errorf("chunkIndex %d 超出范围 [0, %d)", chunkIndex, task.TotalChunk)
	}
	if fh.Size > s.opts.maxChunkBytes() {
		return nil, fmt.Errorf("分片超过单片上限 %d MB", s.opts.maxChunkBytes()/1024/1024)
	}

	// 幂等：已确认的分片直接成功
	exists, err := s.repo.ChunkExists(uploadID, int32(chunkIndex))
	if err != nil {
		return nil, err
	}
	if exists {
		return map[string]any{"chunkIndex": chunkIndex}, nil
	}

	dir := s.chunkDir(uploadID)
	if err := os.MkdirAll(dir, 0o766); err != nil {
		return nil, err
	}
	chunkPath := filepath.Join(dir, strconv.Itoa(chunkIndex))
	dst, err := os.Create(chunkPath)
	if err != nil {
		return nil, err
	}
	src, err := fh.Open()
	if err != nil {
		_ = dst.Close()
		return nil, err
	}
	hash := md5.New()
	written, copyErr := io.Copy(io.MultiWriter(dst, hash), src)
	closeErr := dst.Close()
	_ = src.Close()
	if copyErr != nil {
		return nil, copyErr
	}
	if closeErr != nil {
		return nil, closeErr
	}

	if err := s.repo.CreateChunk(&ChunkInfo{
		UploadId:   uploadID,
		ChunkIndex: int32(chunkIndex),
		ChunkSize:  written,
		ChunkMd5:   fmt.Sprintf("%x", hash.Sum(nil)),
	}); err != nil {
		// 并发下另一请求刚写入同分片：唯一冲突视为幂等成功
		again, e2 := s.repo.ChunkExists(uploadID, int32(chunkIndex))
		if e2 == nil && again {
			return map[string]any{"chunkIndex": chunkIndex}, nil
		}
		return nil, err
	}
	return map[string]any{"chunkIndex": chunkIndex}, nil
}

// ── merge ───────────────────────────────────────────────────────────

type MergeReq struct {
	UploadId string `json:"uploadId" binding:"required"`
}

func (s *service) merge(uploadID string) (any, error) {
	task, err := s.repo.GetTaskByUploadID(uploadID)
	if err != nil {
		return nil, err
	}
	if task == nil {
		return nil, errors.New("上传任务不存在")
	}

	// 幂等：已完成直接返回原地址；合并中拒绝重复触发
	if task.Status == StatusDone && task.StoragePath != "" {
		return map[string]any{"fileUrl": s.fileURL(task.StoragePath)}, nil
	}
	if task.Status == StatusMerging {
		return nil, errors.New("任务正在合并中，请勿重复提交")
	}

	count, err := s.repo.CountChunks(uploadID)
	if err != nil {
		return nil, err
	}
	if count < int(task.TotalChunk) {
		return nil, fmt.Errorf("分片不完整，已上传 %d/%d", count, task.TotalChunk)
	}

	_ = s.repo.SetTaskStatus(task.Id, StatusMerging)

	// 先写 .part，全部成功后改名，避免半成品文件残留
	partPath := filepath.Join(s.chunkDir(uploadID), "merged.part")
	part, err := os.Create(partPath)
	if err != nil {
		_ = s.repo.SetTaskStatus(task.Id, StatusError)
		return nil, err
	}
	hash := md5.New()
	for i := int32(0); i < task.TotalChunk; i++ {
		src, err := os.Open(filepath.Join(s.chunkDir(uploadID), strconv.Itoa(int(i))))
		if err != nil {
			_ = part.Close()
			_ = os.Remove(partPath)
			_ = s.repo.SetTaskStatus(task.Id, StatusError)
			return nil, fmt.Errorf("读取分片 %d 失败: %w", i, err)
		}
		_, copyErr := io.Copy(part, io.TeeReader(src, hash))
		_ = src.Close()
		if copyErr != nil {
			_ = part.Close()
			_ = os.Remove(partPath)
			_ = s.repo.SetTaskStatus(task.Id, StatusError)
			return nil, fmt.Errorf("合并分片 %d 失败: %w", i, copyErr)
		}
	}
	if err := part.Close(); err != nil {
		_ = os.Remove(partPath)
		_ = s.repo.SetTaskStatus(task.Id, StatusError)
		return nil, err
	}

	// 整文件 MD5 校验：保证秒传可信
	if actual := fmt.Sprintf("%x", hash.Sum(nil)); actual != task.FileMd5 {
		_ = os.Remove(partPath)
		_ = s.repo.SetTaskStatus(task.Id, StatusError)
		return nil, fmt.Errorf("合并结果校验失败：实际 MD5 %s 与声称的 %s 不一致", actual, task.FileMd5)
	}

	// 改名进宿主产物目录
	if err := os.MkdirAll(s.opts.CloudDir, 0o766); err != nil {
		_ = os.Remove(partPath)
		_ = s.repo.SetTaskStatus(task.Id, StatusError)
		return nil, err
	}
	storageName := strconv.FormatInt(time.Now().UnixNano(), 10) + "_" + task.FileName
	// storagePath 是 Web 相对路径（含 WebPathPrefix 前缀，如 cloud/xxx），
	// 物理落盘固定在 CloudDir 根下（前缀只是 URL/拼接约定，不是子目录）
	finalPath := filepath.Join(s.opts.CloudDir, storageName)
	storagePath := storageName
	if s.opts.WebPathPrefix != "" {
		storagePath = s.opts.WebPathPrefix + "/" + storageName
	}
	if err := os.Rename(partPath, finalPath); err != nil {
		_ = os.Remove(partPath)
		_ = s.repo.SetTaskStatus(task.Id, StatusError)
		return nil, err
	}

	task.Status = StatusDone
	task.StoragePath = storagePath
	if err := s.repo.UpdateTaskMerged(task.Id, storagePath); err != nil {
		return nil, err
	}

	// 宿主业务衔接（如写云盘记录行）；失败则 merge 整体失败
	if s.opts.OnMerged != nil {
		if err := s.opts.OnMerged(*task); err != nil {
			return nil, err
		}
	}

	// 分片目录异步清理
	go func(dir string) {
		_ = os.RemoveAll(dir)
	}(s.chunkDir(uploadID))

	return map[string]any{"fileUrl": s.fileURL(storagePath)}, nil
}

// ── check / task / abort ────────────────────────────────────────────

func (s *service) check(fileMD5 string) (any, error) {
	if len(fileMD5) != 32 {
		return map[string]any{"status": "none"}, nil
	}
	done, err := s.repo.FindDoneByMD5(fileMD5)
	if err != nil {
		return nil, err
	}
	if done != nil && done.StoragePath != "" {
		return map[string]any{
			"status":  StatusDone,
			"fileUrl": s.fileURL(done.StoragePath),
		}, nil
	}
	resuming, err := s.repo.FindUploadingByMD5(fileMD5)
	if err != nil {
		return nil, err
	}
	if resuming != nil {
		indexes, err := s.repo.ListChunkIndexes(resuming.UploadId)
		if err != nil {
			return nil, err
		}
		return map[string]any{
			"status":         StatusUploading,
			"uploadId":       resuming.UploadId,
			"uploadedChunks": indexes,
		}, nil
	}
	return map[string]any{"status": "none"}, nil
}

func (s *service) taskInfo(uploadID string) (any, error) {
	task, err := s.repo.GetTaskByUploadID(uploadID)
	if err != nil {
		return nil, err
	}
	if task == nil {
		return nil, errors.New("任务不存在")
	}
	indexes, err := s.repo.ListChunkIndexes(uploadID)
	if err != nil {
		return nil, err
	}
	return map[string]any{"task": task, "uploadedChunks": indexes}, nil
}

func (s *service) abort(uploadID string) (any, error) {
	task, err := s.repo.GetTaskByUploadID(uploadID)
	if err != nil {
		return nil, err
	}
	if task == nil {
		return nil, errors.New("上传任务不存在")
	}
	if task.Status == StatusDone {
		return nil, errors.New("任务已完成，无需取消")
	}
	if task.Status == StatusMerging {
		return nil, errors.New("任务正在合并中，无法取消")
	}
	if err := s.repo.DeleteChunks(uploadID); err != nil {
		return nil, err
	}
	if err := s.repo.DeleteTaskByUploadID(uploadID); err != nil {
		return nil, err
	}
	go func(dir string) {
		_ = os.RemoveAll(dir)
	}(s.chunkDir(uploadID))
	return nil, nil
}