122 lines
2.6 KiB
Go
122 lines
2.6 KiB
Go
package storage
|
|
|
|
import (
|
|
"bytes"
|
|
"crypto/md5"
|
|
"encoding/hex"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"mime/multipart"
|
|
"os"
|
|
"path/filepath"
|
|
"time"
|
|
|
|
"github.com/gofiber/fiber/v3"
|
|
"github.com/spf13/afero"
|
|
)
|
|
|
|
type Uploader struct {
|
|
tmpDir string
|
|
chunkPath string
|
|
fileName string
|
|
chunkNumber int
|
|
totalChunks int
|
|
fileMD5 string
|
|
|
|
ext string
|
|
finalPath string
|
|
}
|
|
|
|
type UploadedFile struct {
|
|
ID int64 `json:"id"`
|
|
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("uploads", time.Now().Format("2006/01/02"), fileMD5+filepath.Ext(fileName)),
|
|
}, nil
|
|
}
|
|
|
|
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)
|
|
|
|
// 如果是最后一个分片
|
|
// 生成唯一的文件存储路径
|
|
|
|
// 合并文件
|
|
totalSize, err := up.combineChunks(fs)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("合并文件失败: %w", err)
|
|
}
|
|
|
|
return &UploadedFile{
|
|
Hash: up.fileMD5,
|
|
Name: up.fileName,
|
|
Path: up.finalPath,
|
|
Size: totalSize,
|
|
MimeType: file.Header.Get("Content-Type"),
|
|
}, nil
|
|
}
|
|
|
|
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
|
|
}
|
|
size += int64(len(chunk))
|
|
|
|
if _, err := f.Write(chunk); err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
if _, err := io.Copy(hash, bytes.NewBuffer(chunk)); err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
}
|
|
|
|
md5 := hex.EncodeToString(hash.Sum(nil))
|
|
if md5 != up.fileMD5 {
|
|
return 0, errors.New("文件MD5验证失败")
|
|
}
|
|
|
|
return size, nil
|
|
}
|