diskless-workstation/files/enrol/simtest/tpmsim_test.go

// Package simtest validates tpmkey's TPM ceremony against the go-tpm software
// simulator (the MS TPM 2.0 reference). It lives in its own module so the CGO
// simulator never enters the shipped fleet-enrol build; see go.mod.
package simtest

import (
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"crypto/x509/pkix"
	"testing"

	"fleet-enrol/tpmkey"

	"github.com/google/go-tpm-tools/simulator"
	"github.com/google/go-tpm/tpm2"
	"github.com/google/go-tpm/tpm2/transport"
)

func openSim(t *testing.T) transport.TPM {
	t.Helper()
	sim, err := simulator.Get()
	if err != nil {
		t.Fatalf("simulator: %v", err)
	}
	t.Cleanup(func() { sim.Close() })
	return transport.FromReadWriter(sim)
}

func TestEnsureKeyAndSignCSR(t *testing.T) {
	tpm := openSim(t)
	const h = tpm2.TPMHandle(0x81010001)

	signer, created, err := tpmkey.EnsureKey(tpm, h)
	if err != nil {
		t.Fatalf("EnsureKey: %v", err)
	}
	if !created {
		t.Errorf("first EnsureKey should report created=true")
	}
	if _, ok := signer.Public().(*ecdsa.PublicKey); !ok {
		t.Fatalf("signer public key is not ECDSA")
	}

	// idempotent: second call reuses the persistent key (same pubkey).
	signer2, created2, err := tpmkey.EnsureKey(tpm, h)
	if err != nil || created2 {
		t.Fatalf("second EnsureKey: created=%v err=%v", created2, err)
	}
	if !signer.Public().(*ecdsa.PublicKey).Equal(signer2.Public()) {
		t.Errorf("persistent key changed between calls")
	}

	// the TPM signer must produce a CSR that verifies — this is the exact
	// openssl-tpm2-provider path the Go binary replaces.
	tmpl := &x509.CertificateRequest{Subject: pkix.Name{CommonName: "fleet-host"}}
	der, err := x509.CreateCertificateRequest(rand.Reader, tmpl, signer)
	if err != nil {
		t.Fatalf("CreateCertificateRequest with TPM signer: %v", err)
	}
	csr, err := x509.ParseCertificateRequest(der)
	if err != nil {
		t.Fatalf("parse CSR: %v", err)
	}
	if err := csr.CheckSignature(); err != nil {
		t.Fatalf("CSR signature (TPM-signed) invalid: %v", err)
	}

	// raw signer sanity: sign a digest, verify with the public key.
	dig := sha256.Sum256([]byte("hello"))
	sig, err := signer.Sign(rand.Reader, dig[:], nil)
	if err != nil {
		t.Fatalf("sign: %v", err)
	}
	if !ecdsa.VerifyASN1(signer.Public().(*ecdsa.PublicKey), dig[:], sig) {
		t.Fatalf("TPM ECDSA signature did not verify")
	}
}

func TestNVWriteAndPresence(t *testing.T) {
	tpm := openSim(t)
	const nv = tpm2.TPMHandle(0x01800001)

	if tpmkey.NVPresent(tpm, nv) {
		t.Fatalf("NV should be empty on a fresh TPM")
	}
	// realistic cert-sized payload that also crosses the 512B write chunk.
	want := make([]byte, 700)
	rand.Read(want)
	if err := tpmkey.WriteNVCert(tpm, nv, want); err != nil {
		t.Fatalf("WriteNVCert: %v", err)
	}
	if !tpmkey.NVPresent(tpm, nv) {
		t.Fatalf("NV should be present after write")
	}
	// re-write (re-enrolment) must succeed (undefine + redefine).
	if err := tpmkey.WriteNVCert(tpm, nv, want[:400]); err != nil {
		t.Fatalf("WriteNVCert rewrite: %v", err)
	}
}

// Regression for the bench failure "key at handle is not ECC ... selector value
// was 1": an older provisioning path left an RSA key at 0x81010001, and EnsureKey
// fatally failed trying to read it as ECC. It must now evict the stale key and
// create a fresh ECC one.
func TestEnsureKeyEvictsStaleNonECC(t *testing.T) {
	tpm := openSim(t)
	const h = tpm2.TPMHandle(0x81010001)

	// Pre-populate the handle with an RSA key (mimics tpm2-tools/openssl default).
	rsaTmpl := tpm2.TPMTPublic{
		Type:    tpm2.TPMAlgRSA,
		NameAlg: tpm2.TPMAlgSHA256,
		ObjectAttributes: tpm2.TPMAObject{
			FixedTPM: true, FixedParent: true, SensitiveDataOrigin: true,
			UserWithAuth: true, SignEncrypt: true, NoDA: true,
		},
		Parameters: tpm2.NewTPMUPublicParms(tpm2.TPMAlgRSA, &tpm2.TPMSRSAParms{
			Symmetric: tpm2.TPMTSymDefObject{Algorithm: tpm2.TPMAlgNull},
			Scheme: tpm2.TPMTRSAScheme{
				Scheme: tpm2.TPMAlgRSASSA,
				Details: tpm2.NewTPMUAsymScheme(tpm2.TPMAlgRSASSA, &tpm2.TPMSSigSchemeRSASSA{
					HashAlg: tpm2.TPMAlgSHA256,
				}),
			},
			KeyBits: 2048,
		}),
	}
	cp, err := tpm2.CreatePrimary{
		PrimaryHandle: tpm2.TPMRHOwner, InPublic: tpm2.New2B(rsaTmpl),
	}.Execute(tpm)
	if err != nil {
		t.Fatalf("create RSA primary: %v", err)
	}
	if _, err := (tpm2.EvictControl{
		Auth:             tpm2.TPMRHOwner,
		ObjectHandle:     &tpm2.NamedHandle{Handle: cp.ObjectHandle, Name: cp.Name},
		PersistentHandle: h,
	}).Execute(tpm); err != nil {
		t.Fatalf("persist stale RSA key: %v", err)
	}
	_, _ = (tpm2.FlushContext{FlushHandle: cp.ObjectHandle}).Execute(tpm)

	// EnsureKey must self-heal: evict the RSA key, create + persist a fresh ECC key.
	signer, created, err := tpmkey.EnsureKey(tpm, h)
	if err != nil {
		t.Fatalf("EnsureKey over stale RSA key: %v", err)
	}
	if !created {
		t.Errorf("EnsureKey should have recreated the key (created=true)")
	}
	if _, ok := signer.Public().(*ecdsa.PublicKey); !ok {
		t.Fatalf("recreated signer is not ECDSA")
	}
	// and it must be usable for signing (the whole point of the recovery).
	dig := sha256.Sum256([]byte("post-heal"))
	sig, err := signer.Sign(rand.Reader, dig[:], nil)
	if err != nil {
		t.Fatalf("sign after heal: %v", err)
	}
	if !ecdsa.VerifyASN1(signer.Public().(*ecdsa.PublicKey), dig[:], sig) {
		t.Fatalf("post-heal ECDSA signature did not verify")
	}
}