1// Copyright 2020, 2020 OCI Contributors
2// Copyright 2017 Docker, Inc.
3//
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8//     https://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16package digestset
17
18import (
19	"crypto/sha256"
20	_ "crypto/sha512"
21	"encoding/binary"
22	"math/rand"
23	"testing"
24
25	digest "github.com/opencontainers/go-digest"
26)
27
28func assertEqualDigests(t *testing.T, d1, d2 digest.Digest) {
29	if d1 != d2 {
30		t.Fatalf("Digests do not match:\n\tActual: %s\n\tExpected: %s", d1, d2)
31	}
32}
33
34func TestLookup(t *testing.T) {
35	digests := []digest.Digest{
36		"sha256:1234511111111111111111111111111111111111111111111111111111111111",
37		"sha256:1234111111111111111111111111111111111111111111111111111111111111",
38		"sha256:1234611111111111111111111111111111111111111111111111111111111111",
39		"sha256:5432111111111111111111111111111111111111111111111111111111111111",
40		"sha256:6543111111111111111111111111111111111111111111111111111111111111",
41		"sha256:6432111111111111111111111111111111111111111111111111111111111111",
42		"sha256:6542111111111111111111111111111111111111111111111111111111111111",
43		"sha256:6532111111111111111111111111111111111111111111111111111111111111",
44	}
45
46	dset := NewSet()
47	for i := range digests {
48		if err := dset.Add(digests[i]); err != nil {
49			t.Fatal(err)
50		}
51	}
52
53	dgst, err := dset.Lookup("54")
54	if err != nil {
55		t.Fatal(err)
56	}
57	assertEqualDigests(t, dgst, digests[3])
58
59	_, err = dset.Lookup("1234")
60	if err == nil {
61		t.Fatal("Expected ambiguous error looking up: 1234")
62	}
63	if err != ErrDigestAmbiguous {
64		t.Fatal(err)
65	}
66
67	_, err = dset.Lookup("9876")
68	if err == nil {
69		t.Fatal("Expected not found error looking up: 9876")
70	}
71	if err != ErrDigestNotFound {
72		t.Fatal(err)
73	}
74
75	_, err = dset.Lookup("sha256:1234")
76	if err == nil {
77		t.Fatal("Expected ambiguous error looking up: sha256:1234")
78	}
79	if err != ErrDigestAmbiguous {
80		t.Fatal(err)
81	}
82
83	dgst, err = dset.Lookup("sha256:12345")
84	if err != nil {
85		t.Fatal(err)
86	}
87	assertEqualDigests(t, dgst, digests[0])
88
89	dgst, err = dset.Lookup("sha256:12346")
90	if err != nil {
91		t.Fatal(err)
92	}
93	assertEqualDigests(t, dgst, digests[2])
94
95	dgst, err = dset.Lookup("12346")
96	if err != nil {
97		t.Fatal(err)
98	}
99	assertEqualDigests(t, dgst, digests[2])
100
101	dgst, err = dset.Lookup("12345")
102	if err != nil {
103		t.Fatal(err)
104	}
105	assertEqualDigests(t, dgst, digests[0])
106}
107
108func TestAddDuplication(t *testing.T) {
109	digests := []digest.Digest{
110		"sha256:1234111111111111111111111111111111111111111111111111111111111111",
111		"sha256:1234511111111111111111111111111111111111111111111111111111111111",
112		"sha256:1234611111111111111111111111111111111111111111111111111111111111",
113		"sha256:5432111111111111111111111111111111111111111111111111111111111111",
114		"sha256:6543111111111111111111111111111111111111111111111111111111111111",
115		"sha512:65431111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111",
116		"sha512:65421111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111",
117		"sha512:65321111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111",
118	}
119
120	dset := NewSet()
121	for i := range digests {
122		if err := dset.Add(digests[i]); err != nil {
123			t.Fatal(err)
124		}
125	}
126
127	if len(dset.entries) != 8 {
128		t.Fatal("Invalid dset size")
129	}
130
131	if err := dset.Add(digest.Digest("sha256:1234511111111111111111111111111111111111111111111111111111111111")); err != nil {
132		t.Fatal(err)
133	}
134
135	if len(dset.entries) != 8 {
136		t.Fatal("Duplicate digest insert should not increase entries size")
137	}
138
139	if err := dset.Add(digest.Digest("sha384:123451111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111")); err != nil {
140		t.Fatal(err)
141	}
142
143	if len(dset.entries) != 9 {
144		t.Fatal("Insert with different algorithm should be allowed")
145	}
146}
147
148func TestRemove(t *testing.T) {
149	digests, err := createDigests(10)
150	if err != nil {
151		t.Fatal(err)
152	}
153
154	dset := NewSet()
155	for i := range digests {
156		if err := dset.Add(digests[i]); err != nil {
157			t.Fatal(err)
158		}
159	}
160
161	dgst, err := dset.Lookup(digests[0].String())
162	if err != nil {
163		t.Fatal(err)
164	}
165	if dgst != digests[0] {
166		t.Fatalf("Unexpected digest value:\n\tExpected: %s\n\tActual: %s", digests[0], dgst)
167	}
168
169	if err := dset.Remove(digests[0]); err != nil {
170		t.Fatal(err)
171	}
172
173	if _, err := dset.Lookup(digests[0].String()); err != ErrDigestNotFound {
174		t.Fatalf("Expected error %v when looking up removed digest, got %v", ErrDigestNotFound, err)
175	}
176}
177
178func TestAll(t *testing.T) {
179	digests, err := createDigests(100)
180	if err != nil {
181		t.Fatal(err)
182	}
183
184	dset := NewSet()
185	for i := range digests {
186		if err := dset.Add(digests[i]); err != nil {
187			t.Fatal(err)
188		}
189	}
190
191	all := map[digest.Digest]struct{}{}
192	for _, dgst := range dset.All() {
193		all[dgst] = struct{}{}
194	}
195
196	if len(all) != len(digests) {
197		t.Fatalf("Unexpected number of unique digests found:\n\tExpected: %d\n\tActual: %d", len(digests), len(all))
198	}
199
200	for i, dgst := range digests {
201		if _, ok := all[dgst]; !ok {
202			t.Fatalf("Missing element at position %d: %s", i, dgst)
203		}
204	}
205
206}
207
208func assertEqualShort(t *testing.T, actual, expected string) {
209	if actual != expected {
210		t.Fatalf("Unexpected short value:\n\tExpected: %s\n\tActual: %s", expected, actual)
211	}
212}
213
214func TestShortCodeTable(t *testing.T) {
215	digests := []digest.Digest{
216		"sha256:1234111111111111111111111111111111111111111111111111111111111111",
217		"sha256:1234511111111111111111111111111111111111111111111111111111111111",
218		"sha256:1234611111111111111111111111111111111111111111111111111111111111",
219		"sha256:5432111111111111111111111111111111111111111111111111111111111111",
220		"sha256:6543111111111111111111111111111111111111111111111111111111111111",
221		"sha256:6432111111111111111111111111111111111111111111111111111111111111",
222		"sha256:6542111111111111111111111111111111111111111111111111111111111111",
223		"sha256:6532111111111111111111111111111111111111111111111111111111111111",
224	}
225
226	dset := NewSet()
227	for i := range digests {
228		if err := dset.Add(digests[i]); err != nil {
229			t.Fatal(err)
230		}
231	}
232
233	dump := ShortCodeTable(dset, 2)
234
235	if len(dump) < len(digests) {
236		t.Fatalf("Error unexpected size: %d, expecting %d", len(dump), len(digests))
237	}
238	assertEqualShort(t, dump[digests[0]], "12341")
239	assertEqualShort(t, dump[digests[1]], "12345")
240	assertEqualShort(t, dump[digests[2]], "12346")
241	assertEqualShort(t, dump[digests[3]], "54")
242	assertEqualShort(t, dump[digests[4]], "6543")
243	assertEqualShort(t, dump[digests[5]], "64")
244	assertEqualShort(t, dump[digests[6]], "6542")
245	assertEqualShort(t, dump[digests[7]], "653")
246}
247
248func createDigests(count int) ([]digest.Digest, error) {
249	r := rand.New(rand.NewSource(25823))
250	digests := make([]digest.Digest, count)
251	for i := range digests {
252		h := sha256.New()
253		if err := binary.Write(h, binary.BigEndian, r.Int63()); err != nil {
254			return nil, err
255		}
256		digests[i] = digest.NewDigest("sha256", h)
257	}
258	return digests, nil
259}
260
261func benchAddNTable(b *testing.B, n int) {
262	digests, err := createDigests(n)
263	if err != nil {
264		b.Fatal(err)
265	}
266	b.ResetTimer()
267	for i := 0; i < b.N; i++ {
268		dset := &Set{entries: digestEntries(make([]*digestEntry, 0, n))}
269		for j := range digests {
270			if err = dset.Add(digests[j]); err != nil {
271				b.Fatal(err)
272			}
273		}
274	}
275}
276
277func benchLookupNTable(b *testing.B, n int, shortLen int) {
278	digests, err := createDigests(n)
279	if err != nil {
280		b.Fatal(err)
281	}
282	dset := &Set{entries: digestEntries(make([]*digestEntry, 0, n))}
283	for i := range digests {
284		if err := dset.Add(digests[i]); err != nil {
285			b.Fatal(err)
286		}
287	}
288	shorts := make([]string, 0, n)
289	for _, short := range ShortCodeTable(dset, shortLen) {
290		shorts = append(shorts, short)
291	}
292
293	b.ResetTimer()
294	for i := 0; i < b.N; i++ {
295		if _, err = dset.Lookup(shorts[i%n]); err != nil {
296			b.Fatal(err)
297		}
298	}
299}
300
301func benchRemoveNTable(b *testing.B, n int) {
302	digests, err := createDigests(n)
303	if err != nil {
304		b.Fatal(err)
305	}
306	b.ResetTimer()
307	for i := 0; i < b.N; i++ {
308		dset := &Set{entries: digestEntries(make([]*digestEntry, 0, n))}
309		b.StopTimer()
310		for j := range digests {
311			if err = dset.Add(digests[j]); err != nil {
312				b.Fatal(err)
313			}
314		}
315		b.StartTimer()
316		for j := range digests {
317			if err = dset.Remove(digests[j]); err != nil {
318				b.Fatal(err)
319			}
320		}
321	}
322}
323
324func benchShortCodeNTable(b *testing.B, n int, shortLen int) {
325	digests, err := createDigests(n)
326	if err != nil {
327		b.Fatal(err)
328	}
329	dset := &Set{entries: digestEntries(make([]*digestEntry, 0, n))}
330	for i := range digests {
331		if err := dset.Add(digests[i]); err != nil {
332			b.Fatal(err)
333		}
334	}
335
336	b.ResetTimer()
337	for i := 0; i < b.N; i++ {
338		ShortCodeTable(dset, shortLen)
339	}
340}
341
342func BenchmarkAdd10(b *testing.B) {
343	benchAddNTable(b, 10)
344}
345
346func BenchmarkAdd100(b *testing.B) {
347	benchAddNTable(b, 100)
348}
349
350func BenchmarkAdd1000(b *testing.B) {
351	benchAddNTable(b, 1000)
352}
353
354func BenchmarkRemove10(b *testing.B) {
355	benchRemoveNTable(b, 10)
356}
357
358func BenchmarkRemove100(b *testing.B) {
359	benchRemoveNTable(b, 100)
360}
361
362func BenchmarkRemove1000(b *testing.B) {
363	benchRemoveNTable(b, 1000)
364}
365
366func BenchmarkLookup10(b *testing.B) {
367	benchLookupNTable(b, 10, 12)
368}
369
370func BenchmarkLookup100(b *testing.B) {
371	benchLookupNTable(b, 100, 12)
372}
373
374func BenchmarkLookup1000(b *testing.B) {
375	benchLookupNTable(b, 1000, 12)
376}
377
378func BenchmarkShortCode10(b *testing.B) {
379	benchShortCodeNTable(b, 10, 12)
380}
381func BenchmarkShortCode100(b *testing.B) {
382	benchShortCodeNTable(b, 100, 12)
383}
384func BenchmarkShortCode1000(b *testing.B) {
385	benchShortCodeNTable(b, 1000, 12)
386}
387