From 841a1451740a5e349945d655b5bead1db469cf0a Mon Sep 17 00:00:00 2001
From: Herman Slatman <hermanslatman@hotmail.com>
Date: Tue, 10 Sep 2024 11:04:35 +0200
Subject: [PATCH] Change the simulated TPM options to accept initializers and
 preparers

---
 kms/tpmkms/tpmkms_simulator_test.go | 84 +++++++++++++++++++++++++++--
 tpm/caps_test.go                    |  3 +-
 tpm/tpm.go                          | 24 +++++----
 3 files changed, 96 insertions(+), 15 deletions(-)

diff --git a/kms/tpmkms/tpmkms_simulator_test.go b/kms/tpmkms/tpmkms_simulator_test.go
index 7d8b88e3..a254907a 100644
--- a/kms/tpmkms/tpmkms_simulator_test.go
+++ b/kms/tpmkms/tpmkms_simulator_test.go
@@ -37,9 +37,11 @@ import (
 	"go.step.sm/crypto/tpm/tss2"
 )
 
-type newSimulatedTPMOption func(t *testing.T, tpm *tpmp.TPM)
+type newSimulatedTPMOption any
 
-func withAK(name string) newSimulatedTPMOption {
+type newSimulatedTPMPreparerOption func(t *testing.T, tpm *tpmp.TPM)
+
+func withAK(name string) newSimulatedTPMPreparerOption {
 	return func(t *testing.T, tpm *tpmp.TPM) {
 		t.Helper()
 		_, err := tpm.CreateAK(context.Background(), name)
@@ -47,7 +49,7 @@ func withAK(name string) newSimulatedTPMOption {
 	}
 }
 
-func withKey(name string) newSimulatedTPMOption {
+func withKey(name string) newSimulatedTPMPreparerOption {
 	return func(t *testing.T, tpm *tpmp.TPM) {
 		t.Helper()
 		config := tpmp.CreateKeyConfig{
@@ -59,20 +61,38 @@ func withKey(name string) newSimulatedTPMOption {
 	}
 }
 
+func withCapabilities(caps *tpmp.Capabilities) tpmp.NewTPMOption {
+	return tpmp.WithCapabilities(caps)
+}
+
 func newSimulatedTPM(t *testing.T, opts ...newSimulatedTPMOption) *tpmp.TPM {
 	t.Helper()
+
 	tmpDir := t.TempDir()
 	tpmOpts := []tpmp.NewTPMOption{
 		withSimulator(t),
 		tpmp.WithStore(storage.NewDirstore(tmpDir)),
 	}
 
-	tpm, err := tpmp.New(tpmOpts...)
+	var preparers []newSimulatedTPMPreparerOption
+	for _, opt := range opts {
+		switch o := opt.(type) {
+		case tpmp.NewTPMOption:
+			tpmOpts = append(tpmOpts, o)
+		case newSimulatedTPMPreparerOption:
+			preparers = append(preparers, o)
+		default:
+			require.Fail(t, "invalid TPM option type provided", `TPM option type "%T"`, o)
+		}
+	}
 
+	tpm, err := tpmp.New(tpmOpts...)
 	require.NoError(t, err)
-	for _, applyTo := range opts {
+
+	for _, applyTo := range preparers {
 		applyTo(t, tpm)
 	}
+
 	return tpm
 }
 
@@ -93,6 +113,60 @@ func withSimulator(t *testing.T) tpmp.NewTPMOption {
 	return tpmp.WithSimulator(sim)
 }
 
+func TestTPMKMS_CreateKey_Capabilities(t *testing.T) {
+	tpmWithNoCaps := newSimulatedTPM(t, withCapabilities(&tpmp.Capabilities{}))
+	type fields struct {
+		tpm *tpmp.TPM
+	}
+	type args struct {
+		req *apiv1.CreateKeyRequest
+	}
+	tests := []struct {
+		name       string
+		fields     fields
+		args       args
+		assertFunc assert.ValueAssertionFunc
+		expErr     error
+	}{
+		{
+			name: "fail/unsupported-algorithm",
+			fields: fields{
+				tpm: tpmWithNoCaps,
+			},
+			args: args{
+				req: &apiv1.CreateKeyRequest{
+					Name:               "tpmkms:name=key1",
+					SignatureAlgorithm: apiv1.SHA256WithRSA,
+					Bits:               2048,
+				},
+			},
+			assertFunc: func(tt assert.TestingT, i1 interface{}, i2 ...interface{}) bool {
+				if assert.IsType(t, &apiv1.CreateKeyResponse{}, i1) {
+					r, _ := i1.(*apiv1.CreateKeyResponse)
+					return assert.Nil(t, r)
+				}
+				return false
+			},
+			expErr: errors.New(`signature algorithm "SHA256-RSA" not supported by the TPM device`),
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			k := &TPMKMS{
+				tpm: tt.fields.tpm,
+			}
+			got, err := k.CreateKey(tt.args.req)
+			if tt.expErr != nil {
+				assert.EqualError(t, err, tt.expErr.Error())
+				return
+			}
+
+			assert.NoError(t, err)
+			assert.True(t, tt.assertFunc(t, got))
+		})
+	}
+}
+
 func TestTPMKMS_CreateKey(t *testing.T) {
 	tpmWithAK := newSimulatedTPM(t, withAK("ak1"))
 	type fields struct {
diff --git a/tpm/caps_test.go b/tpm/caps_test.go
index 7cc02383..53aaf1d2 100644
--- a/tpm/caps_test.go
+++ b/tpm/caps_test.go
@@ -3,7 +3,8 @@ package tpm
 import (
 	"testing"
 
-	"github.com/smallstep/assert"
+	"github.com/stretchr/testify/assert"
+
 	"go.step.sm/crypto/tpm/algorithm"
 )
 
diff --git a/tpm/tpm.go b/tpm/tpm.go
index a19a410c..bd0dde01 100644
--- a/tpm/tpm.go
+++ b/tpm/tpm.go
@@ -75,15 +75,6 @@ func WithDisableDownload() NewTPMOption {
 	}
 }
 
-// WithCapabilities explicits sets the capabilities rather
-// than acquiring them from the TPM directly.
-func WithCapabilities(caps *Capabilities) NewTPMOption {
-	return func(o *options) error {
-		o.caps = caps
-		return nil
-	}
-}
-
 // WithSimulator is used to configure a TPM simulator implementation
 // that simulates TPM operations instead of interacting with an actual
 // TPM.
@@ -105,6 +96,21 @@ func WithCommandChannel(commandChannel CommandChannel) NewTPMOption {
 	}
 }
 
+// WithCapabilities explicitly sets the capabilities rather
+// than acquiring them from the TPM directly. The primary use
+// for this option is to ease testing different TPM capabilities.
+//
+// # Experimental
+//
+// Notice: This option is EXPERIMENTAL and may be changed or removed
+// in a later release.
+func WithCapabilities(caps *Capabilities) NewTPMOption {
+	return func(o *options) error {
+		o.caps = caps
+		return nil
+	}
+}
+
 type options struct {
 	deviceName     string
 	attestConfig   *attest.OpenConfig