From 30ec538b47ffdd3d9c4696dbe874c6ab96ea71b3 Mon Sep 17 00:00:00 2001
From: Markus Rudy <mr@edgeless.systems>
Date: Fri, 1 Nov 2024 10:11:21 +0100
Subject: [PATCH] atls: probe for vTPM support

---
 internal/atls/issuer.go      |  35 ++++++++++--
 internal/atls/issuer_test.go | 103 +++++++++++++++++++++++++++++++++++
 2 files changed, 133 insertions(+), 5 deletions(-)
 create mode 100644 internal/atls/issuer_test.go

diff --git a/internal/atls/issuer.go b/internal/atls/issuer.go
index 09ac4a192d..ba89bdd7fc 100644
--- a/internal/atls/issuer.go
+++ b/internal/atls/issuer.go
@@ -6,6 +6,7 @@ package atls
 import (
 	"fmt"
 	"log/slog"
+	"os"
 
 	"github.com/edgelesssys/contrast/internal/attestation/snp"
 	"github.com/edgelesssys/contrast/internal/attestation/tdx"
@@ -15,17 +16,41 @@ import (
 
 // PlatformIssuer creates an attestation issuer for the current platform.
 func PlatformIssuer(log *slog.Logger) (Issuer, error) {
-	cpuid.Detect()
+	var issuer Issuer
 	switch {
 	case cpuid.CPU.Supports(cpuid.SEV_SNP):
-		return snp.NewIssuer(
+		issuer = snp.NewIssuer(
 			logger.NewWithAttrs(logger.NewNamed(log, "issuer"), map[string]string{"tee-type": "snp"}),
-		), nil
+		)
 	case cpuid.CPU.Supports(cpuid.TDX_GUEST):
-		return tdx.NewIssuer(
+		issuer = tdx.NewIssuer(
 			logger.NewWithAttrs(logger.NewNamed(log, "issuer"), map[string]string{"tee-type": "tdx"}),
-		), nil
+		)
 	default:
 		return nil, fmt.Errorf("unsupported platform: %T", cpuid.CPU)
 	}
+
+	if hasTPM() {
+		issuer = &vtpmIssuer{Issuer: issuer}
+	}
+	return issuer, nil
+}
+
+var tpmDevice = "/dev/tpm0"
+
+func hasTPM() bool {
+	f, err := os.Open(tpmDevice)
+	if err == nil {
+		f.Close()
+		return true
+	}
+	// If the device does not exist, we don't have a TPM.
+	// If the device exists but there is no backing TPM, the Open call fails with ENODEV.
+	return false
+}
+
+// vtpmIssuer issues attestation statements for VMs with a TPM.
+// TODO(burgerdev): this is currently a mock that just delegates to an underlying issuer.
+type vtpmIssuer struct {
+	Issuer
 }
diff --git a/internal/atls/issuer_test.go b/internal/atls/issuer_test.go
new file mode 100644
index 0000000000..a600f13391
--- /dev/null
+++ b/internal/atls/issuer_test.go
@@ -0,0 +1,103 @@
+// Copyright 2024 Edgeless Systems GmbH
+// SPDX-License-Identifier: AGPL-3.0-only
+
+package atls
+
+import (
+	"log/slog"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/edgelesssys/contrast/internal/attestation/snp"
+	"github.com/edgelesssys/contrast/internal/attestation/tdx"
+	"github.com/klauspost/cpuid/v2"
+	"github.com/stretchr/testify/require"
+)
+
+func TestPlatformIssuer(t *testing.T) {
+	d := t.TempDir()
+	tpmFile := filepath.Join(d, "tpm0")
+	f, err := os.Create(tpmFile)
+	require.NoError(t, err)
+	t.Cleanup(func() { f.Close() })
+
+	snpCPU := cpuid.CPUInfo{}
+	snpCPU.Enable(cpuid.SEV_SNP)
+
+	tdxCPU := cpuid.CPUInfo{}
+	tdxCPU.Enable(cpuid.TDX_GUEST)
+
+	for _, tc := range []struct {
+		name string
+		tpm  string
+		cpu  *cpuid.CPUInfo
+
+		expectIssuerType any
+		expectError      bool
+	}{
+		{
+			name: "tpm-snp",
+			tpm:  tpmFile,
+			cpu:  &snpCPU,
+
+			expectIssuerType: &vtpmIssuer{},
+		},
+		{
+			name: "tpm-tdx",
+			tpm:  tpmFile,
+			cpu:  &tdxCPU,
+
+			expectIssuerType: &vtpmIssuer{},
+		},
+		{
+			name: "notpm-snp",
+			tpm:  "/invalid",
+			cpu:  &snpCPU,
+
+			expectIssuerType: &snp.Issuer{},
+		},
+		{
+			name: "notpm-tdx",
+			tpm:  "/invalid",
+			cpu:  &tdxCPU,
+
+			expectIssuerType: &tdx.Issuer{},
+		},
+		{
+			name: "notpm-notee",
+			tpm:  "/invalid",
+			cpu:  &cpuid.CPUInfo{},
+
+			expectError: true,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			require := require.New(t)
+			savedCPU := cpuid.CPU
+			t.Cleanup(func() {
+				cpuid.CPU = savedCPU
+			})
+			if tc.cpu != nil {
+				cpuid.CPU = *tc.cpu
+			}
+
+			savedTPMDevice := tpmDevice
+			t.Cleanup(func() {
+				tpmDevice = savedTPMDevice
+			})
+			if tc.tpm != "" {
+				tpmDevice = tc.tpm
+			}
+
+			issuer, err := PlatformIssuer(slog.Default())
+			if tc.expectError {
+				require.Error(err)
+				return
+			}
+			require.NoError(err)
+
+			require.IsType(tc.expectIssuerType, issuer)
+		})
+	}
+}