github.com/machinefi/w3bstream@v1.6.5-rc9.0.20240426031326-b8c7c4876e72/pkg/depends/conf/storage/storage.go (about)

     1  package storage
     2  
     3  import (
     4  	"os"
     5  	"path/filepath"
     6  
     7  	"github.com/pkg/errors"
     8  	"github.com/shirou/gopsutil/v3/disk"
     9  
    10  	"github.com/machinefi/w3bstream/pkg/depends/base/consts"
    11  	"github.com/machinefi/w3bstream/pkg/depends/base/types"
    12  )
    13  
    14  var (
    15  	ErrMissingConfigS3       = errors.New("missing config: s3")
    16  	ErrMissingConfigFS       = errors.New("missing config: fs")
    17  	ErrMissingConfigIPFS     = errors.New("missing config: ipfs")
    18  	ErrUnsupprtedStorageType = errors.New("unsupported storage type")
    19  	ErrEmptyContent          = errors.New("content empty")
    20  	ErrContentSizeExceeded   = errors.New("content size exceeded")
    21  	ErrDiskReservationLimit  = errors.New("disk reservation limit")
    22  )
    23  
    24  type StorageOperations interface {
    25  	Type() StorageType
    26  	Upload(key string, file []byte, chk ...HmacAlgType) error
    27  	Read(key string, chk ...HmacAlgType) (data []byte, sum []byte, err error)
    28  	Delete(key string) error
    29  }
    30  
    31  type StorageOperationsWithValidation interface {
    32  	Validate(data []byte, sum string, chk ...HmacAlgType) bool
    33  }
    34  
    35  type Storage struct {
    36  	Typ             StorageType
    37  	FilesizeLimit   int64
    38  	DiskReserve     int64
    39  	PromiscuousMode bool // PromiscuousMode if support multi storage type
    40  	TempDir         string
    41  
    42  	*S3      `env:"S3"`
    43  	*LocalFs `env:"Fs"`
    44  
    45  	op StorageOperations `env:"-"`
    46  }
    47  
    48  func (s *Storage) Name() string { return "Storage" }
    49  
    50  func (s *Storage) IsZero() bool {
    51  	return s.Typ == STORAGE_TYPE_UNKNOWN || s.S3 == nil && s.LocalFs == nil
    52  }
    53  
    54  func (s *Storage) SetDefault() {
    55  	if s.Typ == STORAGE_TYPE_UNKNOWN {
    56  		s.Typ = STORAGE_TYPE__FILESYSTEM
    57  	}
    58  	if s.FilesizeLimit == 0 {
    59  		s.FilesizeLimit = 1024 * 1024
    60  	}
    61  	if s.DiskReserve == 0 {
    62  		s.DiskReserve = 20 * 1024 * 1024
    63  	}
    64  }
    65  
    66  func (s *Storage) Init() error {
    67  	if s.TempDir == "" {
    68  		tmp := os.Getenv("TMPDIR")
    69  		if tmp == "" {
    70  			tmp = "/tmp"
    71  		}
    72  		service := os.Getenv(consts.EnvProjectName)
    73  		if service == "" {
    74  			service = "service"
    75  		}
    76  		s.TempDir = filepath.Join(tmp, service)
    77  	}
    78  	// overwrite default 'TMPDIR'
    79  	if err := os.Setenv("TMPDIR", s.TempDir); err != nil {
    80  		return err
    81  	}
    82  
    83  	switch s.Typ {
    84  	case STORAGE_TYPE_UNKNOWN, STORAGE_TYPE__FILESYSTEM:
    85  		if s.LocalFs == nil {
    86  			return ErrMissingConfigFS
    87  		}
    88  		s.op = s.LocalFs
    89  		s.Typ = STORAGE_TYPE__FILESYSTEM
    90  	case STORAGE_TYPE__S3:
    91  		if s.S3 == nil || s.S3.IsZero() {
    92  			return ErrMissingConfigS3
    93  		}
    94  		s.op = s.S3
    95  	case STORAGE_TYPE__IPFS:
    96  		return ErrMissingConfigIPFS
    97  	default:
    98  		return ErrUnsupprtedStorageType
    99  	}
   100  
   101  	if canSetDefault, ok := s.op.(types.DefaultSetter); ok {
   102  		canSetDefault.SetDefault()
   103  	}
   104  	if canBeInit, ok := s.op.(types.Initializer); ok {
   105  		canBeInit.Init()
   106  		return nil
   107  	}
   108  	if canBeInit, ok := s.op.(types.ValidatedInitializer); ok {
   109  		return canBeInit.Init()
   110  	}
   111  	return nil
   112  }
   113  
   114  func (s *Storage) WithOperation(op StorageOperations) {
   115  	s.op = op
   116  }
   117  
   118  func (s *Storage) Upload(key string, content []byte, chk ...HmacAlgType) error {
   119  	size := int64(len(content))
   120  	if size == 0 {
   121  		return ErrEmptyContent
   122  	}
   123  	if s.FilesizeLimit != 0 && size > s.FilesizeLimit {
   124  		return ErrContentSizeExceeded
   125  	}
   126  
   127  	free := int64(0)
   128  	stat, err := disk.Usage(s.TempDir)
   129  	if err == nil {
   130  		free = int64(stat.Free) - size
   131  	}
   132  
   133  	if s.DiskReserve != 0 && s.Type() == STORAGE_TYPE__FILESYSTEM && free < s.DiskReserve {
   134  		return ErrDiskReservationLimit
   135  	}
   136  
   137  	if err = s.op.Upload(key, content, chk...); err != nil {
   138  		return err
   139  	}
   140  	if s.Typ == STORAGE_TYPE__S3 {
   141  		_ = os.RemoveAll(s.TempDir)
   142  	}
   143  	return nil
   144  }
   145  
   146  func (s *Storage) Read(key string, chk ...HmacAlgType) ([]byte, []byte, error) {
   147  	return s.op.Read(key, chk...)
   148  }
   149  
   150  func (s *Storage) Delete(key string) error {
   151  	return s.op.Delete(key)
   152  }
   153  
   154  func (s *Storage) Validate(data []byte, sum string, chk ...HmacAlgType) bool {
   155  	if len(data) == 0 || len(sum) == 0 {
   156  		return true
   157  	}
   158  
   159  	t := HMAC_ALG_TYPE__MD5
   160  	if len(chk) > 0 && chk[0] != 0 {
   161  		t = chk[0]
   162  	}
   163  
   164  	return sum == t.HexSum(data)
   165  }
   166  
   167  func (s *Storage) Type() StorageType { return s.op.Type() }