015/backend/internal/controllers/file.go

package controllers

import (
	"backend/internal/services"
	"backend/internal/utils"
	"encoding/json"
	"errors"
	"math"
	"mime/multipart"
	"os"
	"pkg/models"
	u "pkg/utils"
	"time"

	"github.com/hibiken/asynq"
	"github.com/labstack/echo/v5"
)

func CreateUploadTask(c *echo.Context) error {
	// cc := c.(*middleware.CustomContext)
	r := new(models.FileInfo)
	if err := c.Bind(r); err != nil {
		return utils.HTTPErrorHandler(c, err)
	}

	if r.FileSize == 0 || r.MimeType == "" || r.FileHash == "" {
		return utils.HTTPErrorHandler(c, errors.New("调用接口参数错误"))
	}
	fileId := utils.GetFileId(r.FileHash, r.FileSize)
	fileInfo, err := models.GetRedisFileInfo(fileId)
	if err != nil {
		return utils.HTTPErrorHandler(c, err)
	}

	if fileInfo != nil {
		uploadPath, err := utils.GetUploadDirPath()
		if err != nil {
			return utils.HTTPErrorHandler(c, err)
		}
		sliceList, err := services.GetFileSliceList(fileId, uploadPath)
		if err != nil {
			return utils.HTTPErrorHandler(c, err)
		}
		return utils.HTTPSuccessHandler(c, map[string]any{
			"size":       fileInfo.FileSize,
			"mime_type":  fileInfo.MimeType,
			"hash":       fileInfo.FileHash,
			"type":       fileInfo.FileType,
			"expire":     fileInfo.Expire,
			"id":         fileId,
			"chunk_size": fileInfo.ChunkSize,
			"chunks":     sliceList,
		})
	}
	maxStorageSize, err := utils.GetFileSize(u.GetEnv("upload.maximum"))
	if err != nil {
		return utils.HTTPErrorHandler(c, err)
	}
	fileInfoMap, err := models.GetRedisFileInfoAll()
	if err != nil {
		return utils.HTTPErrorHandler(c, err)
	}
	totalSize := int64(0)
	for _, value := range fileInfoMap {
		var fileInfo models.RedisFileInfo
		err := json.Unmarshal([]byte(value), &fileInfo)
		if err != nil {
			return utils.HTTPErrorHandler(c, err)
		}
		totalSize += fileInfo.FileSize
	}
	if totalSize+r.FileSize > int64(maxStorageSize) {
		return utils.HTTPErrorHandler(c, errors.New("存储空间不足"))
	}

	ChunkSize := int64(0.25 * 1024 * 1024)
	// 根据文件大小动态调整块大小
	for r.FileSize/ChunkSize > 1000 {
		ChunkSize *= 2
	}
	uploadTaskExpire := int64(3600)
	newFileInfo := models.RedisFileInfo{
		FileType: models.FileTypeInit,
		FileInfo: models.FileInfo{
			FileSize:  r.FileSize,
			MimeType:  r.MimeType,
			FileHash:  r.FileHash,
			ChunkSize: ChunkSize,
		},
		CreatedAt: time.Now().Unix(),
		Expire:    uploadTaskExpire,
	}
	err = models.SetRedisFileInfo(fileId, newFileInfo)
	if err != nil {
		return utils.HTTPErrorHandler(c, err)
	}

	client := u.GetQueueClient()
	json, err := json.Marshal(map[string]any{
		"file_id": fileId,
	})
	if err != nil {
		return utils.HTTPErrorHandler(c, err)
	}
	_, err = client.Enqueue(asynq.NewTask("file:remove", json), asynq.ProcessIn(time.Duration(uploadTaskExpire)*time.Second))
	if err != nil {
		return utils.HTTPErrorHandler(c, err)
	}

	return utils.HTTPSuccessHandler(c, map[string]any{
		"size":       newFileInfo.FileSize,
		"mime_type":  newFileInfo.MimeType,
		"hash":       newFileInfo.FileHash,
		"type":       newFileInfo.FileType,
		"expire":     newFileInfo.Expire,
		"id":         fileId,
		"chunk_size": newFileInfo.ChunkSize,
	})
}

type UploadFileSliceProps struct {
	FileId    string                `form:"id"`
	FileIndex int64                 `form:"index"`
	FileSlice *multipart.FileHeader `form:"file"`
}

