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
+	KeyScopeArg                  = "key-scope"                 // "machine" or "user"; the keyset that holds the certificate's private key (defaults to store-location)
 )
 
 const (
@@ -83,7 +84,8 @@ var signatureAlgorithmMapping = map[apiv1.SignatureAlgorithm]string{
 }
 
 type uriAttributes struct {
-	containerName             string
+	keyContainerName          string
+	providerName              string
 	hash                      []byte
 	storeLocation             string
 	storeName                 string
@@ -97,9 +99,29 @@ type uriAttributes struct {
 	description               string
 	keySpec                   string
 	skipFindCertificateKey    bool
+	keyScope                  string
 	pin                       string
 }
 
+// isMachineKeyset reports whether the certificate's private key lives in the
+// local machine keyset rather than the current user's. The key scope is taken
+// from the "key-scope" argument; when unset it falls back to the certificate
+// store location ("machine" → machine keyset), since cert and key are usually
+// provisioned in the same scope. This mirrors the tpmkms key-scope resolution.
+func (u *uriAttributes) isMachineKeyset() bool {
+	return u.keyScope == MachineStoreLocation ||
+		(u.keyScope == "" && u.storeLocation == MachineStoreLocation)
+}
+
+// isUserKeyset reports whether the certificate's private key is explicitly
+// scoped to the current user's keyset. Unlike isMachineKeyset it deliberately
+// does not fall back to the store location: when "key-scope" is unset the keyset
+// is left unspecified so key discovery can search both containers, matching
+// CryptFindCertificateKeyProvInfo's historical default.
+func (u *uriAttributes) isUserKeyset() bool {
+	return u.keyScope == UserStoreLocation
+}
+
 func parseURI(rawuri string) (*uriAttributes, error) {
 	u, err := uri.ParseWithScheme(Scheme, rawuri)
 	if err != nil {
@@ -126,7 +148,8 @@ func parseURI(rawuri string) (*uriAttributes, error) {
 	}
 
 	return &uriAttributes{
-		containerName:             u.Get(ContainerNameArg),
+		keyContainerName:          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 +163,7 @@ func parseURI(rawuri string) (*uriAttributes, error) {
 		description:               u.Get(DescriptionArg),
 		keySpec:                   u.Get(KeySpec),
 		skipFindCertificateKey:    u.GetBool(SkipFindCertificateKey),
+		keyScope:                  u.Get(KeyScopeArg),
 		pin:                       u.Pin(),
 	}, nil
 }
@@ -435,10 +459,10 @@ func (k *CAPIKMS) getCertContext(u *uriAttributes) (*windows.CertContext, error)
 				return nil, err
 			}
 		}
-	case u.containerName != "":
+	case u.keyContainerName != "":
 		key, err := k.GetPublicKey(&apiv1.GetPublicKeyRequest{
 			Name: uri.New(Scheme, url.Values{
-				ContainerNameArg: []string{u.containerName},
+				ContainerNameArg: []string{u.keyContainerName},
 			}).String(),
 		})
 		if err != nil {
@@ -532,15 +556,15 @@ func (k *CAPIKMS) CreateSigner(req *apiv1.CreateSignerRequest) (crypto.Signer, e
 		certHandle *windows.CertContext
 	)
 
-	if u.containerName != "" {
+	if u.keyContainerName != "" {
 		keyFlags, err := k.getKeyFlags(u)
 		if err != nil {
 			return nil, err
 		}
 
-		kh, err = nCryptOpenKey(k.providerHandle, u.containerName, 0, keyFlags)
+		kh, err = nCryptOpenKey(k.providerHandle, u.keyContainerName, 0, keyFlags)
 		if err != nil {
-			return nil, fmt.Errorf("unable to open key using %q=%q: %w", ContainerNameArg, u.containerName, err)
+			return nil, fmt.Errorf("unable to open key using %q=%q: %w", ContainerNameArg, u.keyContainerName, err)
 		}
 	} else {
 		// check if a certificate can be located using the URI
@@ -576,7 +600,7 @@ func (k *CAPIKMS) CreateSigner(req *apiv1.CreateSignerRequest) (crypto.Signer, e
 		}
 	}
 
-	return newCAPISigner(kh, u.containerName, u.pin)
+	return newCAPISigner(kh, u.keyContainerName, u.pin)
 }
 
 func setKeySpec(u *uriAttributes) (uint32, error) {
@@ -615,8 +639,8 @@ func (k *CAPIKMS) CreateKey(req *apiv1.CreateKeyRequest) (*apiv1.CreateKeyRespon
 	}
 
 	// generate a random uuid for the container name if it is not present
-	if u.containerName == "" {
-		u.containerName, err = randutil.UUIDv4()
+	if u.keyContainerName == "" {
+		u.keyContainerName, err = randutil.UUIDv4()
 		if err != nil {
 			return nil, fmt.Errorf("failed to generate uuid: %w", err)
 		}
@@ -638,7 +662,7 @@ func (k *CAPIKMS) CreateKey(req *apiv1.CreateKeyRequest) (*apiv1.CreateKeyRespon
 	}
 
 	// TODO: check whether RSA keys require legacyKeySpec set to AT_KEYEXCHANGE
-	kh, err := nCryptCreatePersistedKey(k.providerHandle, u.containerName, alg, keySpec, keyFlags)
+	kh, err := nCryptCreatePersistedKey(k.providerHandle, u.keyContainerName, alg, keySpec, keyFlags)
 	if err != nil {
 		return nil, fmt.Errorf("unable to create persisted key: %w", err)
 	}
@@ -708,7 +732,7 @@ func (k *CAPIKMS) DeleteKey(req *apiv1.DeleteKeyRequest) error {
 		return err
 	}
 
-	if u.containerName == "" {
+	if u.keyContainerName == "" {
 		return fmt.Errorf("%v not specified", ContainerNameArg)
 	}
 
@@ -717,7 +741,7 @@ func (k *CAPIKMS) DeleteKey(req *apiv1.DeleteKeyRequest) error {
 		return err
 	}
 
-	kh, err := nCryptOpenKey(k.providerHandle, u.containerName, 0, keyFlags)
+	kh, err := nCryptOpenKey(k.providerHandle, u.keyContainerName, 0, keyFlags)
 	if err != nil {
 		return fmt.Errorf("unable to open key: %w", err)
 	}
@@ -734,7 +758,7 @@ func (k *CAPIKMS) GetPublicKey(req *apiv1.GetPublicKeyRequest) (crypto.PublicKey
 		return nil, err
 	}
 
-	if u.containerName == "" {
+	if u.keyContainerName == "" {
 		return nil, fmt.Errorf("%v not specified", ContainerNameArg)
 	}
 
@@ -743,7 +767,7 @@ func (k *CAPIKMS) GetPublicKey(req *apiv1.GetPublicKeyRequest) (crypto.PublicKey
 		return nil, err
 	}
 
-	kh, err := nCryptOpenKey(k.providerHandle, u.containerName, 0, keyFlags)
+	kh, err := nCryptOpenKey(k.providerHandle, u.keyContainerName, 0, keyFlags)
 	if err != nil {
 		return nil, fmt.Errorf("unable to open key: %w", err)
 	}
@@ -914,13 +938,40 @@ 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.keyContainerName != "":
+		// 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.isMachineKeyset() {
+			flags |= CRYPT_MACHINE_KEYSET
+		}
+		if err := setCertificateKeyProvInfo(certContext, u.keyContainerName, u.providerName, flags, ncryptKeySpec); err != nil {
+			return fmt.Errorf("failed associating certificate with key %q: %w", u.keyContainerName, 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. With no
+		// keyset flag CryptFindCertificateKeyProvInfo searches both the user and
+		// machine containers (its documented default); restrict the search only
+		// when the key scope is known, so a machine-scoped key is found while
+		// the legacy both-containers behavior is preserved otherwise.
+		var keysetFlags uint32
+		switch {
+		case u.isUserKeyset():
+			keysetFlags = CRYPT_FIND_USER_KEYSET_FLAG
+		case u.isMachineKeyset():
+			keysetFlags = CRYPT_FIND_MACHINE_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 +1020,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.keyContainerName},
+			ProviderNameArg:  []string{u.providerName},
+			KeyScopeArg:      []string{u.keyScope},
 		}).String(),
 		Certificate: leaf,
 	}); err != nil {
@@ -1114,10 +1171,10 @@ func (k *CAPIKMS) DeleteCertificate(req *apiv1.DeleteCertificateRequest) error {
 			}
 			prevCert = certHandle
 		}
-	case u.containerName != "":
+	case u.keyContainerName != "":
 		key, err := k.GetPublicKey(&apiv1.GetPublicKeyRequest{
 			Name: uri.New(Scheme, url.Values{
-				ContainerNameArg: []string{u.containerName},
+				ContainerNameArg: []string{u.keyContainerName},
 			}).String(),
 		})
 		if err != nil {

kms/capi/ncrypt_windows.go

@@ -81,6 +81,20 @@ const (
 	CRYPT_ACQUIRE_PREFER_NCRYPT_KEY_FLAG = uint32(0x00020000)
 	CRYPT_ACQUIRE_ONLY_NCRYPT_KEY_FLAG   = uint32(0x00040000)
 
+	// Keyset selection flags for CryptFindCertificateKeyProvInfo. Each flag
+	// restricts the search to that container; with neither flag the API
+	// searches both. We restrict explicitly so a machine-scoped certificate's
+	// key (in the local machine keyset, e.g. a CNG/PCP key created with
+	// NCRYPT_MACHINE_KEY_FLAG) is matched rather than missed.
+	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 +208,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 +587,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 +610,48 @@ 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, and info (which holds the only
+	// references to container/provider) stays live through the call, so no
+	// runtime.KeepAlive is needed — matching the other cert property helpers.
+	if err := certSetCertificateContextProperty(certContext, CERT_KEY_PROV_INFO_PROP_ID, uintptr(unsafe.Pointer(&info))); 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

@@ -989,16 +989,34 @@ 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 its key scope (machine vs user). 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. CAPI resolves the keyset from key-scope,
+	// falling back to store-location when key-scope is unset.
+	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)
+		if o.keyScope != "" {
+			v.Set("key-scope", o.keyScope)
+		}
+	}
+
 	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,
 	})
 }