Skip to content

feat: Hardware Key Agent - consolidate globally shared PIV service variables #53974

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Joerger merged 3 commits into from
Apr 15, 2025
Merged

feat: Hardware Key Agent - consolidate globally shared PIV service variables #53974

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
84ff27f
Consolidate process-wide shared prompt mutex and yubikey connections …
Joerger Mar 28, 2025
77c414a
Fix interactive piv service tests.
Joerger Apr 14, 2025
0ed5a16
Fix wording on comment.
Joerger Apr 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 40 additions & 25 deletions api/utils/keys/piv/service.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,37 +35,52 @@ import (
"github.com/gravitational/teleport/api/utils/keys/hardwarekey"
)

// TODO(Joerger): Rather than using a global cache and mutexes, clients should be updated
// to create a single YubiKeyService and ensure it is reused across the program execution.
var (
// yubiKeyService is a global YubiKeyService used to share yubikey connections
// and prompt mutex logic across the process in cases where [NewYubiKeyService]
// is called multiple times.
//
// TODO(Joerger): Ensure all clients initialize [NewYubiKeyService] only once so we can
// remove this global variable.
Comment on lines +42 to +43
Copy link
Copy Markdown
Contributor

@rosstimothy rosstimothy Apr 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is the plan for this TODO?

Copy link
Copy Markdown
Contributor Author

@Joerger Joerger Apr 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This PR would address tsh, tctl, and Teleport Connect. For arbitrary API clients or tbot we can do the same type of initialization, or continue to provide a global service through a separate method.

var yubiKeyService *YubiKeyService
var yubiKeyServiceMux sync.Mutex

// YubiKeyService is a YubiKey PIV implementation of [hardwarekey.Service].
type YubiKeyService struct {
prompt hardwarekey.Prompt
promptMux sync.Mutex

// yubiKeys is a shared, thread-safe [YubiKey] cache by serial number. It allows for
// separate goroutines to share a YubiKey connection to work around the single PC/SC
// transaction (connection) per-yubikey limit.
yubiKeys map[uint32]*YubiKey = map[uint32]*YubiKey{}
yubiKeys map[uint32]*YubiKey
yubiKeysMux sync.Mutex

// promptMux is used to prevent over-prompting, especially for back-to-back sign requests
// since touch/PIN from the first signature should be cached for following signatures.
promptMux sync.Mutex
)

// YubiKeyService is a YubiKey PIV implementation of [hardwarekey.Service].
type YubiKeyService struct {
prompt hardwarekey.Prompt
}

// Returns a new [YubiKeyService]. If [customPrompt] is nil, the default CLI prompt will be used.
//
// Only a single service should be created for each process to ensure the cached connections
// are shared and multiple services don't compete for PIV resources.
func NewYubiKeyService(customPrompt hardwarekey.Prompt) *YubiKeyService {
yubiKeyServiceMux.Lock()
defer yubiKeyServiceMux.Unlock()

if yubiKeyService != nil {
// If a prompt is provided, prioritize it over the existing prompt value.
if customPrompt != nil {
yubiKeyService.prompt = customPrompt
}
return yubiKeyService
}
Comment on lines +67 to +73
Copy link
Copy Markdown
Contributor

@espadolini espadolini Apr 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we share the same set of yubikeys but have different prompts? Is that something that we ever need?

Copy link
Copy Markdown
Contributor Author

@Joerger Joerger Apr 14, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's not something we need currently, or should ever need. Right now we are primarily handling cases like:

// pre-action checks in profile/cert store (view keys, expiration,
// active profile, etc.), no prompt expected
piv.NewYubiKeyService(nil) // defaults to CLIPrompt, but no prompt expected (no signatures)

// client and client store initialization in `client.NewClient`. 
// This store is/should be used for all key interactions.
piv.NewYubiKeyService(customPrompt) // overwrite with custom prompt e.g. for Connect

// adhoc `keys.ParsePrivateKey` calls which fall outside of the
// client store's path (no hwKeyService to pass).
piv.NewYubiKeyService(nil) // continue using custom prompt or default cli prompt

Or for non tctl/tsh/connect clients, it's likely they will just call keys.ParsePrivateKey adhoc and we don't have any custom prompt or hwkey service initialization logic preceding it.


if customPrompt == nil {
customPrompt = hardwarekey.NewStdCLIPrompt()
}

return &YubiKeyService{
prompt: customPrompt,
yubiKeyService = &YubiKeyService{
prompt: customPrompt,
yubiKeys: map[uint32]*YubiKey{},
}
return yubiKeyService
}

// NewPrivateKey creates a hardware private key that satisfies the provided [config],
Expand Down Expand Up @@ -170,8 +185,8 @@ func (s *YubiKeyService) Sign(ctx context.Context, ref *hardwarekey.PrivateKeyRe
return nil, trace.Wrap(err)
}

promptMux.Lock()
defer promptMux.Unlock()
s.promptMux.Lock()
defer s.promptMux.Unlock()

return y.sign(ctx, ref, keyInfo, s.prompt, rand, digest, opts)
}
Expand Down Expand Up @@ -227,10 +242,10 @@ func (s *YubiKeyService) GetFullKeyRef(serialNumber uint32, slotKey hardwarekey.
// Get the given YubiKey with the serial number. If the provided serialNumber is "0",
// return the first YubiKey found in the smart card list.
func (s *YubiKeyService) getYubiKey(serialNumber uint32) (*YubiKey, error) {
yubiKeysMux.Lock()
defer yubiKeysMux.Unlock()
s.yubiKeysMux.Lock()
defer s.yubiKeysMux.Unlock()
Comment on lines +245 to +246
Copy link
Copy Markdown
Contributor

@espadolini espadolini Apr 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this mux protecting s.yubiKeys, or is it protecting some global state used by FindYubiKey? Can you run FindYubiKey in parallel in two different *YubiKeyServices? Can you run FindYubiKey in parallel for the same serial number?

Copy link
Copy Markdown
Contributor Author

@Joerger Joerger Apr 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah I see how it could look like that, but no the mutex is just protecting the yubikey map. FindYubiKey is independent.

Copy link
Copy Markdown
Contributor Author

@Joerger Joerger Apr 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FYI this is exercised in TestGetYubiKeyPrivateKey_Interactive, though CI doesn't run it.


if y, ok := yubiKeys[serialNumber]; ok {
if y, ok := s.yubiKeys[serialNumber]; ok {
return y, nil
}

Expand All @@ -239,16 +254,16 @@ func (s *YubiKeyService) getYubiKey(serialNumber uint32) (*YubiKey, error) {
return nil, trace.Wrap(err)
}

yubiKeys[y.serialNumber] = y
s.yubiKeys[y.serialNumber] = y
return y, nil
}

// checkOrSetPIN prompts the user for PIN and verifies it with the YubiKey.
// If the user provides the default PIN, they will be prompted to set a
// non-default PIN and PUK before continuing.
func (s *YubiKeyService) checkOrSetPIN(ctx context.Context, y *YubiKey, keyInfo hardwarekey.ContextualKeyInfo) error {
promptMux.Lock()
defer promptMux.Unlock()
s.promptMux.Lock()
defer s.promptMux.Unlock()

pin, err := s.prompt.AskPIN(ctx, hardwarekey.PINOptional, keyInfo)
if err != nil {
Expand All @@ -270,8 +285,8 @@ func (s *YubiKeyService) checkOrSetPIN(ctx context.Context, y *YubiKey, keyInfo
}

func (s *YubiKeyService) promptOverwriteSlot(ctx context.Context, msg string, keyInfo hardwarekey.ContextualKeyInfo) error {
promptMux.Lock()
defer promptMux.Unlock()
s.promptMux.Lock()
defer s.promptMux.Unlock()

promptQuestion := fmt.Sprintf("%v\nWould you like to overwrite this slot's private key and certificate?", msg)
if confirmed, confirmErr := s.prompt.ConfirmSlotOverwrite(ctx, promptQuestion, keyInfo); confirmErr != nil {
Expand Down
42 changes: 22 additions & 20 deletions api/utils/keys/piv/service_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
package piv_test

import (
"bytes"
"context"
"crypto/x509/pkix"
"fmt"
Expand Down Expand Up @@ -48,7 +49,9 @@ func TestGetYubiKeyPrivateKey_Interactive(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()

s := piv.NewYubiKeyService(hardwarekey.NewStdCLIPrompt())
promptReader := prompt.NewFakeReader()
prompt := hardwarekey.NewCLIPrompt(os.Stderr, promptReader)
s := piv.NewYubiKeyService(prompt)

y, err := piv.FindYubiKey(0)
require.NoError(t, err)
Expand All @@ -64,6 +67,13 @@ func TestGetYubiKeyPrivateKey_Interactive(t *testing.T) {
require.NoError(t, err)
require.NoError(t, priv.WarmupHardwareKey(ctx))

// Set pin and handle expected prompts.
setupPINPrompt := func(t *testing.T) {
const testPIN = "123123"
require.NoError(t, y.SetPIN(pivgo.DefaultPIN, testPIN))
promptReader.AddString(testPIN).AddString(testPIN)
}

for _, policy := range []hardwarekey.PromptPolicy{
hardwarekey.PromptPolicyNone,
hardwarekey.PromptPolicyTouch,
Expand All @@ -73,8 +83,10 @@ func TestGetYubiKeyPrivateKey_Interactive(t *testing.T) {
for _, customSlot := range []bool{true, false} {
t.Run(fmt.Sprintf("policy:%+v", policy), func(t *testing.T) {
t.Run(fmt.Sprintf("custom slot:%v", customSlot), func(t *testing.T) {
resetYubikey(t, y)
setupPINPrompt(t, y)
setupPINPrompt(t)
t.Cleanup(func() {
resetYubikey(t, y)
})

var slot hardwarekey.PIVSlotKeyString = ""
if customSlot {
Expand All @@ -89,7 +101,7 @@ func TestGetYubiKeyPrivateKey_Interactive(t *testing.T) {
require.NoError(t, err)

// test HardwareSigner methods
require.Equal(t, policy, priv.GetPrivateKeyPolicy())
require.Equal(t, policy, priv.GetPrivateKeyPolicy().GetPromptPolicy())
require.NotNil(t, priv.GetAttestationStatement())
require.True(t, priv.IsHardware())

Expand Down Expand Up @@ -122,7 +134,10 @@ func TestOverwritePrompt(t *testing.T) {

ctx := context.Background()

s := piv.NewYubiKeyService(hardwarekey.NewStdCLIPrompt())
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)
Expand All @@ -135,14 +150,14 @@ func TestOverwritePrompt(t *testing.T) {

testOverwritePrompt := func(t *testing.T) {
// Fail to overwrite slot when user denies
prompt.SetStdin(prompt.NewFakeReader().AddString("n"))
promptReader.AddString("n")
_, err := keys.NewHardwarePrivateKey(ctx, s, hardwarekey.PrivateKeyConfig{
Policy: hardwarekey.PromptPolicy{TouchRequired: true},
})
require.True(t, trace.IsCompareFailed(err), "Expected compare failed error but got %v", err)

// Successfully overwrite slot when user accepts
prompt.SetStdin(prompt.NewFakeReader().AddString("y"))
promptReader.AddString("y")
_, err = keys.NewHardwarePrivateKey(ctx, s, hardwarekey.PrivateKeyConfig{
Policy: hardwarekey.PromptPolicy{TouchRequired: true},
})
Expand Down Expand Up @@ -178,16 +193,3 @@ func resetYubikey(t *testing.T, y *piv.YubiKey) {
t.Helper()
require.NoError(t, y.Reset())
}

func setupPINPrompt(t *testing.T, y *piv.YubiKey) {
t.Helper()

// Set pin for tests.
const testPIN = "123123"
require.NoError(t, y.SetPIN(pivgo.DefaultPIN, testPIN))

// Handle PIN prompt.
oldStdin := prompt.Stdin()
t.Cleanup(func() { prompt.SetStdin(oldStdin) })
prompt.SetStdin(prompt.NewFakeReader().AddString(testPIN).AddString(testPIN))
}
Loading