-
Notifications
You must be signed in to change notification settings - Fork 15
/
memstore.go
81 lines (68 loc) · 1.88 KB
/
memstore.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
package contentstore
import (
"bytes"
"fmt"
"io"
"sync"
"github.com/ashirt-ops/ashirt-server/backend"
"github.com/google/uuid"
)
// MemStore is the backing structure needed to interact with local memory -- for unit/integration
// testing purposes only
type MemStore struct {
content map[string][]byte
mutex *sync.Mutex
}
// NewMemStore is the constructor for MemStore
func NewMemStore() (*MemStore, error) {
m := MemStore{
content: make(map[string][]byte),
mutex: new(sync.Mutex),
}
return &m, nil
}
// Upload stores content in memory
func (d *MemStore) Upload(data io.Reader) (key string, err error) {
key = uuid.New().String()
err = d.UploadWithName(key, data)
if err != nil {
err = backend.WrapError("Unable to add to MemStore", err)
}
return
}
// UploadWithName writes the given data to the given memory key -- this
// can allow for re-writing/replacing data if names are not unique
//
// Note: to avoid overwriting random keys, DO NOT use uuids as they key
func (d *MemStore) UploadWithName(key string, data io.Reader) error {
b, err := io.ReadAll(data)
if err != nil {
return backend.WrapError("Unable upload with a given name to MemStore", err)
}
d.mutex.Lock()
d.content[key] = b
d.mutex.Unlock()
return nil
}
func (d *MemStore) Read(key string) (io.Reader, error) {
d.mutex.Lock()
data, ok := d.content[key]
d.mutex.Unlock()
if !ok {
return nil, backend.WrapError("Unable to read from MemStore", fmt.Errorf("No such key"))
}
return bytes.NewReader(data), nil
}
// Delete removes files in in your OS's temp directory
func (d *MemStore) Delete(key string) error {
d.mutex.Lock()
if _, ok := d.content[key]; !ok { // artificial behavior to match other stores
return backend.WrapError("Unable to delete from MemStore", fmt.Errorf("No such key"))
}
delete(d.content, key)
d.mutex.Unlock()
return nil
}
func (d *MemStore) Name() string {
return "memory"
}