From 7c5004737a4a14538b7d51ba03d19576776c5104 Mon Sep 17 00:00:00 2001
From: Joshua Krstic <jkrstic@google.com>
Date: Tue, 5 Sep 2023 13:58:15 -0700
Subject: [PATCH] Add v2 tpm nonce support

---
 launcher/tpm_nonce/tpm_nonce.go      | 59 +++++++++++++++++++
 launcher/tpm_nonce/tpm_nonce_test.go | 87 ++++++++++++++++++++++++++++
 2 files changed, 146 insertions(+)
 create mode 100644 launcher/tpm_nonce/tpm_nonce.go
 create mode 100644 launcher/tpm_nonce/tpm_nonce_test.go

diff --git a/launcher/tpm_nonce/tpm_nonce.go b/launcher/tpm_nonce/tpm_nonce.go
new file mode 100644
index 000000000..25f1abb05
--- /dev/null
+++ b/launcher/tpm_nonce/tpm_nonce.go
@@ -0,0 +1,59 @@
+package internal
+
+import (
+	"bytes"
+	"crypto/sha256"
+	"fmt"
+	"sort"
+	"strings"
+)
+
+const (
+	// ChallengePrefix is the prefix on all challenges generated by the Attestation Service.
+	ChallengePrefix = "GoogAttestV1"
+	// NoncePrefix is the prefix on the nonce given to the TPM before the quote is generated.
+	NoncePrefix = "GoogAttestV2"
+)
+
+// GenerateTpmNonce generates the nonce to give to the TPM before getting a quote.
+// TpmNonce = noncePrefix | SHA256(C | SHA256(sort[applyPrefix(arr)])
+// C == GoogAttestV1 | 128bits entropy
+func GenerateTpmNonce(challenge []byte, nonces [][]byte) ([]byte, error) {
+	if challenge == nil || !strings.HasPrefix(string(challenge), ChallengePrefix) {
+		return nil, fmt.Errorf("expected GoogAttestV1 challenge, got %v", challenge)
+	}
+
+	// No custom nonces specified is OK.
+	if nonces == nil {
+		nonces = [][]byte{}
+	}
+
+	// Prefix the length of each nonce to itself.
+	for i, s := range nonces {
+		l := byte(len(s))
+		nonces[i] = append([]byte{l}, s...)
+	}
+
+	// Sort nonces.
+	sort.Slice(nonces, func(i, j int) bool {
+		return bytes.Compare(nonces[i], nonces[j]) == -1
+	})
+
+	// Concatenate nonces.
+	cat := []byte{}
+	for _, v := range nonces {
+		cat = append(cat, v...)
+	}
+
+	// Take the SHA256 sum of the concatenated nonces.
+	nonceHash := sha256.Sum256(cat)
+
+	// Concatenate the challenge with the hashed nonces.
+	challengeCat := append(challenge, nonceHash[:]...)
+
+	// Hash the newly concatenated challenge.
+	challengeHash := sha256.Sum256(challengeCat)
+
+	// Return the noncePrefix concatenated with the hashed challenge.
+	return append([]byte(NoncePrefix), challengeHash[:]...), nil
+}
diff --git a/launcher/tpm_nonce/tpm_nonce_test.go b/launcher/tpm_nonce/tpm_nonce_test.go
new file mode 100644
index 000000000..a0a4a0e1f
--- /dev/null
+++ b/launcher/tpm_nonce/tpm_nonce_test.go
@@ -0,0 +1,87 @@
+package internal
+
+import (
+	b64 "encoding/base64"
+	"strings"
+	"testing"
+)
+
+func TestGenerateTpmNonce(t *testing.T) {
+	challenge := "GoogAttestV1lcdYd8KX3W3uESVrPbmTjA"
+	nonces := [][]byte{
+		[]byte("GoogAttestV1lddYd8KX3W3uESVrPbmTjA"),
+		[]byte("ThisIsACustomNonce"),
+	}
+
+	tpmNonce, err := GenerateTpmNonce([]byte(challenge), nonces)
+
+	if err != nil {
+		t.Errorf("expected no error, got %v", err)
+		t.Fail()
+	}
+	if len(tpmNonce) != 44 {
+		t.Errorf("expected a nonce 44 bytes long, got %d bytes", len(tpmNonce))
+	}
+	tpmNonceString := string(tpmNonce)
+	if !strings.HasPrefix(tpmNonceString, "GoogAttestV2") {
+		t.Errorf("expected a with prefix 'GoogAttestV2', actual nonce %v", tpmNonceString)
+		t.Fail()
+	}
+}
+
+func TestGenerateTpmNonceNilCustomNonces(t *testing.T) {
+	challenge := "GoogAttestV1lcdYd8KX3W3uESVrPbmTjA"
+
+	tpmNonce, err := GenerateTpmNonce([]byte(challenge), nil)
+
+	if err != nil {
+		t.Errorf("expected no error, got %v", err)
+		t.Fail()
+	}
+	if len(tpmNonce) != 44 {
+		t.Errorf("expected a nonce 44 bytes long, got %d bytes", len(tpmNonce))
+	}
+	tpmNonceString := string(tpmNonce)
+	if !strings.HasPrefix(tpmNonceString, "GoogAttestV2") {
+		t.Errorf("expected a with prefix 'GoogAttestV2', actual nonce %v", tpmNonceString)
+		t.Fail()
+	}
+}
+
+func TestGenerateTpmNonceEmptyCustomNonces(t *testing.T) {
+	challenge := "GoogAttestV1lcdYd8KX3W3uESVrPbmTjA"
+
+	tpmNonce, err := GenerateTpmNonce([]byte(challenge), [][]byte{})
+
+	if err != nil {
+		t.Errorf("expected no error, got %v", err)
+	}
+	if len(tpmNonce) != 44 {
+		t.Errorf("expected a nonce 44 bytes long, got %d bytes", len(tpmNonce))
+	}
+	tpmNonceString := string(tpmNonce)
+	if !strings.HasPrefix(tpmNonceString, "GoogAttestV2") {
+		t.Errorf("expected a with prefix 'GoogAttestV2', actual nonce %v", tpmNonceString)
+	}
+}
+
+func TestGenerateTpmNonceBase64ChallengeError(t *testing.T) {
+	challenge := "GoogAttestV1lcdYd8KX3W3uESVrPbmTjA"
+	base64Challenge := []byte(b64.StdEncoding.EncodeToString([]byte(challenge)))
+
+	_, err := GenerateTpmNonce(base64Challenge, nil)
+
+	if err == nil {
+		t.Errorf("expected error")
+	}
+}
+
+func TestGenerateTpmNonceNoChallengeError(t *testing.T) {
+	challenge := ""
+
+	_, err := GenerateTpmNonce([]byte(challenge), nil)
+
+	if err == nil {
+		t.Errorf("expected error")
+	}
+}