1package storage
  2
  3import (
  4	"bytes"
  5	"fmt"
  6	"io"
  7	"io/fs"
  8	"os"
  9	"path"
 10	"path/filepath"
 11	"strings"
 12	"sync"
 13	"time"
 14
 15	"github.com/picosh/pico/pkg/send/utils"
 16)
 17
 18type seekableReader struct {
 19	*bytes.Reader
 20}
 21
 22func (s *seekableReader) Close() error { return nil }
 23
 24type StorageMemory struct {
 25	storage map[string]map[string]string
 26	mtimes  map[string]map[string]time.Time
 27	mu      sync.RWMutex
 28}
 29
 30var _ StorageServe = &StorageMemory{}
 31var _ StorageServe = (*StorageMemory)(nil)
 32
 33func NewStorageMemory(st map[string]map[string]string) (*StorageMemory, error) {
 34	now := time.Now()
 35	mtimes := map[string]map[string]time.Time{}
 36	for bucket, objects := range st {
 37		mtimes[bucket] = map[string]time.Time{}
 38		for key := range objects {
 39			mtimes[bucket][key] = now
 40		}
 41	}
 42	return &StorageMemory{
 43		storage: st,
 44		mtimes:  mtimes,
 45	}, nil
 46}
 47
 48func memoryKey(fpath string) string {
 49	if !strings.HasPrefix(fpath, "/") {
 50		return "/" + fpath
 51	}
 52	return fpath
 53}
 54
 55// dirKey is the key that keeps a directory listed while it's empty. Keys are
 56// flat like S3's, so a directory is a "dir/" key with no contents.
 57func dirKey(dir string) string {
 58	return strings.TrimSuffix(memoryKey(dir), "/") + "/"
 59}
 60
 61// hasPrefix reports whether any key other than skip starts with prefix.
 62func hasPrefix(objects map[string]string, prefix, skip string) bool {
 63	for key := range objects {
 64		if key != skip && strings.HasPrefix(key, prefix) {
 65			return true
 66		}
 67	}
 68	return false
 69}
 70
 71func (s *StorageMemory) GetBucket(name string) (Bucket, error) {
 72	s.mu.RLock()
 73	defer s.mu.RUnlock()
 74
 75	bucket := Bucket{
 76		Name: name,
 77		Path: name,
 78	}
 79
 80	_, ok := s.storage[name]
 81	if !ok {
 82		return bucket, fmt.Errorf("bucket does not exist")
 83	}
 84
 85	return bucket, nil
 86}
 87
 88func (s *StorageMemory) UpsertBucket(name string) (Bucket, error) {
 89	bucket, err := s.GetBucket(name)
 90	if err == nil {
 91		return bucket, nil
 92	}
 93
 94	s.mu.Lock()
 95	defer s.mu.Unlock()
 96
 97	s.storage[name] = map[string]string{}
 98	s.mtimes[name] = map[string]time.Time{}
 99	return bucket, nil
100}
101
102func (s *StorageMemory) GetBucketQuota(bucket Bucket) (uint64, error) {
103	s.mu.RLock()
104	defer s.mu.RUnlock()
105
106	objects := s.storage[bucket.Path]
107	size := 0
108	for _, val := range objects {
109		size += len([]byte(val))
110	}
111	return uint64(size), nil
112}
113
114func (s *StorageMemory) DeleteBucket(bucket Bucket) error {
115	s.mu.Lock()
116	defer s.mu.Unlock()
117
118	delete(s.storage, bucket.Path)
119	delete(s.mtimes, bucket.Path)
120	return nil
121}
122
123func (s *StorageMemory) GetObject(bucket Bucket, fpath string) (utils.ReadAndReaderAtCloser, *ObjectInfo, error) {
124	s.mu.RLock()
125	defer s.mu.RUnlock()
126
127	fpath = memoryKey(fpath)
128	objInfo := &ObjectInfo{}
129
130	dat, ok := s.storage[bucket.Path][fpath]
131	if !ok || strings.HasSuffix(fpath, "/") {
132		return nil, objInfo, fmt.Errorf("object does not exist: %s: %w", fpath, fs.ErrNotExist)
133	}
134
135	objInfo.Size = int64(len([]byte(dat)))
136	objInfo.LastModified = s.mtimes[bucket.Path][fpath]
137	return &seekableReader{bytes.NewReader([]byte(dat))}, objInfo, nil
138}
139
140func (s *StorageMemory) StatObject(bucket Bucket, fpath string) (*ObjectInfo, error) {
141	s.mu.RLock()
142	defer s.mu.RUnlock()
143
144	fpath = memoryKey(fpath)
145	dat, ok := s.storage[bucket.Path][fpath]
146	if !ok || strings.HasSuffix(fpath, "/") {
147		return nil, fmt.Errorf("object does not exist: %s: %w", fpath, fs.ErrNotExist)
148	}
149	return &ObjectInfo{
150		Size:         int64(len(dat)),
151		LastModified: s.mtimes[bucket.Path][fpath],
152	}, nil
153}
154
155func (s *StorageMemory) PutObject(bucket Bucket, fpath string, contents io.Reader, info *ObjectInfo) (string, int64, error) {
156	s.mu.Lock()
157	defer s.mu.Unlock()
158
159	d, err := io.ReadAll(contents)
160	if err != nil {
161		return "", 0, err
162	}
163
164	fpath = memoryKey(fpath)
165	mtime := info.LastModified
166	if mtime.IsZero() {
167		mtime = time.Now()
168	}
169	s.storage[bucket.Path][fpath] = string(d)
170	if s.mtimes[bucket.Path] == nil {
171		s.mtimes[bucket.Path] = map[string]time.Time{}
172	}
173	s.mtimes[bucket.Path][fpath] = mtime
174	return fmt.Sprintf("%s%s", bucket.Path, fpath), int64(len(d)), nil
175}
176
177func (s *StorageMemory) PutDir(bucket Bucket, dir string) error {
178	s.mu.Lock()
179	defer s.mu.Unlock()
180
181	key := dirKey(dir)
182	if key == "/" {
183		return nil
184	}
185	s.storage[bucket.Path][key] = ""
186	s.mtimes[bucket.Path][key] = time.Now()
187	return nil
188}
189
190func (s *StorageMemory) DeleteObject(bucket Bucket, fpath string) error {
191	s.mu.Lock()
192	defer s.mu.Unlock()
193
194	objects := s.storage[bucket.Path]
195	mtimes := s.mtimes[bucket.Path]
196	fpath = memoryKey(fpath)
197	if _, ok := objects[fpath]; ok && !strings.HasSuffix(fpath, "/") {
198		delete(objects, fpath)
199		delete(mtimes, fpath)
200		if parent := dirKey(path.Dir(fpath)); parent != "/" && !hasPrefix(objects, parent, "") {
201			objects[parent] = ""
202			mtimes[parent] = time.Now()
203		}
204		return nil
205	}
206
207	key := dirKey(fpath)
208	if hasPrefix(objects, key, key) {
209		return fmt.Errorf("%s: directory not empty", fpath)
210	}
211	delete(objects, key)
212	delete(mtimes, key)
213	return nil
214}
215
216func (s *StorageMemory) ListBuckets() ([]string, error) {
217	s.mu.RLock()
218	defer s.mu.RUnlock()
219
220	buckets := []string{}
221	for key := range s.storage {
222		buckets = append(buckets, key)
223	}
224	return buckets, nil
225}
226
227// ListObjects follows StorageFS: a path naming an object lists that object,
228// a directory path without a trailing slash lists only the directory itself,
229// and a path ending in a slash lists its contents relative to it.
230func (s *StorageMemory) ListObjects(bucket Bucket, dir string, recursive bool) ([]os.FileInfo, error) {
231	s.mu.RLock()
232	defer s.mu.RUnlock()
233
234	fileList := []os.FileInfo{}
235	resolved := memoryKey(dir)
236	objects := s.storage[bucket.Path]
237	mtimes := s.mtimes[bucket.Path]
238
239	if val, ok := objects[resolved]; ok && !strings.HasSuffix(resolved, "/") {
240		fileList = append(fileList, &utils.VirtualFile{
241			FName:    filepath.Base(resolved),
242			FSize:    int64(len(val)),
243			FModTime: mtimes[resolved],
244		})
245		return fileList, nil
246	}
247
248	if !strings.HasSuffix(resolved, "/") {
249		for key := range objects {
250			if strings.HasPrefix(key, resolved+"/") {
251				fileList = append(fileList, &utils.VirtualFile{FName: "", FIsDir: true})
252				break
253			}
254		}
255		return fileList, nil
256	}
257
258	seen := map[string]bool{}
259	addDir := func(name string) {
260		if !seen[name] {
261			seen[name] = true
262			fileList = append(fileList, &utils.VirtualFile{FName: name, FIsDir: true})
263		}
264	}
265	for key, val := range objects {
266		rel, ok := strings.CutPrefix(key, resolved)
267		if !ok || rel == "" {
268			continue
269		}
270		if !recursive {
271			if d, _, nested := strings.Cut(rel, "/"); nested {
272				addDir(d)
273				continue
274			}
275		}
276		name, isDir := strings.CutSuffix(rel, "/")
277		start := filepath.Dir(name)
278		if isDir {
279			start = name
280		}
281		for d := start; d != "." && !seen[d]; d = filepath.Dir(d) {
282			addDir(d)
283		}
284		if isDir {
285			continue
286		}
287		fileList = append(fileList, &utils.VirtualFile{
288			FName:    rel,
289			FSize:    int64(len(val)),
290			FModTime: mtimes[key],
291		})
292	}
293
294	return fileList, nil
295}