fix(windows): associate machine-scoped TPM keys with stored certificates

kms/capi/capi.go

@@ -54,6 +54,7 @@ const (
 	IssuerNameArg                = "issuer"
 	KeySpec                      = "key-spec"                  // 0, 1, 2; none/NONE, at_keyexchange/AT_KEYEXCHANGE, at_signature/AT_SIGNATURE
 	SkipFindCertificateKey       = "skip-find-certificate-key" // skips looking up certificate private key when storing a certificate
+	MachineKeysetArg             = "key-machine-keyset"        // when storing a certificate, associate it with a key in the local machine keyset
 )
 
 const (
@@ -84,6 +85,7 @@ var signatureAlgorithmMapping = map[apiv1.SignatureAlgorithm]string{
 
 type uriAttributes struct {
 	containerName             string
+	providerName              string
 	hash                      []byte
 	storeLocation             string
 	storeName                 string
@@ -97,6 +99,7 @@ type uriAttributes struct {
 	description               string
 	keySpec                   string
 	skipFindCertificateKey    bool
+	machineKeyset             bool
 	pin                       string
 }
 
@@ -127,6 +130,7 @@ func parseURI(rawuri string) (*uriAttributes, error) {
 
 	return &uriAttributes{
 		containerName:             u.Get(ContainerNameArg),
+		providerName:              u.Get(ProviderNameArg),
 		hash:                      hashValue,
 		storeLocation:             cmp.Or(u.Get(StoreLocationArg), UserStoreLocation),
 		storeName:                 cmp.Or(u.Get(StoreNameArg), MyStore),
@@ -140,6 +144,7 @@ func parseURI(rawuri string) (*uriAttributes, error) {
 		description:               u.Get(DescriptionArg),
 		keySpec:                   u.Get(KeySpec),
 		skipFindCertificateKey:    u.GetBool(SkipFindCertificateKey),
+		machineKeyset:             u.GetBool(MachineKeysetArg),
 		pin:                       u.Pin(),
 	}, nil
 }
@@ -914,13 +919,35 @@ func (k *CAPIKMS) StoreCertificate(req *apiv1.StoreCertificateRequest) error {
 	}
 	defer windows.CertFreeCertificateContext(certContext)
 
-	// looking up the certificate private key is performed by default, but is made optional,
-	// so that looking up the private key for e.g. intermediate certificates can be skipped.
-	// If not skipped, looking up a private key can prompt the user to insert/select a smart
-	// card, which is usually not what we want to happen.
-	if !u.skipFindCertificateKey {
+	// Associate the certificate with its private key.
+	switch {
+	case u.containerName != "":
+		// The exact key is known (the caller supplied its container name, and
+		// usually its provider). Associate it explicitly rather than by
+		// discovery. This is required for machine-scoped Microsoft Platform
+		// Crypto Provider (TPM) keys: CryptFindCertificateKeyProvInfo does not
+		// search the local machine keyset and so cannot find them. Explicit
+		// association does not enumerate containers and never prompts for a
+		// smart card, so it runs even when skip-find-certificate-key is set.
+		var flags uint32
+		if u.machineKeyset {
+			flags |= CRYPT_MACHINE_KEYSET
+		}
+		if err := setCertificateKeyProvInfo(certContext, u.containerName, u.providerName, flags, ncryptKeySpec); err != nil {
+			return fmt.Errorf("failed associating certificate with key %q: %w", u.containerName, err)
+		}
+	case !u.skipFindCertificateKey:
+		// No specific key was named, so fall back to discovery. Looking up the
+		// private key can prompt the user to insert/select a smart card, which
+		// is why it is skipped for e.g. intermediate certificates. When the
+		// certificate is stored in the machine location its key may live in the
+		// local machine keyset, so search both keysets.
+		var keysetFlags uint32
+		if u.storeLocation == MachineStoreLocation {
+			keysetFlags = CRYPT_FIND_MACHINE_KEYSET_FLAG | CRYPT_FIND_USER_KEYSET_FLAG
+		}
 		// TODO: not finding the associated private key is not a dealbreaker, but maybe a warning should be issued
-		cryptFindCertificateKeyProvInfo(certContext)
+		cryptFindCertificateKeyProvInfo(certContext, keysetFlags)
 	}
 
 	if u.friendlyName != "" {
@@ -969,6 +996,12 @@ func (k *CAPIKMS) StoreCertificateChain(req *apiv1.StoreCertificateChainRequest)
 			FriendlyNameArg:        []string{u.friendlyName},
 			DescriptionArg:         []string{u.description},
 			SkipFindCertificateKey: []string{strconv.FormatBool(u.skipFindCertificateKey)},
+			// Forward the key association hints so the leaf certificate is
+			// linked to its private key (see StoreCertificate). Intermediates
+			// below have no associated key and intentionally omit these.
+			ContainerNameArg: []string{u.containerName},
+			ProviderNameArg:  []string{u.providerName},
+			MachineKeysetArg: []string{strconv.FormatBool(u.machineKeyset)},
 		}).String(),
 		Certificate: leaf,
 	}); err != nil {

kms/capi/ncrypt_windows.go

@@ -10,6 +10,7 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
+	"runtime"
 	"unsafe"
 
 	"golang.org/x/sys/windows"
@@ -81,6 +82,20 @@ const (
 	CRYPT_ACQUIRE_PREFER_NCRYPT_KEY_FLAG = uint32(0x00020000)
 	CRYPT_ACQUIRE_ONLY_NCRYPT_KEY_FLAG   = uint32(0x00040000)
 
+	// Keyset selection flags for CryptFindCertificateKeyProvInfo. When neither
+	// is set the function only searches the current user's key containers, so
+	// a private key stored in the local machine keyset (e.g. a CNG/PCP key
+	// created with NCRYPT_MACHINE_KEY_FLAG) is not found. Setting both searches
+	// the user and machine containers.
+	CRYPT_FIND_USER_KEYSET_FLAG    = uint32(0x00000001)
+	CRYPT_FIND_MACHINE_KEYSET_FLAG = uint32(0x00000002)
+
+	// CRYPT_MACHINE_KEYSET marks a CRYPT_KEY_PROV_INFO as referencing a key in
+	// the local machine keyset rather than the current user's. It must be set
+	// in CRYPT_KEY_PROV_INFO.dwFlags when associating a certificate with a
+	// machine-scoped key (e.g. a TPM key created with NCRYPT_MACHINE_KEY_FLAG).
+	CRYPT_MACHINE_KEYSET = uint32(0x00000020)
+
 	CERT_ID_ISSUER_SERIAL_NUMBER = uint32(1)
 	CERT_ID_KEY_IDENTIFIER       = uint32(2)
 	CERT_ID_SHA1_HASH            = uint32(3)
@@ -194,14 +209,18 @@ type CERT_ID_SERIAL struct {
 	Serial   CERT_ISSUER_SERIAL_NUMBER
 }
 
+// CRYPT_KEY_PROV_INFO mirrors the wincrypt.h structure of the same name. The
+// container/provider names are LPWSTR pointers and the remaining members are
+// DWORDs, so they must be typed as such for the structure to be laid out
+// correctly when passed to CertSetCertificateContextProperty.
 type CRYPT_KEY_PROV_INFO struct {
-	pwszContainerName int
-	pwszProvName      int
-	dwProvType        int
-	dwFlags           int
-	cProvParam        int
-	rgProvParam       int
-	dwKeySpec         int
+	pwszContainerName *uint16
+	pwszProvName      *uint16
+	dwProvType        uint32
+	dwFlags           uint32
+	cProvParam        uint32
+	rgProvParam       uintptr
+	dwKeySpec         uint32
 }
 
 func errNoToStr(e uint32) string {
@@ -569,10 +588,15 @@ func findCertificateInStore(store windows.Handle, enc, findFlags, findType uint3
 	return (*windows.CertContext)(unsafe.Pointer(h)), nil
 }
 
-func cryptFindCertificateKeyProvInfo(certContext *windows.CertContext) error {
+// cryptFindCertificateKeyProvInfo locates the private key matching the
+// certificate and records the association (CERT_KEY_PROV_INFO_PROP_ID) on the
+// certificate context. keysetFlags selects which key containers are searched
+// (CRYPT_FIND_USER_KEYSET_FLAG / CRYPT_FIND_MACHINE_KEYSET_FLAG); when zero the
+// API defaults to the current user's containers only.
+func cryptFindCertificateKeyProvInfo(certContext *windows.CertContext, keysetFlags uint32) error {
 	r, _, err := procCryptFindCertificateKeyProvInfo.Call(
 		uintptr(unsafe.Pointer(certContext)),
-		uintptr(CRYPT_ACQUIRE_PREFER_NCRYPT_KEY_FLAG),
+		uintptr(CRYPT_ACQUIRE_PREFER_NCRYPT_KEY_FLAG|keysetFlags),
 		0,
 	)
 
@@ -587,6 +611,51 @@ func cryptFindCertificateKeyProvInfo(certContext *windows.CertContext) error {
 	return nil
 }
 
+// setCertificateKeyProvInfo explicitly associates the certificate with a named
+// private key by attaching a CERT_KEY_PROV_INFO_PROP_ID property to the
+// certificate context. Unlike cryptFindCertificateKeyProvInfo, it does not
+// enumerate key containers to discover the key, so it works for keys that
+// discovery cannot locate — notably machine-scoped Microsoft Platform Crypto
+// Provider (TPM) keys, which live in the local machine keyset
+// (CRYPT_MACHINE_KEYSET) that CryptFindCertificateKeyProvInfo does not search.
+//
+// containerName is the CNG key (container) name, provName the storage provider
+// (e.g. "Microsoft Platform Crypto Provider"), dwFlags carries keyset flags
+// such as CRYPT_MACHINE_KEYSET, and dwKeySpec is the key spec (CNG keys use
+// CERT_NCRYPT_KEY_SPEC).
+func setCertificateKeyProvInfo(certContext *windows.CertContext, containerName, provName string, dwFlags, dwKeySpec uint32) error {
+	container, err := windows.UTF16PtrFromString(containerName)
+	if err != nil {
+		return fmt.Errorf("invalid key container name %q: %w", containerName, err)
+	}
+	var provider *uint16
+	if provName != "" {
+		if provider, err = windows.UTF16PtrFromString(provName); err != nil {
+			return fmt.Errorf("invalid provider name %q: %w", provName, err)
+		}
+	}
+
+	info := CRYPT_KEY_PROV_INFO{
+		pwszContainerName: container,
+		pwszProvName:      provider,
+		dwProvType:        0, // 0 selects a CNG (NCrypt) storage provider
+		dwFlags:           dwFlags,
+		dwKeySpec:         dwKeySpec,
+	}
+
+	// CertSetCertificateContextProperty copies the structure and the strings it
+	// references into the certificate context, so the Go-allocated memory only
+	// needs to stay alive for the duration of the call.
+	err = certSetCertificateContextProperty(certContext, CERT_KEY_PROV_INFO_PROP_ID, uintptr(unsafe.Pointer(&info)))
+	runtime.KeepAlive(container)
+	runtime.KeepAlive(provider)
+	runtime.KeepAlive(&info)
+	if err != nil {
+		return fmt.Errorf("CertSetCertificateContextProperty(CERT_KEY_PROV_INFO_PROP_ID) failed: %w", err)
+	}
+	return nil
+}
+
 func cryptFindCertificatePrivateKey(certContext *windows.CertContext) (uintptr, error) {
 	var (
 		kh      windows.Handle

kms/tpmkms/tpmkms.go

@@ -20,6 +20,7 @@ import (
 	"os"
 	"path/filepath"
 	"runtime"
+	"strconv"
 	"time"
 
 	"go.step.sm/crypto/kms/apiv1"
@@ -989,16 +990,31 @@ func (k *TPMKMS) storeCertificateChainToWindowsCertificateStore(req *apiv1.Store
 		intermediateCAStore = o.intermediateStore
 	}
 
+	// Associate the stored certificate with the TPM key explicitly. The agent
+	// stores certificates with skip-find-certificate-key set (to avoid a smart
+	// card prompt during discovery), and CryptFindCertificateKeyProvInfo cannot
+	// discover machine-scoped Platform Crypto Provider keys anyway, so we hand
+	// the CAPI layer the exact key: its CNG container name, the TPM provider,
+	// and whether it lives in the local machine keyset. The container name is
+	// the key name prefixed with "app-" (see prefixKey / go-attestation), which
+	// is how the key was persisted in the PCP KSP.
+	v := url.Values{
+		"store-location":              []string{location},
+		"store":                       []string{store},
+		"friendly-name":               []string{o.friendlyName},
+		"description":                 []string{o.description},
+		"skip-find-certificate-key":   []string{skipFindCertificateKey},
+		"intermediate-store-location": []string{intermediateCAStoreLocation},
+		"intermediate-store":          []string{intermediateCAStore},
+	}
+	if o.name != "" {
+		v.Set("key", tpm.ApplicationKeyName(o.name))
+		v.Set("provider", microsoftPCP)
+		v.Set("key-machine-keyset", strconv.FormatBool(o.isMachineKey()))
+	}
+
 	return k.windowsCertificateManager.StoreCertificateChain(&apiv1.StoreCertificateChainRequest{
-		Name: uri.New("capi", url.Values{
-			"store-location":              []string{location},
-			"store":                       []string{store},
-			"friendly-name":               []string{o.friendlyName},
-			"description":                 []string{o.description},
-			"skip-find-certificate-key":   []string{skipFindCertificateKey},
-			"intermediate-store-location": []string{intermediateCAStoreLocation},
-			"intermediate-store":          []string{intermediateCAStore},
-		}).String(),
+		Name:             uri.New("capi", v).String(),
 		CertificateChain: req.CertificateChain,
 	})
 }

tpm/names.go

@@ -36,3 +36,13 @@ func prefixAK(name string) string {
 func prefixKey(name string) string {
 	return fmt.Sprintf("app-%s", name)
 }
+
+// ApplicationKeyName returns the name an application key is persisted under by
+// the underlying provider, given its logical name. It prefixes `app-`, matching
+// go-attestation's default. On Windows this is the CNG container name of the
+// key, which is needed to associate a certificate with it in the certificate
+// store. It is the single source of truth shared with callers (e.g. tpmkms)
+// that must reference the persisted key name.
+func ApplicationKeyName(name string) string {
+	return prefixKey(name)
+}