backend/cloudprint-backend/pkg/upload/upload.go

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)
}