From bae4f6d779fd730496416304ba2a8f34ea2eda16 Mon Sep 17 00:00:00 2001
From: Philipp Winter <phw@nymity.ch>
Date: Sun, 20 Oct 2024 09:32:49 -0500
Subject: [PATCH] Initialize NSM session on first call to `Attest`.

This makes it possible to use the nitro attester on non-nitro platforms
to verify attestation documents.
---
 cmd/veil/main.go                        |  7 +------
 internal/enclave/attester_nitro.go      | 18 +++++++++---------
 internal/enclave/attester_nitro_test.go |  7 ++-----
 internal/enclave/pcr.go                 |  7 +------
 4 files changed, 13 insertions(+), 26 deletions(-)

diff --git a/cmd/veil/main.go b/cmd/veil/main.go
index 0c9f370..5d90eb3 100644
--- a/cmd/veil/main.go
+++ b/cmd/veil/main.go
@@ -101,14 +101,9 @@ func run(ctx context.Context, out io.Writer, args []string) (err error) {
 	}
 
 	// Initialize dependencies and start the service.
-	var attester enclave.Attester
+	var attester enclave.Attester = enclave.NewNitroAttester()
 	if cfg.Testing {
 		attester = enclave.NewNoopAttester()
-	} else {
-		attester, err = enclave.NewNitroAttester()
-		if err != nil {
-			return err
-		}
 	}
 	service.Run(ctx, cfg, attester, tunnel.NewNoop())
 	return nil
diff --git a/internal/enclave/attester_nitro.go b/internal/enclave/attester_nitro.go
index f3b25e3..bf70432 100644
--- a/internal/enclave/attester_nitro.go
+++ b/internal/enclave/attester_nitro.go
@@ -19,15 +19,8 @@ type NitroAttester struct {
 }
 
 // NewNitroAttester returns a new nitroAttester.
-func NewNitroAttester() (attester Attester, err error) {
-	defer errs.Wrap(&err, "failed to create nitro attester")
-	a := new(NitroAttester)
-
-	// Open a session to the Nitro Secure Module.
-	if a.session, err = nsm.OpenDefaultSession(); err != nil {
-		return nil, err
-	}
-	return a, nil
+func NewNitroAttester() Attester {
+	return new(NitroAttester)
 }
 
 func (*NitroAttester) Type() string {
@@ -37,6 +30,13 @@ func (*NitroAttester) Type() string {
 func (a *NitroAttester) Attest(aux *AuxInfo) (_ *AttestationDoc, err error) {
 	defer errs.Wrap(&err, "failed to create attestation document")
 
+	if a.session == nil {
+		// Open a session to the Nitro Secure Module.
+		if a.session, err = nsm.OpenDefaultSession(); err != nil {
+			return nil, err
+		}
+	}
+
 	if aux == nil {
 		return nil, errors.New("aux info is nil")
 	}
diff --git a/internal/enclave/attester_nitro_test.go b/internal/enclave/attester_nitro_test.go
index a9c6784..4f15140 100644
--- a/internal/enclave/attester_nitro_test.go
+++ b/internal/enclave/attester_nitro_test.go
@@ -18,8 +18,7 @@ func TestNitroAttest(t *testing.T) {
 	if !IsEnclave() {
 		t.Skip("skipping test; not running in an enclave")
 	}
-	attester, err := NewNitroAttester()
-	require.NoError(t, err)
+	attester := NewNitroAttester()
 
 	cases := []struct {
 		name    string
@@ -59,9 +58,7 @@ func TestNitroVerify(t *testing.T) {
 		t.Skip("skipping test; not running in an enclave")
 	}
 
-	attester, err := NewNitroAttester()
-	require.NoError(t, err)
-
+	attester := NewNitroAttester()
 	getDoc := func(t *testing.T, n *nonce.Nonce) *AttestationDoc {
 		doc, err := attester.Attest(&AuxInfo{Nonce: ToAuxField(n.ToSlice())})
 		require.NoError(t, err)
diff --git a/internal/enclave/pcr.go b/internal/enclave/pcr.go
index f469cad..561be4e 100644
--- a/internal/enclave/pcr.go
+++ b/internal/enclave/pcr.go
@@ -14,12 +14,7 @@ type pcr map[uint][]byte
 func getPCRs() (_ pcr, err error) {
 	defer errs.Wrap(&err, "failed to get PCRs")
 
-	attester, err := NewNitroAttester()
-	if err != nil {
-		return nil, err
-	}
-
-	attestation, err := attester.Attest(&AuxInfo{})
+	attestation, err := NewNitroAttester().Attest(&AuxInfo{})
 	if err != nil {
 		return nil, err
 	}