68d33d58创建于 2025年12月18日历史提交
package storage

import (
	"context"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/goodrain/rainbond/config/configs"
	"github.com/goodrain/rainbond/event"
	"github.com/sirupsen/logrus"
	"mime/multipart"
	"net/http"
)

// StorageComponent -
type StorageComponent struct {
	StorageCli    InterfaceStorage
	storageConfig *configs.StorageConfig
}

var defaultStorageComponent *StorageComponent

// New -
func New() *StorageComponent {
	storageConfig := configs.Default().StorageConfig

	defaultStorageComponent = &StorageComponent{
		storageConfig: storageConfig,
	}
	return defaultStorageComponent
}

// Start -
func (s *StorageComponent) Start(ctx context.Context) error {
	var storageCli InterfaceStorage
	logrus.Infof("create s3 client %v,----%v,----%v", s.storageConfig.StorageType, s.storageConfig.S3AccessKeyID, s.storageConfig.S3SecretAccessKey)
	if s.storageConfig.StorageType == "s3" {
		sess, err := session.NewSession(&aws.Config{
			Endpoint:         aws.String(s.storageConfig.S3Endpoint),
			Region:           aws.String("rainbond"), // 可以根据需要选择区域
			Credentials:      credentials.NewStaticCredentials(s.storageConfig.S3AccessKeyID, s.storageConfig.S3SecretAccessKey, ""),
			S3ForcePathStyle: aws.Bool(true), // 使用路径风格
		})
		if err != nil {
			logrus.Errorf("failed to create session: %v", err)
			return err
		}
		s3Client := s3.New(sess)
		s3Storage := &S3Storage{s3Client: s3Client}

		// API 启动时主动初始化 bucket 生命周期策略
		logrus.Info("Initializing S3 bucket lifecycle policies on startup...")
		if err := s3Storage.InitBucketLifecycle(); err != nil {
			logrus.Warnf("Failed to initialize bucket lifecycle policies: %v (non-fatal, continuing startup)", err)
			// 不返回错误,允许 API 继续启动,生命周期策略会在后续操作中自动创建
		} else {
			logrus.Info("Successfully initialized S3 bucket lifecycle policies")
		}

		storageCli = s3Storage
	} else {
		storageCli = &LocalStorage{}
	}
	s.StorageCli = storageCli
	return nil
}

// CloseHandle -
func (s *StorageComponent) CloseHandle() {
}

// Default -
func Default() *StorageComponent {
	return defaultStorageComponent
}

type InterfaceStorage interface {
	MkdirAll(path string) error
	Unzip(archive, target string, currentDirectory bool) error
	ReadDir(dirName string) ([]string, error)
	ServeFile(w http.ResponseWriter, r *http.Request, filePath string)
	SaveFile(fileName string, reader multipart.File) error
	UploadFileToFile(src string, dst string, logger event.Logger) error
	DownloadDirToDir(srcDir, dstDir string) error
	DownloadFileToDir(srcFile, dstDir string) error
	// ReadFile reads a file directly from storage and returns a reader
	ReadFile(filePath string) (ReadCloser, error)

	// 分片上传相关方法
	SaveChunk(sessionID string, chunkIndex int, reader multipart.File) (string, error)
	MergeChunks(sessionID string, outputPath string, totalChunks int) error
	ChunkExists(sessionID string, chunkIndex int) bool
	CleanupChunks(sessionID string) error
	GetChunkDir(sessionID string) string
}

type ReadCloser interface {
	Read(p []byte) (n int, err error)
	Close() error
}

type SrcFile interface {
	Read([]byte) (int, error)
}

// DstFile 目标文件接口
type DstFile interface {
	Write([]byte) (int, error)
}