Skip to content

Commit 3706061

Browse files
ktarpleektarplee
authored
allow pkg/v1/random to accept a RNG source (#1675)
* allow pkg/v1/random to accept a RNG source * added a unit test for the new options
1 parent b7c6e9d commit 3706061

File tree

4 files changed

+113
-9
lines changed

4 files changed

+113
-9
lines changed

‎pkg/v1/random/image.go‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,10 @@ import (
1818
"archive/tar"
1919
"bytes"
2020
"crypto"
21-
"crypto/rand"
2221
"encoding/hex"
2322
"fmt"
2423
"io"
25-
mrand "math/rand"
24+
"math/rand"
2625
"time"
2726

2827
v1 "github.com/google/go-containerregistry/pkg/v1"
@@ -57,10 +56,10 @@ func (ul *uncompressedLayer) MediaType() (types.MediaType, error) {
5756
var _ partial.UncompressedLayer = (*uncompressedLayer)(nil)
5857

5958
// Image returns a pseudo-randomly generated Image.
60-
func Image(byteSize, layers int64) (v1.Image, error) {
59+
func Image(byteSize, layers int64, options ...Option) (v1.Image, error) {
6160
adds := make([]mutate.Addendum, 0, 5)
6261
for i := int64(0); i < layers; i++ {
63-
layer, err := Layer(byteSize, types.DockerLayer)
62+
layer, err := Layer(byteSize, types.DockerLayer, options...)
6463
if err != nil {
6564
return nil, err
6665
}
@@ -79,8 +78,11 @@ func Image(byteSize, layers int64) (v1.Image, error) {
7978
}
8079

8180
// Layer returns a layer with pseudo-randomly generated content.
82-
func Layer(byteSize int64, mt types.MediaType) (v1.Layer, error) {
83-
fileName := fmt.Sprintf("random_file_%d.txt", mrand.Int()) //nolint: gosec
81+
func Layer(byteSize int64, mt types.MediaType, options ...Option) (v1.Layer, error) {
82+
o := getOptions(options)
83+
rng := rand.New(o.source) //nolint:gosec
84+
85+
fileName := fmt.Sprintf("random_file_%d.txt", rng.Int())
8486

8587
// Hash the contents as we write it out to the buffer.
8688
var b bytes.Buffer
@@ -96,7 +98,7 @@ func Layer(byteSize int64, mt types.MediaType) (v1.Layer, error) {
9698
}); err != nil {
9799
return nil, err
98100
}
99-
if _, err := io.CopyN(tw, rand.Reader, byteSize); err != nil {
101+
if _, err := io.CopyN(tw, rng, byteSize); err != nil {
100102
return nil, err
101103
}
102104
if err := tw.Close(); err != nil {

‎pkg/v1/random/image_test.go‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,10 @@ package random
1616

1717
import (
1818
"archive/tar"
19+
"bytes"
1920
"errors"
2021
"io"
22+
"math/rand"
2123
"testing"
2224

2325
"github.com/google/go-containerregistry/pkg/v1/types"
@@ -127,3 +129,43 @@ func TestRandomLayer(t *testing.T) {
127129
t.Errorf("Layer contained more files; got %v, want EOF", err)
128130
}
129131
}
132+
133+
func TestRandomLayerSource(t *testing.T) {
134+
layerData := func(o ...Option) []byte {
135+
l, err := Layer(1024, types.DockerLayer, o...)
136+
if err != nil {
137+
t.Fatalf("Layer: %v", err)
138+
}
139+
140+
rc, err := l.Compressed()
141+
if err != nil {
142+
t.Fatalf("Compressed(): %v", err)
143+
}
144+
defer rc.Close()
145+
146+
data, err := io.ReadAll(rc)
147+
if err != nil {
148+
t.Fatalf("Read: %v", err)
149+
}
150+
return data
151+
}
152+
153+
data0a := layerData(WithSource(rand.NewSource(0)))
154+
data0b := layerData(WithSource(rand.NewSource(0)))
155+
data1 := layerData(WithSource(rand.NewSource(1)))
156+
157+
if !bytes.Equal(data0a, data0b) {
158+
t.Error("Expected the layer data to be the same with the same seed")
159+
}
160+
161+
if bytes.Equal(data0a, data1) {
162+
t.Error("Expected the layer data to be different with different seeds")
163+
}
164+
165+
dataA := layerData()
166+
dataB := layerData()
167+
168+
if bytes.Equal(dataA, dataB) {
169+
t.Error("Expected the layer data to be different with different random seeds")
170+
}
171+
}

‎pkg/v1/random/index.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ type randomIndex struct {
3131

3232
// Index returns a pseudo-randomly generated ImageIndex with count images, each
3333
// having the given number of layers of size byteSize.
34-
func Index(byteSize, layers, count int64) (v1.ImageIndex, error) {
34+
func Index(byteSize, layers, count int64, options ...Option) (v1.ImageIndex, error) {
3535
manifest := v1.IndexManifest{
3636
SchemaVersion: 2,
3737
MediaType: types.OCIImageIndex,
@@ -40,7 +40,7 @@ func Index(byteSize, layers, count int64) (v1.ImageIndex, error) {
4040

4141
images := make(map[v1.Hash]v1.Image)
4242
for i := int64(0); i < count; i++ {
43-
img, err := Image(byteSize, layers)
43+
img, err := Image(byteSize, layers, options...)
4444
if err != nil {
4545
return nil, err
4646
}

‎pkg/v1/random/options.go‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
// Copyright 2018 Google LLC All Rights Reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package random
16+
17+
import "math/rand"
18+
19+
// Option is an optional parameter to the random functions
20+
type Option func(opts *options)
21+
22+
type options struct {
23+
source rand.Source
24+
25+
// TODO opens the door to add this in the future
26+
// algorithm digest.Algorithm
27+
}
28+
29+
func getOptions(opts []Option) *options {
30+
// get a random seed
31+
32+
// TODO in go 1.20 this is fine (it will be random)
33+
seed := rand.Int63() //nolint:gosec
34+
/*
35+
// in prior go versions this needs to come from crypto/rand
36+
var b [8]byte
37+
_, err := crypto_rand.Read(b[:])
38+
if err != nil {
39+
panic("cryptographically secure random number generator is not working")
40+
}
41+
seed := int64(binary.LittleEndian.Int64(b[:]))
42+
*/
43+
44+
// defaults
45+
o := &options{
46+
source: rand.NewSource(seed),
47+
}
48+
49+
for _, opt := range opts {
50+
opt(o)
51+
}
52+
return o
53+
}
54+
55+
// WithSource sets the random number generator source
56+
func WithSource(source rand.Source) Option {
57+
return func(opts *options) {
58+
opts.source = source
59+
}
60+
}

0 commit comments

Comments
 (0)