fix: don't cache an incorrect PIV PIN (#54614)

* Verify PIV PIN before we cache it; Fix lint for pincache.go with build tag.

* Add local test for PIN caching in the YubiKey service.

* Minor fixes.

Author: Joerger

api/utils/keys/hardwarekey/cliprompt.go

@@ -73,8 +73,16 @@ func (c *cliPrompt) AskPIN(ctx context.Context, requirement PINPromptRequirement
 		msg = fmt.Sprintf("%v to continue with command %q", msg, keyInfo.Command)
 	}
 
-	password, err := prompt.Password(ctx, c.writer, c.reader, msg)
-	return password, trace.Wrap(err)
+	pin, err := prompt.Password(ctx, c.writer, c.reader, msg)
+	if err != nil {
+		return "", nil
+	}
+
+	if pin == "" {
+		pin = DefaultPIN
+	}
+
+	return pin, trace.Wrap(err)
 }
 
 // Touch prompts the user to touch the hardware key.

api/utils/keys/piv/pincache.go

@@ -1,3 +1,5 @@
+//go:build piv || pivtest
+
 // Copyright 2025 Gravitational, Inc.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,14 +17,10 @@
 package piv
 
 import (
-	"context"
 	"sync"
 	"time"
 
-	"github.com/gravitational/trace"
 	"github.com/jonboulle/clockwork"
-
-	"github.com/gravitational/teleport/api/utils/keys/hardwarekey"
 )
 
 // pinCache is a PIN cache that supports consumers with varying required TTLs.
@@ -46,38 +44,14 @@ func newPINCache() *pinCache { //nolint:unused // used in yubikey.go with piv bu
 	}
 }
 
