From 927a3b3a86cfa87909b06ea7f13d597551452fa4 Mon Sep 17 00:00:00 2001
From: Mariano Cano <mariano@smallstep.com>
Date: Wed, 15 Jan 2020 17:27:21 -0800
Subject: [PATCH] Return crypto.PublicKey on kms.GetPublicKey.

---
 kms/apiv1/requests.go    |  5 -----
 kms/cloudkms/cloudkms.go |  7 ++-----
 kms/kms.go               |  2 +-
 kms/softkms/softkms.go   | 24 ++++++++++++++++--------
 4 files changed, 19 insertions(+), 19 deletions(-)

diff --git a/kms/apiv1/requests.go b/kms/apiv1/requests.go
index ddcbb108..40187d4f 100644
--- a/kms/apiv1/requests.go
+++ b/kms/apiv1/requests.go
@@ -94,11 +94,6 @@ type GetPublicKeyRequest struct {
 	Name string
 }
 
-type GetPublicKeyResponse struct {
-	Name      string
-	PublicKey crypto.PublicKey
-}
-
 type CreateKeyRequest struct {
 	Name               string
 	SignatureAlgorithm SignatureAlgorithm
diff --git a/kms/cloudkms/cloudkms.go b/kms/cloudkms/cloudkms.go
index f734a470..78a734f0 100644
--- a/kms/cloudkms/cloudkms.go
+++ b/kms/cloudkms/cloudkms.go
@@ -223,7 +223,7 @@ func (k *CloudKMS) createKeyRingIfNeeded(name string) error {
 // GetPublicKey gets from Google's Cloud KMS a public key by name. Key names
 // follow the pattern:
 //   projects/([^/]+)/locations/([a-zA-Z0-9_-]{1,63})/keyRings/([a-zA-Z0-9_-]{1,63})/cryptoKeys/([a-zA-Z0-9_-]{1,63})/cryptoKeyVersions/([a-zA-Z0-9_-]{1,63})
-func (k *CloudKMS) GetPublicKey(req *apiv1.GetPublicKeyRequest) (*apiv1.GetPublicKeyResponse, error) {
+func (k *CloudKMS) GetPublicKey(req *apiv1.GetPublicKeyRequest) (crypto.PublicKey, error) {
 	ctx, cancel := defaultContext()
 	defer cancel()
 
@@ -239,10 +239,7 @@ func (k *CloudKMS) GetPublicKey(req *apiv1.GetPublicKeyRequest) (*apiv1.GetPubli
 		return nil, err
 	}
 
-	return &apiv1.GetPublicKeyResponse{
-		Name:      req.Name,
-		PublicKey: pk,
-	}, nil
+	return pk, nil
 }
 
 func defaultContext() (context.Context, context.CancelFunc) {
diff --git a/kms/kms.go b/kms/kms.go
index 524afae7..209783e5 100644
--- a/kms/kms.go
+++ b/kms/kms.go
@@ -13,7 +13,7 @@ import (
 
 // KeyManager is the interface implemented by all the KMS.
 type KeyManager interface {
-	GetPublicKey(req *apiv1.GetPublicKeyRequest) (*apiv1.GetPublicKeyResponse, error)
+	GetPublicKey(req *apiv1.GetPublicKeyRequest) (crypto.PublicKey, error)
 	CreateKey(req *apiv1.CreateKeyRequest) (*apiv1.CreateKeyResponse, error)
 	CreateSigner(req *apiv1.CreateSignerRequest) (crypto.Signer, error)
 	Close() error
diff --git a/kms/softkms/softkms.go b/kms/softkms/softkms.go
index ec79daeb..b5ec6468 100644
--- a/kms/softkms/softkms.go
+++ b/kms/softkms/softkms.go
@@ -19,6 +19,9 @@ type algorithmAttributes struct {
 	Curve string
 }
 
+// DefaultRSAKeySize is the default size for RSA keys.
+const DefaultRSAKeySize = 3072
+
 var signatureAlgorithmMapping = map[apiv1.SignatureAlgorithm]algorithmAttributes{
 	apiv1.UnspecifiedSignAlgorithm: algorithmAttributes{"EC", "P-256"},
 	apiv1.SHA256WithRSA:            algorithmAttributes{"RSA", ""},
@@ -33,6 +36,14 @@ var signatureAlgorithmMapping = map[apiv1.SignatureAlgorithm]algorithmAttributes
 	apiv1.PureEd25519:              algorithmAttributes{"OKP", "Ed25519"},
 }
 
+// generateKey is used for testing purposes.
+var generateKey = func(kty, crv string, size int) (interface{}, interface{}, error) {
+	if kty == "RSA" && size == 0 {
+		size = DefaultRSAKeySize
+	}
+	return keys.GenerateKeyPair(kty, crv, size)
+}
+
 // SoftKSM is a key manager that uses keys stored in disk.
 type SoftKMS struct{}
 
@@ -87,7 +98,7 @@ func (k *SoftKMS) CreateKey(req *apiv1.CreateKeyRequest) (*apiv1.CreateKeyRespon
 		return nil, errors.Errorf("softKMS does not support signature algorithm '%s'", req.SignatureAlgorithm)
 	}
 
-	pub, priv, err := keys.GenerateKeyPair(v.Type, v.Curve, req.Bits)
+	pub, priv, err := generateKey(v.Type, v.Curve, req.Bits)
 	if err != nil {
 		return nil, err
 	}
@@ -106,21 +117,18 @@ func (k *SoftKMS) CreateKey(req *apiv1.CreateKeyRequest) (*apiv1.CreateKeyRespon
 	}, nil
 }
 
-func (k *SoftKMS) GetPublicKey(req *apiv1.GetPublicKeyRequest) (*apiv1.GetPublicKeyResponse, error) {
+func (k *SoftKMS) GetPublicKey(req *apiv1.GetPublicKeyRequest) (crypto.PublicKey, error) {
 	v, err := pemutil.Read(req.Name)
 	if err != nil {
 		return nil, err
 	}
 
-	switch v.(type) {
+	switch vv := v.(type) {
 	case *x509.Certificate:
+		return vv.PublicKey, nil
 	case *rsa.PublicKey, *ecdsa.PublicKey, ed25519.PublicKey:
+		return vv, nil
 	default:
 		return nil, errors.Errorf("unsupported public key type %T", v)
 	}
-
-	return &apiv1.GetPublicKeyResponse{
-		Name:      req.Name,
-		PublicKey: v,
-	}, nil
 }