Add support for imported keys

This commit adds support for imported keys in YubiKey KMS. Now, we
attempt to retrieve a key using go-piv's KeyInfo that supports both
imported and generated keys. This functionality is only available from
YubiKey firmware  5.3.0. We will  fallback to use Attest and Certificate
methods that are available in older versions.

Fixes #655

kms/yubikey/yubikey.go

@@ -42,6 +42,7 @@ type pivKey interface {
 	Certificate(slot piv.Slot) (*x509.Certificate, error)
 	SetCertificate(key []byte, slot piv.Slot, cert *x509.Certificate) error
 	GenerateKey(key []byte, slot piv.Slot, opts piv.Key) (crypto.PublicKey, error)
+	KeyInfo(slot piv.Slot) (piv.KeyInfo, error)
 	PrivateKey(slot piv.Slot, public crypto.PublicKey, auth piv.KeyAuth) (crypto.PrivateKey, error)
 	Attest(slot piv.Slot) (*x509.Certificate, error)
 	Serial() (uint32, error)
@@ -381,17 +382,33 @@ func (k *YubiKey) Close() error {
 	return nil
 }
 
-// getPublicKey returns the public key on a slot. First it attempts to do
-// attestation to get a certificate with the public key in it, if this succeeds
-// means that the key was generated in the device. If not we'll try to get the
-// key from a stored certificate in the same slot.
+// getPublicKey returns the public key on a slot. First it attempts to use
+// KeyInfo to get the public key, then tries to do attestation to get a
+// certificate with the public key in it, if this succeeds means that the key
+// was generated in the device. If not we'll try to get the key from a stored
+// certificate in the same slot.
 func (k *YubiKey) getPublicKey(slot piv.Slot) (crypto.PublicKey, error) {
-	cert, err := k.yk.Attest(slot)
+	// YubiKey >= 5.3.0 (generated and imported keys)
+	if ki, err := k.yk.KeyInfo(slot); err == nil && ki.PublicKey != nil {
+		return ki.PublicKey, nil
+	}
+
+	// YubiKey >= 4.3.0 (generated keys)
+	if cert, err := k.yk.Attest(slot); err == nil {
+		return cert.PublicKey, nil
+	}
+
+	// Fallback to certificate in slot (generated and imported)
+	cert, err := k.yk.Certificate(slot)
 	if err != nil {
-		if cert, err = k.yk.Certificate(slot); err != nil {
-			return nil, errors.Wrap(err, "error retrieving public key")
+		if errors.Is(err, piv.ErrNotFound) {
+			return nil, apiv1.NotFoundError{
+				Message: err.Error(),
+			}
 		}
+		return nil, fmt.Errorf("error retrieving public key: %w", err)
 	}
+
 	return cert.PublicKey, nil
 }
 

kms/yubikey/yubikey_test.go

@@ -17,6 +17,7 @@ import (
 	"crypto/x509/pkix"
 	"encoding/asn1"
 	"errors"
+	"fmt"
 	"reflect"
 	"sync"
 	"testing"
@@ -33,6 +34,7 @@ type stubPivKey struct {
 	attestCA      *minica.CA
 	attestSigner  privateKey
 	userCA        *minica.CA
+	keyInfoMap    map[piv.Slot]piv.KeyInfo
 	attestMap     map[piv.Slot]*x509.Certificate
 	certMap       map[piv.Slot]*x509.Certificate
 	signerMap     map[piv.Slot]interface{}
@@ -73,8 +75,10 @@ func newStubPivKey(t *testing.T, alg symmetricAlgorithm) *stubPivKey {
 		t.Fatal(err)
 	}
 
+	var keyInfoAlgo piv.Algorithm
 	switch alg {
 	case ECDSA:
+		keyInfoAlgo = piv.AlgorithmEC256
 		attSigner, err = ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
 		if err != nil {
 			t.Fatal(err)
@@ -84,6 +88,7 @@ func newStubPivKey(t *testing.T, alg symmetricAlgorithm) *stubPivKey {
 			t.Fatal(err)
 		}
 	case RSA:
+		keyInfoAlgo = piv.AlgorithmRSA2048
 		attSigner, err = rsa.GenerateKey(rand.Reader, rsaKeySize)
 		if err != nil {
 			t.Fatal(err)
@@ -124,6 +129,15 @@ func newStubPivKey(t *testing.T, alg symmetricAlgorithm) *stubPivKey {
 		attestCA:     attestCA,
 		attestSigner: attSigner,
 		userCA:       userCA,
+		keyInfoMap: map[piv.Slot]piv.KeyInfo{
+			piv.SlotKeyManagement: {
+				PublicKey:   attSigner.Public(),
+				Algorithm:   keyInfoAlgo,
+				PINPolicy:   piv.PINPolicyOnce,
+				TouchPolicy: piv.TouchPolicyCached,
+				Origin:      piv.OriginGenerated,
+			}, // 9d
+		},
 		attestMap: map[piv.Slot]*x509.Certificate{
 			piv.SlotAuthentication: attCert, // 9a
 		},
@@ -140,10 +154,21 @@ func newStubPivKey(t *testing.T, alg symmetricAlgorithm) *stubPivKey {
 	}
 }
 
+func (s *stubPivKey) KeyInfo(slot piv.Slot) (piv.KeyInfo, error) {
+	keyInfo, ok := s.keyInfoMap[slot]
+	if !ok {
+		return piv.KeyInfo{}, errors.New("public key not found")
+	}
+	return keyInfo, nil
+}
+
 func (s *stubPivKey) Certificate(slot piv.Slot) (*x509.Certificate, error) {
 	cert, ok := s.certMap[slot]
 	if !ok {
-		return nil, errors.New("certificate not found")
+		if slot == slotMapping["82"] {
+			return nil, errors.New("command failed: some error")
+		}
+		return nil, fmt.Errorf("command failed: %w", piv.ErrNotFound)
 	}
 	return cert, nil
 }
@@ -523,13 +548,22 @@ func TestYubiKey_GetPublicKey(t *testing.T) {
 		want    crypto.PublicKey
 		wantErr bool
 	}{
-		{"ok", fields{yk, "123456", piv.DefaultManagementKey}, args{&apiv1.GetPublicKeyRequest{
+		{"ok with keyInfo", fields{yk, "123456", piv.DefaultManagementKey}, args{&apiv1.GetPublicKeyRequest{
+			Name: "yubikey:slot-id=9d",
+		}}, yk.keyInfoMap[piv.SlotKeyManagement].PublicKey, false},
+		{"ok with Attest", fields{yk, "123456", piv.DefaultManagementKey}, args{&apiv1.GetPublicKeyRequest{
+			Name: "yubikey:slot-id=9a",
+		}}, yk.attestMap[piv.SlotAuthentication].PublicKey, false},
+		{"ok with certificate", fields{yk, "123456", piv.DefaultManagementKey}, args{&apiv1.GetPublicKeyRequest{
 			Name: "yubikey:slot-id=9c",
 		}}, yk.certMap[piv.SlotSignature].PublicKey, false},
 		{"fail getSlot", fields{yk, "123456", piv.DefaultManagementKey}, args{&apiv1.GetPublicKeyRequest{
 			Name: "slot-id=9c",
 		}}, nil, true},
 		{"fail getPublicKey", fields{yk, "123456", piv.DefaultManagementKey}, args{&apiv1.GetPublicKeyRequest{
+			Name: "yubikey:slot-id=82",
+		}}, nil, true},
+		{"fail getPublicKey not found", fields{yk, "123456", piv.DefaultManagementKey}, args{&apiv1.GetPublicKeyRequest{
 			Name: "yubikey:slot-id=85",
 		}}, nil, true},
 	}