From efb4281d28ec3e2c144ebdc54b0a7abd434a5c7e Mon Sep 17 00:00:00 2001
From: Juerg Wullschleger <juerg@google.com>
Date: Wed, 13 Mar 2024 07:54:11 -0700
Subject: [PATCH] Add Benchmark tests for PRFs in Golang.

BenchmarkComputePRF/HMAC_SHA256_PRF_16-8                  285105              3884 ns/op             528 B/op          5 allocs/op
BenchmarkComputePRF/HMAC_SHA256_PRF_16k-8                  19966             59463 ns/op             528 B/op          5 allocs/op
BenchmarkComputePRF/HMAC_SHA512_PRF_16-8                  256862              4397 ns/op             768 B/op          5 allocs/op
BenchmarkComputePRF/HMAC_SHA512_PRF_16k-8                  28039             42464 ns/op             768 B/op          5 allocs/op
BenchmarkComputePRF/HKDF_SHA256_16-8                      133785              8502 ns/op            1361 B/op         15 allocs/op
BenchmarkComputePRF/HKDF_SHA256_16k-8                      18424             65048 ns/op            1361 B/op         15 allocs/op
BenchmarkComputePRF/AES_CMAC_PRF_16-8                    3682260               323.0 ns/op            48 B/op          3 allocs/op
BenchmarkComputePRF/AES_CMAC_PRF_16k-8                     10000            114050 ns/op              48 B/op          3 allocs/op

PiperOrigin-RevId: 615416650
Change-Id: Ifbf696f49c94544f5e51a1edd86a0acf06fbf4ee
---
 prf/BUILD.bazel           |  1 +
 prf/prf_benchmark_test.go | 95 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+)
 create mode 100644 prf/prf_benchmark_test.go

diff --git a/prf/BUILD.bazel b/prf/BUILD.bazel
index 022cfbd..53e5250 100644
--- a/prf/BUILD.bazel
+++ b/prf/BUILD.bazel
@@ -41,6 +41,7 @@ go_test(
         "aes_cmac_prf_key_manager_test.go",
         "hkdf_prf_key_manager_test.go",
         "hmac_prf_key_manager_test.go",
+        "prf_benchmark_test.go",
         "prf_key_templates_test.go",
         "prf_set_factory_test.go",
         "prf_test.go",
diff --git a/prf/prf_benchmark_test.go b/prf/prf_benchmark_test.go
new file mode 100644
index 0000000..b093e25
--- /dev/null
+++ b/prf/prf_benchmark_test.go
@@ -0,0 +1,95 @@
+// Copyright 2024 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 prf_test
+
+import (
+	"testing"
+
+	"github.com/tink-crypto/tink-go/v2/keyset"
+	"github.com/tink-crypto/tink-go/v2/prf"
+	"github.com/tink-crypto/tink-go/v2/subtle/random"
+	tinkpb "github.com/tink-crypto/tink-go/v2/proto/tink_go_proto"
+)
+
+// Benchmarks for PRF algorithms.
+
+func BenchmarkComputePRF(b *testing.B) {
+	const (
+		outputLength = 16
+	)
+	testCases := []struct {
+		name     string
+		template *tinkpb.KeyTemplate
+		dataSize uint32
+	}{
+		{
+			name:     "HMAC_SHA256_PRF_16",
+			template: prf.HMACSHA256PRFKeyTemplate(),
+			dataSize: 16,
+		}, {
+			name:     "HMAC_SHA256_PRF_16k",
+			template: prf.HMACSHA256PRFKeyTemplate(),
+			dataSize: 16 * 1024,
+		}, {
+			name:     "HMAC_SHA512_PRF_16",
+			template: prf.HMACSHA512PRFKeyTemplate(),
+			dataSize: 16,
+		}, {
+			name:     "HMAC_SHA512_PRF_16k",
+			template: prf.HMACSHA512PRFKeyTemplate(),
+			dataSize: 16 * 1024,
+		}, {
+			name:     "HKDF_SHA256_16",
+			template: prf.HKDFSHA256PRFKeyTemplate(),
+			dataSize: 16,
+		}, {
+			name:     "HKDF_SHA256_16k",
+			template: prf.HKDFSHA256PRFKeyTemplate(),
+			dataSize: 16 * 1024,
+		}, {
+			name:     "AES_CMAC_PRF_16",
+			template: prf.AESCMACPRFKeyTemplate(),
+			dataSize: 16,
+		}, {
+			name:     "AES_CMAC_PRF_16k",
+			template: prf.AESCMACPRFKeyTemplate(),
+			dataSize: 16 * 1024,
+		},
+	}
+	for _, tc := range testCases {
+		b.Run(tc.name, func(b *testing.B) {
+			b.ReportAllocs()
+
+			handle, err := keyset.NewHandle(tc.template)
+			if err != nil {
+				b.Fatal(err)
+			}
+			primitive, err := prf.NewPRFSet(handle)
+			if err != nil {
+				b.Fatal(err)
+			}
+			data := random.GetRandomBytes(tc.dataSize)
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				_, err := primitive.ComputePrimaryPRF(data, outputLength)
+				if err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}