From ffbd58fd572659ee9bccd990f5bfbc20d55f07a6 Mon Sep 17 00:00:00 2001
From: Andrew Ayer <agwa@andrewayer.name>
Date: Tue, 1 Oct 2024 17:44:28 -0400
Subject: [PATCH] Add ECDSAPrivateKey.ECDH which takes a *ecdh.PublicKey

The current ECDSAPrivateKey.SharedKey method takes a *ecdsa.PublicKey.
However, using crypto/ecdsa and crypto/elliptic for ECDH has been
deprecated in the standard library in favor of crypto/ecdh.

This commit adds a new ECDH method to ECDSAPrivateKey which takes
a *ecdh.PublicKey.  This method has the same signature
as ecdh.PrivateKey.ECDH, meaning the following interface can be be used
to do ECDH with both standard library private keys and piv-go keys, providing
the same flexibility as crypto.Signer and crypto.Decrypter:

interface {
	ECDH(*ecdh.PublicKey) ([]byte, error)
}

ECDSAPrivateKey.SharedKey has been re-implemented as a small wrapper around
ECDSAPrivateKey.ECDH.
---
 v2/piv/key.go      | 26 ++++++++++++++++--
 v2/piv/key_test.go | 66 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 90 insertions(+), 2 deletions(-)

diff --git a/v2/piv/key.go b/v2/piv/key.go
index aa82d74..69eeecb 100644
--- a/v2/piv/key.go
+++ b/v2/piv/key.go
@@ -17,6 +17,7 @@ package piv
 import (
 	"bytes"
 	"crypto"
+	"crypto/ecdh"
 	"crypto/ecdsa"
 	"crypto/ed25519"
 	"crypto/elliptic"
@@ -1151,10 +1152,31 @@ func (k *ECDSAPrivateKey) Sign(rand io.Reader, digest []byte, opts crypto.Signer
 // used for the operation. Callers should use a cryptographic key
 // derivation function to extract the amount of bytes they need.
 func (k *ECDSAPrivateKey) SharedKey(peer *ecdsa.PublicKey) ([]byte, error) {
-	if peer.Curve.Params().BitSize != k.pub.Curve.Params().BitSize {
+	peerECDH, err := peer.ECDH()
+	if err != nil {
+		return nil, unsupportedCurveError{curve: peer.Params().BitSize}
+	}
+	return k.ECDH(peerECDH)
+}
+
+// ECDH performs a Diffie-Hellman key agreement with the peer
+// to produce a shared secret key.
+//
+// Peer's public key must use the same algorithm as the key in
+// this slot, or an error will be returned.
+//
+// Length of the result depends on the types and sizes of the keys
+// used for the operation. Callers should use a cryptographic key
+// derivation function to extract the amount of bytes they need.
+func (k *ECDSAPrivateKey) ECDH(peer *ecdh.PublicKey) ([]byte, error) {
+	ourECDH, err := k.pub.ECDH()
+	if err != nil {
+		return nil, unsupportedCurveError{curve: k.pub.Params().BitSize}
+	}
+	if peer.Curve() != ourECDH.Curve() {
 		return nil, errMismatchingAlgorithms
 	}
-	msg := elliptic.Marshal(peer.Curve, peer.X, peer.Y)
+	msg := peer.Bytes()
 	return k.auth.do(k.yk, k.pp, func(tx *scTx) ([]byte, error) {
 		var alg byte
 		size := k.pub.Params().BitSize
diff --git a/v2/piv/key_test.go b/v2/piv/key_test.go
index fec6739..35d4a43 100644
--- a/v2/piv/key_test.go
+++ b/v2/piv/key_test.go
@@ -17,6 +17,7 @@ package piv
 import (
 	"bytes"
 	"crypto"
+	"crypto/ecdh"
 	"crypto/ecdsa"
 	"crypto/ed25519"
 	"crypto/elliptic"
@@ -84,6 +85,71 @@ func TestYubiKeySignECDSA(t *testing.T) {
 	}
 }
 
+func TestYubiKeyECDSAECDH(t *testing.T) {
+	yk, close := newTestYubiKey(t)
+	defer close()
+
+	slot := SlotAuthentication
+
+	key := Key{
+		Algorithm:   AlgorithmEC256,
+		TouchPolicy: TouchPolicyNever,
+		PINPolicy:   PINPolicyNever,
+	}
+	pubKey, err := yk.GenerateKey(DefaultManagementKey, slot, key)
+	if err != nil {
+		t.Fatalf("generating key: %v", err)
+	}
+	pub, ok := pubKey.(*ecdsa.PublicKey)
+	if !ok {
+		t.Fatalf("public key is not an ecdsa key")
+	}
+	pubECDH, err := pub.ECDH()
+	if err != nil {
+		t.Fatalf("converting pubkey to ECDH key: %v", err)
+	}
+	priv, err := yk.PrivateKey(slot, pub, KeyAuth{})
+	if err != nil {
+		t.Fatalf("getting private key: %v", err)
+	}
+	privECDSA, ok := priv.(*ECDSAPrivateKey)
+	if !ok {
+		t.Fatalf("expected private key to be ECDSA private key")
+	}
+
+	t.Run("good", func(t *testing.T) {
+		privECDH, err := ecdh.P256().GenerateKey(rand.Reader)
+		if err != nil {
+			t.Fatalf("cannot generate key: %v", err)
+		}
+		secret1, err := privECDH.ECDH(pubECDH)
+		if err != nil {
+			t.Fatalf("key agreement 1 failed: %v", err)
+		}
+
+		secret2, err := privECDSA.ECDH(privECDH.PublicKey())
+		if err != nil {
+			t.Fatalf("key agreement 2 failed: %v", err)
+		}
+		if !bytes.Equal(secret1, secret2) {
+			t.Errorf("key agreement didn't match")
+		}
+	})
+
+	t.Run("bad", func(t *testing.T) {
+		t.Run("size", func(t *testing.T) {
+			privECDH, err := ecdh.P384().GenerateKey(rand.Reader)
+			if err != nil {
+				t.Fatalf("cannot generate key: %v", err)
+			}
+			_, err = privECDSA.ECDH(privECDH.PublicKey())
+			if !errors.Is(err, errMismatchingAlgorithms) {
+				t.Fatalf("unexpected error value: wanted errMismatchingAlgorithms: %v", err)
+			}
+		})
+	})
+}
+
 func TestYubiKeyECDSASharedKey(t *testing.T) {
 	yk, close := newTestYubiKey(t)
 	defer close()