upkg/controller/file_manager/local.go

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
}