-// PromptOrGetPIN retrieves the cached PIN if set. Otherwise it prompts for the PIN and caches it.
-func (p *pinCache) PromptOrGetPIN(ctx context.Context, prompt hardwarekey.Prompt, requirement hardwarekey.PINPromptRequirement, keyInfo hardwarekey.ContextualKeyInfo, pinCacheTTL time.Duration) (string, error) {
-	// If the provided ttl is 0, it doesn't support caching, so we just prompt.
-	if pinCacheTTL == 0 {
-		return prompt.AskPIN(ctx, requirement, keyInfo)
-	}
-
-	p.mu.Lock()
-	defer p.mu.Unlock()
-
-	if pin := p.getPIN(pinCacheTTL); pin != "" {
-		return pin, nil
-	}
-
-	// Add a timeout to prevent an unanswered PIN prompt from holding the lock.
-	const pinPromptTimeout = time.Minute
-	ctx, cancel := context.WithTimeout(ctx, pinPromptTimeout)
-	defer cancel()
-
-	pin, err := prompt.AskPIN(ctx, requirement, keyInfo)
-	if err != nil {
-		return "", trace.Wrap(err)
-	}
-
-	p.setPIN(pin, pinCacheTTL)
-	return pin, nil
-}
-
 // getPIN retrieves the cached PIN. If the PIN was cached before by an amount of
 // time equal to the provided TTL, the PIN will not be returned.
 // Must be called under [p.mu] lock.
 func (p *pinCache) getPIN(ttl time.Duration) string {
+	if ttl == 0 {
+		return ""
+	}
+
 	if p.pin == "" {
 		return ""
 	}
@@ -105,6 +79,10 @@ func (p *pinCache) getPIN(ttl time.Duration) string {
 // TTL would exceed that expiration.
 // Must be called under [p.mu] lock.
 func (p *pinCache) setPIN(pin string, ttl time.Duration) {
+	if ttl == 0 {
+		return
+	}
+
 	now := p.clock.Now()
 	expiry := now.Add(ttl)
 

api/utils/keys/piv/pincache_test.go

@@ -1,3 +1,5 @@
+//go:build pivtest
+
 // Copyright 2025 Gravitational, Inc.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");

api/utils/keys/piv/service_test.go

@@ -23,6 +23,7 @@ import (
 	"fmt"
 	"os"
 	"testing"
+	"time"
 
 	pivgo "github.com/go-piv/piv-go/piv"
 	"github.com/gravitational/trace"
@@ -217,6 +218,79 @@ func TestOverwritePrompt(t *testing.T) {
 	})
 }
 
+func TestPINCaching(t *testing.T) {
+	// This test will overwrite any PIV data on the yubiKey.
+	if os.Getenv("TELEPORT_TEST_YUBIKEY_PIV") == "" {
+		t.Skipf("Skipping TestGenerateYubiKeyPrivateKey because TELEPORT_TEST_YUBIKEY_PIV is not set")
+	}
+
+	ctx := context.Background()
+
+	promptWriter := bytes.NewBuffer([]byte{})
+	promptReader := prompt.NewFakeReader()
+	prompt := hardwarekey.NewCLIPrompt(promptWriter, promptReader)
+	s := piv.NewYubiKeyService(prompt)
+
+	y, err := piv.FindYubiKey(0)
+	require.NoError(t, err)
+
+	resetYubikey(t, y)
+	t.Cleanup(func() { resetYubikey(t, y) })
+
+	// Set pin.
+	const testPIN = "123123"
+	require.NoError(t, y.SetPIN(pivgo.DefaultPIN, testPIN))
+	const wrongPIN = "123321"
+
+	// Generate a key with PINPolicyAlways so that the PIN isn't cached internally on the YubiKey.
+	pivSlot := pivgo.SlotAuthentication
+	err = y.GenerateKey(pivSlot, pivgo.Key{
+		Algorithm:   pivgo.AlgorithmEC384,
+		PINPolicy:   pivgo.PINPolicyAlways,
+		TouchPolicy: pivgo.TouchPolicyNever,
+	})
+	require.NoError(t, err)
+
+	// Providing the wrong PIN should fail without caching it.
+	promptReader.AddString(wrongPIN)
+	_, err = keys.NewHardwarePrivateKey(ctx, s, hardwarekey.PrivateKeyConfig{
+		Policy:      hardwarekey.PromptPolicyPIN,
+		PINCacheTTL: time.Second,
+		CustomSlot:  hardwarekey.PIVSlotKeyString(pivSlot.String()),
+	})
+	require.Error(t, err)
+
+	// Retrieve the key with the right PIN and cache it.
+	promptReader.AddString(testPIN)
+	priv, err := keys.NewHardwarePrivateKey(ctx, s, hardwarekey.PrivateKeyConfig{
+		Policy:      hardwarekey.PromptPolicyPIN,
+		PINCacheTTL: time.Second,
+		CustomSlot:  hardwarekey.PIVSlotKeyString(pivSlot.String()),
+	})
+	require.NoError(t, err)
+
+	// The PIN is cached, no prompt needed.
+	err = priv.WarmupHardwareKey(ctx)
+	require.NoError(t, err)
+
+	// Wait for the cache to expire.
+	time.Sleep(time.Second)
+
+	// Signing should fail with the wrong PIN without caching it.
+	promptReader.AddString(wrongPIN)
+	err = priv.WarmupHardwareKey(ctx)
+	require.Error(t, err)
+
+	// Signing with the right PIN should cache it.
+	promptReader.AddString(testPIN)
+	err = priv.WarmupHardwareKey(ctx)
+	require.Error(t, err)
+
+	// The PIN is cached, no prompt needed.
+	err = priv.WarmupHardwareKey(ctx)
+	require.Error(t, err)
+}
+
 // resetYubikey connects to the first yubiKey and resets it to defaults.
 func resetYubikey(t *testing.T, y *piv.YubiKey) {
 	t.Helper()

api/utils/keys/piv/yubikey.go

@@ -238,8 +238,8 @@ func (y *YubiKey) sign(ctx context.Context, ref *hardwarekey.PrivateKeyRef, keyI
 				defer touchPromptDelayTimer.Reset(signTouchPromptDelay)
 			}
 		}
-		pin, err := y.pinCache.PromptOrGetPIN(ctx, prompt, hardwarekey.PINRequired, keyInfo, ref.PINCacheTTL)
-		return pin, trace.Wrap(err)
+
+		return y.promptPIN(ctx, prompt, hardwarekey.PINRequired, keyInfo, ref.PINCacheTTL)
 	}
 
 	pinPolicy := piv.PINPolicyNever
@@ -371,7 +371,7 @@ func (y *YubiKey) generatePrivateKey(slot piv.Slot, policy hardwarekey.PromptPol
 		TouchPolicy: touchPolicy,
 	}
 
-	if _, err := y.conn.generateKey(piv.DefaultManagementKey, slot, opts); err != nil {
+	if err := y.GenerateKey(slot, opts); err != nil {
 		return nil, trace.Wrap(err)
 	}
 
@@ -385,6 +385,12 @@ func (y *YubiKey) generatePrivateKey(slot piv.Slot, policy hardwarekey.PromptPol
 	return y.getKeyRef(slot, pinCacheTTL)
 }
 
+// GenerateKey generates a new private key in the given PIV slot.
+func (y *YubiKey) GenerateKey(slot piv.Slot, opts piv.Key) error {
+	_, err := y.conn.generateKey(piv.DefaultManagementKey, slot, opts)
+	return trace.Wrap(err)
+}
+
 // SetMetadataCertificate creates a self signed certificate and stores it in the YubiKey's
 // PIV certificate slot. This certificate is purely used as metadata to determine when a
 // slot is in used by a Teleport Client and is not fit to be used in cryptographic operations.
@@ -468,27 +474,62 @@ func (y *YubiKey) SetPIN(oldPin, newPin string) error {
 // If the user provides the default PIN, they will be prompted to set a
 // non-default PIN and PUK before continuing.
 func (y *YubiKey) checkOrSetPIN(ctx context.Context, prompt hardwarekey.Prompt, keyInfo hardwarekey.ContextualKeyInfo, pinCacheTTL time.Duration) error {
-	pin, err := y.pinCache.PromptOrGetPIN(ctx, prompt, hardwarekey.PINOptional, keyInfo, pinCacheTTL)
+	pin, err := y.promptPIN(ctx, prompt, hardwarekey.PINOptional, keyInfo, pinCacheTTL)
 	if err != nil {
 		return trace.Wrap(err)
 	}
 
-	switch pin {
-	case piv.DefaultPIN:
+	if pin == piv.DefaultPIN {
 		fmt.Fprintf(os.Stderr, "The default PIN %q is not supported.\n", piv.DefaultPIN)
-		fallthrough
-	case "":
-		pin, err = y.setPINAndPUKFromDefault(ctx, prompt, keyInfo)
+
+		pin, err = y.setPINAndPUKFromDefault(ctx, prompt, keyInfo, pinCacheTTL)
 		if err != nil {
 			return trace.Wrap(err)
 		}
-		y.pinCache.setPIN(pin, pinCacheTTL)
 	}
 
-	return trace.Wrap(y.verifyPIN(pin))
+	return nil
 }
 
-func (y *YubiKey) setPINAndPUKFromDefault(ctx context.Context, prompt hardwarekey.Prompt, keyInfo hardwarekey.ContextualKeyInfo) (string, error) {
+// PIN (or PUK) prompts time out after 1 minute to prevent an indefinite hold of
+// the pin cache mutex or the exclusive PC/SC transaction.
+const pinPromptTimeout = time.Minute
+
+func (y *YubiKey) promptPIN(ctx context.Context, prompt hardwarekey.Prompt, requirement hardwarekey.PINPromptRequirement, keyInfo hardwarekey.ContextualKeyInfo, pinCacheTTL time.Duration) (string, error) {
+	y.pinCache.mu.Lock()
+	defer y.pinCache.mu.Unlock()
+
+	pin := y.pinCache.getPIN(pinCacheTTL)
+	if pin != "" {
+		return pin, nil
+	}
+
+	ctx, cancel := context.WithTimeout(ctx, pinPromptTimeout)
+	defer cancel()
+
+	pin, err := prompt.AskPIN(ctx, requirement, keyInfo)
+	if err != nil {
+		return "", trace.Wrap(err)
+	}
+
+	// Verify that the PIN is correct before we cache it. This also caches it internally in the PC/SC transaction.
+	// TODO(Joerger): In the signature pin prompt logic, we unfortunately repeat this verification
+	// due to the way the upstream piv-go library handles PIN prompts.
+	if err := y.verifyPIN(pin); err != nil {
+		return "", trace.Wrap(err)
+	}
+
+	y.pinCache.setPIN(pin, pinCacheTTL)
+	return pin, nil
+}
+
+func (y *YubiKey) setPINAndPUKFromDefault(ctx context.Context, prompt hardwarekey.Prompt, keyInfo hardwarekey.ContextualKeyInfo, pinCacheTTL time.Duration) (string, error) {
+	y.pinCache.mu.Lock()
+	defer y.pinCache.mu.Unlock()
+
+	ctx, cancel := context.WithTimeout(ctx, pinPromptTimeout)
+	defer cancel()
+
 	pinAndPUK, err := prompt.ChangePIN(ctx, keyInfo)
 	if err != nil {
 		return "", trace.Wrap(err)
@@ -504,10 +545,12 @@ func (y *YubiKey) setPINAndPUKFromDefault(ctx context.Context, prompt hardwareke
 		}
 	}
 
+	// unblock caches the new PIN the same way verify does.
 	if err := y.conn.unblock(pinAndPUK.PUK, pinAndPUK.PIN); err != nil {
 		return "", trace.Wrap(err)
 	}
 
+	y.pinCache.setPIN(pinAndPUK.PIN, pinCacheTTL)
 	return pinAndPUK.PIN, nil
 }
 

lib/teleterm/daemon/hardwarekeyprompt.go

@@ -95,7 +95,13 @@ func (h *hardwareKeyPrompter) AskPIN(ctx context.Context, requirement hardwareke
 	if err != nil {
 		return "", trace.Wrap(err)
 	}
-	return res.Pin, nil
+
+	pin := res.Pin
+	if pin == "" {
+		pin = hardwarekey.DefaultPIN
+	}
+
+	return pin, nil
 }
 
 // ChangePIN asks for a new PIN.