Repository: google/exposure-notifications-server
Pull request #1461 (2021-03-16 23:24:16)

Clean up duplicate and unused code in util
It's my goal to eventually get rid of util packages, but this is a start.

**Release Note**

<!-- Enter your extended release note in the below block. If the PR requires
additional action from users switching to the new release, include the string
"action required". If no release note is required, write "NONE". -->

```release-note
NONE
```


Patch:
diff --git a/pkg/util/generate.go b/pkg/util/generate.go
index 04bbf79..da956cc 100644
--- a/pkg/util/generate.go
+++ b/pkg/util/generate.go
@@ -23,8 +23,8 @@ import (
 	"math/big"
 	"time"
 
+	"github.com/google/exposure-notifications-server/internal/project"
 	v1 "github.com/google/exposure-notifications-server/pkg/api/v1"
-	"github.com/google/exposure-notifications-server/pkg/base64util"
 	"github.com/google/exposure-notifications-server/pkg/timeutils"
 )
 
@@ -52,15 +52,6 @@ func RandomInt(maxValue int) (int, error) {
 	return int(n.Int64()), nil
 }
 
-// RandomIntWithMin is inclusive, [min:max].
-func RandomIntWithMin(min, max int) (int, error) {
-	n, err := rand.Int(rand.Reader, big.NewInt(int64(max-min+1)))
-	if err != nil {
-		return 0, err
-	}
-	return int(n.Int64()) + min, nil
-}
-
 func RandomRevisedReportType() (string, error) {
 	n, err := RandomInt(2)
 	if err != nil {
@@ -97,15 +88,6 @@ func RandomTransmissionRisk() (int, error) {
 	return n + 1, err
 }
 
-// RandomArrValue chooses a random element from the array.
-func RandomArrValue(arr []string) (string, error) {
-	n, err := RandomInt(len(arr))
-	if err != nil {
-		return "", err
-	}
-	return arr[n], nil
-}
-
 // GenerateExposuresForIntervals generates a key for each interval start passed in
 func GenerateExposuresForIntervals(intervals []int32) ([]v1.ExposureKey, error) {
 	exposureKeys := make([]v1.ExposureKey, len(intervals))
@@ -161,10 +143,12 @@ func GenerateExposureKeys(numKeys, tr int, randomInterval bool) []v1.ExposureKey
 
 // RandomExposureKey creates a random exposure key.
 func RandomExposureKey(intervalNumber int32, intervalCount int32, transmissionRisk int) (v1.ExposureKey, error) {
-	key, err := GenerateKey()
+	b, err := RandomTEK()
 	if err != nil {
 		return v1.ExposureKey{}, err
 	}
+	key := base64.StdEncoding.EncodeToString(b)
+
 	return v1.ExposureKey{
 		Key:              key,
 		IntervalNumber:   intervalNumber,
@@ -173,35 +157,11 @@ func RandomExposureKey(intervalNumber int32, intervalCount int32, transmissionRi
 	}, nil
 }
 
-// RandomBytes generates a random byte sequence.
-func RandomBytes(arrLen int) ([]byte, error) {
-	padding := make([]byte, arrLen)
-	_, err := rand.Read(padding)
+// RandomTEK generates a new random 16-byte TEK.
+func RandomTEK() ([]byte, error) {
+	b, err := project.RandomBytes(dkLen)
 	if err != nil {
 		return nil, err
 	}
-	return padding, nil
-}
-
-// GenerateKey generates a key.
-func GenerateKey() (string, error) {
-	b, err := RandomBytes(dkLen)
-	if err != nil {
-		return "", err
-	}
-	return ToBase64(b), nil
-}
-
-// ToBase64 encodes bytes array to base64.
-func ToBase64(key []byte) string {
-	return base64.StdEncoding.EncodeToString(key)
-}
-
-// DecodeKey decodes base64 string to []byte.
-func DecodeKey(b64key string) []byte {
-	k, err := base64util.DecodeString(b64key)
-	if err != nil {
-		log.Fatalf("unable to decode key: %v", err)
-	}
-	return k
+	return b, nil
 }
diff --git a/tools/exposure-client/main.go b/tools/exposure-client/main.go
index 15bce5c..51b9031 100644
--- a/tools/exposure-client/main.go
+++ b/tools/exposure-client/main.go
@@ -30,6 +30,7 @@ import (
 	"time"
 
 	"github.com/google/exposure-notifications-server/internal/buildinfo"
+	"github.com/google/exposure-notifications-server/internal/project"
 	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1"
 	"github.com/google/exposure-notifications-server/pkg/logging"
 	"github.com/google/exposure-notifications-server/pkg/util"
@@ -76,7 +77,7 @@ func realMain(ctx context.Context) error {
 		return fmt.Errorf("failed to get random int: %w", err)
 	}
 
-	padding, err := util.RandomBytes(i + 1000)
+	padding, err := project.RandomBytes(i + 1000)
 	if err != nil {
 		return fmt.Errorf("failed to get random padding: %w", err)
 	}


Test patch:
diff --git a/internal/export/worker_test.go b/internal/export/worker_test.go
index c4778b6..8b3241a 100644
--- a/internal/export/worker_test.go
+++ b/internal/export/worker_test.go
@@ -99,11 +99,7 @@ func TestEnsureMinExposures(t *testing.T) {
 	}
 
 	for _, k := range exposures {
-		b64, err := util.GenerateKey()
-		if err != nil {
-			t.Fatalf("unable to generate exposure keys: %v", err)
-		}
-		k.ExposureKey = util.DecodeKey(b64)
+		k.ExposureKey = randomTEK(t)
 	}
 
 	numKeys := 100
@@ -136,15 +132,6 @@ func TestEnsureMinExposures(t *testing.T) {
 	}
 }
 
-func getKey(t *testing.T) []byte {
-	t.Helper()
-	eKey, err := util.RandomBytes(verifyapi.KeyLength)
-	if err != nil {
-		t.Fatal(err)
-	}
-	return eKey
-}
-
 func TestBatchExposures(t *testing.T) {
 	t.Parallel()
 
@@ -192,7 +179,7 @@ func TestBatchExposures(t *testing.T) {
 			for i := 0; i < tc.count; i++ {
 				// home country non-traveler
 				exposures[i*3] = &publishmodel.Exposure{
-					ExposureKey:     getKey(t),
+					ExposureKey:     randomTEK(t),
 					Regions:         []string{"US"},
 					IntervalNumber:  100,
 					IntervalCount:   144,
@@ -203,7 +190,7 @@ func TestBatchExposures(t *testing.T) {
 				}
 				// foreign country traveler
 				exposures[i*3+1] = &publishmodel.Exposure{
-					ExposureKey:     getKey(t),
+					ExposureKey:     randomTEK(t),
 					Regions:         []string{"CA"},
 					IntervalNumber:  100,
 					IntervalCount:   144,
@@ -214,7 +201,7 @@ func TestBatchExposures(t *testing.T) {
 				}
 				// foreign country non-traveler
 				exposures[i*3+2] = &publishmodel.Exposure{
-					ExposureKey:     getKey(t),
+					ExposureKey:     randomTEK(t),
 					Regions:         []string{"CA"},
 					IntervalNumber:  100,
 					IntervalCount:   144,
@@ -401,7 +388,7 @@ func TestVariableBatchMaxSize(t *testing.T) {
 	for i := 0; i < 20; i++ {
 		// All keys aligned to the same hour.
 		exposures[i] = &publishmodel.Exposure{
-			ExposureKey:     getKey(t),
+			ExposureKey:     randomTEK(t),
 			Regions:         []string{"US"},
 			IntervalNumber:  100,
 			IntervalCount:   144,
@@ -467,6 +454,17 @@ func TestVariableBatchMaxSize(t *testing.T) {
 	}
 }
 
+// randomTEK is like util.RandomTEK, but handles the error from tb.
+func randomTEK(tb testing.TB) []byte {
+	tb.Helper()
+
+	b, err := util.RandomTEK()
+	if err != nil {
+		tb.Fatalf("failed to generate tek: %v", err)
+	}
+	return b
+}
+
 func TestExportFilename(t *testing.T) {
 	t.Parallel()