153 lines
3.5 KiB
Go
153 lines
3.5 KiB
Go
package storage
|
|
|
|
import (
|
|
"crypto/md5"
|
|
"encoding/hex"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"mime/multipart"
|
|
"os"
|
|
"path/filepath"
|
|
"time"
|
|
|
|
"github.com/gofiber/fiber/v3"
|
|
)
|
|
|
|
type Uploader struct {
|
|
tmpDir string
|
|
chunkPath string
|
|
fileName string
|
|
chunkNumber int
|
|
totalChunks int
|
|
fileMD5 string
|
|
|
|
dst string
|
|
ext string
|
|
finalPath string
|
|
}
|
|
|
|
type UploadedFile struct {
|
|
Hash string `json:"hash"`
|
|
Name string `json:"name"`
|
|
Size int64 `json:"size"`
|
|
MimeType string `json:"type"`
|
|
Path string `json:"path"`
|
|
Preview string `json:"preview"`
|
|
}
|
|
|
|
func NewUploader(fileName string, chunkNumber, totalChunks int, fileMD5 string) (*Uploader, error) {
|
|
// 使用MD5创建唯一的临时目录
|
|
tempDir := filepath.Join(os.TempDir(), fileMD5)
|
|
if err := os.MkdirAll(tempDir, 0o755); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &Uploader{
|
|
tmpDir: filepath.Join(os.TempDir(), fileMD5),
|
|
chunkPath: filepath.Join(os.TempDir(), fileMD5, fmt.Sprintf("chunk_%d", chunkNumber)),
|
|
fileName: fileName,
|
|
chunkNumber: chunkNumber,
|
|
totalChunks: totalChunks,
|
|
fileMD5: fileMD5,
|
|
ext: filepath.Ext(fileName),
|
|
finalPath: filepath.Join(os.TempDir(), fileMD5+filepath.Ext(fileName)),
|
|
}, nil
|
|
}
|
|
|
|
func (up *Uploader) Save(ctx fiber.Ctx, file *multipart.FileHeader) (*UploadedFile, error) {
|
|
if up.chunkNumber != up.totalChunks-1 {
|
|
return nil, ctx.SaveFile(file, up.chunkPath)
|
|
}
|
|
|
|
// 如果是最后一个分片
|
|
// 生成唯一的文件存储路径
|
|
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)
|
|
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,
|
|
Name: up.fileName,
|
|
Path: up.finalPath,
|
|
Size: totalSize,
|
|
MimeType: file.Header.Get("Content-Type"),
|
|
}, 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)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
totalSize += info.Size()
|
|
}
|
|
return totalSize, nil
|
|
}
|
|
|
|
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 := finalFile.Write(chunk); err != nil {
|
|
return 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
|
|
}
|