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}