feat: fix upload file

This commit is contained in:
Rogee
2025-01-17 15:48:53 +08:00
parent b5583bb34a
commit 5b511a5ea1
5 changed files with 99 additions and 76 deletions

View File

@@ -1,6 +1,7 @@
package storage
import (
"bytes"
"crypto/md5"
"encoding/hex"
"errors"
@@ -12,6 +13,7 @@ import (
"time"
"github.com/gofiber/fiber/v3"
"github.com/spf13/afero"
)
type Uploader struct {
@@ -22,7 +24,6 @@ type Uploader struct {
totalChunks int
fileMD5 string
dst string
ext string
finalPath string
}
@@ -52,49 +53,27 @@ func NewUploader(fileName string, chunkNumber, totalChunks int, fileMD5 string)
totalChunks: totalChunks,
fileMD5: fileMD5,
ext: filepath.Ext(fileName),
finalPath: filepath.Join(os.TempDir(), fileMD5+filepath.Ext(fileName)),
finalPath: filepath.Join("uploads", time.Now().Format("2006/01/02"), fileMD5+filepath.Ext(fileName)),
}, nil
}
func (up *Uploader) Save(ctx fiber.Ctx, file *multipart.FileHeader) (*UploadedFile, error) {
func (up *Uploader) Save(ctx fiber.Ctx, fs afero.Fs, file *multipart.FileHeader) (*UploadedFile, error) {
if up.chunkNumber != up.totalChunks-1 {
return nil, ctx.SaveFile(file, up.chunkPath)
}
defer os.RemoveAll(up.tmpDir)
// 如果是最后一个分片
// 生成唯一的文件存储路径
storageDir := filepath.Join(up.dst, time.Now().Format("2006/01/02"))
if err := os.MkdirAll(storageDir, 0o755); err != nil {
os.RemoveAll(filepath.Join(os.TempDir(), up.fileMD5))
return nil, err
}
// 计算所有分片的实际大小总和
totalSize, err := calculateTotalSize(up.tmpDir, up.totalChunks)
if err != nil {
os.RemoveAll(up.tmpDir)
return nil, fmt.Errorf("计算文件大小失败: %w", err)
}
// 合并文件
if err := combineChunks(up.tmpDir, up.finalPath, up.totalChunks); err != nil {
os.RemoveAll(up.tmpDir)
totalSize, err := up.combineChunks(fs)
if err != nil {
return nil, fmt.Errorf("合并文件失败: %w", err)
}
// 验证MD5
calculatedMD5, err := calculateFileMD5(up.finalPath)
if err != nil || calculatedMD5 != up.fileMD5 {
os.RemoveAll(up.tmpDir)
os.Remove(up.finalPath)
return nil, errors.New("文件MD5验证失败")
}
// 清理临时目录
os.RemoveAll(up.tmpDir)
return &UploadedFile{
Hash: calculatedMD5,
Hash: up.fileMD5,
Name: up.fileName,
Path: up.finalPath,
Size: totalSize,
@@ -102,52 +81,41 @@ func (up *Uploader) Save(ctx fiber.Ctx, file *multipart.FileHeader) (*UploadedFi
}, nil
}
// 计算所有分片的实际大小总和
func calculateTotalSize(tempDir string, totalChunks int) (int64, error) {
var totalSize int64
for i := 0; i < totalChunks; i++ {
chunkPath := filepath.Join(tempDir, fmt.Sprintf("chunk_%d", i))
info, err := os.Stat(chunkPath)
func (up *Uploader) combineChunks(fs afero.Fs) (int64, error) {
if err := fs.MkdirAll(filepath.Dir(up.finalPath), os.ModePerm); err != nil {
return 0, err
}
f, err := fs.Create(up.finalPath)
if err != nil {
return 0, err
}
defer f.Close()
hash := md5.New()
size := int64(0)
for i := 0; i < up.totalChunks; i++ {
chunkPath := fmt.Sprintf("%s/chunk_%d", up.tmpDir, i)
chunk, err := os.ReadFile(chunkPath)
if err != nil {
return 0, err
}
totalSize += info.Size()
}
return totalSize, nil
}
size += int64(len(chunk))
func combineChunks(tempDir, finalPath string, totalChunks int) error {
finalFile, err := os.Create(finalPath)
if err != nil {
return err
}
defer finalFile.Close()
for i := 0; i < totalChunks; i++ {
chunkPath := fmt.Sprintf("%s/chunk_%d", tempDir, i)
chunk, err := os.ReadFile(chunkPath)
if err != nil {
return err
if _, err := f.Write(chunk); err != nil {
return 0, err
}
if _, err := finalFile.Write(chunk); err != nil {
return err
if _, err := io.Copy(hash, bytes.NewBuffer(chunk)); err != nil {
return 0, err
}
}
return nil
}
func calculateFileMD5(filePath string) (string, error) {
file, err := os.Open(filePath)
if err != nil {
return "", err
}
defer file.Close()
hash := md5.New()
if _, err := io.Copy(hash, file); err != nil {
return "", err
}
return hex.EncodeToString(hash.Sum(nil)), nil
md5 := hex.EncodeToString(hash.Sum(nil))
if md5 != up.fileMD5 {
return 0, errors.New("文件MD5验证失败")
}
return size, nil
}