Define limitFromHash() for use in streamingprf_test.

PiperOrigin-RevId: 476248334

go/keyderivation/internal/streamingprf/BUILD.bazel

@@ -31,6 +31,7 @@ go_test(
     srcs = [
         "hkdf_streaming_prf_key_manager_test.go",
         "hkdf_streaming_prf_test.go",
+        "streaming_prf_test.go",
     ],
     data = ["@wycheproof//testvectors:all"],
     embed = [":streamingprf"],

go/keyderivation/internal/streamingprf/hkdf_streaming_prf_key_manager_test.go

@@ -18,8 +18,6 @@ package streamingprf_test
 
 import (
 	"bytes"
-	"crypto/sha256"
-	"crypto/sha512"
 	"testing"
 
 	"google.golang.org/protobuf/proto"
@@ -58,7 +56,6 @@ func TestHKDFStreamingPRFKeyManagerPrimitive(t *testing.T) {
 				},
 				KeyValue: random.GetRandomBytes(32),
 			}
-			data := random.GetRandomBytes(32)
 			serializedKey, err := proto.Marshal(key)
 			if err != nil {
 				t.Fatalf("proto.Marshal(%v) err = %v, want nil", key, err)
@@ -67,26 +64,18 @@ func TestHKDFStreamingPRFKeyManagerPrimitive(t *testing.T) {
 			if err != nil {
 				t.Fatalf("Primitive() err = %v, want nil", err)
 			}
-			hkdf, ok := p.(streamingprf.StreamingPRF)
+			prf, ok := p.(streamingprf.StreamingPRF)
 			if !ok {
 				t.Fatal("primitive is not StreamingPRF")
 			}
-			r, err := hkdf.Compute(data)
+			r, err := prf.Compute(random.GetRandomBytes(32))
 			if err != nil {
 				t.Fatalf("Compute() err = %v, want nil", err)
 			}
-			var outLimit int
-			switch test.hash {
-			case commonpb.HashType_SHA256:
-				outLimit = sha256.Size * 255
-			case commonpb.HashType_SHA512:
-				outLimit = sha512.Size * 255
-			default:
-				outLimit = 0
-			}
-			out := make([]byte, outLimit)
+			limit := limitFromHash(t, test.hash)
+			out := make([]byte, limit)
 			n, err := r.Read(out)
-			if n != outLimit || err != nil {
+			if n != limit || err != nil {
 				t.Errorf("Read() not enough bytes: %d, %v", n, err)
 			}
 		})

go/keyderivation/internal/streamingprf/streaming_prf_test.go

@@ -0,0 +1,38 @@
+// Copyright 2022 Google LLC
+//
+// 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 streamingprf_test
+
+import (
+	"crypto/sha256"
+	"crypto/sha512"
+	"testing"
+
+	commonpb "github.com/google/tink/go/proto/common_go_proto"
+)
+
+func limitFromHash(t *testing.T, hash commonpb.HashType) (limit int) {
+	t.Helper()
+	switch hash {
+	case commonpb.HashType_SHA256:
+		limit = sha256.Size * 255
+	case commonpb.HashType_SHA512:
+		limit = sha512.Size * 255
+	default:
+		t.Fatalf("unsupported hash type: %s", hash.String())
+	}
+	return
+}