From 3706061a683bf461511f2691ceec78cab4757093 Mon Sep 17 00:00:00 2001 From: Kyle Tarplee Date: Wed, 26 Apr 2023 15:09:30 -0400 Subject: [PATCH] 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 --- pkg/v1/random/image.go | 16 +++++----- pkg/v1/random/image_test.go | 42 ++++++++++++++++++++++++++ pkg/v1/random/index.go | 4 +-- pkg/v1/random/options.go | 60 +++++++++++++++++++++++++++++++++++++ 4 files changed, 113 insertions(+), 9 deletions(-) create mode 100644 pkg/v1/random/options.go diff --git a/pkg/v1/random/image.go b/pkg/v1/random/image.go index 4b28913db..0986e204c 100644 --- a/pkg/v1/random/image.go +++ b/pkg/v1/random/image.go @@ -18,11 +18,10 @@ import ( "archive/tar" "bytes" "crypto" - "crypto/rand" "encoding/hex" "fmt" "io" - mrand "math/rand" + "math/rand" "time" v1 "github.com/google/go-containerregistry/pkg/v1" @@ -57,10 +56,10 @@ func (ul *uncompressedLayer) MediaType() (types.MediaType, error) { var _ partial.UncompressedLayer = (*uncompressedLayer)(nil) // Image returns a pseudo-randomly generated Image. -func Image(byteSize, layers int64) (v1.Image, error) { +func Image(byteSize, layers int64, options ...Option) (v1.Image, error) { adds := make([]mutate.Addendum, 0, 5) for i := int64(0); i < layers; i++ { - layer, err := Layer(byteSize, types.DockerLayer) + layer, err := Layer(byteSize, types.DockerLayer, options...) if err != nil { return nil, err } @@ -79,8 +78,11 @@ func Image(byteSize, layers int64) (v1.Image, error) { } // Layer returns a layer with pseudo-randomly generated content. -func Layer(byteSize int64, mt types.MediaType) (v1.Layer, error) { - fileName := fmt.Sprintf("random_file_%d.txt", mrand.Int()) //nolint: gosec +func Layer(byteSize int64, mt types.MediaType, options ...Option) (v1.Layer, error) { + o := getOptions(options) + rng := rand.New(o.source) //nolint:gosec + + fileName := fmt.Sprintf("random_file_%d.txt", rng.Int()) // Hash the contents as we write it out to the buffer. var b bytes.Buffer @@ -96,7 +98,7 @@ func Layer(byteSize int64, mt types.MediaType) (v1.Layer, error) { }); err != nil { return nil, err } - if _, err := io.CopyN(tw, rand.Reader, byteSize); err != nil { + if _, err := io.CopyN(tw, rng, byteSize); err != nil { return nil, err } if err := tw.Close(); err != nil { diff --git a/pkg/v1/random/image_test.go b/pkg/v1/random/image_test.go index 8f30bc7e1..3acbc65af 100644 --- a/pkg/v1/random/image_test.go +++ b/pkg/v1/random/image_test.go @@ -16,8 +16,10 @@ package random import ( "archive/tar" + "bytes" "errors" "io" + "math/rand" "testing" "github.com/google/go-containerregistry/pkg/v1/types" @@ -127,3 +129,43 @@ func TestRandomLayer(t *testing.T) { t.Errorf("Layer contained more files; got %v, want EOF", err) } } + +func TestRandomLayerSource(t *testing.T) { + layerData := func(o ...Option) []byte { + l, err := Layer(1024, types.DockerLayer, o...) + if err != nil { + t.Fatalf("Layer: %v", err) + } + + rc, err := l.Compressed() + if err != nil { + t.Fatalf("Compressed(): %v", err) + } + defer rc.Close() + + data, err := io.ReadAll(rc) + if err != nil { + t.Fatalf("Read: %v", err) + } + return data + } + + data0a := layerData(WithSource(rand.NewSource(0))) + data0b := layerData(WithSource(rand.NewSource(0))) + data1 := layerData(WithSource(rand.NewSource(1))) + + if !bytes.Equal(data0a, data0b) { + t.Error("Expected the layer data to be the same with the same seed") + } + + if bytes.Equal(data0a, data1) { + t.Error("Expected the layer data to be different with different seeds") + } + + dataA := layerData() + dataB := layerData() + + if bytes.Equal(dataA, dataB) { + t.Error("Expected the layer data to be different with different random seeds") + } +} diff --git a/pkg/v1/random/index.go b/pkg/v1/random/index.go index 89a88438f..4368bddff 100644 --- a/pkg/v1/random/index.go +++ b/pkg/v1/random/index.go @@ -31,7 +31,7 @@ type randomIndex struct { // Index returns a pseudo-randomly generated ImageIndex with count images, each // having the given number of layers of size byteSize. -func Index(byteSize, layers, count int64) (v1.ImageIndex, error) { +func Index(byteSize, layers, count int64, options ...Option) (v1.ImageIndex, error) { manifest := v1.IndexManifest{ SchemaVersion: 2, MediaType: types.OCIImageIndex, @@ -40,7 +40,7 @@ func Index(byteSize, layers, count int64) (v1.ImageIndex, error) { images := make(map[v1.Hash]v1.Image) for i := int64(0); i < count; i++ { - img, err := Image(byteSize, layers) + img, err := Image(byteSize, layers, options...) if err != nil { return nil, err } diff --git a/pkg/v1/random/options.go b/pkg/v1/random/options.go new file mode 100644 index 000000000..af1d2f969 --- /dev/null +++ b/pkg/v1/random/options.go @@ -0,0 +1,60 @@ +// Copyright 2018 Google LLC All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package random + +import "math/rand" + +// Option is an optional parameter to the random functions +type Option func(opts *options) + +type options struct { + source rand.Source + + // TODO opens the door to add this in the future + // algorithm digest.Algorithm +} + +func getOptions(opts []Option) *options { + // get a random seed + + // TODO in go 1.20 this is fine (it will be random) + seed := rand.Int63() //nolint:gosec + /* + // in prior go versions this needs to come from crypto/rand + var b [8]byte + _, err := crypto_rand.Read(b[:]) + if err != nil { + panic("cryptographically secure random number generator is not working") + } + seed := int64(binary.LittleEndian.Int64(b[:])) + */ + + // defaults + o := &options{ + source: rand.NewSource(seed), + } + + for _, opt := range opts { + opt(o) + } + return o +} + +// WithSource sets the random number generator source +func WithSource(source rand.Source) Option { + return func(opts *options) { + opts.source = source + } +}