package upload import ( "context" "fmt" "io" "mime/multipart" "os" "path/filepath" "strings" "sync" "github.com/google/uuid" ) var ( // AllowedFileTypes 文件类型白名单 AllowedFileTypes = map[string]bool{ "pdf": true, "doc": true, "docx": true, "xls": true, "xlsx": true, "ppt": true, "pptx": true, "txt": true, "jpg": true, "jpeg": true, "png": true, "bmp": true, "gif": true, "zip": true, "rar": true, } ) type FileUploader struct { basePath string chunkPath string maxSize int64 mu sync.RWMutex chunks map[string]*ChunkInfo } type ChunkInfo struct { FileName string TotalParts int Parts map[int][]byte CreatedAt int64 } func NewFileUploader(basePath string, chunkPath string, maxSize int64) *FileUploader { uploader := &FileUploader{ basePath: basePath, chunkPath: chunkPath, maxSize: maxSize, chunks: make(map[string]*ChunkInfo), } // Ensure directories exist os.MkdirAll(basePath, 0755) os.MkdirAll(chunkPath, 0755) return uploader } func (u *FileUploader) getExt(filename string) string { ext := strings.ToLower(filepath.Ext(filename)) return strings.TrimPrefix(ext, ".") } func (u *FileUploader) ValidateFileType(filename string) bool { ext := u.getExt(filename) return AllowedFileTypes[ext] } func (u *FileUploader) GenerateFilePath(filename string) string { ext := u.getExt(filename) newName := fmt.Sprintf("%s.%s", uuid.New().String(), ext) // 按日期分目录 // dateDir := time.Now().Format("2006/01/02") // return filepath.Join(dateDir, newName) return newName } func (u *FileUploader) Upload(ctx context.Context, file *multipart.FileHeader) (string, error) { if u.maxSize > 0 && file.Size > u.maxSize { return "", fmt.Errorf("file size exceeds limit: %d > %d", file.Size, u.maxSize) } if !u.ValidateFileType(file.Filename) { return "", fmt.Errorf("file type not allowed: %s", file.Filename) } src, err := file.Open() if err != nil { return "", fmt.Errorf("failed to open file: %w", err) } defer src.Close() filePath := u.GenerateFilePath(file.Filename) fullPath := filepath.Join(u.basePath, filePath) // Ensure parent directory exists if err := os.MkdirAll(filepath.Dir(fullPath), 0755); err != nil { return "", fmt.Errorf("failed to create directory: %w", err) } dst, err := os.Create(fullPath) if err != nil { return "", fmt.Errorf("failed to create file: %w", err) } defer dst.Close() if _, err := io.Copy(dst, src); err != nil { return "", fmt.Errorf("failed to write file: %w", err) } return filePath, nil } // Chunk upload support func (u *FileUploader) InitChunkUpload(chunkID, fileName string, totalParts int) error { u.mu.Lock() defer u.mu.Unlock() u.chunks[chunkID] = &ChunkInfo{ FileName: fileName, TotalParts: totalParts, Parts: make(map[int][]byte), } return nil } func (u *FileUploader) SaveChunk(chunkID string, partIndex int, data []byte) error { u.mu.Lock() defer u.mu.Unlock() chunk, ok := u.chunks[chunkID] if !ok { return fmt.Errorf("chunk upload not initialized: %s", chunkID) } if partIndex < 0 || partIndex >= chunk.TotalParts { return fmt.Errorf("invalid part index: %d", partIndex) } chunk.Parts[partIndex] = data return nil } func (u *FileUploader) CompleteChunkUpload(ctx context.Context, chunkID string) (string, error) { u.mu.Lock() chunk, ok := u.chunks[chunkID] if !ok { u.mu.Unlock() return "", fmt.Errorf("chunk upload not found: %s", chunkID) } if len(chunk.Parts) != chunk.TotalParts { u.mu.Unlock() return "", fmt.Errorf("incomplete chunks: %d/%d", len(chunk.Parts), chunk.TotalParts) } // Delete chunk info early to allow other operations delete(u.chunks, chunkID) u.mu.Unlock() filePath := u.GenerateFilePath(chunk.FileName) fullPath := filepath.Join(u.basePath, filePath) if err := os.MkdirAll(filepath.Dir(fullPath), 0755); err != nil { return "", fmt.Errorf("failed to create directory: %w", err) } f, err := os.Create(fullPath) if err != nil { return "", fmt.Errorf("failed to create file: %w", err) } defer f.Close() for i := 0; i < chunk.TotalParts; i++ { if _, err := f.Write(chunk.Parts[i]); err != nil { return "", fmt.Errorf("failed to write chunk %d: %w", i, err) } } return filePath, nil } func (u *FileUploader) DeleteFile(filePath string) error { fullPath := filepath.Join(u.basePath, filePath) return os.Remove(fullPath) } func (u *FileUploader) GetFilePath(filePath string) string { return filepath.Join(u.basePath, filePath) }