func UploadFileSlice(c *echo.Context) error {
	r := new(UploadFileSliceProps)
	if err := c.Bind(r); err != nil {
		return utils.HTTPErrorHandler(c, err)
	}

	if r.FileId == "" || r.FileIndex == 0 || r.FileSlice == nil {
		return utils.HTTPErrorHandler(c, errors.New("调用接口参数错误"))
	}
	fileInfo, err := models.GetRedisFileInfo(r.FileId)
	if err != nil {
		return utils.HTTPErrorHandler(c, err)
	}

	now := time.Now().Unix()
	if fileInfo.CreatedAt+fileInfo.Expire < now {
		return utils.HTTPErrorHandler(c, errors.New("上传任务已过期"))
	}

	if fileInfo.FileType != models.FileTypeInit {
		return utils.HTTPErrorHandler(c, errors.New("上传任务状态错误"))
	}
	if r.FileIndex > ((fileInfo.FileSize / fileInfo.ChunkSize) + 1) {
		return utils.HTTPErrorHandler(c, errors.New("文件切片索引错误"))
	}

	if r.FileSlice.Size > fileInfo.ChunkSize {
		return utils.HTTPErrorHandler(c, errors.New("文件切片大小错误"))
	}

	// 打开文件
	file, err := r.FileSlice.Open()
	if err != nil {
		return utils.HTTPErrorHandler(c, err)
	}
	defer file.Close()

	uploadPath, err := utils.GetUploadDirPath()
	if err != nil {
		return utils.HTTPErrorHandler(c, err)
	}

	if _, err := services.CreateFileSlice(r.FileId, uploadPath, file, r.FileIndex); err != nil {
		return utils.HTTPErrorHandler(c, err)
	}

	return utils.HTTPSuccessHandler(c, map[string]any{
		"message": "成功上传",
	})
}

type FinishUploadTaskProps struct {
	FileId string `json:"id"`
}

func FinishUploadTask(c *echo.Context) error {
	r := new(FinishUploadTaskProps)
	if err := c.Bind(r); err != nil {
		return utils.HTTPErrorHandler(c, err)
	}

	if r.FileId == "" {
		return utils.HTTPErrorHandler(c, errors.New("文件ID不能为空"))
	}

	fileInfo, err := models.GetRedisFileInfo(r.FileId)
	if err != nil {
		return utils.HTTPErrorHandler(c, err)
	}

	if fileInfo.FileType != models.FileTypeInit {
		return utils.HTTPErrorHandler(c, errors.New("上传任务状态错误"))
	}

	now := time.Now().Unix()
	if fileInfo.CreatedAt+fileInfo.Expire < now {
		return utils.HTTPErrorHandler(c, errors.New("上传任务已过期"))
	}

	uploadPath, err := utils.GetUploadDirPath()
	if err != nil {
		return utils.HTTPErrorHandler(c, err)
	}

	fileSliceList, err := services.GetFileSliceList(r.FileId, uploadPath)
	if err != nil {
		return utils.HTTPErrorHandler(c, err)
	}

	if len(fileSliceList) != int(math.Ceil(float64(fileInfo.FileSize)/float64(fileInfo.ChunkSize))) {
		return utils.HTTPErrorHandler(c, errors.New("文件切片不完整"))
	}

	// 最终合并后的文件路径
	mergeFilePath, err := services.MergeFileSlices(r.FileId, uploadPath)
	if err != nil {
		return utils.HTTPErrorHandler(c, err)
	}

	// 计算文件MD5
	file, err := os.Open(mergeFilePath)
	if err != nil {
		file.Close()
		os.Remove(mergeFilePath)
		return utils.HTTPErrorHandler(c, err)
	}

	file_hash, err := utils.GetFileMd5(file)

	if err != nil {
		file.Close()
		os.Remove(mergeFilePath)
		return utils.HTTPErrorHandler(c, err)
	}

	if file_hash != fileInfo.FileHash {
		file.Close()
		os.Remove(mergeFilePath)
		return utils.HTTPErrorHandler(c, errors.New("文件MD5不一致"))
	}
	defer file.Close()
	// 更新文件信息
	models.SetRedisFileInfo(r.FileId, models.RedisFileInfo{
		FileType: models.FileTypeUpload,
	})
	// 统计
	currentDate := time.Now().Format("2006-01-02")
	statData, _ := models.GetRedisStat(currentDate)
	if statData == nil {
		statData = &models.StatData{
			FileSize:    0,
			FileNum:     0,
			ShareNum:    0,
			DownloadNum: 0,
		}
	}
	statData.FileSize += fileInfo.FileSize
	statData.FileNum += 1
	models.SetRedisStat(currentDate, *statData)

	return utils.HTTPSuccessHandler(c, map[string]any{
		"size":      fileInfo.FileSize,
		"mime_type": fileInfo.MimeType,
		"hash":      fileInfo.FileHash,
		"type":      models.FileTypeUpload,
		"id":        r.FileId,
	})
}