From 476469cd4979f91f920ed265e4ba64e23714834a Mon Sep 17 00:00:00 2001
From: joerger <bjoerger@goteleport.com>
Date: Fri, 9 May 2025 13:24:56 -0700
Subject: [PATCH 1/8] * Have hardware key agent validate known keys instead of
 always validating by the PIV slot cert.

* Refactor cert check logic with a custom error.
---
 api/utils/keys/hardwarekey/hardwarekey.go     | 13 ++--
 api/utils/keys/hardwarekeyagent/agent.go      | 25 ++++++-
 .../keys/hardwarekeyagent/service_test.go     |  2 +-
 api/utils/keys/piv/service.go                 | 14 ++--
 api/utils/keys/piv/yubikey.go                 | 72 ++++++++++---------
 lib/client/client_store.go                    | 29 ++++++++
 lib/hardwarekey/agent.go                      |  4 +-
 lib/hardwarekey/agent_test.go                 |  8 +--
 lib/teleterm/teleterm.go                      |  2 +-
 tool/tsh/common/piv.go                        |  6 +-
 10 files changed, 117 insertions(+), 58 deletions(-)

diff --git a/api/utils/keys/hardwarekey/hardwarekey.go b/api/utils/keys/hardwarekey/hardwarekey.go
index 7b3d2f77864f5..f6b18c677ad6b 100644
--- a/api/utils/keys/hardwarekey/hardwarekey.go
+++ b/api/utils/keys/hardwarekey/hardwarekey.go
@@ -254,11 +254,14 @@ type ContextualKeyInfo struct {
 	Username string
 	// ClusterName is a Teleport cluster name that the key is associated with.
 	ClusterName string
-	// AgentKey specifies whether this key is being utilized through an agent.
-	// The hardware key service may impose additional restrictions in this case,
-	// such as checking that the PIV slot certificate matches the Teleport client
-	// metadata certificate format, to ensure the agent doesn't provide access to
-	// non teleport client PIV keys.
+	// AgentKey specifies whether this key is being utilized through an agent and
+	// the key is unknown to the client running the agent, probably because the
+	// client on the other side of the agent is using a different Teleport Home
+	// directory.
+	//
+	// When true, the hardware key service will check that the certificate in the
+	// same slot as the key matches a Teleport client metadata certificate in order
+	// to ensure the agent doesn't provide access to non teleport client PIV keys.
 	AgentKey bool
 	// Command is the running command utilizing this key.
 	Command string
diff --git a/api/utils/keys/hardwarekeyagent/agent.go b/api/utils/keys/hardwarekeyagent/agent.go
index fdcdb8b2f525e..06233b05da032 100644
--- a/api/utils/keys/hardwarekeyagent/agent.go
+++ b/api/utils/keys/hardwarekeyagent/agent.go
@@ -61,19 +61,30 @@ func NewClient(ctx context.Context, socketPath string, creds credentials.Transpo
 }
 
 // NewServer returns a new hardware key agent server.
-func NewServer(ctx context.Context, s hardwarekey.Service, creds credentials.TransportCredentials) *grpc.Server {
+//
+// If [knownKeyFn] is not provided, all PIV keys are treated as agent keys.
+func NewServer(_ context.Context, s hardwarekey.Service, creds credentials.TransportCredentials, knownKeyFn KnownHardwareKeyFn) *grpc.Server {
 	grpcServer := grpc.NewServer(
 		grpc.Creds(creds),
 		grpc.UnaryInterceptor(interceptors.GRPCServerUnaryErrorInterceptor),
 	)
-	hardwarekeyagentv1.RegisterHardwareKeyAgentServiceServer(grpcServer, &agentService{s: s})
+	hardwarekeyagentv1.RegisterHardwareKeyAgentServiceServer(grpcServer, &agentService{s: s, knownKeyFn: knownKeyFn})
 	return grpcServer
 }
 
+type KnownHardwareKeyFn func(ref *hardwarekey.PrivateKeyRef, keyInfo hardwarekey.ContextualKeyInfo) (bool, error)
+
 // agentService implements [hardwarekeyagentv1.HardwareKeyAgentServiceServer].
 type agentService struct {
 	hardwarekeyagentv1.UnimplementedHardwareKeyAgentServiceServer
 	s hardwarekey.Service
+
+	// knownKeyFn dictates whether a PIV key referenced in a [agentService.Sign] request is
+	// treated as a known key or an unknown agent key. Unknown agent keys are treated with
+	// additional restrictions to ensure the PIV slot is intended for Teleport client usage.
+	//
+	// If knownKeyFn is nil, all keys are treated as agent keys.
+	knownKeyFn KnownHardwareKeyFn
 }
 
 // Sign the given digest with the specified hardware private key.
@@ -108,10 +119,18 @@ func (s *agentService) Sign(ctx context.Context, req *hardwarekeyagentv1.SignReq
 		ProxyHost:   req.KeyInfo.ProxyHost,
 		Username:    req.KeyInfo.Username,
 		ClusterName: req.KeyInfo.ClusterName,
-		AgentKey:    true,
 		Command:     req.Command,
 	}
 
+	var knownKey bool
+	if s.knownKeyFn != nil {
+		knownKey, err = s.knownKeyFn(keyRef, keyInfo)
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+	}
+	keyInfo.AgentKey = !knownKey
+
 	var signerOpts crypto.SignerOpts
 	switch req.Hash {
 	case hardwarekeyagentv1.Hash_HASH_NONE:
diff --git a/api/utils/keys/hardwarekeyagent/service_test.go b/api/utils/keys/hardwarekeyagent/service_test.go
index a78895103bbea..7a572543ecd75 100644
--- a/api/utils/keys/hardwarekeyagent/service_test.go
+++ b/api/utils/keys/hardwarekeyagent/service_test.go
@@ -37,7 +37,7 @@ func TestHardwareKeyAgentService(t *testing.T) {
 
 	// Prepare the agent server
 	mockService := hardwarekey.NewMockHardwareKeyService(nil /*prompt*/)
-	server := hardwarekeyagent.NewServer(ctx, mockService, insecure.NewCredentials())
+	server := hardwarekeyagent.NewServer(ctx, mockService, insecure.NewCredentials(), nil /*knownKeyFn*/)
 	t.Cleanup(server.Stop)
 
 	agentDir := t.TempDir()
diff --git a/api/utils/keys/piv/service.go b/api/utils/keys/piv/service.go
index b0c5daa0544bd..8f2f09b81cfed 100644
--- a/api/utils/keys/piv/service.go
+++ b/api/utils/keys/piv/service.go
@@ -135,20 +135,20 @@ func (s *YubiKeyService) NewPrivateKey(ctx context.Context, config hardwarekey.P
 	// If a custom slot was not specified, check for a key in the
 	// default slot for the given policy and generate a new one if needed.
 	if config.CustomSlot == "" {
-		switch cert, err := y.getCertificate(pivSlot); {
-		case errors.Is(err, piv.ErrNotFound):
+		switch err := y.checkCertificate(pivSlot); {
+		case trace.IsNotFound(err):
 			return generatePrivateKey()
 
-		case err != nil:
-			return nil, trace.Wrap(err)
-
 		// Unknown cert found, this slot could be in use by a non-teleport client.
 		// Prompt the user before we overwrite the slot.
-		case !isTeleportMetadataCertificate(cert):
-			if err := s.promptOverwriteSlot(ctx, nonTeleportCertificateMessage(pivSlot, cert), config.ContextualKeyInfo); err != nil {
+		case errors.As(err, &nonTeleportCertError{}):
+			if err := s.promptOverwriteSlot(ctx, err.Error(), config.ContextualKeyInfo); err != nil {
 				return nil, trace.Wrap(err)
 			}
 			return generatePrivateKey()
+
+		case err != nil:
+			return nil, trace.Wrap(err)
 		}
 	}
 
diff --git a/api/utils/keys/piv/yubikey.go b/api/utils/keys/piv/yubikey.go
index 5eaaefa0599eb..ba45f86ecfc0a 100644
--- a/api/utils/keys/piv/yubikey.go
+++ b/api/utils/keys/piv/yubikey.go
@@ -147,16 +147,6 @@ const (
 	signTouchPromptDelay = time.Millisecond * 200
 )
 
-var (
-	ErrMissingTeleportCert = trace.BadParameterError{
-		Message: "hardware key agent cannot perform signatures on PIV slots that aren't configured for Teleport. " +
-			"The PIV slot should be configured automatically by the Teleport client during login. If you are " +
-			"are configuring the PIV slot manually, you must also generate a certificate in the slot with " +
-			"\"teleport\" as the organization name: " +
-			"e.g. \"ykman piv keys generate -a ECCP256 9a pub.pem && ykman piv certificate generate 9a pub.pem -s O=teleport\"",
-	}
-)
-
 func (y *YubiKey) sign(ctx context.Context, ref *hardwarekey.PrivateKeyRef, keyInfo hardwarekey.ContextualKeyInfo, prompt hardwarekey.Prompt, rand io.Reader, digest []byte, opts crypto.SignerOpts) ([]byte, error) {
 	pivSlot, err := parsePIVSlot(ref.SlotKey)
 	if err != nil {
@@ -177,19 +167,14 @@ func (y *YubiKey) sign(ctx context.Context, ref *hardwarekey.PrivateKeyRef, keyI
 		return nil, trace.CompareFailed("public key mismatch on PIV slot 0x%x", pivSlot.Key)
 	}
 
-	// If this sign request is coming from the hardware key agent, ensure that the requested PIV
-	// slot was configured by a Teleport client, or manually configured by the user / hardware key
-	// administrator. Manual configuration is used in cases where the default PIV management key
-	// is not used, e.g. when the hardware key is managed by a third party provider by an admin.
+	// If the sign request is for an unknown agent key, ensure that the requested PIV slot was
+	// configured with a self-signed Teleport metadata certificate.
 	if keyInfo.AgentKey {
-		cert, err := y.getCertificate(pivSlot)
-		switch {
-		case errors.Is(err, piv.ErrNotFound):
-			return nil, trace.Wrap(&ErrMissingTeleportCert, "certificate not found in PIV slot 0x%x", pivSlot.Key)
+		switch err := y.checkCertificate(pivSlot); {
+		case trace.IsNotFound(err), errors.As(err, &nonTeleportCertError{}):
+			return nil, trace.Wrap(err, agentRequiresTeleportCertMessage)
 		case err != nil:
 			return nil, trace.Wrap(err)
-		case !isTeleportMetadataCertificate(cert):
-			return nil, trace.Wrap(&ErrMissingTeleportCert, nonTeleportCertificateMessage(pivSlot, cert))
 		}
 	}
 
@@ -404,10 +389,22 @@ func (y *YubiKey) SetMetadataCertificate(slot piv.Slot, subject pkix.Name) error
 	return trace.Wrap(err)
 }
 
-// getCertificate gets a certificate from the given PIV slot.
-func (y *YubiKey) getCertificate(slot piv.Slot) (*x509.Certificate, error) {
+// checkCertificate checks for a certificate on the PIV slot matching a Teleport client
+// metadata certificate. Expected errors include [trace.NotFoundError] and [nonTeleportCertError].
+func (y *YubiKey) checkCertificate(slot piv.Slot) error {
 	cert, err := y.conn.certificate(slot)
-	return cert, trace.Wrap(err)
+	switch {
+	case errors.Is(err, piv.ErrNotFound):
+		return trace.NotFound("certificate not found in PIV slot 0x%x", slot.Key)
+	case err != nil:
+		return trace.Wrap(err)
+	case !isTeleportMetadataCertificate(cert):
+		return nonTeleportCertError{
+			slot: slot,
+			cert: cert,
+		}
+	}
+	return nil
 }
 
 // attestKey attests the key in the given PIV slot.
@@ -810,9 +807,14 @@ func isTeleportMetadataCertificate(cert *x509.Certificate) bool {
 	return len(cert.Subject.Organization) > 0 && cert.Subject.Organization[0] == certOrgName
 }
 
-func nonTeleportCertificateMessage(slot piv.Slot, cert *x509.Certificate) string {
+type nonTeleportCertError struct {
+	slot piv.Slot
+	cert *x509.Certificate
+}
+
+func (e nonTeleportCertError) Error() string {
 	// Gather a small list of user-readable x509 certificate fields to display to the user.
-	sum := sha256.Sum256(cert.Raw)
+	sum := sha256.Sum256(e.cert.Raw)
 	fingerPrint := hex.EncodeToString(sum[:])
 	return fmt.Sprintf(`Certificate in YubiKey PIV slot %q is not a Teleport client cert:
 Slot %s:
@@ -824,13 +826,19 @@ Slot %s:
 	Not before:		%v
 	Not after:		%v
 `,
-		slot, slot,
-		cert.SignatureAlgorithm,
-		cert.Subject,
-		cert.Issuer,
-		cert.SerialNumber,
+		e.slot, e.slot,
+		e.cert.SignatureAlgorithm,
+		e.cert.Subject,
+		e.cert.Issuer,
+		e.cert.SerialNumber,
 		fingerPrint,
-		cert.NotBefore,
-		cert.NotAfter,
+		e.cert.NotBefore,
+		e.cert.NotAfter,
 	)
 }
+
+var agentRequiresTeleportCertMessage = "hardware key agent cannot perform signatures on PIV slots that aren't configured for Teleport. " +
+	"The PIV slot should be configured automatically by the Teleport client during login. If you are " +
+	"are configuring the PIV slot manually, you must also generate a certificate in the slot with " +
+	"\"teleport\" as the organization name: " +
+	"e.g. \"ykman piv keys generate -a ECCP256 9a pub.pem && ykman piv certificate generate 9a pub.pem -s O=teleport\""
diff --git a/lib/client/client_store.go b/lib/client/client_store.go
index 8a58be144589d..e0735698d42c7 100644
--- a/lib/client/client_store.go
+++ b/lib/client/client_store.go
@@ -118,6 +118,35 @@ func (s *Store) NewHardwarePrivateKey(ctx context.Context, config hardwarekey.Pr
 	return keys.NewHardwarePrivateKey(ctx, s.HardwareKeyService, config)
 }
 
+func (s *Store) KnownHardwareKey(ref *hardwarekey.PrivateKeyRef, keyInfo hardwarekey.ContextualKeyInfo) (bool, error) {
+	keyRing, err := s.GetKeyRing(KeyRingIndex{
+		ProxyHost:   keyInfo.ProxyHost,
+		Username:    keyInfo.Username,
+		ClusterName: keyInfo.ClusterName,
+	})
+	if trace.IsNotFound(err) {
+		return false, nil
+	} else if err != nil {
+		return false, trace.Wrap(err)
+	}
+
+	// There is a known key matching the key info from the agent client, now check
+	// if it is the same key, with the same hardware key reference.
+	hwSigner, ok := keyRing.TLSPrivateKey.Signer.(*hardwarekey.Signer)
+	if !ok {
+		return false, nil
+	}
+
+	// We only need to compare the serial number and slot key. Other values, like the
+	// public key and prompt policy, will be validated against the hardware key directly
+	// when needed.
+	sameKeyRef :=
+		hwSigner.Ref.SerialNumber == ref.SerialNumber &&
+			hwSigner.Ref.SlotKey == ref.SlotKey
+
+	return sameKeyRef, nil
+}
+
 // AddKeyRing adds the given key ring to the key store. The key's trusted certificates are
 // added to the trusted certs store.
 func (s *Store) AddKeyRing(keyRing *KeyRing) error {
diff --git a/lib/hardwarekey/agent.go b/lib/hardwarekey/agent.go
index e86cc4099c743..e03ccf98a6e2f 100644
--- a/lib/hardwarekey/agent.go
+++ b/lib/hardwarekey/agent.go
@@ -75,7 +75,7 @@ type Server struct {
 // The given directory will be created when the server is served and destroyed with the server is stopped.
 //
 // [DefaultAgentDir] should be used for [keyAgentDir] outside of tests.
-func NewAgentServer(ctx context.Context, s hardwarekey.Service, keyAgentDir string) (*Server, error) {
+func NewAgentServer(ctx context.Context, s hardwarekey.Service, keyAgentDir string, knownKeyFn hardwarekeyagent.KnownHardwareKeyFn) (*Server, error) {
 	if err := os.MkdirAll(keyAgentDir, 0o700); err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -91,7 +91,7 @@ func NewAgentServer(ctx context.Context, s hardwarekey.Service, keyAgentDir stri
 		return nil, trace.Wrap(err)
 	}
 
-	grpcServer := hardwarekeyagent.NewServer(ctx, s, credentials.NewServerTLSFromCert(&cert))
+	grpcServer := hardwarekeyagent.NewServer(ctx, s, credentials.NewServerTLSFromCert(&cert), knownKeyFn)
 	return &Server{
 		grpcServer: grpcServer,
 		listener:   l,
diff --git a/lib/hardwarekey/agent_test.go b/lib/hardwarekey/agent_test.go
index 99a551fa9db8d..0ee4c74b0d0aa 100644
--- a/lib/hardwarekey/agent_test.go
+++ b/lib/hardwarekey/agent_test.go
@@ -36,7 +36,7 @@ func TestHardwareKeyAgent_Server(t *testing.T) {
 
 	// Prepare the agent server
 	mockService := hardwarekey.NewMockHardwareKeyService(nil /*prompt*/)
-	server, err := libhwk.NewAgentServer(ctx, mockService, agentDir)
+	server, err := libhwk.NewAgentServer(ctx, mockService, agentDir, nil /*knownKeyFn*/)
 	require.NoError(t, err)
 	t.Cleanup(server.Stop)
 
@@ -46,7 +46,7 @@ func TestHardwareKeyAgent_Server(t *testing.T) {
 	}()
 
 	// Should fail to open a new server in the same directory.
-	_, err = libhwk.NewAgentServer(ctx, mockService, agentDir)
+	_, err = libhwk.NewAgentServer(ctx, mockService, agentDir, nil /*knownKeyFn*/)
 	require.Error(t, err)
 
 	// Existing server should be unaffected.
@@ -61,7 +61,7 @@ func TestHardwareKeyAgent_Server(t *testing.T) {
 		_, err := os.Stat(agentDir)
 		return errors.Is(err, os.ErrNotExist)
 	}, 5*time.Second, 100*time.Millisecond)
-	server, err = libhwk.NewAgentServer(ctx, mockService, agentDir)
+	server, err = libhwk.NewAgentServer(ctx, mockService, agentDir, nil /*knownKeyFn*/)
 	require.NoError(t, err)
 	t.Cleanup(server.Stop)
 
@@ -69,7 +69,7 @@ func TestHardwareKeyAgent_Server(t *testing.T) {
 	// Use a timeoutCtx so that the failed Ping request fails quickly.
 	timeoutCtx, cancel := context.WithTimeout(ctx, 500*time.Millisecond)
 	defer cancel()
-	server, err = libhwk.NewAgentServer(timeoutCtx, mockService, agentDir)
+	server, err = libhwk.NewAgentServer(timeoutCtx, mockService, agentDir, nil /*knownKeyFn*/)
 	require.NoError(t, err)
 	t.Cleanup(server.Stop)
 }
diff --git a/lib/teleterm/teleterm.go b/lib/teleterm/teleterm.go
index 91a5b49b66e72..9425ceda6a791 100644
--- a/lib/teleterm/teleterm.go
+++ b/lib/teleterm/teleterm.go
@@ -106,7 +106,7 @@ func Serve(ctx context.Context, cfg Config) error {
 
 	var hardwareKeyAgentServer *libhwk.Server
 	if cfg.HardwareKeyAgent {
-		hardwareKeyAgentServer, err = libhwk.NewAgentServer(ctx, hwks, libhwk.DefaultAgentDir())
+		hardwareKeyAgentServer, err = libhwk.NewAgentServer(ctx, hwks, libhwk.DefaultAgentDir(), storage.ClientStore.KnownHardwareKey)
 		if err != nil {
 			slog.WarnContext(ctx, "failed to create the hardware key agent server", "err", err)
 		} else {
diff --git a/tool/tsh/common/piv.go b/tool/tsh/common/piv.go
index 16d6097e6e7c6..d99fa9e7f64bc 100644
--- a/tool/tsh/common/piv.go
+++ b/tool/tsh/common/piv.go
@@ -22,7 +22,6 @@ import (
 	"github.com/alecthomas/kingpin/v2"
 	"github.com/gravitational/trace"
 
-	"github.com/gravitational/teleport/api/utils/keys/piv"
 	libhwk "github.com/gravitational/teleport/lib/hardwarekey"
 )
 
@@ -50,8 +49,9 @@ func newPIVAgentCommand(parent *kingpin.CmdClause) *pivAgentCommand {
 }
 
 func (c *pivAgentCommand) run(cf *CLIConf) error {
-	hwKeyService := piv.NewYubiKeyService(nil /*prompt*/)
-	s, err := libhwk.NewAgentServer(cf.Context, hwKeyService, libhwk.DefaultAgentDir())
+	cf.disableHardwareKeyAgentClient = true
+	store := cf.getClientStore()
+	s, err := libhwk.NewAgentServer(cf.Context, store.HardwareKeyService, libhwk.DefaultAgentDir(), store.KnownHardwareKey)
 	if err != nil {
 		return trace.Wrap(err)
 	}

From a6b5d03cba8999e6557eb196554f2b9461ce2980 Mon Sep 17 00:00:00 2001
From: joerger <bjoerger@goteleport.com>
Date: Fri, 9 May 2025 14:17:56 -0700
Subject: [PATCH 2/8] Add test.

---
 api/utils/keys/hardwarekey/service_mock.go    | 21 ++++++++++--
 api/utils/keys/hardwarekeyagent/service.go    |  8 ++++-
 .../keys/hardwarekeyagent/service_test.go     | 34 +++++++++++++++++--
 lib/client/client_store.go                    |  2 ++
 4 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/api/utils/keys/hardwarekey/service_mock.go b/api/utils/keys/hardwarekey/service_mock.go
index 0ec1a1e291133..0151a1000686c 100644
--- a/api/utils/keys/hardwarekey/service_mock.go
+++ b/api/utils/keys/hardwarekey/service_mock.go
@@ -51,6 +51,9 @@ type MockHardwareKeyService struct {
 
 	fakeHardwarePrivateKeys    map[hardwareKeySlot]*fakeHardwarePrivateKey
 	fakeHardwarePrivateKeysMux sync.Mutex
+
+	// mock a PIV slot with a key but no teleport metadata cert.
+	unknownAgentKey map[hardwareKeySlot]bool
 }
 
 // NewMockHardwareKeyService returns a [mockHardwareKeyService] for use in tests.
@@ -60,6 +63,7 @@ func NewMockHardwareKeyService(prompt Prompt) *MockHardwareKeyService {
 		prompt:                  prompt,
 		mockTouch:               make(chan struct{}),
 		fakeHardwarePrivateKeys: map[hardwareKeySlot]*fakeHardwarePrivateKey{},
+		unknownAgentKey:         map[hardwareKeySlot]bool{},
 	}
 }
 
@@ -141,10 +145,16 @@ func (s *MockHardwareKeyService) Sign(ctx context.Context, ref *PrivateKeyRef, k
 	s.fakeHardwarePrivateKeysMux.Lock()
 	defer s.fakeHardwarePrivateKeysMux.Unlock()
 
-	priv, ok := s.fakeHardwarePrivateKeys[hardwareKeySlot{
+	slot := hardwareKeySlot{
 		serialNumber: serialNumber,
 		slot:         ref.SlotKey,
-	}]
+	}
+
+	if keyInfo.AgentKey && s.unknownAgentKey[slot] {
+		return nil, trace.BadParameter("unknown agent key")
+	}
+
+	priv, ok := s.fakeHardwarePrivateKeys[slot]
 	if !ok {
 		return nil, trace.NotFound("key not found in slot 0x%x", ref.SlotKey)
 	}
@@ -193,6 +203,13 @@ func (s *MockHardwareKeyService) SetPrompt(prompt Prompt) {
 	s.prompt = prompt
 }
 
+func (s *MockHardwareKeyService) AddUnknownAgentKey(ref *PrivateKeyRef) {
+	s.unknownAgentKey[hardwareKeySlot{
+		serialNumber: ref.SerialNumber,
+		slot:         ref.SlotKey,
+	}] = true
+}
+
 // TODO(Joerger): DELETE IN v19.0.0
 func (s *MockHardwareKeyService) GetFullKeyRef(serialNumber uint32, slotKey PIVSlotKey) (*PrivateKeyRef, error) {
 	s.fakeHardwarePrivateKeysMux.Lock()
diff --git a/api/utils/keys/hardwarekeyagent/service.go b/api/utils/keys/hardwarekeyagent/service.go
index 20859733e10fd..610431a5c8163 100644
--- a/api/utils/keys/hardwarekeyagent/service.go
+++ b/api/utils/keys/hardwarekeyagent/service.go
@@ -56,6 +56,9 @@ func NewService(agentClient hardwarekeyagentv1.HardwareKeyAgentServiceClient, fa
 
 // NewPrivateKey creates or retrieves a hardware private key for the given config.
 func (s *Service) NewPrivateKey(ctx context.Context, config hardwarekey.PrivateKeyConfig) (*hardwarekey.Signer, error) {
+	if s.fallbackService == nil {
+		return nil, trace.BadParameter("hardware key agent client cannot create a new key without a fallback service")
+	}
 	return s.fallbackService.NewPrivateKey(ctx, config)
 }
 
@@ -64,7 +67,7 @@ func (s *Service) NewPrivateKey(ctx context.Context, config hardwarekey.PrivateK
 func (s *Service) Sign(ctx context.Context, ref *hardwarekey.PrivateKeyRef, keyInfo hardwarekey.ContextualKeyInfo, rand io.Reader, digest []byte, opts crypto.SignerOpts) ([]byte, error) {
 	// First try to sign with the agent, then fallback to the direct service if needed.
 	signature, err := s.agentSign(ctx, ref, keyInfo, rand, digest, opts)
-	if err != nil {
+	if err != nil && s.fallbackService != nil {
 		slog.ErrorContext(ctx, "Failed to perform signature over hardware key agent, falling back to fallback service", "agent_err", err)
 		signature, err = s.fallbackService.Sign(ctx, ref, keyInfo, rand, digest, opts)
 	}
@@ -161,5 +164,8 @@ func (s *Service) agentSign(ctx context.Context, ref *hardwarekey.PrivateKeyRef,
 
 // TODO(Joerger): DELETE IN v19.0.0
 func (s *Service) GetFullKeyRef(serialNumber uint32, slotKey hardwarekey.PIVSlotKey) (*hardwarekey.PrivateKeyRef, error) {
+	if s.fallbackService == nil {
+		return nil, trace.BadParameter("hardware key agent client cannot get missing key ref info without a fallback service")
+	}
 	return s.fallbackService.GetFullKeyRef(serialNumber, slotKey)
 }
diff --git a/api/utils/keys/hardwarekeyagent/service_test.go b/api/utils/keys/hardwarekeyagent/service_test.go
index 7a572543ecd75..02e887d9cf90c 100644
--- a/api/utils/keys/hardwarekeyagent/service_test.go
+++ b/api/utils/keys/hardwarekeyagent/service_test.go
@@ -23,11 +23,14 @@ import (
 	"crypto/rsa"
 	"net"
 	"path/filepath"
+	"slices"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc/credentials/insecure"
 
+	"github.com/gravitational/trace"
+
 	"github.com/gravitational/teleport/api/utils/keys/hardwarekey"
 	"github.com/gravitational/teleport/api/utils/keys/hardwarekeyagent"
 )
@@ -35,9 +38,17 @@ import (
 func TestHardwareKeyAgentService(t *testing.T) {
 	ctx := context.Background()
 
+	// Mock known keys. Usually the server's login session storage would be used to check for known keys.
+	var serverKnownKeySlots []hardwarekey.PIVSlotKey
+	knownKeyFn := func(ref *hardwarekey.PrivateKeyRef, _ hardwarekey.ContextualKeyInfo) (bool, error) {
+		return slices.ContainsFunc(serverKnownKeySlots, func(s hardwarekey.PIVSlotKey) bool {
+			return ref.SlotKey == s
+		}), nil
+	}
+
 	// Prepare the agent server
 	mockService := hardwarekey.NewMockHardwareKeyService(nil /*prompt*/)
-	server := hardwarekeyagent.NewServer(ctx, mockService, insecure.NewCredentials(), nil /*knownKeyFn*/)
+	server := hardwarekeyagent.NewServer(ctx, mockService, insecure.NewCredentials(), knownKeyFn)
 	t.Cleanup(server.Stop)
 
 	agentDir := t.TempDir()
@@ -54,7 +65,7 @@ func TestHardwareKeyAgentService(t *testing.T) {
 	agentClient, err := hardwarekeyagent.NewClient(ctx, socketPath, insecure.NewCredentials())
 	require.NoError(t, err)
 
-	agentService := hardwarekeyagent.NewService(agentClient, hardwarekey.NewMockHardwareKeyService(nil /*prompt*/))
+	agentServiceNoFallback := hardwarekeyagent.NewService(agentClient, nil)
 	agentServiceWithFallback := hardwarekeyagent.NewService(agentClient, mockService)
 
 	for _, tc := range []struct {
@@ -189,7 +200,7 @@ func TestHardwareKeyAgentService(t *testing.T) {
 			}
 
 			// Perform a signature over the agent.
-			_, err = agentService.Sign(ctx, hwSigner.Ref, hwSigner.KeyInfo, rand.Reader, digest, tc.opts)
+			_, err = agentServiceNoFallback.Sign(ctx, hwSigner.Ref, hwSigner.KeyInfo, rand.Reader, digest, tc.opts)
 			if tc.expectErr {
 				require.Error(t, err)
 			} else {
@@ -198,6 +209,23 @@ func TestHardwareKeyAgentService(t *testing.T) {
 		})
 	}
 
+	t.Run("unknown agent key", func(t *testing.T) {
+		mockService.Reset()
+
+		hwSigner, err := mockService.NewPrivateKey(ctx, hardwarekey.PrivateKeyConfig{})
+		require.NoError(t, err)
+
+		// Mark the hardware key as unknown by the Hardware Key Service.
+		mockService.AddUnknownAgentKey(hwSigner.Ref)
+		_, err = agentServiceNoFallback.Sign(ctx, hwSigner.Ref, hwSigner.KeyInfo, rand.Reader, []byte{}, crypto.Hash(0))
+		require.True(t, trace.IsBadParameter(err), "expected trace.BadParameter but got %v", err)
+
+		// Make the hardware key as known by the  Hardware Key Agent Server.
+		serverKnownKeySlots = append(serverKnownKeySlots, hwSigner.Ref.SlotKey)
+		_, err = agentServiceNoFallback.Sign(ctx, hwSigner.Ref, hwSigner.KeyInfo, rand.Reader, []byte{}, crypto.Hash(0))
+		require.NoError(t, err)
+	})
+
 	t.Run("fallback", func(t *testing.T) {
 		mockService.Reset()
 
diff --git a/lib/client/client_store.go b/lib/client/client_store.go
index e0735698d42c7..8dc5a3163d49f 100644
--- a/lib/client/client_store.go
+++ b/lib/client/client_store.go
@@ -118,6 +118,8 @@ func (s *Store) NewHardwarePrivateKey(ctx context.Context, config hardwarekey.Pr
 	return keys.NewHardwarePrivateKey(ctx, s.HardwareKeyService, config)
 }
 
+// KnownHardwareKey returns whether the given hardware key ref and info corresponds to a hardware key known
+// to this client store.
 func (s *Store) KnownHardwareKey(ref *hardwarekey.PrivateKeyRef, keyInfo hardwarekey.ContextualKeyInfo) (bool, error) {
 	keyRing, err := s.GetKeyRing(KeyRingIndex{
 		ProxyHost:   keyInfo.ProxyHost,

From bb543576e31263465f7aa2938db46b3f5e2f8659 Mon Sep 17 00:00:00 2001
From: joerger <bjoerger@goteleport.com>
Date: Fri, 9 May 2025 14:21:52 -0700
Subject: [PATCH 3/8] Update rfd.

---
 rfd/0199-hardware-key-agent.md | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/rfd/0199-hardware-key-agent.md b/rfd/0199-hardware-key-agent.md
index e11c55c0926d2..b7954bf1df266 100644
--- a/rfd/0199-hardware-key-agent.md
+++ b/rfd/0199-hardware-key-agent.md
@@ -680,7 +680,8 @@ the hardware key agent for non Teleport client keys. e.g PIV keys traditionally
 used for email encryption.
 
 The hardware key agent will determine whether a key in a PIV slot is a Teleport
-client key by checking for the [self-signed metadata certificate](./0080-hardware-key-support.md#piv-slot-logic)
+client key by checking for a recent login session matching the key or checking
+for the [self-signed metadata certificate](./0080-hardware-key-support.md#piv-slot-logic)
 generated by Teleport clients on hardware key login.
 
 Note: When user hardware keys are externally managed, administrators are currently

From 857a155f45118c67924a27c73a4f746a297f0c1c Mon Sep 17 00:00:00 2001
From: joerger <bjoerger@goteleport.com>
Date: Mon, 12 May 2025 10:29:12 -0700
Subject: [PATCH 4/8] Fix lint.

---
 api/utils/keys/hardwarekeyagent/service_test.go | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/api/utils/keys/hardwarekeyagent/service_test.go b/api/utils/keys/hardwarekeyagent/service_test.go
index 02e887d9cf90c..ab0035c617fd4 100644
--- a/api/utils/keys/hardwarekeyagent/service_test.go
+++ b/api/utils/keys/hardwarekeyagent/service_test.go
@@ -26,11 +26,10 @@ import (
 	"slices"
 	"testing"
 
+	"github.com/gravitational/trace"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc/credentials/insecure"
 
-	"github.com/gravitational/trace"
-
 	"github.com/gravitational/teleport/api/utils/keys/hardwarekey"
 	"github.com/gravitational/teleport/api/utils/keys/hardwarekeyagent"
 )

From 250ef271f3b0dac659430ea2a9422b1571f38860 Mon Sep 17 00:00:00 2001
From: joerger <bjoerger@goteleport.com>
Date: Tue, 13 May 2025 17:27:49 -0700
Subject: [PATCH 5/8] Address comments.

---
 api/utils/keys/hardwarekeyagent/agent.go   |  6 ++++--
 api/utils/keys/hardwarekeyagent/service.go | 14 +++++++++++---
 api/utils/keys/piv/yubikey.go              |  2 +-
 lib/client/client_store.go                 |  4 +---
 4 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/api/utils/keys/hardwarekeyagent/agent.go b/api/utils/keys/hardwarekeyagent/agent.go
index 06233b05da032..495197e8758fb 100644
--- a/api/utils/keys/hardwarekeyagent/agent.go
+++ b/api/utils/keys/hardwarekeyagent/agent.go
@@ -62,7 +62,7 @@ func NewClient(ctx context.Context, socketPath string, creds credentials.Transpo
 
 // NewServer returns a new hardware key agent server.
 //
-// If [knownKeyFn] is not provided, all PIV keys are treated as agent keys.
+// If a [KnownHardwareKeyFn] is not provided, all PIV keys are treated as agent keys (used in tests).
 func NewServer(_ context.Context, s hardwarekey.Service, creds credentials.TransportCredentials, knownKeyFn KnownHardwareKeyFn) *grpc.Server {
 	grpcServer := grpc.NewServer(
 		grpc.Creds(creds),
@@ -72,6 +72,8 @@ func NewServer(_ context.Context, s hardwarekey.Service, creds credentials.Trans
 	return grpcServer
 }
 
+// KnownHardwareKeyFn is a function to determine if the hardware private key described by the given key ref
+// and info is known or not, usually based on whether a matching key is found in the client's key store.
 type KnownHardwareKeyFn func(ref *hardwarekey.PrivateKeyRef, keyInfo hardwarekey.ContextualKeyInfo) (bool, error)
 
 // agentService implements [hardwarekeyagentv1.HardwareKeyAgentServiceServer].
@@ -83,7 +85,7 @@ type agentService struct {
 	// treated as a known key or an unknown agent key. Unknown agent keys are treated with
 	// additional restrictions to ensure the PIV slot is intended for Teleport client usage.
 	//
-	// If knownKeyFn is nil, all keys are treated as agent keys.
+	// When nil, all keys are treated as agent keys (used in tests).
 	knownKeyFn KnownHardwareKeyFn
 }
 
diff --git a/api/utils/keys/hardwarekeyagent/service.go b/api/utils/keys/hardwarekeyagent/service.go
index 610431a5c8163..b4736674fd754 100644
--- a/api/utils/keys/hardwarekeyagent/service.go
+++ b/api/utils/keys/hardwarekeyagent/service.go
@@ -40,13 +40,21 @@ type Service struct {
 	agentClient hardwarekeyagentv1.HardwareKeyAgentServiceClient
 	// Used for non signature methods and as a fallback for signatures if the
 	// agent client fails to handle a sign request.
+	//
+	// When nil, the service will not fallback on failed sign requests and methods
+	// that aren't supported through the agent service, like NewPrivateKey, will
+	// return an error (used in tests).
 	fallbackService hardwarekey.Service
 }
 
 // NewService creates a new hardware key agent service from the given
-// agent client and fallback service. The fallback service is used for
-// non-signature methods of [hardwarekey.Service] which are not implemented
-// by the agent. Generally this fallback service is only used during login.
+// agent client and fallback service.
+//
+// The fallback service is used for methods unsupported by the agent service,
+// such as [Service.NewPrivateKey], and as a fallback for failed agent signatures.
+//
+// If the fallback service is not provided, the service will return an error
+// for unsupported services and skip the signature fallback logic (used in tests).
 func NewService(agentClient hardwarekeyagentv1.HardwareKeyAgentServiceClient, fallbackService hardwarekey.Service) *Service {
 	return &Service{
 		agentClient:     agentClient,
diff --git a/api/utils/keys/piv/yubikey.go b/api/utils/keys/piv/yubikey.go
index ba45f86ecfc0a..7d46c599e2ca0 100644
--- a/api/utils/keys/piv/yubikey.go
+++ b/api/utils/keys/piv/yubikey.go
@@ -837,7 +837,7 @@ Slot %s:
 	)
 }
 
-var agentRequiresTeleportCertMessage = "hardware key agent cannot perform signatures on PIV slots that aren't configured for Teleport. " +
+const agentRequiresTeleportCertMessage = "hardware key agent cannot perform signatures on PIV slots that aren't configured for Teleport. " +
 	"The PIV slot should be configured automatically by the Teleport client during login. If you are " +
 	"are configuring the PIV slot manually, you must also generate a certificate in the slot with " +
 	"\"teleport\" as the organization name: " +
diff --git a/lib/client/client_store.go b/lib/client/client_store.go
index 8dc5a3163d49f..6899f855f9be5 100644
--- a/lib/client/client_store.go
+++ b/lib/client/client_store.go
@@ -142,9 +142,7 @@ func (s *Store) KnownHardwareKey(ref *hardwarekey.PrivateKeyRef, keyInfo hardwar
 	// We only need to compare the serial number and slot key. Other values, like the
 	// public key and prompt policy, will be validated against the hardware key directly
 	// when needed.
-	sameKeyRef :=
-		hwSigner.Ref.SerialNumber == ref.SerialNumber &&
-			hwSigner.Ref.SlotKey == ref.SlotKey
+	sameKeyRef := hwSigner.Ref.SerialNumber == ref.SerialNumber && hwSigner.Ref.SlotKey == ref.SlotKey
 
 	return sameKeyRef, nil
 }

From 746635a3f6f0140b307e245f070ccde8ed87c604 Mon Sep 17 00:00:00 2001
From: joerger <bjoerger@goteleport.com>
Date: Wed, 14 May 2025 10:41:20 -0700
Subject: [PATCH 6/8] Require knownKeyFn to be provided.

---
 api/utils/keys/hardwarekeyagent/agent.go | 13 +++----------
 lib/hardwarekey/agent.go                 |  4 ++++
 lib/hardwarekey/agent_test.go            | 16 +++++++++++-----
 3 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/api/utils/keys/hardwarekeyagent/agent.go b/api/utils/keys/hardwarekeyagent/agent.go
index 495197e8758fb..dea3f327bacb0 100644
--- a/api/utils/keys/hardwarekeyagent/agent.go
+++ b/api/utils/keys/hardwarekeyagent/agent.go
@@ -61,8 +61,6 @@ func NewClient(ctx context.Context, socketPath string, creds credentials.Transpo
 }
 
 // NewServer returns a new hardware key agent server.
-//
-// If a [KnownHardwareKeyFn] is not provided, all PIV keys are treated as agent keys (used in tests).
 func NewServer(_ context.Context, s hardwarekey.Service, creds credentials.TransportCredentials, knownKeyFn KnownHardwareKeyFn) *grpc.Server {
 	grpcServer := grpc.NewServer(
 		grpc.Creds(creds),
@@ -84,8 +82,6 @@ type agentService struct {
 	// knownKeyFn dictates whether a PIV key referenced in a [agentService.Sign] request is
 	// treated as a known key or an unknown agent key. Unknown agent keys are treated with
 	// additional restrictions to ensure the PIV slot is intended for Teleport client usage.
-	//
-	// When nil, all keys are treated as agent keys (used in tests).
 	knownKeyFn KnownHardwareKeyFn
 }
 
@@ -124,12 +120,9 @@ func (s *agentService) Sign(ctx context.Context, req *hardwarekeyagentv1.SignReq
 		Command:     req.Command,
 	}
 
-	var knownKey bool
-	if s.knownKeyFn != nil {
-		knownKey, err = s.knownKeyFn(keyRef, keyInfo)
-		if err != nil {
-			return nil, trace.Wrap(err)
-		}
+	knownKey, err := s.knownKeyFn(keyRef, keyInfo)
+	if err != nil {
+		return nil, trace.Wrap(err)
 	}
 	keyInfo.AgentKey = !knownKey
 
diff --git a/lib/hardwarekey/agent.go b/lib/hardwarekey/agent.go
index e03ccf98a6e2f..4dc60fb03dfed 100644
--- a/lib/hardwarekey/agent.go
+++ b/lib/hardwarekey/agent.go
@@ -76,6 +76,10 @@ type Server struct {
 //
 // [DefaultAgentDir] should be used for [keyAgentDir] outside of tests.
 func NewAgentServer(ctx context.Context, s hardwarekey.Service, keyAgentDir string, knownKeyFn hardwarekeyagent.KnownHardwareKeyFn) (*Server, error) {
+	if knownKeyFn == nil {
+		return nil, trace.BadParameter("knownKeyFn must be provided")
+	}
+
 	if err := os.MkdirAll(keyAgentDir, 0o700); err != nil {
 		return nil, trace.Wrap(err)
 	}
diff --git a/lib/hardwarekey/agent_test.go b/lib/hardwarekey/agent_test.go
index 0ee4c74b0d0aa..91ce8e3823d9e 100644
--- a/lib/hardwarekey/agent_test.go
+++ b/lib/hardwarekey/agent_test.go
@@ -34,9 +34,15 @@ func TestHardwareKeyAgent_Server(t *testing.T) {
 	ctx := context.Background()
 	agentDir := t.TempDir()
 
-	// Prepare the agent server
 	mockService := hardwarekey.NewMockHardwareKeyService(nil /*prompt*/)
-	server, err := libhwk.NewAgentServer(ctx, mockService, agentDir, nil /*knownKeyFn*/)
+
+	// treat all keys as unknown (agent) keys.
+	knownKeyFn := func(_ *hardwarekey.PrivateKeyRef, _ hardwarekey.ContextualKeyInfo) (bool, error) {
+		return false, nil
+	}
+
+	// Prepare the agent server
+	server, err := libhwk.NewAgentServer(ctx, mockService, agentDir, knownKeyFn)
 	require.NoError(t, err)
 	t.Cleanup(server.Stop)
 
@@ -46,7 +52,7 @@ func TestHardwareKeyAgent_Server(t *testing.T) {
 	}()
 
 	// Should fail to open a new server in the same directory.
-	_, err = libhwk.NewAgentServer(ctx, mockService, agentDir, nil /*knownKeyFn*/)
+	_, err = libhwk.NewAgentServer(ctx, mockService, agentDir, knownKeyFn)
 	require.Error(t, err)
 
 	// Existing server should be unaffected.
@@ -61,7 +67,7 @@ func TestHardwareKeyAgent_Server(t *testing.T) {
 		_, err := os.Stat(agentDir)
 		return errors.Is(err, os.ErrNotExist)
 	}, 5*time.Second, 100*time.Millisecond)
-	server, err = libhwk.NewAgentServer(ctx, mockService, agentDir, nil /*knownKeyFn*/)
+	server, err = libhwk.NewAgentServer(ctx, mockService, agentDir, knownKeyFn)
 	require.NoError(t, err)
 	t.Cleanup(server.Stop)
 
@@ -69,7 +75,7 @@ func TestHardwareKeyAgent_Server(t *testing.T) {
 	// Use a timeoutCtx so that the failed Ping request fails quickly.
 	timeoutCtx, cancel := context.WithTimeout(ctx, 500*time.Millisecond)
 	defer cancel()
-	server, err = libhwk.NewAgentServer(timeoutCtx, mockService, agentDir, nil /*knownKeyFn*/)
+	server, err = libhwk.NewAgentServer(timeoutCtx, mockService, agentDir, knownKeyFn)
 	require.NoError(t, err)
 	t.Cleanup(server.Stop)
 }

From f7bc36e264bb38794eca9356bf3817d7be804dda Mon Sep 17 00:00:00 2001
From: joerger <bjoerger@goteleport.com>
Date: Wed, 14 May 2025 10:44:15 -0700
Subject: [PATCH 7/8] Require fallbackService to be provided to agent service.

---
 api/utils/keys/hardwarekeyagent/service.go      | 15 +--------------
 api/utils/keys/hardwarekeyagent/service_test.go |  6 +++---
 2 files changed, 4 insertions(+), 17 deletions(-)

diff --git a/api/utils/keys/hardwarekeyagent/service.go b/api/utils/keys/hardwarekeyagent/service.go
index b4736674fd754..b3505409778f7 100644
--- a/api/utils/keys/hardwarekeyagent/service.go
+++ b/api/utils/keys/hardwarekeyagent/service.go
@@ -40,10 +40,6 @@ type Service struct {
 	agentClient hardwarekeyagentv1.HardwareKeyAgentServiceClient
 	// Used for non signature methods and as a fallback for signatures if the
 	// agent client fails to handle a sign request.
-	//
-	// When nil, the service will not fallback on failed sign requests and methods
-	// that aren't supported through the agent service, like NewPrivateKey, will
-	// return an error (used in tests).
 	fallbackService hardwarekey.Service
 }
 
@@ -52,9 +48,6 @@ type Service struct {
 //
 // The fallback service is used for methods unsupported by the agent service,
 // such as [Service.NewPrivateKey], and as a fallback for failed agent signatures.
-//
-// If the fallback service is not provided, the service will return an error
-// for unsupported services and skip the signature fallback logic (used in tests).
 func NewService(agentClient hardwarekeyagentv1.HardwareKeyAgentServiceClient, fallbackService hardwarekey.Service) *Service {
 	return &Service{
 		agentClient:     agentClient,
@@ -64,9 +57,6 @@ func NewService(agentClient hardwarekeyagentv1.HardwareKeyAgentServiceClient, fa
 
 // NewPrivateKey creates or retrieves a hardware private key for the given config.
 func (s *Service) NewPrivateKey(ctx context.Context, config hardwarekey.PrivateKeyConfig) (*hardwarekey.Signer, error) {
-	if s.fallbackService == nil {
-		return nil, trace.BadParameter("hardware key agent client cannot create a new key without a fallback service")
-	}
 	return s.fallbackService.NewPrivateKey(ctx, config)
 }
 
@@ -75,7 +65,7 @@ func (s *Service) NewPrivateKey(ctx context.Context, config hardwarekey.PrivateK
 func (s *Service) Sign(ctx context.Context, ref *hardwarekey.PrivateKeyRef, keyInfo hardwarekey.ContextualKeyInfo, rand io.Reader, digest []byte, opts crypto.SignerOpts) ([]byte, error) {
 	// First try to sign with the agent, then fallback to the direct service if needed.
 	signature, err := s.agentSign(ctx, ref, keyInfo, rand, digest, opts)
-	if err != nil && s.fallbackService != nil {
+	if err != nil {
 		slog.ErrorContext(ctx, "Failed to perform signature over hardware key agent, falling back to fallback service", "agent_err", err)
 		signature, err = s.fallbackService.Sign(ctx, ref, keyInfo, rand, digest, opts)
 	}
@@ -172,8 +162,5 @@ func (s *Service) agentSign(ctx context.Context, ref *hardwarekey.PrivateKeyRef,
 
 // TODO(Joerger): DELETE IN v19.0.0
 func (s *Service) GetFullKeyRef(serialNumber uint32, slotKey hardwarekey.PIVSlotKey) (*hardwarekey.PrivateKeyRef, error) {
-	if s.fallbackService == nil {
-		return nil, trace.BadParameter("hardware key agent client cannot get missing key ref info without a fallback service")
-	}
 	return s.fallbackService.GetFullKeyRef(serialNumber, slotKey)
 }
diff --git a/api/utils/keys/hardwarekeyagent/service_test.go b/api/utils/keys/hardwarekeyagent/service_test.go
index ab0035c617fd4..960be30711731 100644
--- a/api/utils/keys/hardwarekeyagent/service_test.go
+++ b/api/utils/keys/hardwarekeyagent/service_test.go
@@ -26,7 +26,6 @@ import (
 	"slices"
 	"testing"
 
-	"github.com/gravitational/trace"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc/credentials/insecure"
 
@@ -64,7 +63,8 @@ func TestHardwareKeyAgentService(t *testing.T) {
 	agentClient, err := hardwarekeyagent.NewClient(ctx, socketPath, insecure.NewCredentials())
 	require.NoError(t, err)
 
-	agentServiceNoFallback := hardwarekeyagent.NewService(agentClient, nil)
+	unusedService := hardwarekey.NewMockHardwareKeyService(nil /*prompt*/)
+	agentServiceNoFallback := hardwarekeyagent.NewService(agentClient, unusedService)
 	agentServiceWithFallback := hardwarekeyagent.NewService(agentClient, mockService)
 
 	for _, tc := range []struct {
@@ -217,7 +217,7 @@ func TestHardwareKeyAgentService(t *testing.T) {
 		// Mark the hardware key as unknown by the Hardware Key Service.
 		mockService.AddUnknownAgentKey(hwSigner.Ref)
 		_, err = agentServiceNoFallback.Sign(ctx, hwSigner.Ref, hwSigner.KeyInfo, rand.Reader, []byte{}, crypto.Hash(0))
-		require.True(t, trace.IsBadParameter(err), "expected trace.BadParameter but got %v", err)
+		require.Error(t, err)
 
 		// Make the hardware key as known by the  Hardware Key Agent Server.
 		serverKnownKeySlots = append(serverKnownKeySlots, hwSigner.Ref.SlotKey)

From d83ac46dd12c27978bd243d3c766617908a9be51 Mon Sep 17 00:00:00 2001
From: joerger <bjoerger@goteleport.com>
Date: Fri, 16 May 2025 11:10:48 -0700
Subject: [PATCH 8/8] Address comments.

---
 api/utils/keys/hardwarekey/cliprompt.go       |  8 ++---
 api/utils/keys/hardwarekey/hardwarekey.go     | 23 ++++++++-----
 api/utils/keys/hardwarekey/service_mock.go    |  2 +-
 api/utils/keys/hardwarekeyagent/agent.go      | 33 +++++++++++++------
 .../keys/hardwarekeyagent/service_test.go     |  5 +--
 api/utils/keys/piv/yubikey.go                 |  2 +-
 lib/hardwarekey/agent.go                      |  9 +++--
 lib/teleterm/daemon/hardwarekeyprompt.go      | 12 +++----
 8 files changed, 59 insertions(+), 35 deletions(-)

diff --git a/api/utils/keys/hardwarekey/cliprompt.go b/api/utils/keys/hardwarekey/cliprompt.go
index 8a44973e063f0..72095abd7531c 100644
--- a/api/utils/keys/hardwarekey/cliprompt.go
+++ b/api/utils/keys/hardwarekey/cliprompt.go
@@ -69,8 +69,8 @@ func (c *cliPrompt) AskPIN(ctx context.Context, requirement PINPromptRequirement
 
 	// If this is a hardware key agent request with command context info,
 	// include the command in the prompt.
-	if keyInfo.Command != "" {
-		msg = fmt.Sprintf("%v to continue with command %q", msg, keyInfo.Command)
+	if keyInfo.AgentKeyInfo.Command != "" {
+		msg = fmt.Sprintf("%v to continue with command %q", msg, keyInfo.AgentKeyInfo.Command)
 	}
 
 	pin, err := prompt.Password(ctx, c.writer, c.reader, msg)
@@ -88,8 +88,8 @@ func (c *cliPrompt) AskPIN(ctx context.Context, requirement PINPromptRequirement
 // Touch prompts the user to touch the hardware key.
 func (c *cliPrompt) Touch(_ context.Context, keyInfo ContextualKeyInfo) error {
 	msg := "Tap your YubiKey"
-	if keyInfo.Command != "" {
-		msg = fmt.Sprintf("%v to continue with command %q", msg, keyInfo.Command)
+	if keyInfo.AgentKeyInfo.Command != "" {
+		msg = fmt.Sprintf("%v to continue with command %q", msg, keyInfo.AgentKeyInfo.Command)
 	}
 
 	_, err := fmt.Fprintln(c.writer, msg)
diff --git a/api/utils/keys/hardwarekey/hardwarekey.go b/api/utils/keys/hardwarekey/hardwarekey.go
index f6b18c677ad6b..c152fe746ff68 100644
--- a/api/utils/keys/hardwarekey/hardwarekey.go
+++ b/api/utils/keys/hardwarekey/hardwarekey.go
@@ -254,16 +254,21 @@ type ContextualKeyInfo struct {
 	Username string
 	// ClusterName is a Teleport cluster name that the key is associated with.
 	ClusterName string
-	// AgentKey specifies whether this key is being utilized through an agent and
-	// the key is unknown to the client running the agent, probably because the
-	// client on the other side of the agent is using a different Teleport Home
-	// directory.
+	// AgentKeyInfo contains info associated with an hardware key agent signature request.
+	AgentKeyInfo AgentKeyInfo
+}
+
+// AgentKeyInfo contains info associated with an hardware key agent signature request.
+type AgentKeyInfo struct {
+	// UnknownAgentKey indicates whether this hardware private key is known to the hardware key agent
+	// process, usually based on whether a matching key is found in the process's client key store.
 	//
-	// When true, the hardware key service will check that the certificate in the
-	// same slot as the key matches a Teleport client metadata certificate in order
-	// to ensure the agent doesn't provide access to non teleport client PIV keys.
-	AgentKey bool
-	// Command is the running command utilizing this key.
+	// For unknown agent keys, the hardware key service will check that the certificate in the same
+	// slot as the key matches a Teleport client metadata certificate in order to ensure the agent
+	// doesn't provide access to non teleport client PIV keys.
+	UnknownAgentKey bool
+	// Command is the command reported by the agent client which this agent key is being utilized to
+	// complete, e.g. `tsh ssh server01`.
 	Command string
 }
 
diff --git a/api/utils/keys/hardwarekey/service_mock.go b/api/utils/keys/hardwarekey/service_mock.go
index 0151a1000686c..9a200c9dde025 100644
--- a/api/utils/keys/hardwarekey/service_mock.go
+++ b/api/utils/keys/hardwarekey/service_mock.go
@@ -150,7 +150,7 @@ func (s *MockHardwareKeyService) Sign(ctx context.Context, ref *PrivateKeyRef, k
 		slot:         ref.SlotKey,
 	}
 
-	if keyInfo.AgentKey && s.unknownAgentKey[slot] {
+	if keyInfo.AgentKeyInfo.UnknownAgentKey && s.unknownAgentKey[slot] {
 		return nil, trace.BadParameter("unknown agent key")
 	}
 
diff --git a/api/utils/keys/hardwarekeyagent/agent.go b/api/utils/keys/hardwarekeyagent/agent.go
index dea3f327bacb0..e419e701c06ee 100644
--- a/api/utils/keys/hardwarekeyagent/agent.go
+++ b/api/utils/keys/hardwarekeyagent/agent.go
@@ -38,7 +38,7 @@ import (
 )
 
 // NewClient creates a new hardware key agent client.
-func NewClient(ctx context.Context, socketPath string, creds credentials.TransportCredentials) (hardwarekeyagentv1.HardwareKeyAgentServiceClient, error) {
+func NewClient(socketPath string, creds credentials.TransportCredentials) (hardwarekeyagentv1.HardwareKeyAgentServiceClient, error) {
 	if _, err := os.Stat(socketPath); err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -61,17 +61,22 @@ func NewClient(ctx context.Context, socketPath string, creds credentials.Transpo
 }
 
 // NewServer returns a new hardware key agent server.
-func NewServer(_ context.Context, s hardwarekey.Service, creds credentials.TransportCredentials, knownKeyFn KnownHardwareKeyFn) *grpc.Server {
+func NewServer(s hardwarekey.Service, creds credentials.TransportCredentials, knownKeyFn KnownHardwareKeyFn) (*grpc.Server, error) {
+	if knownKeyFn == nil {
+		return nil, trace.BadParameter("knownKeyFn must be provided")
+	}
+
 	grpcServer := grpc.NewServer(
 		grpc.Creds(creds),
 		grpc.UnaryInterceptor(interceptors.GRPCServerUnaryErrorInterceptor),
 	)
 	hardwarekeyagentv1.RegisterHardwareKeyAgentServiceServer(grpcServer, &agentService{s: s, knownKeyFn: knownKeyFn})
-	return grpcServer
+	return grpcServer, nil
 }
 
-// KnownHardwareKeyFn is a function to determine if the hardware private key described by the given key ref
-// and info is known or not, usually based on whether a matching key is found in the client's key store.
+// KnownHardwareKeyFn is a function to determine if the hardware private key, described by the given
+// key ref and key info, is known by this process. This is usually based on whether a matching key
+// is found in the process's client key store.
 type KnownHardwareKeyFn func(ref *hardwarekey.PrivateKeyRef, keyInfo hardwarekey.ContextualKeyInfo) (bool, error)
 
 // agentService implements [hardwarekeyagentv1.HardwareKeyAgentServiceServer].
@@ -79,9 +84,14 @@ type agentService struct {
 	hardwarekeyagentv1.UnimplementedHardwareKeyAgentServiceServer
 	s hardwarekey.Service
 
-	// knownKeyFn dictates whether a PIV key referenced in a [agentService.Sign] request is
-	// treated as a known key or an unknown agent key. Unknown agent keys are treated with
-	// additional restrictions to ensure the PIV slot is intended for Teleport client usage.
+	// knownKeyFn is a function to determine if the hardware private key, described by the given
+	// key ref and key info, is known by this process. This is usually based on whether a matching key
+	// is found in the process's client key store.
+	//
+	// Unknown keys will treated with additional restrictions in [agentService.Sign] requests to
+	// ensure the PIV slot is intended for Teleport client usage, e.g. the agent will require that
+	// the PIV slot has a self-signed metadata certificate used to identify PIV keys generated
+	// specifically for Teleport use.
 	knownKeyFn KnownHardwareKeyFn
 }
 
@@ -117,14 +127,17 @@ func (s *agentService) Sign(ctx context.Context, req *hardwarekeyagentv1.SignReq
 		ProxyHost:   req.KeyInfo.ProxyHost,
 		Username:    req.KeyInfo.Username,
 		ClusterName: req.KeyInfo.ClusterName,
-		Command:     req.Command,
 	}
 
 	knownKey, err := s.knownKeyFn(keyRef, keyInfo)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
-	keyInfo.AgentKey = !knownKey
+
+	keyInfo.AgentKeyInfo = hardwarekey.AgentKeyInfo{
+		UnknownAgentKey: !knownKey,
+		Command:         req.Command,
+	}
 
 	var signerOpts crypto.SignerOpts
 	switch req.Hash {
diff --git a/api/utils/keys/hardwarekeyagent/service_test.go b/api/utils/keys/hardwarekeyagent/service_test.go
index 960be30711731..9cf66daabcb53 100644
--- a/api/utils/keys/hardwarekeyagent/service_test.go
+++ b/api/utils/keys/hardwarekeyagent/service_test.go
@@ -46,7 +46,8 @@ func TestHardwareKeyAgentService(t *testing.T) {
 
 	// Prepare the agent server
 	mockService := hardwarekey.NewMockHardwareKeyService(nil /*prompt*/)
-	server := hardwarekeyagent.NewServer(ctx, mockService, insecure.NewCredentials(), knownKeyFn)
+	server, err := hardwarekeyagent.NewServer(mockService, insecure.NewCredentials(), knownKeyFn)
+	require.NoError(t, err)
 	t.Cleanup(server.Stop)
 
 	agentDir := t.TempDir()
@@ -60,7 +61,7 @@ func TestHardwareKeyAgentService(t *testing.T) {
 	}()
 
 	// Prepare the agent client
-	agentClient, err := hardwarekeyagent.NewClient(ctx, socketPath, insecure.NewCredentials())
+	agentClient, err := hardwarekeyagent.NewClient(socketPath, insecure.NewCredentials())
 	require.NoError(t, err)
 
 	unusedService := hardwarekey.NewMockHardwareKeyService(nil /*prompt*/)
diff --git a/api/utils/keys/piv/yubikey.go b/api/utils/keys/piv/yubikey.go
index 7d46c599e2ca0..a1ba9570f6b56 100644
--- a/api/utils/keys/piv/yubikey.go
+++ b/api/utils/keys/piv/yubikey.go
@@ -169,7 +169,7 @@ func (y *YubiKey) sign(ctx context.Context, ref *hardwarekey.PrivateKeyRef, keyI
 
 	// If the sign request is for an unknown agent key, ensure that the requested PIV slot was
 	// configured with a self-signed Teleport metadata certificate.
-	if keyInfo.AgentKey {
+	if keyInfo.AgentKeyInfo.UnknownAgentKey {
 		switch err := y.checkCertificate(pivSlot); {
 		case trace.IsNotFound(err), errors.As(err, &nonTeleportCertError{}):
 			return nil, trace.Wrap(err, agentRequiresTeleportCertMessage)
diff --git a/lib/hardwarekey/agent.go b/lib/hardwarekey/agent.go
index 4dc60fb03dfed..08a9e86907fb9 100644
--- a/lib/hardwarekey/agent.go
+++ b/lib/hardwarekey/agent.go
@@ -61,7 +61,7 @@ func NewAgentClient(ctx context.Context, keyAgentDir string) (hardwarekeyagentv1
 		return nil, err
 	}
 
-	return hardwarekeyagent.NewClient(ctx, socketPath, creds)
+	return hardwarekeyagent.NewClient(socketPath, creds)
 }
 
 // Server implementation [hardwarekeyagentv1.HardwareKeyAgentServiceServer].
@@ -95,7 +95,12 @@ func NewAgentServer(ctx context.Context, s hardwarekey.Service, keyAgentDir stri
 		return nil, trace.Wrap(err)
 	}
 
-	grpcServer := hardwarekeyagent.NewServer(ctx, s, credentials.NewServerTLSFromCert(&cert), knownKeyFn)
+	grpcServer, err := hardwarekeyagent.NewServer(s, credentials.NewServerTLSFromCert(&cert), knownKeyFn)
+	if err != nil {
+		l.Close()
+		return nil, trace.Wrap(err)
+	}
+
 	return &Server{
 		grpcServer: grpcServer,
 		listener:   l,
diff --git a/lib/teleterm/daemon/hardwarekeyprompt.go b/lib/teleterm/daemon/hardwarekeyprompt.go
index d7e8248a8665f..4059725bcf99e 100644
--- a/lib/teleterm/daemon/hardwarekeyprompt.go
+++ b/lib/teleterm/daemon/hardwarekeyprompt.go
@@ -56,8 +56,8 @@ type hardwareKeyPrompter struct {
 // Touch prompts the user to touch the hardware key.
 func (h *hardwareKeyPrompter) Touch(ctx context.Context, keyInfo hardwarekey.ContextualKeyInfo) error {
 	// Don't include "tsh daemon" commands.
-	if strings.Contains(keyInfo.Command, "tsh daemon") {
-		keyInfo.Command = ""
+	if strings.Contains(keyInfo.AgentKeyInfo.Command, "tsh daemon") {
+		keyInfo.AgentKeyInfo.Command = ""
 	}
 
 	clt, err := h.c.GetClient(ctx)
@@ -67,7 +67,7 @@ func (h *hardwareKeyPrompter) Touch(ctx context.Context, keyInfo hardwarekey.Con
 
 	_, err = clt.PromptHardwareKeyTouch(ctx, &api.PromptHardwareKeyTouchRequest{
 		ProxyHostname: keyInfo.ProxyHost,
-		Command:       keyInfo.Command,
+		Command:       keyInfo.AgentKeyInfo.Command,
 	})
 	if err != nil {
 		return trace.Wrap(err)
@@ -78,8 +78,8 @@ func (h *hardwareKeyPrompter) Touch(ctx context.Context, keyInfo hardwarekey.Con
 // AskPIN prompts the user for a PIN.
 func (h *hardwareKeyPrompter) AskPIN(ctx context.Context, requirement hardwarekey.PINPromptRequirement, keyInfo hardwarekey.ContextualKeyInfo) (string, error) {
 	// Don't include "tsh daemon" commands.
-	if strings.Contains(keyInfo.Command, "tsh daemon") {
-		keyInfo.Command = ""
+	if strings.Contains(keyInfo.AgentKeyInfo.Command, "tsh daemon") {
+		keyInfo.AgentKeyInfo.Command = ""
 	}
 
 	clt, err := h.c.GetClient(ctx)
@@ -90,7 +90,7 @@ func (h *hardwareKeyPrompter) AskPIN(ctx context.Context, requirement hardwareke
 	res, err := clt.PromptHardwareKeyPIN(ctx, &api.PromptHardwareKeyPINRequest{
 		ProxyHostname: keyInfo.ProxyHost,
 		PinOptional:   requirement == hardwarekey.PINOptional,
-		Command:       keyInfo.Command,
+		Command:       keyInfo.AgentKeyInfo.Command,
 	})
 	if err != nil {
 		return "", trace.Wrap(err)