Initial commit: file_manager package with local and S3 support

2026-01-17 15:19:50 +08:00
commit f7160ce416
10 changed files with 1526 additions and 0 deletions
--- a/controller/file_manager/local.go
+++ b/controller/file_manager/local.go
@@ -0,0 +1,465 @@
+package file_manager
+
+import (
+	"context"
+	"crypto/rand"
+	"crypto/sha256"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"sync"
+	"time"
+)
+
+type HashMismatchError struct {
+	Expected string
+	Actual   string
+}
+
+func (e *HashMismatchError) Error() string {
+	return fmt.Sprintf("SHA256 hash mismatch: expected %s, got %s", e.Expected, e.Actual)
+}
+
+func NewHashMismatchError(expected, actual string) error {
+	return &HashMismatchError{
+		Expected: expected,
+		Actual:   actual,
+	}
+}
+
+type FileNotReadyError struct {
+	Code         string
+	CurrentSize  int64
+	ExpectedSize int64
+}
+
+func (e *FileNotReadyError) Error() string {
+	return fmt.Sprintf("file %s not ready: current size %d, expected %d", e.Code, e.CurrentSize, e.ExpectedSize)
+}
+
+func NewFileNotReadyError(code string, currentSize, expectedSize int64) error {
+	return &FileNotReadyError{
+		Code:         code,
+		CurrentSize:  currentSize,
+		ExpectedSize: expectedSize,
+	}
+}
+
+type DeleteError struct {
+	Code    string
+	Message string
+	Errors  []string
+}
+
+func (e *DeleteError) Error() string {
+	if len(e.Errors) == 0 {
+		return e.Message
+	}
+	return fmt.Sprintf("%s: %v", e.Message, e.Errors)
+}
+
+type LocalFileManager struct {
+	dir         string
+	fileHandles sync.Map
+	ctx         context.Context
+	cancel      context.CancelFunc
+	timeout     time.Duration
+	expire      time.Duration
+	verifyHash  bool
+	mu          sync.RWMutex
+}
+
+func NewLocalFileManager(opts *option) *LocalFileManager {
+	defaultDir := "./uploads"
+	lfm := &LocalFileManager{
+		dir:        defaultDir,
+		timeout:    time.Minute,
+		expire:     24 * time.Hour,
+		verifyHash: true,
+	}
+
+	if opts != nil {
+		if opts.dir != nil {
+			lfm.dir = *opts.dir
+		}
+		if opts.timeout > 0 {
+			lfm.timeout = opts.timeout
+		}
+		if opts.expire > 0 {
+			lfm.expire = opts.expire
+		}
+	}
+
+	lfm.ctx, lfm.cancel = context.WithCancel(context.Background())
+
+	go lfm.startCleaner()
+
+	return lfm
+}
+
+func (lfm *LocalFileManager) manifestPath(code string) string {
+	return filepath.Join(lfm.dir, "."+code+".manifest")
+}
+
+func (lfm *LocalFileManager) readManifest(code string) (*FileInfo, error) {
+	manifestPath := lfm.manifestPath(code)
+	data, err := os.ReadFile(manifestPath)
+	if err != nil {
+		return nil, err
+	}
+	var fileInfo FileInfo
+	if err := json.Unmarshal(data, &fileInfo); err != nil {
+		return nil, err
+	}
+	return &fileInfo, nil
+}
+
+func (lfm *LocalFileManager) writeManifest(code string, fileInfo *FileInfo) error {
+	manifestPath := lfm.manifestPath(code)
+	data, err := json.Marshal(fileInfo)
+	if err != nil {
+		return err
+	}
+	return os.WriteFile(manifestPath, data, 0644)
+}
+
+func (lfm *LocalFileManager) deleteManifest(code string) error {
+	return os.Remove(lfm.manifestPath(code))
+}
+
+func (lfm *LocalFileManager) scanManifests() ([]string, error) {
+	var codes []string
+	entries, err := os.ReadDir(lfm.dir)
+	if err != nil {
+		return nil, err
+	}
+	for _, entry := range entries {
+		if entry.IsDir() {
+			continue
+		}
+		name := entry.Name()
+		if len(name) > 10 && name[:1] == "." && name[len(name)-9:] == ".manifest" {
+			code := name[1 : len(name)-9]
+			codes = append(codes, code)
+		}
+	}
+	return codes, nil
+}
+
+func (lfm *LocalFileManager) startCleaner() {
+	ticker := time.NewTicker(time.Minute)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-lfm.ctx.Done():
+			return
+		case <-ticker.C:
+			lfm.cleanup()
+		}
+	}
+}
+
+func (lfm *LocalFileManager) cleanup() {
+	codes, err := lfm.scanManifests()
+	if err != nil {
+		return
+	}
+	now := time.Now()
+	for _, code := range codes {
+		fileInfo, err := lfm.readManifest(code)
+		if err != nil {
+			continue
+		}
+		if fileInfo.Complete {
+			if now.Sub(fileInfo.CreateTime) > lfm.expire {
+				lfm.cleanupFile(code, fileInfo)
+			}
+		} else {
+			if now.Sub(fileInfo.CreateTime) > lfm.timeout {
+				lfm.cleanupFile(code, fileInfo)
+			}
+		}
+	}
+}
+
+func (lfm *LocalFileManager) cleanupFile(code string, fileInfo *FileInfo) {
+	if fileHandleValue, exists := lfm.fileHandles.LoadAndDelete(code); exists {
+		file := fileHandleValue.(*os.File)
+		file.Close()
+	}
+	if fileInfo.Path != "" {
+		os.Remove(fileInfo.Path)
+	}
+	lfm.deleteManifest(code)
+}
+
+func generateRandomCode(length int) (string, error) {
+	bytes := make([]byte, length/2)
+	if _, err := rand.Read(bytes); err != nil {
+		return "", err
+	}
+	return hex.EncodeToString(bytes), nil
+}
+
+func (lfm *LocalFileManager) CloseManager() {
+	lfm.cancel()
+	lfm.fileHandles.Range(func(key, value any) bool {
+		code := key.(string)
+		file := value.(*os.File)
+		file.Close()
+		lfm.fileHandles.Delete(code)
+		return true
+	})
+}
+
+func (lfm *LocalFileManager) Create(ctx context.Context, filename string, size int64, sha256 string) (*CreateResult, error) {
+	code, err := generateRandomCode(16)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := os.MkdirAll(lfm.dir, 0755); err != nil {
+		return nil, err
+	}
+
+	filePath := filepath.Join(lfm.dir, code)
+
+	fileInfo := &FileInfo{
+		Filename:   filename,
+		Size:       size,
+		SHA256:     sha256,
+		Path:       filePath,
+		CreateTime: time.Now(),
+		Complete:   false,
+	}
+
+	if err := lfm.writeManifest(code, fileInfo); err != nil {
+		return nil, err
+	}
+
+	file, err := os.Create(filePath)
+	if err != nil {
+		lfm.deleteManifest(code)
+		return nil, err
+	}
+
+	lfm.fileHandles.Store(code, file)
+
+	return &CreateResult{
+		Code:   code,
+		SHA256: sha256,
+	}, nil
+}
+
+func (lfm *LocalFileManager) Upload(ctx context.Context, code string, start int64, end int64, reader io.Reader) (int64, int64, error) {
+	fileHandleValue, exists := lfm.fileHandles.Load(code)
+	if !exists {
+		return 0, 0, os.ErrNotExist
+	}
+
+	file := fileHandleValue.(*os.File)
+
+	currentInfo, err := file.Stat()
+	if err != nil {
+		return 0, 0, err
+	}
+
+	currentSize := currentInfo.Size()
+
+	if start != currentSize {
+		return currentSize, 0, nil
+	}
+
+	written, err := io.Copy(file, reader)
+	if err != nil {
+		return currentSize, 0, err
+	}
+
+	newInfo, err := file.Stat()
+	if err != nil {
+		return currentSize, written, err
+	}
+
+	totalSize := newInfo.Size()
+
+	fileInfo, err := lfm.readManifest(code)
+	if err != nil {
+		return totalSize, written, os.ErrNotExist
+	}
+
+	if end > 0 && totalSize == end && totalSize == fileInfo.Size {
+		if fileInfo.SHA256 != "" {
+			if _, err := file.Seek(0, 0); err != nil {
+				return totalSize, written, err
+			}
+
+			hasher := sha256.New()
+			if _, err := io.Copy(hasher, file); err != nil {
+				return totalSize, written, err
+			}
+
+			calculatedHash := hex.EncodeToString(hasher.Sum(nil))
+
+			if calculatedHash != fileInfo.SHA256 {
+				return totalSize, written, NewHashMismatchError(fileInfo.SHA256, calculatedHash)
+			}
+		}
+
+		if fileInfo.SHA256 == "" {
+			if _, err := file.Seek(0, 0); err != nil {
+				return totalSize, written, err
+			}
+
+			hasher := sha256.New()
+			if _, err := io.Copy(hasher, file); err != nil {
+				return totalSize, written, err
+			}
+
+			calculatedHash := hex.EncodeToString(hasher.Sum(nil))
+			fileInfo.SHA256 = calculatedHash
+		}
+
+		fileInfo.Complete = true
+		if err := lfm.writeManifest(code, fileInfo); err != nil {
+			return totalSize, written, err
+		}
+
+		file.Close()
+		lfm.fileHandles.Delete(code)
+	}
+
+	return totalSize, written, nil
+}
+
+func (lfm *LocalFileManager) verifyFileHash(code string, fileInfo *FileInfo) error {
+	if fileInfo.SHA256 == "" {
+		return nil
+	}
+
+	file, err := os.Open(fileInfo.Path)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	hasher := sha256.New()
+	if _, err := io.Copy(hasher, file); err != nil {
+		return err
+	}
+
+	calculatedHash := hex.EncodeToString(hasher.Sum(nil))
+
+	if calculatedHash != fileInfo.SHA256 {
+		lfm.cleanupFile(code, fileInfo)
+		return NewHashMismatchError(fileInfo.SHA256, calculatedHash)
+	}
+
+	return nil
+}
+
+func (lfm *LocalFileManager) Get(ctx context.Context, code string) ([]byte, error) {
+	fileInfo, err := lfm.readManifest(code)
+	if err != nil {
+		return nil, os.ErrNotExist
+	}
+
+	fileHandleValue, handleExists := lfm.fileHandles.Load(code)
+	if handleExists {
+		file := fileHandleValue.(*os.File)
+		fileInfoStat, err := file.Stat()
+		if err != nil {
+			return nil, err
+		}
+
+		if fileInfoStat.Size() != fileInfo.Size {
+			return nil, NewFileNotReadyError(code, fileInfoStat.Size(), fileInfo.Size)
+		}
+
+		if file, exists := lfm.fileHandles.LoadAndDelete(code); exists {
+			f := file.(*os.File)
+			f.Close()
+		}
+	}
+
+	if err := lfm.verifyFileHash(code, fileInfo); err != nil {
+		return nil, err
+	}
+
+	return os.ReadFile(fileInfo.Path)
+}
+
+func (lfm *LocalFileManager) GetInfo(ctx context.Context, code string) (*FileInfo, error) {
+	fileInfo, err := lfm.readManifest(code)
+	if err != nil {
+		return nil, os.ErrNotExist
+	}
+
+	fileHandleValue, handleExists := lfm.fileHandles.Load(code)
+	if handleExists {
+		file := fileHandleValue.(*os.File)
+		fileInfoStat, err := file.Stat()
+		if err != nil {
+			return nil, err
+		}
+
+		if fileInfoStat.Size() != fileInfo.Size {
+			return nil, NewFileNotReadyError(code, fileInfoStat.Size(), fileInfo.Size)
+		}
+
+		if file, exists := lfm.fileHandles.LoadAndDelete(code); exists {
+			f := file.(*os.File)
+			f.Close()
+		}
+	}
+
+	if err := lfm.verifyFileHash(code, fileInfo); err != nil {
+		return nil, err
+	}
+
+	return fileInfo, nil
+}
+
+func (lfm *LocalFileManager) Delete(ctx context.Context, code string) error {
+	fileInfo, err := lfm.readManifest(code)
+	if err != nil {
+		return os.ErrNotExist
+	}
+
+	var errors []string
+
+	if fileHandleValue, handleExists := lfm.fileHandles.LoadAndDelete(code); handleExists {
+		file := fileHandleValue.(*os.File)
+		if err := file.Close(); err != nil {
+			errors = append(errors, fmt.Sprintf("failed to close file handle: %v", err))
+		}
+	}
+
+	if err := os.Remove(fileInfo.Path); err != nil && !os.IsNotExist(err) {
+		errors = append(errors, fmt.Sprintf("failed to delete file: %v", err))
+	}
+
+	if err := lfm.deleteManifest(code); err != nil && !os.IsNotExist(err) {
+		errors = append(errors, fmt.Sprintf("failed to delete manifest: %v", err))
+	}
+
+	if len(errors) > 0 {
+		return &DeleteError{
+			Code:    code,
+			Message: "partial file deletion completed with errors",
+			Errors:  errors,
+		}
+	}
+
+	return nil
+}
+
+func (lfm *LocalFileManager) Close(code string) error {
+	if fileHandleValue, exists := lfm.fileHandles.LoadAndDelete(code); exists {
+		file := fileHandleValue.(*os.File)
+		return file.Close()
+	}
+	return nil
+}