1package merkletree2
2
3import (
4	"errors"
5	"fmt"
6	"math/bits"
7	"math/rand"
8	"sort"
9
10	"github.com/keybase/client/go/logger"
11)
12
13// In memory StorageEngine implementation, used for tests. It ignores
14// Transaction arguments, so it can't be used for concurrency tests.
15type InMemoryStorageEngine struct {
16	Roots            map[Seqno]RootMetadata
17	SortedKVPRs      []*KVPRecord
18	Nodes            map[string]*NodeRecord
19	MasterSecretsMap map[Seqno]MasterSecret
20	ReverseRootMap   map[string]Seqno
21
22	// used to make prefix queries efficient. Not otherwise necessary
23	//PositionToKeys map[string](map[string]bool)
24	cfg Config
25}
26
27var _ StorageEngine = &InMemoryStorageEngine{}
28var _ StorageEngineWithBlinding = &InMemoryStorageEngine{}
29
30type SortedKVPR []*KVPRecord
31
32func (s SortedKVPR) Len() int {
33	return len(s)
34}
35
36func (s SortedKVPR) Less(i, j int) bool {
37	return s[i].kevp.Key.Cmp(s[j].kevp.Key) < 0
38}
39
40func (s SortedKVPR) Swap(i, j int) {
41	s[i], s[j] = s[j], s[i]
42}
43
44var _ sort.Interface = SortedKVPR{}
45
46type NodeRecord struct {
47	p    Position
48	s    Seqno
49	h    Hash
50	next *NodeRecord
51}
52
53type KVPRecord struct {
54	kevp KeyEncodedValuePair
55	s    Seqno
56	next *KVPRecord
57}
58
59func NewInMemoryStorageEngine(cfg Config) *InMemoryStorageEngine {
60	i := InMemoryStorageEngine{}
61	i.MasterSecretsMap = make(map[Seqno]MasterSecret)
62	i.Roots = make(map[Seqno]RootMetadata)
63	i.Nodes = make(map[string]*NodeRecord)
64	i.ReverseRootMap = make(map[string]Seqno)
65	i.cfg = cfg
66	return &i
67}
68
69func (i *InMemoryStorageEngine) ExecTransaction(ctx logger.ContextInterface, txFn func(logger.ContextInterface, Transaction) error) error {
70	return fmt.Errorf("Transactions are not supported by this engine!")
71}
72
73func (i *InMemoryStorageEngine) findKVPR(k Key) *KVPRecord {
74	j := sort.Search(len(i.SortedKVPRs), func(n int) bool {
75		return i.SortedKVPRs[n].kevp.Key.Cmp(k) >= 0
76	})
77	if j < len(i.SortedKVPRs) && i.SortedKVPRs[j].kevp.Key.Cmp(k) == 0 {
78		return i.SortedKVPRs[j]
79	}
80	return nil
81}
82
83func (i *InMemoryStorageEngine) StoreKEVPairs(c logger.ContextInterface, t Transaction, s Seqno, kevps []KeyEncodedValuePair) error {
84	newSKVPR := make([]*KVPRecord, len(kevps))
85
86	for j, kevp := range kevps {
87		oldKvp := i.findKVPR(kevp.Key)
88		newSKVPR[j] = &KVPRecord{kevp: kevp, s: s, next: oldKvp}
89		if oldKvp != nil && oldKvp.s >= s {
90			return errors.New("Engine does not support out of order insertions")
91		}
92	}
93
94	sort.Sort(SortedKVPR(newSKVPR))
95	i.SortedKVPRs = newSKVPR
96	return nil
97}
98
99func (i *InMemoryStorageEngine) StoreNodes(c logger.ContextInterface, t Transaction, s Seqno, phps []PositionHashPair) error {
100	for _, php := range phps {
101		err := i.StoreNode(c, t, s, &php.Position, php.Hash)
102		if err != nil {
103			return err
104		}
105	}
106	return nil
107}
108
109func (i *InMemoryStorageEngine) StoreNode(c logger.ContextInterface, t Transaction, s Seqno, p *Position, h Hash) error {
110	strKey := p.AsString()
111	oldNodeRec := i.Nodes[strKey]
112	i.Nodes[strKey] = &NodeRecord{s: s, p: *p, h: h, next: oldNodeRec}
113	if oldNodeRec != nil && oldNodeRec.s >= s {
114		return errors.New("Engine does not support out of order insertions")
115	}
116	return nil
117}
118
119func (i *InMemoryStorageEngine) StoreRootMetadata(c logger.ContextInterface, t Transaction, r RootMetadata, h Hash) error {
120	i.Roots[r.Seqno] = r
121	i.ReverseRootMap[string(h)] = r.Seqno
122	return nil
123}
124
125func (i *InMemoryStorageEngine) LookupLatestRoot(c logger.ContextInterface, t Transaction) (Seqno, RootMetadata, error) {
126	if len(i.Roots) == 0 {
127		return 0, RootMetadata{}, NewNoLatestRootFoundError()
128	}
129	max := Seqno(0)
130	for k := range i.Roots {
131		if k > max {
132			max = k
133		}
134	}
135	return max, i.Roots[max], nil
136}
137
138func (i *InMemoryStorageEngine) LookupRoot(c logger.ContextInterface, t Transaction, s Seqno) (RootMetadata, error) {
139	r, found := i.Roots[s]
140	if found {
141		return r, nil
142	}
143	return RootMetadata{}, NewInvalidSeqnoError(s, fmt.Errorf("No root at seqno %v", s))
144}
145
146func (i *InMemoryStorageEngine) LookupRootFromHash(c logger.ContextInterface, t Transaction, h Hash) (RootMetadata, error) {
147	s, found := i.ReverseRootMap[string(h)]
148	if found {
149		return i.Roots[s], nil
150	}
151	return RootMetadata{}, NewInvalidSeqnoError(s, fmt.Errorf("No root at seqno %v", s))
152}
153
154func (i *InMemoryStorageEngine) LookupRoots(c logger.ContextInterface, t Transaction, seqnos []Seqno) (roots []RootMetadata, err error) {
155	seqnosSorted := make([]Seqno, len(seqnos))
156	copy(seqnosSorted, seqnos)
157	sort.Sort(SeqnoSortedAsInt(seqnosSorted))
158
159	roots = make([]RootMetadata, len(seqnosSorted))
160
161	for j, s := range seqnosSorted {
162		root, found := i.Roots[s]
163		if !found {
164			return nil, NewInvalidSeqnoError(s, fmt.Errorf("No root at seqno %v", s))
165		}
166		roots[j] = root
167	}
168
169	return roots, nil
170}
171
172func (i *InMemoryStorageEngine) LookupRootHashes(c logger.ContextInterface, t Transaction, seqnos []Seqno) (hashes []Hash, err error) {
173	if len(seqnos) == 0 {
174		return nil, fmt.Errorf("No seqnos requested")
175	}
176	seqnosSorted := make([]Seqno, len(seqnos))
177	copy(seqnosSorted, seqnos)
178	sort.Sort(SeqnoSortedAsInt(seqnosSorted))
179
180	hashes = make([]Hash, len(seqnosSorted))
181
182	for j, s := range seqnosSorted {
183		r, found := i.Roots[s]
184		if !found {
185			return nil, NewInvalidSeqnoError(s, fmt.Errorf("No root at seqno %v", s))
186		}
187		_, hashes[j], err = i.cfg.Encoder.EncodeAndHashGeneric(r)
188		if err != nil {
189			panic(fmt.Sprintf("Error encoding %+v", r))
190		}
191	}
192
193	return hashes, nil
194}
195
196func (i *InMemoryStorageEngine) LookupNode(c logger.ContextInterface, t Transaction, s Seqno, p *Position) (Hash, error) {
197	node, found := i.Nodes[string(p.GetBytes())]
198	if !found {
199		return nil, NewNodeNotFoundError()
200	}
201	for ; node != nil; node = node.next {
202		if node.s <= s {
203			return node.h, nil
204		}
205	}
206	return nil, NewNodeNotFoundError()
207}
208
209func (i *InMemoryStorageEngine) LookupNodes(c logger.ContextInterface, t Transaction, s Seqno, positions []Position) (res []PositionHashPair, err error) {
210	for _, p := range positions {
211		h, err := i.LookupNode(c, t, s, &p)
212		if err == nil {
213			res = append(res, PositionHashPair{Position: p, Hash: h})
214		}
215	}
216
217	// Shuffle the result to catch bugs that happen when ordering is different.
218	rand.Shuffle(len(res), func(i, j int) { res[i], res[j] = res[j], res[i] })
219
220	return res, nil
221}
222
223func (i *InMemoryStorageEngine) LookupKEVPair(c logger.ContextInterface, t Transaction, s Seqno, k Key) (EncodedValue, Seqno, error) {
224	kvpr := i.findKVPR(k)
225	if kvpr == nil {
226		return nil, 0, NewKeyNotFoundError()
227	}
228	for ; kvpr != nil; kvpr = kvpr.next {
229		if kvpr.s <= s {
230			return kvpr.kevp.Value, kvpr.s, nil
231		}
232	}
233	return nil, 0, NewKeyNotFoundError()
234}
235
236func (i *InMemoryStorageEngine) LookupKEVPairsUnderPosition(ctx logger.ContextInterface, t Transaction, s Seqno, p *Position) (kvps []KeyEncodedValuePair, seqnos []Seqno, err error) {
237	pBytes := p.GetBytes()
238
239	minIndex := sort.Search(len(i.SortedKVPRs), func(n int) bool {
240		k := i.SortedKVPRs[n].kevp.Key
241		pLeadingZeros := bits.LeadingZeros8(pBytes[0])
242		for j := 1; j < 8*len(pBytes)-pLeadingZeros; j++ {
243			jthBitP := (pBytes[(pLeadingZeros+j)/8] & byte(1<<uint(7-((pLeadingZeros+j)%8)))) != 0
244			jthBitK := (k[(j-1)/8] & byte(1<<uint(7-((j-1)%8)))) != 0
245			if !jthBitK && jthBitP {
246				return false
247			} else if jthBitK == jthBitP {
248				continue
249			} else if jthBitK && !jthBitP {
250				return true
251			}
252		}
253		return true
254	})
255
256	maxIndex := sort.Search(len(i.SortedKVPRs), func(n int) bool {
257		k := i.SortedKVPRs[n].kevp.Key
258		pLeadingZeros := bits.LeadingZeros8(pBytes[0])
259		for j := 1; j < 8*len(pBytes)-pLeadingZeros; j++ {
260			jthBitP := (pBytes[(pLeadingZeros+j)/8] & byte(1<<uint(7-((pLeadingZeros+j)%8)))) != 0
261			jthBitK := (k[(j-1)/8] & byte(1<<uint(7-((j-1)%8)))) != 0
262			if !jthBitK && jthBitP {
263				return false
264			} else if jthBitK == jthBitP {
265				continue
266			} else if jthBitK && !jthBitP {
267				return true
268			}
269		}
270		return false
271	})
272
273	kvps = make([]KeyEncodedValuePair, 0, maxIndex-minIndex)
274	seqnos = make([]Seqno, 0, maxIndex-minIndex)
275	for j := minIndex; j < maxIndex; j++ {
276		kvpr := i.SortedKVPRs[j]
277		for ; kvpr != nil; kvpr = kvpr.next {
278			if kvpr.s <= s {
279				kvps = append(kvps, kvpr.kevp)
280				seqnos = append(seqnos, kvpr.s)
281				break
282			}
283		}
284	}
285	if len(kvps) == 0 {
286		return nil, nil, NewKeyNotFoundError()
287	}
288	return kvps, seqnos, nil
289}
290
291// LookupAllKEVPairs returns all the keys and encoded values at the specified Seqno.
292func (i *InMemoryStorageEngine) LookupAllKEVPairs(ctx logger.ContextInterface, t Transaction, s Seqno) (kevps []KeyEncodedValuePair, err error) {
293	kevps = make([]KeyEncodedValuePair, 0, len(i.SortedKVPRs))
294	for _, kvpr := range i.SortedKVPRs {
295		for ; kvpr != nil; kvpr = kvpr.next {
296			if kvpr.s <= s {
297				kevps = append(kevps, kvpr.kevp)
298				break
299			}
300		}
301	}
302	return kevps, nil
303}
304
305func (i *InMemoryStorageEngine) StoreMasterSecret(ctx logger.ContextInterface, t Transaction, s Seqno, ms MasterSecret) (err error) {
306	i.MasterSecretsMap[s] = MasterSecret(append([]byte{}, []byte(ms)...))
307	return nil
308}
309
310func (i *InMemoryStorageEngine) LookupMasterSecrets(ctx logger.ContextInterface, t Transaction, seqnos []Seqno) (msMap map[Seqno]MasterSecret, err error) {
311	if i.MasterSecretsMap == nil {
312		i.MasterSecretsMap = make(map[Seqno]MasterSecret)
313	}
314	msMap = make(map[Seqno]MasterSecret)
315	for _, s := range seqnos {
316		ms, found := i.MasterSecretsMap[s]
317		if found {
318			msMap[s] = ms
319		} else {
320			return nil, fmt.Errorf("MasterSecret for Seqno %v not found", s)
321		}
322	}
323	return msMap, nil
324}
325