From a947779795c83ab0e2902994d96d259a7b3528a5 Mon Sep 17 00:00:00 2001
From: Mariano Cano <mariano.cano@gmail.com>
Date: Tue, 16 Feb 2021 13:11:47 -0800
Subject: [PATCH] Add uri support initializing cloudkms.

---
 kms/cloudkms/cloudkms.go      | 22 +++++++++++++-
 kms/cloudkms/cloudkms_test.go | 57 +++++++++++++++++++++++++++++------
 2 files changed, 68 insertions(+), 11 deletions(-)

diff --git a/kms/cloudkms/cloudkms.go b/kms/cloudkms/cloudkms.go
index 547bfc62..98fe375c 100644
--- a/kms/cloudkms/cloudkms.go
+++ b/kms/cloudkms/cloudkms.go
@@ -14,11 +14,15 @@ import (
 	gax "github.com/googleapis/gax-go/v2"
 	"github.com/pkg/errors"
 	"github.com/smallstep/certificates/kms/apiv1"
+	"github.com/smallstep/certificates/kms/uri"
 	"go.step.sm/crypto/pemutil"
 	"google.golang.org/api/option"
 	kmspb "google.golang.org/genproto/googleapis/cloud/kms/v1"
 )
 
+// Scheme is the scheme used in uris.
+const Scheme = "cloudkms"
+
 const pendingGenerationRetries = 10
 
 // protectionLevelMapping maps step protection levels with cloud kms ones.
@@ -71,6 +75,10 @@ type KeyManagementClient interface {
 	CreateCryptoKeyVersion(ctx context.Context, req *kmspb.CreateCryptoKeyVersionRequest, opts ...gax.CallOption) (*kmspb.CryptoKeyVersion, error)
 }
 
+var newKeyManagementClient = func(ctx context.Context, opts ...option.ClientOption) (KeyManagementClient, error) {
+	return cloudkms.NewKeyManagementClient(ctx, opts...)
+}
+
 // CloudKMS implements a KMS using Google's Cloud apiv1.
 type CloudKMS struct {
 	client KeyManagementClient
@@ -79,11 +87,23 @@ type CloudKMS struct {
 // New creates a new CloudKMS configured with a new client.
 func New(ctx context.Context, opts apiv1.Options) (*CloudKMS, error) {
 	var cloudOpts []option.ClientOption
+
+	if opts.URI != "" {
+		u, err := uri.ParseWithScheme(Scheme, opts.URI)
+		if err != nil {
+			return nil, err
+		}
+		if f := u.Get("credentials-file"); f != "" {
+			cloudOpts = append(cloudOpts, option.WithCredentialsFile(f))
+		}
+	}
+
+	// Deprecated way to setting configuration parameters.
 	if opts.CredentialsFile != "" {
 		cloudOpts = append(cloudOpts, option.WithCredentialsFile(opts.CredentialsFile))
 	}
 
-	client, err := cloudkms.NewKeyManagementClient(ctx, cloudOpts...)
+	client, err := newKeyManagementClient(ctx, cloudOpts...)
 	if err != nil {
 		return nil, err
 	}
diff --git a/kms/cloudkms/cloudkms_test.go b/kms/cloudkms/cloudkms_test.go
index 1038432a..e04e0198 100644
--- a/kms/cloudkms/cloudkms_test.go
+++ b/kms/cloudkms/cloudkms_test.go
@@ -5,13 +5,13 @@ import (
 	"crypto"
 	"fmt"
 	"io/ioutil"
-	"os"
 	"reflect"
 	"testing"
 
 	gax "github.com/googleapis/gax-go/v2"
 	"github.com/smallstep/certificates/kms/apiv1"
 	"go.step.sm/crypto/pemutil"
+	"google.golang.org/api/option"
 	kmspb "google.golang.org/genproto/googleapis/cloud/kms/v1"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -50,26 +50,63 @@ func TestParent(t *testing.T) {
 }
 
 func TestNew(t *testing.T) {
+	tmp := newKeyManagementClient
+	t.Cleanup(func() {
+		newKeyManagementClient = tmp
+	})
+	newKeyManagementClient = func(ctx context.Context, opts ...option.ClientOption) (KeyManagementClient, error) {
+		if len(opts) > 0 {
+			return nil, fmt.Errorf("test error")
+		}
+		return &MockClient{}, nil
+	}
+
 	type args struct {
 		ctx  context.Context
 		opts apiv1.Options
 	}
 	tests := []struct {
-		name     string
-		skipOnCI bool
-		args     args
-		want     *CloudKMS
-		wantErr  bool
+		name    string
+		args    args
+		want    *CloudKMS
+		wantErr bool
 	}{
-		{"fail authentication", true, args{context.Background(), apiv1.Options{}}, nil, true},
-		{"fail credentials", false, args{context.Background(), apiv1.Options{CredentialsFile: "testdata/missing"}}, nil, true},
+		{"ok", args{context.Background(), apiv1.Options{}}, &CloudKMS{client: &MockClient{}}, false},
+		{"ok with uri", args{context.Background(), apiv1.Options{URI: "cloudkms:"}}, &CloudKMS{client: &MockClient{}}, false},
+		{"fail credentials", args{context.Background(), apiv1.Options{CredentialsFile: "testdata/missing"}}, nil, true},
+		{"fail with uri", args{context.Background(), apiv1.Options{URI: "cloudkms:credentials-file=testdata/missing"}}, nil, true},
+		{"fail schema", args{context.Background(), apiv1.Options{URI: "pkcs11:"}}, nil, true},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			if tt.skipOnCI && os.Getenv("CI") == "true" {
-				t.SkipNow()
+			got, err := New(tt.args.ctx, tt.args.opts)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("New() error = %v, wantErr %v", err, tt.wantErr)
+				return
 			}
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("New() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
 
+func TestNew_real(t *testing.T) {
+	type args struct {
+		ctx  context.Context
+		opts apiv1.Options
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    *CloudKMS
+		wantErr bool
+	}{
+		{"fail credentials", args{context.Background(), apiv1.Options{CredentialsFile: "testdata/missing"}}, nil, true},
+		{"fail with uri", args{context.Background(), apiv1.Options{URI: "cloudkms:credentials-file=testdata/missing"}}, nil, true},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			got, err := New(tt.args.ctx, tt.args.opts)
 			if (err != nil) != tt.wantErr {
 				t.Errorf("New() error = %v, wantErr %v", err, tt.wantErr)