From db516b7e0121c5800e78b2828651a9adab1dc62a Mon Sep 17 00:00:00 2001
From: James Pickett <James-Pickett@users.noreply.github.com>
Date: Fri, 29 Dec 2023 08:28:26 -0800
Subject: [PATCH] use secure enclave public key, tweaks (#35)

* use secure enclave public key, tweaks
---
 pkg/challenge/challenge.go              |  2 +-
 pkg/challenge/response.go               |  2 +-
 pkg/echelper/helpers.go                 | 10 +++++-----
 pkg/secureenclave/secureenclave.go      | 23 ++++++++++++++---------
 pkg/secureenclave/secureenclave_test.go |  2 +-
 pkg/tpm/tpm_test.go                     |  2 +-
 6 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/pkg/challenge/challenge.go b/pkg/challenge/challenge.go
index 041aad0..8d0e644 100644
--- a/pkg/challenge/challenge.go
+++ b/pkg/challenge/challenge.go
@@ -26,7 +26,7 @@ type OuterChallenge struct {
 }
 
 func (o *OuterChallenge) Verify(counterParty ecdsa.PublicKey) error {
-	if err := echelper.VerifySignature(counterParty, o.Msg, o.Sig); err != nil {
+	if err := echelper.VerifySignature(&counterParty, o.Msg, o.Sig); err != nil {
 		return err
 	}
 
diff --git a/pkg/challenge/response.go b/pkg/challenge/response.go
index 3ed78c7..d994084 100644
--- a/pkg/challenge/response.go
+++ b/pkg/challenge/response.go
@@ -72,7 +72,7 @@ func verifyWithKeyBytes(keyBytes []byte, msg []byte, sig []byte) error {
 		return fmt.Errorf("parsing public key: %w", err)
 	}
 
-	return echelper.VerifySignature(*key, msg, sig)
+	return echelper.VerifySignature(key, msg, sig)
 }
 
 type InnerResponse struct {
diff --git a/pkg/echelper/helpers.go b/pkg/echelper/helpers.go
index 0c74695..edf254b 100644
--- a/pkg/echelper/helpers.go
+++ b/pkg/echelper/helpers.go
@@ -18,7 +18,7 @@ import (
 )
 
 func Sign(signer crypto.Signer, data []byte) ([]byte, error) {
-	digest, err := hashForSignature(data)
+	digest, err := HashForSignature(data)
 	if err != nil {
 		return nil, fmt.Errorf("hashing data: %w", err)
 	}
@@ -31,13 +31,13 @@ func Sign(signer crypto.Signer, data []byte) ([]byte, error) {
 	return signature, nil
 }
 
-func VerifySignature(counterParty ecdsa.PublicKey, data []byte, signature []byte) error {
-	digest, err := hashForSignature(data)
+func VerifySignature(counterParty *ecdsa.PublicKey, data []byte, signature []byte) error {
+	digest, err := HashForSignature(data)
 	if err != nil {
 		return fmt.Errorf("hashing inner box: %w", err)
 	}
 
-	if !ecdsa.VerifyASN1(&counterParty, digest, signature) {
+	if !ecdsa.VerifyASN1(counterParty, digest, signature) {
 		return fmt.Errorf("invalid signature")
 	}
 
@@ -135,7 +135,7 @@ func SignWithTimeout(signer crypto.Signer, data []byte, duration, interval time.
 	}
 }
 
-func hashForSignature(data []byte) ([]byte, error) {
+func HashForSignature(data []byte) ([]byte, error) {
 	hash := sha256.New()
 	_, err := hash.Write(data)
 	if err != nil {
diff --git a/pkg/secureenclave/secureenclave.go b/pkg/secureenclave/secureenclave.go
index ff3c28f..0991f51 100644
--- a/pkg/secureenclave/secureenclave.go
+++ b/pkg/secureenclave/secureenclave.go
@@ -29,8 +29,17 @@ type SecureEnclaveSigner struct {
 
 // New verifies that the provided public key already exists in the secure enclave.
 // Then returns a new Secure Enclave Keyer using the provided public key.
-func New(publicKeySha1 []byte) (*SecureEnclaveSigner, error) {
-	pubKey, err := findKey(publicKeySha1)
+func New(pubKey *ecdsa.PublicKey) (*SecureEnclaveSigner, error) {
+	if pubKey == nil {
+		return nil, errors.New("nil public key")
+	}
+
+	lookUp, err := publicKeyLookUpHash(pubKey)
+	if err != nil {
+		return nil, err
+	}
+
+	pubKey, err = findKey(lookUp)
 	if err != nil {
 		return nil, fmt.Errorf("finding existing public key: %w", err)
 	}
@@ -72,19 +81,15 @@ func (s *SecureEnclaveSigner) Sign(rand io.Reader, digest []byte, opts crypto.Si
 	return result, nil
 }
 
-// CreateKey creates a new secure enclave key and returns the hash used to access it.
-func CreateKey() ([]byte, error) {
+// CreateKey creates a new secure enclave key and returns the public key.
+func CreateKey() (*ecdsa.PublicKey, error) {
 	wrapper := C.wrapCreateKey()
 	result, err := unwrap(wrapper)
 	if err != nil {
 		return nil, err
 	}
 
-	sha1 := sha1.New()
-	if _, err := sha1.Write(result); err != nil {
-		return nil, fmt.Errorf("hashing secure enclave create key result to sha1: %w", err)
-	}
-	return sha1.Sum(nil), nil
+	return rawToEcdsa(result), nil
 }
 
 // unwrap a Wrapper struct to a Go byte slice
diff --git a/pkg/secureenclave/secureenclave_test.go b/pkg/secureenclave/secureenclave_test.go
index 9f07101..a5e25aa 100644
--- a/pkg/secureenclave/secureenclave_test.go
+++ b/pkg/secureenclave/secureenclave_test.go
@@ -95,7 +95,7 @@ func TestSecureEnclaveSigning(t *testing.T) {
 
 	publicKey := seSigner.Public().(*ecdsa.PublicKey)
 
-	require.NoError(t, echelper.VerifySignature(*publicKey, dataToSign, signature))
+	require.NoError(t, echelper.VerifySignature(publicKey, dataToSign, signature))
 }
 
 func TestSecureEnclaveErrors(t *testing.T) {
diff --git a/pkg/tpm/tpm_test.go b/pkg/tpm/tpm_test.go
index 4a83673..64fc7b2 100644
--- a/pkg/tpm/tpm_test.go
+++ b/pkg/tpm/tpm_test.go
@@ -35,7 +35,7 @@ func TestTpmSigning(t *testing.T) {
 
 	publicKey := tpmSigner.Public().(*ecdsa.PublicKey)
 
-	require.NoError(t, echelper.VerifySignature(*publicKey, dataToSign, signature))
+	require.NoError(t, echelper.VerifySignature(publicKey, dataToSign, signature))
 }
 
 func TestTpmErrors(t *testing.T) {