From 3013be3e6a55b35503c46b48e988087eaca3e7f0 Mon Sep 17 00:00:00 2001
From: Chris Fenner <cfenn@google.com>
Date: Fri, 9 May 2025 16:52:45 +0000
Subject: [PATCH] add test-vector-based unit test for KDFa and KDFe

---
 tpm2/test/kdfa_test.go               | 24 +++++++++++++
 tpm2/test/kdfe_test.go               | 24 +++++++++++++
 tpm2/test/testvectors/testvectors.go | 52 ++++++++++++++++++++++++++++
 3 files changed, 100 insertions(+)
 create mode 100644 tpm2/test/kdfa_test.go
 create mode 100644 tpm2/test/kdfe_test.go

diff --git a/tpm2/test/kdfa_test.go b/tpm2/test/kdfa_test.go
new file mode 100644
index 00000000..4969a23b
--- /dev/null
+++ b/tpm2/test/kdfa_test.go
@@ -0,0 +1,24 @@
+package tpm2test
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/google/go-tpm/tpm2"
+	"github.com/google/go-tpm/tpm2/test/testvectors"
+)
+
+func TestKDFa(t *testing.T) {
+	for _, testcase := range testvectors.KDFa(t) {
+		h, err := tpm2.TPMIAlgHash(testcase.HashAlg).Hash()
+		if err != nil {
+			t.Fatalf("%v", err)
+		}
+		t.Run(testcase.Name, func(t *testing.T) {
+			result := tpm2.KDFa(h, testcase.Key, testcase.Label, testcase.ContextU, testcase.ContextV, testcase.Bits)
+			if !bytes.Equal(result, testcase.Result) {
+				t.Errorf("KDFa() = %x\nwant %x", result, testcase.Result)
+			}
+		})
+	}
+}
diff --git a/tpm2/test/kdfe_test.go b/tpm2/test/kdfe_test.go
new file mode 100644
index 00000000..db06d0ec
--- /dev/null
+++ b/tpm2/test/kdfe_test.go
@@ -0,0 +1,24 @@
+package tpm2test
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/google/go-tpm/tpm2"
+	"github.com/google/go-tpm/tpm2/test/testvectors"
+)
+
+func TestKDFe(t *testing.T) {
+	for _, testcase := range testvectors.KDFe(t) {
+		h, err := tpm2.TPMIAlgHash(testcase.HashAlg).Hash()
+		if err != nil {
+			t.Fatalf("%v", err)
+		}
+		t.Run(testcase.Name, func(t *testing.T) {
+			result := tpm2.KDFe(h, testcase.Z, testcase.Label, testcase.ContextU, testcase.ContextV, testcase.Bits)
+			if !bytes.Equal(result, testcase.Result) {
+				t.Errorf("KDFe() = %x\nwant %x", result, testcase.Result)
+			}
+		})
+	}
+}
diff --git a/tpm2/test/testvectors/testvectors.go b/tpm2/test/testvectors/testvectors.go
index 3742a5ed..105bb33c 100644
--- a/tpm2/test/testvectors/testvectors.go
+++ b/tpm2/test/testvectors/testvectors.go
@@ -71,3 +71,55 @@ func RSALabeledEncapsulation(t *testing.T) []RSALabeledEncapsTestCase {
 
 	return testCases
 }
+
+//go:embed kdfa.json
+var kdfaJSON []byte
+
+// KDFaTestCase is a test case for KDFa.
+type KDFaTestCase struct {
+	Name     string
+	HashAlg  uint16
+	Key      hexBytes
+	Label    string
+	ContextU hexBytes
+	ContextV hexBytes
+	Bits     int
+	Result   hexBytes
+}
+
+func KDFa(t *testing.T) []KDFaTestCase {
+	t.Helper()
+
+	var testCases []KDFaTestCase
+	if err := json.Unmarshal(kdfaJSON, &testCases); err != nil {
+		t.Fatalf("could not unmarshal JSON: %v", err)
+	}
+
+	return testCases
+}
+
+//go:embed kdfe.json
+var kdfeJSON []byte
+
+// KDFeTestCase is a test case for KDFe.
+type KDFeTestCase struct {
+	Name     string
+	HashAlg  uint16
+	Z        hexBytes
+	Label    string
+	ContextU hexBytes
+	ContextV hexBytes
+	Bits     int
+	Result   hexBytes
+}
+
+func KDFe(t *testing.T) []KDFeTestCase {
+	t.Helper()
+
+	var testCases []KDFeTestCase
+	if err := json.Unmarshal(kdfeJSON, &testCases); err != nil {
+		t.Fatalf("could not unmarshal JSON: %v", err)
+	}
+
+	return testCases
+}