Go · 4642 bytes Raw Blame History
1 // SPDX-License-Identifier: AGPL-3.0-or-later
2
3 package storage
4
5 import (
6 "bytes"
7 "context"
8 "crypto/md5" //nolint:gosec // ETag, not security-sensitive — matches S3 etag derivation.
9 "encoding/hex"
10 "fmt"
11 "io"
12 "sort"
13 "strings"
14 "sync"
15 "time"
16 )
17
18 // MemoryStore is an in-process ObjectStore implementation. Used by tests.
19 // Honors the same If-None-Match semantics as the s3 backend.
20 type MemoryStore struct {
21 mu sync.RWMutex
22 objects map[string]memObject
23 // signedURLBase is the prefix of generated SignedURLs so tests can
24 // assert their shape. Defaults to "mem://".
25 signedURLBase string
26 }
27
28 type memObject struct {
29 body []byte
30 etag string
31 contentType string
32 lastModified time.Time
33 }
34
35 // NewMemoryStore constructs an empty in-memory store.
36 func NewMemoryStore() *MemoryStore {
37 return &MemoryStore{
38 objects: make(map[string]memObject),
39 signedURLBase: "mem://",
40 }
41 }
42
43 // Put implements ObjectStore.
44 func (m *MemoryStore) Put(_ context.Context, key string, body io.Reader, opts PutOpts) (PutResult, error) {
45 if key == "" {
46 return PutResult{}, fmt.Errorf("storage: put: %w: empty key", ErrInvalidPath)
47 }
48 buf, err := io.ReadAll(body)
49 if err != nil {
50 return PutResult{}, fmt.Errorf("storage: put: read body: %w", err)
51 }
52 m.mu.Lock()
53 defer m.mu.Unlock()
54 if opts.IfNoneMatch == "*" {
55 if _, exists := m.objects[key]; exists {
56 return PutResult{}, ErrPreconditionFailed
57 }
58 }
59 sum := md5.Sum(buf) //nolint:gosec // not security-sensitive.
60 etag := hex.EncodeToString(sum[:])
61 m.objects[key] = memObject{
62 body: buf,
63 etag: etag,
64 contentType: opts.ContentType,
65 lastModified: time.Now().UTC(),
66 }
67 return PutResult{ETag: etag, Size: int64(len(buf))}, nil
68 }
69
70 // Get implements ObjectStore.
71 func (m *MemoryStore) Get(_ context.Context, key string) (io.ReadCloser, ObjectMeta, error) {
72 m.mu.RLock()
73 defer m.mu.RUnlock()
74 o, ok := m.objects[key]
75 if !ok {
76 return nil, ObjectMeta{}, ErrNotFound
77 }
78 return io.NopCloser(bytes.NewReader(o.body)), m.metaOf(key, o), nil
79 }
80
81 // Stat implements ObjectStore.
82 func (m *MemoryStore) Stat(_ context.Context, key string) (ObjectMeta, error) {
83 m.mu.RLock()
84 defer m.mu.RUnlock()
85 o, ok := m.objects[key]
86 if !ok {
87 return ObjectMeta{}, ErrNotFound
88 }
89 return m.metaOf(key, o), nil
90 }
91
92 // Delete implements ObjectStore. Idempotent.
93 func (m *MemoryStore) Delete(_ context.Context, key string) error {
94 m.mu.Lock()
95 defer m.mu.Unlock()
96 delete(m.objects, key)
97 return nil
98 }
99
100 // List implements ObjectStore. ContinuationToken is the last key returned
101 // in the previous page; results are sorted lexicographically.
102 func (m *MemoryStore) List(_ context.Context, prefix string, opts ListOpts) (ListResult, error) {
103 m.mu.RLock()
104 defer m.mu.RUnlock()
105
106 keys := make([]string, 0, len(m.objects))
107 for k := range m.objects {
108 if strings.HasPrefix(k, prefix) {
109 keys = append(keys, k)
110 }
111 }
112 sort.Strings(keys)
113
114 if opts.ContinuationToken != "" {
115 i := sort.SearchStrings(keys, opts.ContinuationToken)
116 if i < len(keys) && keys[i] == opts.ContinuationToken {
117 i++
118 }
119 keys = keys[i:]
120 }
121
122 maxKeys := opts.MaxKeys
123 if maxKeys <= 0 {
124 maxKeys = 1000
125 }
126
127 var (
128 out []ObjectMeta
129 prefixes []string
130 )
131 seenPrefix := map[string]struct{}{}
132
133 for _, k := range keys {
134 if !opts.Recursive {
135 rest := strings.TrimPrefix(k, prefix)
136 if i := strings.Index(rest, "/"); i >= 0 {
137 cp := prefix + rest[:i+1]
138 if _, ok := seenPrefix[cp]; !ok {
139 seenPrefix[cp] = struct{}{}
140 prefixes = append(prefixes, cp)
141 }
142 continue
143 }
144 }
145 o := m.objects[k]
146 out = append(out, m.metaOf(k, o))
147 if len(out) >= maxKeys {
148 break
149 }
150 }
151
152 res := ListResult{Objects: out, CommonPrefixes: prefixes}
153 if len(out) >= maxKeys && len(out) > 0 {
154 res.IsTruncated = true
155 res.NextContinuationToken = out[len(out)-1].Key
156 }
157 return res, nil
158 }
159
160 // SignedURL implements ObjectStore. Tests can rely on the prefix to
161 // distinguish memory-backed URLs from real ones.
162 func (m *MemoryStore) SignedURL(_ context.Context, key string, ttl time.Duration, method string) (string, error) {
163 switch method {
164 case "GET", "PUT":
165 default:
166 return "", fmt.Errorf("storage: signed url: unsupported method %q", method)
167 }
168 if key == "" {
169 return "", fmt.Errorf("storage: signed url: %w: empty key", ErrInvalidPath)
170 }
171 return fmt.Sprintf("%s%s?method=%s&ttl=%s", m.signedURLBase, key, method, ttl), nil
172 }
173
174 func (m *MemoryStore) metaOf(key string, o memObject) ObjectMeta {
175 return ObjectMeta{
176 Key: key,
177 Size: int64(len(o.body)),
178 ETag: o.etag,
179 ContentType: o.contentType,
180 LastModified: o.lastModified,
181 }
182 }
183