Skip to content

Commit d2531cb

Browse files
jvdmjvdm
committed
jsonblob: Support iteration with rangefunc style
1 parent f4731b3 commit d2531cb

File tree

2 files changed

+211
-0
lines changed

2 files changed

+211
-0
lines changed

‎libvuln/jsonblob/jsonblob.go

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,64 @@ type Store struct {
4444
latest map[driver.UpdateKind]uuid.UUID
4545
}
4646

47+
type iter2[X, Y any] func(yield func(X, Y) bool)
48+
49+
// RecordIter iterates over records of an update operation.
50+
type RecordIter iter2[*claircore.Vulnerability, *driver.EnrichmentRecord]
51+
52+
// OperationIter iterates over operations, offering a nested iterator for records.
53+
type OperationIter iter2[*driver.UpdateOperation, RecordIter]
54+
55+
// Iterate iterates over each record serialized in the [io.Reader] grouping by
56+
// update operations. It returns an OperationIter, which is an iterator over each
57+
// update operation with a nested iterator for the associated vulnerability
58+
// entries, and an error function, to check for iteration errors.
59+
func Iterate(r io.Reader) (OperationIter, func() error) {
60+
var err error
61+
var de diskEntry
62+
63+
d := json.NewDecoder(r)
64+
err = d.Decode(&de)
65+
66+
it := func(yield func(*driver.UpdateOperation, RecordIter) bool) {
67+
for err == nil {
68+
op := &driver.UpdateOperation{
69+
Ref: de.Ref,
70+
Updater: de.Updater,
71+
Fingerprint: de.Fingerprint,
72+
Date: de.Date,
73+
Kind: de.Kind,
74+
}
75+
it := func(yield func(*claircore.Vulnerability, *driver.EnrichmentRecord) bool) {
76+
var vuln *claircore.Vulnerability
77+
var en *driver.EnrichmentRecord
78+
for err == nil && op.Ref == de.Ref {
79+
vuln, en, err = de.Unmarshal()
80+
if err != nil || !yield(vuln, en) {
81+
break
82+
}
83+
err = d.Decode(&de)
84+
}
85+
}
86+
if !yield(op, it) {
87+
break
88+
}
89+
for err == nil && op.Ref == de.Ref {
90+
err = d.Decode(&de)
91+
}
92+
}
93+
}
94+
95+
errF := func() error {
96+
if errors.Is(err, io.EOF) {
97+
return nil
98+
}
99+
return err
100+
}
101+
102+
return it, errF
103+
}
104+
47105
// Load reads in all the records serialized in the provided [io.Reader].
48106
func Load(ctx context.Context, r io.Reader) (*Loader, error) {
49107
l := Loader{
@@ -252,6 +310,25 @@ type diskEntry struct {
252310
Kind driver.UpdateKind
253311
}
254312

313+
// Unmarshal parses the JSON-encoded vulnerability or enrichment record encoded
314+
// in the disk entry, based on the update kind.
315+
func (de *diskEntry) Unmarshal() (v *claircore.Vulnerability, e *driver.EnrichmentRecord, err error) {
316+
switch de.Kind {
317+
case driver.VulnerabilityKind:
318+
v = &claircore.Vulnerability{}
319+
if err = json.Unmarshal(de.Vuln.buf, v); err != nil {
320+
return
321+
}
322+
case driver.EnrichmentKind:
323+
e = &driver.EnrichmentRecord{}
324+
err = json.Unmarshal(de.Enrichment.buf, e)
325+
if err != nil {
326+
return
327+
}
328+
}
329+
return
330+
}
331+
255332
// Entries returns a map containing all the Entries stored by calls to
256333
// UpdateVulnerabilities.
257334
//
@@ -455,6 +532,14 @@ func (s *Store) DeltaUpdateVulnerabilities(ctx context.Context, updater string,
455532
return uuid.Nil, nil
456533
}
457534

535+
func (s *Store) UpdateEnrichmentsIter(_ context.Context, _ string, _ driver.Fingerprint, _ datastore.EnrichmentIter) (uuid.UUID, error) {
536+
return uuid.Nil, errors.ErrUnsupported
537+
}
538+
539+
func (s *Store) UpdateVulnerabilitiesIter(_ context.Context, _ string, _ driver.Fingerprint, _ datastore.VulnerabilityIter) (uuid.UUID, error) {
540+
return uuid.Nil, errors.ErrUnsupported
541+
}
542+
458543
var bufPool sync.Pool
459544

460545
func getBuf() []byte {

‎libvuln/jsonblob/jsonblob_test.go

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,3 +118,129 @@ func TestEnrichments(t *testing.T) {
118118
}
119119
t.Logf("wrote:\n%s", buf.String())
120120
}
121+
122+
func TestIterationWithBreak(t *testing.T) {
123+
ctx := context.Background()
124+
a, err := New()
125+
if err != nil {
126+
t.Fatal(err)
127+
}
128+
129+
var want, got struct {
130+
V []*claircore.Vulnerability
131+
E []driver.EnrichmentRecord
132+
}
133+
134+
want.V = test.GenUniqueVulnerabilities(10, "test")
135+
ref, err := a.UpdateVulnerabilities(ctx, "test", "", want.V)
136+
if err != nil {
137+
t.Error(err)
138+
}
139+
t.Logf("ref: %v", ref)
140+
141+
// We will break after getting vulnerabilities.
142+
test.GenEnrichments(15)
143+
ref, err = a.UpdateEnrichments(ctx, "test", "", want.E)
144+
if err != nil {
145+
t.Error(err)
146+
}
147+
t.Logf("ref: %v", ref)
148+
149+
var buf bytes.Buffer
150+
defer func() {
151+
t.Logf("wrote:\n%s", buf.String())
152+
}()
153+
r, w := io.Pipe()
154+
eg, ctx := errgroup.WithContext(ctx)
155+
eg.Go(func() error { defer w.Close(); return a.Store(w) })
156+
eg.Go(func() error {
157+
i, iErr := Iterate(io.TeeReader(r, &buf))
158+
i(func(o *driver.UpdateOperation, i RecordIter) bool {
159+
i(func(v *claircore.Vulnerability, e *driver.EnrichmentRecord) bool {
160+
switch o.Kind {
161+
case driver.VulnerabilityKind:
162+
got.V = append(got.V, v)
163+
case driver.EnrichmentKind:
164+
got.E = append(got.E, *e)
165+
default:
166+
t.Errorf("unnexpected kind: %s", o.Kind)
167+
}
168+
return true
169+
})
170+
// Stop the operation iter, effectively skipping enrichments.
171+
return false
172+
})
173+
return iErr()
174+
})
175+
if err := eg.Wait(); err != nil {
176+
t.Error(err)
177+
}
178+
if !cmp.Equal(got, want) {
179+
t.Error(cmp.Diff(got, want))
180+
}
181+
}
182+
183+
func TestIterationWithSkip(t *testing.T) {
184+
ctx := context.Background()
185+
a, err := New()
186+
if err != nil {
187+
t.Fatal(err)
188+
}
189+
190+
var want, got struct {
191+
V []*claircore.Vulnerability
192+
E []driver.EnrichmentRecord
193+
}
194+
195+
want.V = test.GenUniqueVulnerabilities(10, "test")
196+
ref, err := a.UpdateVulnerabilities(ctx, "test", "", want.V)
197+
if err != nil {
198+
t.Error(err)
199+
}
200+
t.Logf("ref: %v", ref)
201+
202+
// We will skip the updater "skip this".
203+
test.GenUniqueVulnerabilities(10, "skip this")
204+
205+
want.E = test.GenEnrichments(15)
206+
ref, err = a.UpdateEnrichments(ctx, "test", "", want.E)
207+
if err != nil {
208+
t.Error(err)
209+
}
210+
t.Logf("ref: %v", ref)
211+
212+
var buf bytes.Buffer
213+
defer func() {
214+
t.Logf("wrote:\n%s", buf.String())
215+
}()
216+
r, w := io.Pipe()
217+
eg, ctx := errgroup.WithContext(ctx)
218+
eg.Go(func() error { defer w.Close(); return a.Store(w) })
219+
eg.Go(func() error {
220+
i, iErr := Iterate(io.TeeReader(r, &buf))
221+
i(func(o *driver.UpdateOperation, i RecordIter) bool {
222+
if o.Updater == "skip this" {
223+
return true
224+
}
225+
i(func(v *claircore.Vulnerability, e *driver.EnrichmentRecord) bool {
226+
switch o.Kind {
227+
case driver.VulnerabilityKind:
228+
got.V = append(got.V, v)
229+
case driver.EnrichmentKind:
230+
got.E = append(got.E, *e)
231+
default:
232+
t.Errorf("unnexpected kind: %s", o.Kind)
233+
}
234+
return true
235+
})
236+
return true
237+
})
238+
return iErr()
239+
})
240+
if err := eg.Wait(); err != nil {
241+
t.Error(err)
242+
}
243+
if !cmp.Equal(got, want) {
244+
t.Error(cmp.Diff(got, want))
245+
}
246+
}

0 commit comments

Comments
 (0)