kra-new/internal/data/storage/local.go

package storage

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io"
	"mime"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"kra/internal/biz"
	"kra/internal/conf"
)

type fileStorage struct {
	root      string
	urlPrefix string
}

func New(config *conf.AdminBackend) (biz.FileStorage, error) {
	storageType := "local"
	if config != nil && config.Storage != nil && config.Storage.Type != "" {
		storageType = strings.ToLower(config.Storage.Type)
	}
	if storageType == "qiniu" {
		return newQiniuStorage(config.Storage.Qiniu)
	}
	if storageType != "local" {
		var object *conf.AdminBackend_ObjectStore
		switch storageType {
		case "aliyun-oss":
			return newAliyunStorage(config.Storage.AliyunOss)
		case "huawei-obs":
			return newHuaweiStorage(config.Storage.HuaweiObs)
		case "tencent-cos":
			return newTencentStorage(config.Storage.TencentCos)
		case "aws-s3":
			return newAWSStorage(storageType, config.Storage.AwsS3)
		case "cloudflare-r2":
			return newAWSStorage(storageType, config.Storage.CloudflareR2)
		case "minio":
			object = config.Storage.Minio
		default:
			return nil, fmt.Errorf("unsupported storage type %q", storageType)
		}
		return newS3Storage(storageType, object)
	}
	root, prefix := "uploads/file", "uploads/file"
	if config != nil && config.Local != nil {
		if config.Local.StorePath != "" {
			root = config.Local.StorePath
		}
		if config.Local.PathPrefix != "" {
			prefix = config.Local.PathPrefix
		}
	}
	root = filepath.Clean(root)
	if err := os.MkdirAll(root, 0o755); err != nil {
		return nil, fmt.Errorf("create upload directory: %w", err)
	}
	return &fileStorage{root: root, urlPrefix: "/" + strings.Trim(prefix, "/")}, nil
}

func composeFiles(ctx context.Context, storage biz.FileStorage, names []string, destination string) (*biz.StoredFile, string, error) {
	reader, writer := io.Pipe()
	hash := md5.New()
	errCh := make(chan error, 1)
	go func() {
		defer writer.Close()
		for _, name := range names {
			if ctx.Err() != nil {
				errCh <- ctx.Err()
				return
			}
			file, err := storage.Open(ctx, name)
			if err != nil {
				errCh <- err
				return
			}
			_, err = io.Copy(io.MultiWriter(writer, hash), file)
			_ = file.Close()
			if err != nil {
				errCh <- err
				return
			}
		}
		errCh <- nil
	}()
	stored, err := storage.Put(ctx, destination, reader)
	composeErr := <-errCh
	if err != nil {
		return nil, "", err
	}
	if composeErr != nil {
		_ = storage.Delete(ctx, destination)
		return nil, "", composeErr
	}
	return stored, hex.EncodeToString(hash.Sum(nil)), nil
}

func (s *fileStorage) resolve(name string) (string, error) {
	clean := filepath.Clean(strings.TrimPrefix(name, "/"))
	if clean == "." || filepath.IsAbs(clean) || clean == ".." || strings.HasPrefix(clean, ".."+string(filepath.Separator)) {
		return "", fmt.Errorf("invalid storage path")
	}
	return filepath.Join(s.root, clean), nil
}

func (s *fileStorage) Put(ctx context.Context, name string, reader io.Reader) (*biz.StoredFile, error) {
	path, err := s.resolve(name)
	if err != nil {
		return nil, err
	}
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return nil, err
	}
	file, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644)
	if err != nil {
		return nil, err
	}
	size, copyErr := io.Copy(file, reader)
	closeErr := file.Close()
	if copyErr != nil {
		return nil, copyErr
	}
	if closeErr != nil {
		return nil, closeErr
	}
	return &biz.StoredFile{Name: filepath.Base(name), Path: name, URL: s.urlPrefix + "/" + strings.TrimPrefix(filepath.ToSlash(name), "/"), Size: size}, nil
}

func (s *fileStorage) Open(ctx context.Context, name string) (io.ReadCloser, error) {
	path, err := s.resolve(name)
	if err != nil {
		return nil, err
	}
	return os.Open(path)
}

func (s *fileStorage) Delete(ctx context.Context, name string) error {
	path, err := s.resolve(name)
	if err != nil {
		return err
	}
	if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
		return err
	}
	return nil
}

func (s *fileStorage) Compose(ctx context.Context, names []string, destination string) (*biz.StoredFile, string, error) {
	return composeFiles(ctx, s, names, destination)
}
func (s *fileStorage) DeletePrefix(ctx context.Context, prefix string) error {
	path, err := s.resolve(strings.TrimSuffix(prefix, "/") + "/placeholder")
	if err != nil {
		return err
	}
	dir := filepath.Dir(path)
	if !strings.HasPrefix(dir, s.root+string(filepath.Separator)) {
		return fmt.Errorf("invalid storage path")
	}
	if err = os.RemoveAll(dir); err != nil {
		return err
	}
	return nil
}
func (s *fileStorage) List(ctx context.Context, prefix, cursor string, limit int) ([]*biz.StoredFile, string, bool, error) {
	if limit < 1 || limit > 1000 {
		limit = 100
	}
	root, err := s.resolve(strings.TrimSuffix(prefix, "/") + "/placeholder")
	if err != nil {
		return nil, "", false, err
	}
	walkRoot := filepath.Dir(root)
	items := []*biz.StoredFile{}
	err = filepath.WalkDir(walkRoot, func(path string, entry os.DirEntry, walkErr error) error {
		if walkErr != nil {
			if os.IsNotExist(walkErr) {
				return nil
			}
			return walkErr
		}
		if ctx.Err() != nil {
			return ctx.Err()
		}
		if entry.IsDir() || entry.Type()&os.ModeSymlink != 0 {
			return nil
		}
		relative, _ := filepath.Rel(s.root, path)
		relative = filepath.ToSlash(relative)
		if relative <= cursor {
			return nil
		}
		info, infoErr := entry.Info()
		if infoErr != nil {
			return infoErr
		}
		items = append(items, &biz.StoredFile{Name: entry.Name(), Path: relative, URL: s.urlPrefix + "/" + relative, Size: info.Size(), LastModified: info.ModTime(), ContentType: mime.TypeByExtension(filepath.Ext(entry.Name()))})
		return nil
	})
	if err != nil {
		return nil, "", false, err
	}
	sort.Slice(items, func(i, j int) bool { return items[i].Path < items[j].Path })
	hasMore := len(items) > limit
	if hasMore {
		items = items[:limit]
	}
	next := ""
	if len(items) > 0 {
		next = items[len(items)-1].Path
	}
	return items, next, hasMore, nil
}