feat: Hardware Key Agent - Enrich the PEM encoded hardware private key file

api/utils/keys/hardwarekey/hardwarekey.go

@@ -19,6 +19,7 @@ import (
 	"context"
 	"crypto"
 	"crypto/rand"
+	"crypto/x509"
 	"encoding/json"
 	"io"
 
@@ -33,6 +34,7 @@ type Service interface {
 	// Sign performs a cryptographic signature using the specified hardware
 	// private key and provided signature parameters.
 	Sign(ctx context.Context, ref *PrivateKeyRef, rand io.Reader, digest []byte, opts crypto.SignerOpts) (signature []byte, err error)
+	// TODO(Joerger): DELETE IN v19.0.0
 	// GetFullKeyRef gets the full [PrivateKeyRef] for an existing hardware private
 	// key in the given slot of the hardware key with the given serial number.
 	GetFullKeyRef(serialNumber uint32, slotKey PIVSlotKey) (*PrivateKeyRef, error)
@@ -96,17 +98,8 @@ func EncodeSigner(p *Signer) ([]byte, error) {
 }
 
 // DecodeSigner decodes an encoded hardware key signer for the given service.
-func DecodeSigner(s Service, encodedKey []byte) (*Signer, error) {
-	partialRef, err := decodeKeyRef(encodedKey)
-	if err != nil {
-		return nil, trace.Wrap(err)
-	}
-
-	// TODO(Joerger): all fields should be encoded to the key's PEM file
-	// rather than being retrieved from the hardware key each time. This
-	// will result in a massive performance boost by avoiding re-attesting
-	// the key for every client call.
-	ref, err := s.GetFullKeyRef(partialRef.SerialNumber, partialRef.SlotKey)
+func DecodeSigner(encodedKey []byte, s Service) (*Signer, error) {
+	ref, err := decodeKeyRef(encodedKey, s)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -121,16 +114,21 @@ type PrivateKeyRef struct {
 	// SlotKey is the key name for the hardware key PIV slot, e.g. "9a".
 	SlotKey PIVSlotKey `json:"slot_key"`
 	// PublicKey is the public key paired with the hardware private key.
-	PublicKey crypto.PublicKey `json:"-"`
+	PublicKey crypto.PublicKey `json:"-"` // uses custom JSON marshaling in PKIX, ASN.1 DER form
 	// Policy specifies the hardware private key's PIN/touch prompt policies.
-	Policy PromptPolicy `json:"-"`
+	Policy PromptPolicy `json:"policy"`
 	// AttestationStatement contains the hardware private key's attestation statement, which is
 	// to attest the touch and pin requirements for this hardware private key during login.
-	AttestationStatement *AttestationStatement `json:"-"`
+	AttestationStatement *AttestationStatement `json:"attestation_statement"`
 }
 
 // encode encodes a [PrivateKeyRef] to JSON.
 func (r *PrivateKeyRef) encode() ([]byte, error) {
+	// Ensure that all required fields are provided to encode.
+	if err := r.Validate(); err != nil {
+		return nil, trace.Wrap(err)
+	}
+
 	keyRefBytes, err := json.Marshal(r)
 	if err != nil {
 		return nil, trace.Wrap(err)
@@ -139,13 +137,86 @@ func (r *PrivateKeyRef) encode() ([]byte, error) {
 }
 
 // decodeKeyRef decodes a [PrivateKeyRef] from JSON.
-func decodeKeyRef(encodedKeyRef []byte) (*PrivateKeyRef, error) {
-	keyRef := &PrivateKeyRef{}
-	if err := json.Unmarshal(encodedKeyRef, keyRef); err != nil {
+func decodeKeyRef(encodedKeyRef []byte, s Service) (*PrivateKeyRef, error) {
+	ref := &PrivateKeyRef{}
+	if err := json.Unmarshal(encodedKeyRef, ref); err != nil {
 		return nil, trace.Wrap(err)
 	}
 
-	return keyRef, nil
+	// Ensure that all required fields are decoded.
+	if err := ref.Validate(); err != nil {
+		// If some fields are missing, this is likely an old login key with only
+		// the serial number and slot. Fetch missing data from the hardware key.
+		// This data will be saved to the login key on next login
+		// TODO(Joerger): DELETE IN v19.0.0
+		if ref.SerialNumber != 0 && ref.SlotKey != 0 {
+			return s.GetFullKeyRef(ref.SerialNumber, ref.SlotKey)
+		}
+
+		return nil, trace.Wrap(err)
+	}
+
+	return ref, nil
+}
+
+func (r *PrivateKeyRef) Validate() error {
+	if r.SerialNumber == 0 {
+		return trace.BadParameter("private key ref missing SerialNumber")
+	}
+	if r.SlotKey == 0 {
+		return trace.BadParameter("private key ref missing SlotKey")
+	}
+	if r.PublicKey == nil {
+		return trace.BadParameter("private key ref missing PublicKey")
+	}
+	if r.AttestationStatement == nil {
+		return trace.BadParameter("private key ref missing AttestationStatement")
+	}
+	return nil
+}
+
+// These types are used for custom marshaling of the crypto.PublicKey field in [PrivateKeyRef].
+type rawPrivateKeyRef PrivateKeyRef
+type hardwarePrivateKeyRefJSON struct {
+	// embedding an alias type instead of [HardwarePrivateKeyRef] prevents the custom marshaling
+	// from recursively applying, which would result in a stack overflow.
+	rawPrivateKeyRef
+	PublicKeyDER []byte `json:"public_key,omitempty"`
+}
+
+// MarshalJSON marshals [PrivateKeyRef] with custom logic for the public key.
+func (r PrivateKeyRef) MarshalJSON() ([]byte, error) {
+	var pubDER []byte
+	if r.PublicKey != nil {
+		var err error
+		if pubDER, err = x509.MarshalPKIXPublicKey(r.PublicKey); err != nil {
+			return nil, trace.Wrap(err)
+		}
+	}
+
+	return json.Marshal(&hardwarePrivateKeyRefJSON{
+		rawPrivateKeyRef: rawPrivateKeyRef(r),
+		PublicKeyDER:     pubDER,
+	})
+}
+
+// UnmarshalJSON unmarshals [PrivateKeyRef] with custom logic for the public key.
+func (r *PrivateKeyRef) UnmarshalJSON(b []byte) error {
+	var ref hardwarePrivateKeyRefJSON
+	err := json.Unmarshal(b, &ref)
+	if err != nil {
+		return trace.Wrap(err)
+	}
+
+	if ref.PublicKeyDER != nil {
+		ref.rawPrivateKeyRef.PublicKey, err = x509.ParsePKIXPublicKey(ref.PublicKeyDER)
+		if err != nil {
+			return trace.Wrap(err)
+		}
+	}
+
+	*r = PrivateKeyRef(ref.rawPrivateKeyRef)
+	return nil
 }
 
 // PrivateKeyConfig contains config for creating a new hardware private key.

api/utils/keys/hardwarekey/hardwarekey_test.go

@@ -21,6 +21,7 @@ import (
 	"crypto"
 	"crypto/rand"
 	"crypto/sha512"
+	"encoding/json"
 	"fmt"
 	"io"
 	"testing"
@@ -50,9 +51,32 @@ func TestPrivateKey_EncodeDecode(t *testing.T) {
 	encoded, err := hardwarekey.EncodeSigner(priv)
 	require.NoError(t, err)
 
-	decodedPriv, err := hardwarekey.DecodeSigner(s, encoded)
+	decodedSigner, err := hardwarekey.DecodeSigner(encoded, s)
 	require.NoError(t, err)
-	require.Equal(t, hwSigner, decodedPriv)
+	require.Equal(t, hwSigner, decodedSigner)
+}
+
+// Old client logins would only have encoded the serial number and slot key.
+// TODO(Joerger): DELETE IN v19.0.0
+func TestPrivateKey_DecodePartialKeyRef(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	s := hardwarekey.NewMockHardwareKeyService(nil /*prompt*/)
+	hwSigner, err := s.NewPrivateKey(ctx, hardwarekey.PrivateKeyConfig{
+		Policy: hardwarekey.PromptPolicyNone,
+	})
+	require.NoError(t, err)
+
+	partialKeyRefJSON, err := json.Marshal(&hardwarekey.PrivateKeyRef{
+		SerialNumber: hwSigner.Ref.SerialNumber,
+		SlotKey:      hwSigner.Ref.SlotKey,
+	})
+	require.NoError(t, err)
+
+	decodedSigner, err := hardwarekey.DecodeSigner(partialKeyRefJSON, s)
+	require.NoError(t, err)
+	require.Equal(t, hwSigner, decodedSigner)
 }
 
 // TestPrivateKey_Prompt tests hardware key service PIN/Touch logic with a mocked service.

api/utils/keys/hardwarekey/service_mock.go

@@ -106,6 +106,10 @@ func (s *MockHardwareKeyService) NewPrivateKey(ctx context.Context, config Priva
 		AttestationStatement: &AttestationStatement{},
 	}
 
+	if err := ref.Validate(); err != nil {
+		return nil, trace.Wrap(err)
+	}
+
 	s.fakeHardwarePrivateKeys[keySlot] = &fakeHardwarePrivateKey{
 		Signer: priv,
 		ref:    ref,
@@ -172,6 +176,7 @@ func (s *MockHardwareKeyService) SetPrompt(prompt Prompt) {
 	s.prompt = prompt
 }
 
+// TODO(Joerger): DELETE IN v19.0.0
 func (s *MockHardwareKeyService) GetFullKeyRef(serialNumber uint32, slotKey PIVSlotKey) (*PrivateKeyRef, error) {
 	s.fakeHardwarePrivateKeysMux.Lock()
 	defer s.fakeHardwarePrivateKeysMux.Unlock()

api/utils/keys/piv/service.go

@@ -193,6 +193,9 @@ type baseKeyRef struct {
 
 // GetFullKeyRef gets the full [PrivateKeyRef] for an existing hardware private
 // key in the given slot of the hardware key with the given serial number.
+//
+// Used for backwards compatibility with old logins.
+// TODO(Joerger): DELETE IN v19.0.0
 func (s *YubiKeyService) GetFullKeyRef(serialNumber uint32, slotKey hardwarekey.PIVSlotKey) (*hardwarekey.PrivateKeyRef, error) {
 	keyRefsMux.Lock()
 	defer keyRefsMux.Unlock()

api/utils/keys/piv/yubikey.go

@@ -368,7 +368,7 @@ func (y *YubiKey) getKeyRef(slot piv.Slot) (*hardwarekey.PrivateKeyRef, error) {
 		return nil, trace.Wrap(err)
 	}
 
-	return &hardwarekey.PrivateKeyRef{
+	ref := &hardwarekey.PrivateKeyRef{
 		SerialNumber: y.serialNumber,
 		SlotKey:      hardwarekey.PIVSlotKey(slot.Key),
 		PublicKey:    slotCert.PublicKey,
@@ -384,7 +384,13 @@ func (y *YubiKey) getKeyRef(slot piv.Slot) (*hardwarekey.PrivateKeyRef, error) {
 				},
 			},
 		},
-	}, nil
+	}
+
+	if err := ref.Validate(); err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	return ref, nil
 }
 
 // SetPIN sets the YubiKey PIV PIN. This doesn't require user interaction like touch, just the correct old PIN.

api/utils/keys/privatekey.go

@@ -286,12 +286,12 @@ func ParsePrivateKey(keyPEM []byte, opts ...ParsePrivateKeyOpt) (*PrivateKey, er
 		// it in the client store. This allows the process to properly share PIV connections
 		// and prompt logic (pin caching, etc.).
 		hwKeyService := piv.NewYubiKeyService(appliedOpts.CustomHardwareKeyPrompt)
-		hwPrivateKey, err := hardwarekey.DecodeSigner(hwKeyService, block.Bytes)
+		hwSigner, err := hardwarekey.DecodeSigner(block.Bytes, hwKeyService)
 		if err != nil {
 			return nil, trace.Wrap(err, "failed to parse hardware key signer")
 		}
 
-		return newPrivateKeyWithKeyPEM(hwPrivateKey, keyPEM)
+		return newPrivateKeyWithKeyPEM(hwSigner, keyPEM)
 	case OpenSSHPrivateKeyType:
 		priv, err := ssh.ParseRawPrivateKey(keyPEM)
 		if err != nil {