feat: Hardware Key Agent validates known keys

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)

api/utils/keys/hardwarekey/hardwarekey.go

@@ -254,13 +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.
-	// 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 bool
-	// Command is the running command utilizing this key.
+	// 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.
+	//
+	// 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
 }
 

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.AgentKeyInfo.UnknownAgentKey && 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()

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,19 +61,38 @@ 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 {
+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})
-	return grpcServer
+	hardwarekeyagentv1.RegisterHardwareKeyAgentServiceServer(grpcServer, &agentService{s: s, knownKeyFn: knownKeyFn})
+	return grpcServer, nil
 }
 
+// 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].
 type agentService struct {
 	hardwarekeyagentv1.UnimplementedHardwareKeyAgentServiceServer
 	s hardwarekey.Service
+
+	// 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
 }
 
 // Sign the given digest with the specified hardware private key.
@@ -108,8 +127,16 @@ 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,
+	}
+
+	knownKey, err := s.knownKeyFn(keyRef, keyInfo)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	keyInfo.AgentKeyInfo = hardwarekey.AgentKeyInfo{
+		UnknownAgentKey: !knownKey,
+		Command:         req.Command,
 	}
 
 	var signerOpts crypto.SignerOpts

api/utils/keys/hardwarekeyagent/service.go

@@ -44,9 +44,10 @@ type Service struct {
 }
 
 // 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.
 func NewService(agentClient hardwarekeyagentv1.HardwareKeyAgentServiceClient, fallbackService hardwarekey.Service) *Service {
 	return &Service{
 		agentClient:     agentClient,

api/utils/keys/hardwarekeyagent/service_test.go

@@ -23,6 +23,7 @@ import (
 	"crypto/rsa"
 	"net"
 	"path/filepath"
+	"slices"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -35,9 +36,18 @@ 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())
+	server, err := hardwarekeyagent.NewServer(mockService, insecure.NewCredentials(), knownKeyFn)
+	require.NoError(t, err)
 	t.Cleanup(server.Stop)
 
 	agentDir := t.TempDir()
@@ -51,10 +61,11 @@ 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)
 
-	agentService := hardwarekeyagent.NewService(agentClient, hardwarekey.NewMockHardwareKeyService(nil /*prompt*/))
+	unusedService := hardwarekey.NewMockHardwareKeyService(nil /*prompt*/)
+	agentServiceNoFallback := hardwarekeyagent.NewService(agentClient, unusedService)
 	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.Error(t, 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()
 

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)
 		}
 	}