repos / pico

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

pico / pkg / storage
Eric Bower  ·  2026-04-20

memory.go

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