repos / pico

pico services mono repo
git clone https://github.com/picosh/pico.git

pico / pkg / pobj / storage
Antonio Mika  ·  2025-03-12

memory.go

  1package storage
  2
  3import (
  4	"fmt"
  5	"io"
  6	"os"
  7	"path/filepath"
  8	"strings"
  9	"sync"
 10	"time"
 11
 12	"github.com/picosh/pico/pkg/send/utils"
 13)
 14
 15type StorageMemory struct {
 16	storage map[string]map[string]string
 17	mu      sync.RWMutex
 18}
 19
 20var _ ObjectStorage = &StorageMemory{}
 21var _ ObjectStorage = (*StorageMemory)(nil)
 22
 23func NewStorageMemory(st map[string]map[string]string) (*StorageMemory, error) {
 24	return &StorageMemory{
 25		storage: st,
 26	}, nil
 27}
 28
 29func (s *StorageMemory) GetBucket(name string) (Bucket, error) {
 30	s.mu.RLock()
 31	defer s.mu.RUnlock()
 32
 33	bucket := Bucket{
 34		Name: name,
 35		Path: name,
 36	}
 37
 38	_, ok := s.storage[name]
 39	if !ok {
 40		return bucket, fmt.Errorf("bucket does not exist")
 41	}
 42
 43	return bucket, nil
 44}
 45
 46func (s *StorageMemory) UpsertBucket(name string) (Bucket, error) {
 47	bucket, err := s.GetBucket(name)
 48	if err == nil {
 49		return bucket, nil
 50	}
 51
 52	s.mu.Lock()
 53	defer s.mu.Unlock()
 54
 55	s.storage[name] = map[string]string{}
 56	return bucket, nil
 57}
 58
 59func (s *StorageMemory) GetBucketQuota(bucket Bucket) (uint64, error) {
 60	s.mu.RLock()
 61	defer s.mu.RUnlock()
 62
 63	objects := s.storage[bucket.Path]
 64	size := 0
 65	for _, val := range objects {
 66		size += len([]byte(val))
 67	}
 68	return uint64(size), nil
 69}
 70
 71func (s *StorageMemory) DeleteBucket(bucket Bucket) error {
 72	s.mu.Lock()
 73	defer s.mu.Unlock()
 74
 75	delete(s.storage, bucket.Path)
 76	return nil
 77}
 78
 79func (s *StorageMemory) GetObject(bucket Bucket, fpath string) (utils.ReadAndReaderAtCloser, *ObjectInfo, error) {
 80	s.mu.RLock()
 81	defer s.mu.RUnlock()
 82
 83	if !strings.HasPrefix(fpath, "/") {
 84		fpath = "/" + fpath
 85	}
 86
 87	objInfo := &ObjectInfo{
 88		LastModified: time.Time{},
 89		Metadata:     nil,
 90		UserMetadata: map[string]string{},
 91	}
 92
 93	dat, ok := s.storage[bucket.Path][fpath]
 94	if !ok {
 95		return nil, objInfo, fmt.Errorf("object does not exist: %s", fpath)
 96	}
 97
 98	objInfo.Size = int64(len([]byte(dat)))
 99	reader := utils.NopReadAndReaderAtCloser(strings.NewReader(dat))
100	return reader, objInfo, nil
101}
102
103func (s *StorageMemory) PutObject(bucket Bucket, fpath string, contents io.Reader, entry *utils.FileEntry) (string, int64, error) {
104	s.mu.Lock()
105	defer s.mu.Unlock()
106
107	d, err := io.ReadAll(contents)
108	if err != nil {
109		return "", 0, err
110	}
111
112	s.storage[bucket.Path][fpath] = string(d)
113	return fmt.Sprintf("%s%s", bucket.Path, fpath), int64(len(d)), nil
114}
115
116func (s *StorageMemory) DeleteObject(bucket Bucket, fpath string) error {
117	s.mu.Lock()
118	defer s.mu.Unlock()
119
120	delete(s.storage[bucket.Path], fpath)
121	return nil
122}
123
124func (s *StorageMemory) ListBuckets() ([]string, error) {
125	s.mu.RLock()
126	defer s.mu.RUnlock()
127
128	buckets := []string{}
129	for key := range s.storage {
130		buckets = append(buckets, key)
131	}
132	return buckets, nil
133}
134
135func (s *StorageMemory) ListObjects(bucket Bucket, dir string, recursive bool) ([]os.FileInfo, error) {
136	s.mu.RLock()
137	defer s.mu.RUnlock()
138
139	var fileList []os.FileInfo
140
141	resolved := dir
142
143	if !strings.HasPrefix(resolved, "/") {
144		resolved = "/" + resolved
145	}
146
147	objects := s.storage[bucket.Path]
148	// dir is actually an object
149	oval, ok := objects[resolved]
150	if ok {
151		fileList = append(fileList, &utils.VirtualFile{
152			FName:    filepath.Base(resolved),
153			FIsDir:   false,
154			FSize:    int64(len([]byte(oval))),
155			FModTime: time.Time{},
156		})
157		return fileList, nil
158	}
159
160	for key, val := range objects {
161		if !strings.HasPrefix(key, resolved) {
162			continue
163		}
164
165		rep := strings.Replace(key, resolved, "", 1)
166		fdir := filepath.Dir(rep)
167		fname := filepath.Base(rep)
168		paths := strings.Split(fdir, "/")
169
170		if fdir == "/" {
171			ffname := filepath.Base(resolved)
172			fileList = append(fileList, &utils.VirtualFile{
173				FName:  ffname,
174				FIsDir: true,
175			})
176		}
177
178		for _, p := range paths {
179			if p == "" || p == "/" || p == "." {
180				continue
181			}
182			fileList = append(fileList, &utils.VirtualFile{
183				FName:  p,
184				FIsDir: true,
185			})
186		}
187
188		trimRes := strings.TrimSuffix(resolved, "/")
189		dirKey := filepath.Dir(key)
190		if recursive {
191			fileList = append(fileList, &utils.VirtualFile{
192				FName:    fname,
193				FIsDir:   false,
194				FSize:    int64(len([]byte(val))),
195				FModTime: time.Time{},
196			})
197		} else if resolved == dirKey || trimRes == dirKey {
198			fileList = append(fileList, &utils.VirtualFile{
199				FName:    fname,
200				FIsDir:   false,
201				FSize:    int64(len([]byte(val))),
202				FModTime: time.Time{},
203			})
204		}
205	}
206
207	return fileList, nil
208}