From be89459524a4d74edd7e2375ae98d00095bd55f5 Mon Sep 17 00:00:00 2001
From: Gary Belvin <gdbelvin@gmail.com>
Date: Thu, 17 Jun 2021 09:06:35 -0400
Subject: [PATCH] Set key export bit

---
 cmd/step-pkcs11-init/main.go |  4 ++++
 go.mod                       |  1 +
 kms/apiv1/requests.go        |  4 ++++
 kms/pkcs11/pkcs11.go         | 24 ++++++++++++++++++++----
 4 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/cmd/step-pkcs11-init/main.go b/cmd/step-pkcs11-init/main.go
index 34f9f8f8..b190c261 100644
--- a/cmd/step-pkcs11-init/main.go
+++ b/cmd/step-pkcs11-init/main.go
@@ -50,6 +50,7 @@ type Config struct {
 	NoCerts          bool
 	EnableSSH        bool
 	Force            bool
+	Extractable      bool
 }
 
 // Validate checks the flags in the config.
@@ -117,6 +118,7 @@ func main() {
 	flag.BoolVar(&c.EnableSSH, "ssh", false, "Enable the creation of ssh keys.")
 	flag.BoolVar(&c.NoCerts, "no-certs", false, "Do not store certificates in the module.")
 	flag.BoolVar(&c.Force, "force", false, "Force the delete of previous keys.")
+	flag.BoolVar(&c.Extractable, "extractable", false, "Allow export of private keys under wrap.")
 	flag.Usage = usage
 	flag.Parse()
 
@@ -286,6 +288,7 @@ func createPKI(k kms.KeyManager, c Config) error {
 		resp, err := k.CreateKey(&apiv1.CreateKeyRequest{
 			Name:               c.RootKeyObject,
 			SignatureAlgorithm: apiv1.ECDSAWithSHA256,
+			Extractable:        c.Extractable,
 		})
 		if err != nil {
 			return err
@@ -366,6 +369,7 @@ func createPKI(k kms.KeyManager, c Config) error {
 		resp, err := k.CreateKey(&apiv1.CreateKeyRequest{
 			Name:               c.CrtKeyObject,
 			SignatureAlgorithm: apiv1.ECDSAWithSHA256,
+			Extractable:        c.Extractable,
 		})
 		if err != nil {
 			return err
diff --git a/go.mod b/go.mod
index ac3d54fb..74aa56e9 100644
--- a/go.mod
+++ b/go.mod
@@ -15,6 +15,7 @@ require (
 	github.com/googleapis/gax-go/v2 v2.0.5
 	github.com/konsorten/go-windows-terminal-sequences v1.0.2 // indirect
 	github.com/micromdm/scep/v2 v2.0.0
+	github.com/miekg/pkcs11 v1.0.3-0.20190429190417-a667d056470f
 	github.com/newrelic/go-agent v2.15.0+incompatible
 	github.com/pkg/errors v0.9.1
 	github.com/rs/xid v1.2.1
diff --git a/kms/apiv1/requests.go b/kms/apiv1/requests.go
index f6fe7dd2..321b308b 100644
--- a/kms/apiv1/requests.go
+++ b/kms/apiv1/requests.go
@@ -112,6 +112,10 @@ type CreateKeyRequest struct {
 	// ProtectionLevel specifies how cryptographic operations are performed.
 	// Used by: cloudkms
 	ProtectionLevel ProtectionLevel
+
+	// Whether the key may be exported from the HSM under a wrap key.
+	// Sets the CKA_EXTRACTABLE bit.
+	Extractable bool
 }
 
 // CreateKeyResponse is the response value of the kms.CreateKey method.
diff --git a/kms/pkcs11/pkcs11.go b/kms/pkcs11/pkcs11.go
index 47c298a5..f6a86a0d 100644
--- a/kms/pkcs11/pkcs11.go
+++ b/kms/pkcs11/pkcs11.go
@@ -14,6 +14,7 @@ import (
 	"sync"
 
 	"github.com/ThalesIgnite/crypto11"
+	"github.com/miekg/pkcs11"
 	"github.com/pkg/errors"
 	"github.com/smallstep/certificates/kms/apiv1"
 	"github.com/smallstep/certificates/kms/uri"
@@ -35,6 +36,7 @@ type P11 interface {
 	DeleteCertificate(id, label []byte, serial *big.Int) error
 	GenerateRSAKeyPairWithLabel(id, label []byte, bits int) (crypto11.SignerDecrypter, error)
 	GenerateECDSAKeyPairWithLabel(id, label []byte, curve elliptic.Curve) (crypto11.Signer, error)
+	GenerateECDSAKeyPairWithAttributes(public, private crypto11.AttributeSet, curve elliptic.Curve) (crypto11.Signer, error)
 	Close() error
 }
 
@@ -291,17 +293,17 @@ func generateKey(ctx P11, req *apiv1.CreateKeyRequest) (crypto11.Signer, error)
 
 	switch req.SignatureAlgorithm {
 	case apiv1.UnspecifiedSignAlgorithm:
-		return ctx.GenerateECDSAKeyPairWithLabel(id, object, elliptic.P256())
+		return GenerateECDSAKeyPairWithLabel(ctx, id, object, elliptic.P256(), req.Extractable)
 	case apiv1.SHA256WithRSA, apiv1.SHA384WithRSA, apiv1.SHA512WithRSA:
 		return ctx.GenerateRSAKeyPairWithLabel(id, object, bits)
 	case apiv1.SHA256WithRSAPSS, apiv1.SHA384WithRSAPSS, apiv1.SHA512WithRSAPSS:
 		return ctx.GenerateRSAKeyPairWithLabel(id, object, bits)
 	case apiv1.ECDSAWithSHA256:
-		return ctx.GenerateECDSAKeyPairWithLabel(id, object, elliptic.P256())
+		return GenerateECDSAKeyPairWithLabel(ctx, id, object, elliptic.P256(), req.Extractable)
 	case apiv1.ECDSAWithSHA384:
-		return ctx.GenerateECDSAKeyPairWithLabel(id, object, elliptic.P384())
+		return GenerateECDSAKeyPairWithLabel(ctx, id, object, elliptic.P384(), req.Extractable)
 	case apiv1.ECDSAWithSHA512:
-		return ctx.GenerateECDSAKeyPairWithLabel(id, object, elliptic.P521())
+		return GenerateECDSAKeyPairWithLabel(ctx, id, object, elliptic.P521(), req.Extractable)
 	case apiv1.PureEd25519:
 		return nil, fmt.Errorf("signature algorithm %s is not supported", req.SignatureAlgorithm)
 	default:
@@ -309,6 +311,20 @@ func generateKey(ctx P11, req *apiv1.CreateKeyRequest) (crypto11.Signer, error)
 	}
 }
 
+func GenerateECDSAKeyPairWithLabel(ctx P11, id, label []byte, curve elliptic.Curve, extractable bool) (crypto11.Signer, error) {
+	public, err := crypto11.NewAttributeSetWithIDAndLabel(id, label)
+	if err != nil {
+		return nil, err
+	}
+	// Copy the AttributeSet to allow modifications.
+	private := public.Copy()
+	private.AddIfNotPresent([]*pkcs11.Attribute{
+		pkcs11.NewAttribute(pkcs11.CKA_EXTRACTABLE, extractable),
+	})
+
+	return ctx.GenerateECDSAKeyPairWithAttributes(public, private, curve)
+}
+
 func findSigner(ctx P11, rawuri string) (crypto11.Signer, error) {
 	id, object, err := parseObject(rawuri)
 	if err != nil {