Update GCP KMS tests and implementation

Signed-off-by: Daniel Weiße <dw@edgeless.systems>

coordinator/kms/setup.go

@@ -113,7 +113,7 @@ func getKMS(ctx context.Context, kmsURI string, store kms.Storage) (kms.CloudKMS
 		if err != nil {
 			return nil, err
 		}
-		return gcp.New(project, location, keyRing, store, kmspb.ProtectionLevel(protectionLvl)), nil
+		return gcp.New(ctx, project, location, keyRing, store, kmspb.ProtectionLevel(protectionLvl))
 
 	case "cluster-kms":
 		return &ClusterKMS{}, nil

kms/kms/gcp/gcp.go

@@ -7,6 +7,7 @@ import (
 	"crypto/sha1"
 	"errors"
 	"fmt"
+	"io"
 	"strings"
 	"time"
 
@@ -18,23 +19,36 @@ import (
 	"github.com/googleapis/gax-go/v2"
 	"google.golang.org/api/option"
 	kmspb "google.golang.org/genproto/googleapis/cloud/kms/v1"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
+type clientAPI interface {
+	io.Closer
+	CreateCryptoKey(context.Context, *kmspb.CreateCryptoKeyRequest, ...gax.CallOption) (*kmspb.CryptoKey, error)
+	CreateImportJob(context.Context, *kmspb.CreateImportJobRequest, ...gax.CallOption) (*kmspb.ImportJob, error)
+	Decrypt(context.Context, *kmspb.DecryptRequest, ...gax.CallOption) (*kmspb.DecryptResponse, error)
+	Encrypt(context.Context, *kmspb.EncryptRequest, ...gax.CallOption) (*kmspb.EncryptResponse, error)
+	GetKeyRing(context.Context, *kmspb.GetKeyRingRequest, ...gax.CallOption) (*kmspb.KeyRing, error)
+	ImportCryptoKeyVersion(context.Context, *kmspb.ImportCryptoKeyVersionRequest, ...gax.CallOption) (*kmspb.CryptoKeyVersion, error)
+	UpdateCryptoKeyPrimaryVersion(context.Context, *kmspb.UpdateCryptoKeyPrimaryVersionRequest, ...gax.CallOption) (*kmspb.CryptoKey, error)
+	GetImportJob(context.Context, *kmspb.GetImportJobRequest, ...gax.CallOption) (*kmspb.ImportJob, error)
+}
+
 // KMSClient implements the CloudKMS interface for Google Cloud Platform.
 type KMSClient struct {
 	projectID        string
 	locationID       string
 	keyRingID        string
+	newClient        func(ctx context.Context, opts ...option.ClientOption) (clientAPI, error)
 	waitBackoffLimit int
 	storage          kmsInterface.Storage
 	protectionLevel  kmspb.ProtectionLevel
 	opts             []gax.CallOption
-	// Used for testing purposes
-	clientOpts []option.ClientOption
 }
 
 // New initializes a KMS client for Google Cloud Platform.
-func New(projectID, locationID, keyRingID string, store kmsInterface.Storage, protectionLvl kmspb.ProtectionLevel, opts ...gax.CallOption) *KMSClient {
+func New(ctx context.Context, projectID, locationID, keyRingID string, store kmsInterface.Storage, protectionLvl kmspb.ProtectionLevel, opts ...gax.CallOption) (*KMSClient, error) {
 	if store == nil {
 		store = storage.NewMemMapStorage()
 	}
@@ -43,22 +57,30 @@ func New(projectID, locationID, keyRingID string, store kmsInterface.Storage, pr
 		protectionLvl = kmspb.ProtectionLevel_SOFTWARE
 	}
 
-	return &KMSClient{
+	c := &KMSClient{
 		projectID:        projectID,
 		locationID:       locationID,
 		keyRingID:        keyRingID,
+		newClient:        keyManagementClientFactory,
 		waitBackoffLimit: 10,
 		storage:          store,
 		protectionLevel:  protectionLvl,
 		opts:             opts,
 	}
+
+	// test if the KMS can be reached with the given configuration
+	if err := c.testConnection(ctx); err != nil {
+		return nil, fmt.Errorf("testing connection to GCP KMS: %w", err)
+	}
+
+	return c, nil
 }
 
 // CreateKEK creates a new Key Encryption Key using Google Key Management System.
 //
 // If no key material is provided, a new key is generated by Google's KMS, otherwise the key material is used to import the key.
 func (c *KMSClient) CreateKEK(ctx context.Context, keyID string, key []byte) error {
-	client, err := kms.NewKeyManagementClient(ctx, c.clientOpts...)
+	client, err := c.newClient(ctx)
 	if err != nil {
 		return err
 	}
@@ -81,6 +103,12 @@ func (c *KMSClient) CreateKEK(ctx context.Context, keyID string, key []byte) err
 
 // GetDEK fetches an encrypted Data Encryption Key from storage and decrypts it using a KEK stored in Google's KMS.
 func (c *KMSClient) GetDEK(ctx context.Context, kekID, keyID string, dekSize int) ([]byte, error) {
+	client, err := c.newClient(ctx)
+	if err != nil {
+		return nil, err
+	}
+	defer client.Close()
+
 	encryptedDEK, err := c.storage.Get(ctx, keyID)
 	if err != nil {
 		if !errors.Is(err, storage.ErrDEKUnset) {
@@ -92,15 +120,9 @@ func (c *KMSClient) GetDEK(ctx context.Context, kekID, keyID string, dekSize int
 		if err != nil {
 			return nil, fmt.Errorf("could not generate key: %w", err)
 		}
-		return newDEK, c.putDEK(ctx, kekID, keyID, newDEK)
+		return newDEK, c.putDEK(ctx, client, kekID, keyID, newDEK)
 	}
 
-	client, err := kms.NewKeyManagementClient(ctx, c.clientOpts...)
-	if err != nil {
-		return nil, err
-	}
-	defer client.Close()
-
 	request := &kmspb.DecryptRequest{
 		Name:       fmt.Sprintf("projects/%s/locations/%s/keyRings/%s/cryptoKeys/%s", c.projectID, c.locationID, c.keyRingID, kekID),
 		Ciphertext: encryptedDEK,
@@ -108,7 +130,7 @@ func (c *KMSClient) GetDEK(ctx context.Context, kekID, keyID string, dekSize int
 
 	res, err := client.Decrypt(ctx, request, c.opts...)
 	if err != nil {
-		if strings.Contains(err.Error(), "code = NotFound") {
+		if status.Convert(err).Code() == codes.NotFound {
 			return nil, kmsInterface.ErrKEKUnknown
 		}
 		return nil, fmt.Errorf("decrypting DEK: %w", err)
@@ -118,13 +140,7 @@ func (c *KMSClient) GetDEK(ctx context.Context, kekID, keyID string, dekSize int
 }
 
 // putDEK encrypts a Data Encryption Key using a KEK stored in Google's KMS and saves it to storage.
-func (c *KMSClient) putDEK(ctx context.Context, kekID, keyID string, plainDEK []byte) error {
+func (c *KMSClient) putDEK(ctx context.Context, client clientAPI, kekID, keyID string, plainDEK []byte) error {
-	client, err := kms.NewKeyManagementClient(ctx, c.clientOpts...)
-	if err != nil {
-		return err
-	}
-	defer client.Close()
-
 	request := &kmspb.EncryptRequest{
 		Name:      fmt.Sprintf("projects/%s/locations/%s/keyRings/%s/cryptoKeys/%s", c.projectID, c.locationID, c.keyRingID, kekID),
 		Plaintext: plainDEK,
@@ -132,7 +148,7 @@ func (c *KMSClient) putDEK(ctx context.Context, kekID, keyID string, plainDEK []
 
 	res, err := client.Encrypt(ctx, request, c.opts...)
 	if err != nil {
-		if strings.Contains(err.Error(), "code = NotFound") {
+		if status.Convert(err).Code() == codes.NotFound {
 			return kmsInterface.ErrKEKUnknown
 		}
 		return fmt.Errorf("encrypting DEK: %w", err)
@@ -142,7 +158,7 @@ func (c *KMSClient) putDEK(ctx context.Context, kekID, keyID string, plainDEK []
 }
 
 // createNewKEK creates a new symmetric Crypto Key in Google's KMS.
-func (c *KMSClient) createNewKEK(ctx context.Context, keyID string, client *kms.KeyManagementClient, importOnly bool) (*kmspb.CryptoKey, error) {
+func (c *KMSClient) createNewKEK(ctx context.Context, keyID string, client clientAPI, importOnly bool) (*kmspb.CryptoKey, error) {
 	request := &kmspb.CreateCryptoKeyRequest{
 		Parent:      fmt.Sprintf("projects/%s/locations/%s/keyRings/%s", c.projectID, c.locationID, c.keyRingID),
 		CryptoKeyId: keyID,
@@ -168,7 +184,7 @@ func (c *KMSClient) createNewKEK(ctx context.Context, keyID string, client *kms.
 //
 // Keys in the Google KMS can not be removed, only disabled and/or key material destroyed.
 // Since we create the initial key with `SkipInitialVersionCreation=true`, no key material is created and we do not perform any cleanup on failure.
-func (c *KMSClient) importKEK(ctx context.Context, keyID string, key []byte, client *kms.KeyManagementClient) (*kmspb.CryptoKey, error) {
+func (c *KMSClient) importKEK(ctx context.Context, keyID string, key []byte, client clientAPI) (*kmspb.CryptoKey, error) {
 	// we need an empty crypto key to import into
 	parentKey, err := c.createNewKEK(ctx, keyID, client, true)
 	if err != nil {
@@ -228,7 +244,7 @@ func (c *KMSClient) importKEK(ctx context.Context, keyID string, key []byte, cli
 }
 
 // waitBackoff is a utility function to wait for the creation of an import job.
-func (c *KMSClient) waitBackoff(ctx context.Context, jobName string, client *kms.KeyManagementClient) (*kmspb.ImportJob, bool) {
+func (c *KMSClient) waitBackoff(ctx context.Context, jobName string, client clientAPI) (*kmspb.ImportJob, bool) {
 	for i := 0; i < c.waitBackoffLimit; i++ {
 		res, err := client.GetImportJob(ctx, &kmspb.GetImportJobRequest{
 			Name: fmt.Sprintf("projects/%s/locations/%s/keyRings/%s/importJobs/%s", c.projectID, c.locationID, c.keyRingID, jobName),
@@ -243,6 +259,26 @@ func (c *KMSClient) waitBackoff(ctx context.Context, jobName string, client *kms
 	return nil, false
 }
 
+// testConnection checks if the KMS is reachable with the given configuration.
+func (c *KMSClient) testConnection(ctx context.Context) error {
+	client, err := c.newClient(ctx)
+	if err != nil {
+		return err
+	}
+	defer client.Close()
+
+	if _, err := client.GetKeyRing(ctx, &kmspb.GetKeyRingRequest{
+		Name: fmt.Sprintf("projects/%s/locations/%s/keyRings/%s", c.projectID, c.locationID, c.keyRingID),
+	}); err != nil {
+		return fmt.Errorf("GCP KMS not reachable: %w", err)
+	}
+	return nil
+}
+
+func keyManagementClientFactory(ctx context.Context, opts ...option.ClientOption) (clientAPI, error) {
+	return kms.NewKeyManagementClient(ctx, opts...)
+}
+
 // wrapCryptoKey wraps a key for import using a public RSA key, see: https://cloud.google.com/kms/docs/wrapping-a-key
 func wrapCryptoKey(key []byte, wrapKeyRSA *rsa.PublicKey) ([]byte, error) {
 	// Enforce 256bit key length

kms/kms/gcp/gcp_test.go

@@ -3,28 +3,20 @@ package gcp
 import (
 	"context"
 	"errors"
-	"fmt"
-	"net"
-	"strings"
 	"testing"
 
-	"github.com/edgelesssys/constellation/kms/config"
 	kmsInterface "github.com/edgelesssys/constellation/kms/kms"
 	"github.com/edgelesssys/constellation/kms/kms/util"
 	"github.com/edgelesssys/constellation/kms/storage"
+	"github.com/googleapis/gax-go/v2"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 	"google.golang.org/api/option"
 	kmspb "google.golang.org/genproto/googleapis/cloud/kms/v1"
-	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/status"
-	"google.golang.org/grpc/metadata"
-	"google.golang.org/protobuf/proto"
 )
 
-var (
+var testKeyRSA = `-----BEGIN PUBLIC KEY-----
-	testKey    = []byte{0x52, 0xFD, 0xFC, 0x07, 0x21, 0x82, 0x65, 0x4F, 0x16, 0x3F, 0x5F, 0x0F, 0x9A, 0x62, 0x1D, 0x72, 0x95, 0x66, 0xC7, 0x4D, 0x10, 0x03, 0x7C, 0x4D, 0x7B, 0xBB, 0x04, 0x07, 0xD1, 0xE2, 0xC6, 0x49}
-	testKeyRSA = `-----BEGIN PUBLIC KEY-----
 MIICIjANBgkqhkiG9w0BAQEFAAOCAg8AMIICCgKCAgEAu+OepfHCTiTi27nkTGke
 dn+AIkiM1AIWWDwqfqG85aNulcj60mGQGXIYV8LoEVkyKOhYBIUmJUaVczB4ltqq
 ZhR7l46RQw2vnv+XiUmfK555d4ZDInyjTusO69hE6tkuYKdXLlG1HzcrhJ254LE2
@@ -38,310 +30,346 @@ Ubhn4tvjy/q5XzVqZtBeoseW2TyyrsAN53LBkSqag5tG/264CQDigQ6Y/OADOE2x
 n08MyrFHIL/wFMscOvJo7c2Eo4EW1yXkEkAy5tF5PZgnfRObakj4gdqPeq18FNzc
 Y+t5OxL3kL15VzY1Ob0d5cMCAwEAAQ==
 -----END PUBLIC KEY-----`
-)
 
-// Google KMS testing implementation taken from: https://github.com/googleapis/google-cloud-go/blob/kms/v1.1.0/kms/apiv1/mock_test.go
+type stubGCPClient struct {
-//
+	createErr                           error
-// To keep the tests simple this only implements the methods required by our Google KMS client.
+	createCryptoKeyCalled               bool
-// More methods can be added as needed.
+	createCryptoKeyErr                  error
-type mockKeyManagementServer struct {
+	createImportJobErr                  error
-	// Embed for forward compatibility.
+	decryptResponse                     []byte
-	// Tests will keep working if more methods are added
+	decryptErr                          error
-	// in the future.
+	encryptErr                          error
-	kmspb.KeyManagementServiceServer
+	getKeyRingErr                       error
-
+	importCryptoKeyVersionErr           error
-	reqs []proto.Message
+	updateCryptoKeyPrimaryVersionCalled bool
-
+	updateCryptoKeyPrimaryVersionErr    error
-	// If set, all calls return this error.
+	getImportJobErr                     error
-	err error
+	getImportJobResponse                *kmspb.ImportJob
-
-	// responses to return if err == nil
-	resps []proto.Message
 }
 
-// CreateCryptoKey creates a new KEK.
+func newStubGCPClientFactory(stub *stubGCPClient) func(ctx context.Context, opts ...option.ClientOption) (clientAPI, error) {
-func (s *mockKeyManagementServer) CreateCryptoKey(ctx context.Context, req *kmspb.CreateCryptoKeyRequest) (*kmspb.CryptoKey, error) {
+	return func(ctx context.Context, opts ...option.ClientOption) (clientAPI, error) {
-	md, _ := metadata.FromIncomingContext(ctx)
+		return stub, stub.createErr
-	if xg := md["x-goog-api-client"]; len(xg) == 0 || !strings.Contains(xg[0], "gl-go/") {
-		return nil, fmt.Errorf("x-goog-api-client = %v, expected gl-go key", xg)
 	}
-	s.reqs = append(s.reqs, req)
-	if s.err != nil {
-		return nil, s.err
-	}
-	return s.popResponse().(*kmspb.CryptoKey), nil
 }
 
-// Decrypt performs decryption.
+func (s *stubGCPClient) Close() error {
-func (s *mockKeyManagementServer) Decrypt(ctx context.Context, req *kmspb.DecryptRequest) (*kmspb.DecryptResponse, error) {
+	return nil
-	md, _ := metadata.FromIncomingContext(ctx)
-	if xg := md["x-goog-api-client"]; len(xg) == 0 || !strings.Contains(xg[0], "gl-go/") {
-		return nil, fmt.Errorf("x-goog-api-client = %v, expected gl-go key", xg)
-	}
-	s.reqs = append(s.reqs, req)
-	if s.err != nil {
-		return nil, s.err
-	}
-	res := s.popResponse().(*kmspb.DecryptResponse)
-	res.Plaintext = make([]byte, len(req.Ciphertext))
-	for i, v := range req.Ciphertext {
-		res.Plaintext[len(res.Plaintext)-1-i] = v
-	}
-	return res, nil
 }
 
-// Encrypt performs encryption.
+func (s *stubGCPClient) CreateCryptoKey(ctx context.Context, req *kmspb.CreateCryptoKeyRequest, opts ...gax.CallOption) (*kmspb.CryptoKey, error) {
-func (s *mockKeyManagementServer) Encrypt(ctx context.Context, req *kmspb.EncryptRequest) (*kmspb.EncryptResponse, error) {
+	s.createCryptoKeyCalled = true
-	md, _ := metadata.FromIncomingContext(ctx)
+	return &kmspb.CryptoKey{}, s.createCryptoKeyErr
-	if xg := md["x-goog-api-client"]; len(xg) == 0 || !strings.Contains(xg[0], "gl-go/") {
-		return nil, fmt.Errorf("x-goog-api-client = %v, expected gl-go key", xg)
-	}
-	s.reqs = append(s.reqs, req)
-	if s.err != nil {
-		return nil, s.err
-	}
-	res := s.popResponse().(*kmspb.EncryptResponse)
-	// Reverse the input string to generate a ciphertext
-	res.Ciphertext = make([]byte, len(req.Plaintext))
-	for i, v := range req.Plaintext {
-		res.Ciphertext[len(res.Ciphertext)-1-i] = v
-	}
-	return res, nil
 }
 
-// CreateImportJob creates a new import job.
+func (s *stubGCPClient) CreateImportJob(ctx context.Context, req *kmspb.CreateImportJobRequest, opts ...gax.CallOption) (*kmspb.ImportJob, error) {
-func (s *mockKeyManagementServer) CreateImportJob(ctx context.Context, req *kmspb.CreateImportJobRequest) (*kmspb.ImportJob, error) {
+	return &kmspb.ImportJob{}, s.createImportJobErr
-	md, _ := metadata.FromIncomingContext(ctx)
-	if xg := md["x-goog-api-client"]; len(xg) == 0 || !strings.Contains(xg[0], "gl-go/") {
-		return nil, fmt.Errorf("x-goog-api-client = %v, expected gl-go key", xg)
-	}
-	s.reqs = append(s.reqs, req)
-	if s.err != nil {
-		return nil, s.err
-	}
-	return s.popResponse().(*kmspb.ImportJob), nil
 }
 
-// ImportCryptoKeyVersion imports a KEK using an import job.
+func (s *stubGCPClient) Decrypt(ctx context.Context, req *kmspb.DecryptRequest, opts ...gax.CallOption) (*kmspb.DecryptResponse, error) {
-func (s *mockKeyManagementServer) ImportCryptoKeyVersion(ctx context.Context, req *kmspb.ImportCryptoKeyVersionRequest) (*kmspb.CryptoKeyVersion, error) {
+	return &kmspb.DecryptResponse{Plaintext: s.decryptResponse}, s.decryptErr
-	md, _ := metadata.FromIncomingContext(ctx)
-	if xg := md["x-goog-api-client"]; len(xg) == 0 || !strings.Contains(xg[0], "gl-go/") {
-		return nil, fmt.Errorf("x-goog-api-client = %v, expected gl-go key", xg)
-	}
-	s.reqs = append(s.reqs, req)
-	if s.err != nil {
-		return nil, s.err
-	}
-	return s.popResponse().(*kmspb.CryptoKeyVersion), nil
 }
 
-// UpdateCryptoKeyPrimaryVersion sets the primary version of a KEK.
+func (s *stubGCPClient) Encrypt(ctx context.Context, req *kmspb.EncryptRequest, opts ...gax.CallOption) (*kmspb.EncryptResponse, error) {
-func (s *mockKeyManagementServer) UpdateCryptoKeyPrimaryVersion(ctx context.Context, req *kmspb.UpdateCryptoKeyPrimaryVersionRequest) (*kmspb.CryptoKey, error) {
+	return &kmspb.EncryptResponse{}, s.encryptErr
-	md, _ := metadata.FromIncomingContext(ctx)
-	if xg := md["x-goog-api-client"]; len(xg) == 0 || !strings.Contains(xg[0], "gl-go/") {
-		return nil, fmt.Errorf("x-goog-api-client = %v, expected gl-go key", xg)
-	}
-	s.reqs = append(s.reqs, req)
-	if s.err != nil {
-		return nil, s.err
-	}
-	return s.popResponse().(*kmspb.CryptoKey), nil
 }
 
-// GetImportJob returns information about a running import job.
+func (s *stubGCPClient) GetKeyRing(ctx context.Context, req *kmspb.GetKeyRingRequest, opts ...gax.CallOption) (*kmspb.KeyRing, error) {
-func (s *mockKeyManagementServer) GetImportJob(ctx context.Context, req *kmspb.GetImportJobRequest) (*kmspb.ImportJob, error) {
+	return &kmspb.KeyRing{}, s.getKeyRingErr
-	md, _ := metadata.FromIncomingContext(ctx)
-	if xg := md["x-goog-api-client"]; len(xg) == 0 || !strings.Contains(xg[0], "gl-go/") {
-		return nil, fmt.Errorf("x-goog-api-client = %v, expected gl-go key", xg)
-	}
-	s.reqs = append(s.reqs, req)
-	if s.err != nil {
-		return nil, s.err
-	}
-	return s.popResponse().(*kmspb.ImportJob), nil
 }
 
-func (s *mockKeyManagementServer) popResponse() proto.Message {
+func (s *stubGCPClient) ImportCryptoKeyVersion(ctx context.Context, req *kmspb.ImportCryptoKeyVersionRequest, opts ...gax.CallOption) (*kmspb.CryptoKeyVersion, error) {
-	resp := s.resps[0]
+	return &kmspb.CryptoKeyVersion{}, s.importCryptoKeyVersionErr
-	if len(s.resps) > 1 {
-		s.resps = s.resps[1:]
-	}
-	return resp
 }
 
-func TestGoogleKMS(t *testing.T) {
+func (s *stubGCPClient) UpdateCryptoKeyPrimaryVersion(ctx context.Context, req *kmspb.UpdateCryptoKeyPrimaryVersionRequest, opts ...gax.CallOption) (*kmspb.CryptoKey, error) {
-	assert := assert.New(t)
+	s.updateCryptoKeyPrimaryVersionCalled = true
-	require := require.New(t)
+	return &kmspb.CryptoKey{}, s.updateCryptoKeyPrimaryVersionErr
+}
 
-	serv := grpc.NewServer()
+func (s *stubGCPClient) GetImportJob(ctx context.Context, req *kmspb.GetImportJobRequest, opts ...gax.CallOption) (*kmspb.ImportJob, error) {
-	defer serv.GracefulStop()
+	return s.getImportJobResponse, s.getImportJobErr
-	var mockKeyManagement mockKeyManagementServer
+}
-	kmspb.RegisterKeyManagementServiceServer(serv, &mockKeyManagement)
-	lis, err := net.Listen("tcp", "localhost:0")
-	require.NoError(err)
-	go serv.Serve(lis)
 
-	project := "test-project"
+type stubStorage struct {
-	location := "global"
+	key    []byte
-	keyRing := "test-key-ring"
+	getErr error
-	kekName := "test-kek"
+	putErr error
-	dekName := "test-dek"
+}
-	plainDEK := []byte("plain DEK")
 
-	// load responses
+func (s *stubStorage) Get(context.Context, string) ([]byte, error) {
-	mockKeyManagement.resps = []proto.Message{
+	return s.key, s.getErr
-		&kmspb.CryptoKey{
+}
-			Name: fmt.Sprintf("projects/%s/locations/%s/keyRings/%s/cryptoKeys/%s", project, location, keyRing, kekName),
+
+func (s *stubStorage) Put(context.Context, string, []byte) error {
+	return s.putErr
+}
+
+func TestCreateKEK(t *testing.T) {
+	someErr := errors.New("error")
+	importKey := []byte("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA")
+
+	testCases := map[string]struct {
+		client      *stubGCPClient
+		importKey   []byte
+		errExpected bool
+	}{
+		"create new kek successful": {
+			client: &stubGCPClient{},
 		},
-		&kmspb.EncryptResponse{
+		"import kek successful": {
-			Name: dekName,
+			client: &stubGCPClient{
-		},
+				getImportJobResponse: &kmspb.ImportJob{
-		&kmspb.DecryptResponse{
+					PublicKey: &kmspb.ImportJob_WrappingPublicKey{
-			Plaintext: plainDEK,
+						Pem: testKeyRSA,
-		},
+					},
-		&kmspb.CryptoKey{
+					State: kmspb.ImportJob_ACTIVE,
-			Name: fmt.Sprintf("projects/%s/locations/%s/keyRings/%s/cryptoKeys/%s", project, location, keyRing, kekName),
+				},
-		},
-		&kmspb.ImportJob{
-			Name: "import-job",
-		},
-		&kmspb.ImportJob{
-			Name:  "import-job",
-			State: kmspb.ImportJob_ACTIVE,
-			PublicKey: &kmspb.ImportJob_WrappingPublicKey{
-				Pem: testKeyRSA,
 			},
+			importKey: importKey,
 		},
-		&kmspb.CryptoKeyVersion{
+		"CreateCryptoKey fails": {
-			Name: fmt.Sprintf("projects/%s/locations/%s/keyRings/%s/cryptoKeys/%s/cryptoKeyVersions/1", project, location, keyRing, kekName),
+			client:      &stubGCPClient{createCryptoKeyErr: someErr},
+			errExpected: true,
 		},
-		&kmspb.CryptoKey{
+		"CreatCryptoKey fails on import": {
-			Name: fmt.Sprintf("projects/%s/locations/%s/keyRings/%s/cryptoKeys/%s", project, location, keyRing, kekName),
+			client: &stubGCPClient{
+				createCryptoKeyErr: someErr,
+				getImportJobResponse: &kmspb.ImportJob{
+					PublicKey: &kmspb.ImportJob_WrappingPublicKey{
+						Pem: testKeyRSA,
+					},
+					State: kmspb.ImportJob_ACTIVE,
+				},
+			},
+			importKey:   importKey,
+			errExpected: true,
+		},
+		"CreateImportJob fails": {
+			client:      &stubGCPClient{createImportJobErr: someErr},
+			importKey:   importKey,
+			errExpected: true,
+		},
+		"ImportCryptoKeyVersion fails": {
+			client: &stubGCPClient{
+				getImportJobResponse: &kmspb.ImportJob{
+					PublicKey: &kmspb.ImportJob_WrappingPublicKey{
+						Pem: testKeyRSA,
+					},
+					State: kmspb.ImportJob_ACTIVE,
+				},
+				importCryptoKeyVersionErr: someErr,
+			},
+			importKey:   importKey,
+			errExpected: true,
+		},
+		"UpdateCryptoKeyPrimaryVersion fails": {
+			client: &stubGCPClient{
+				getImportJobResponse: &kmspb.ImportJob{
+					PublicKey: &kmspb.ImportJob_WrappingPublicKey{
+						Pem: testKeyRSA,
+					},
+					State: kmspb.ImportJob_ACTIVE,
+				},
+				updateCryptoKeyPrimaryVersionErr: someErr,
+			},
+			importKey:   importKey,
+			errExpected: true,
+		},
+		"GetImportJob fails during waitBackoff": {
+			client:      &stubGCPClient{getImportJobErr: someErr},
+			importKey:   importKey,
+			errExpected: true,
+		},
+		"GetImportJob returns no key": {
+			client: &stubGCPClient{
+				getImportJobResponse: &kmspb.ImportJob{
+					State: kmspb.ImportJob_ACTIVE,
+				},
+			},
+			importKey:   importKey,
+			errExpected: true,
+		},
+		"waitBackoff times out": {
+			client: &stubGCPClient{
+				getImportJobResponse: &kmspb.ImportJob{
+					PublicKey: &kmspb.ImportJob_WrappingPublicKey{
+						Pem: testKeyRSA,
+					},
+					State: kmspb.ImportJob_PENDING_GENERATION,
+				},
+			},
+			importKey:   importKey,
+			errExpected: true,
+		},
+		"creating client fails": {
+			client:      &stubGCPClient{createErr: someErr},
+			errExpected: true,
 		},
 	}
 
-	store := storage.NewMemMapStorage()
+	for name, tc := range testCases {
-	client := New(project, location, keyRing, store, kmspb.ProtectionLevel_SOFTWARE)
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
 
-	// redirect client calls to mock kms
+			client := &KMSClient{
-	// since the connection is closed after each call, we need to reset this option every time
+				projectID:        "test-project",
-	client.clientOpts = []option.ClientOption{getConnection(lis.Addr().String(), require)}
+				locationID:       "global",
-	ctx := context.Background()
+				keyRingID:        "test-ring",
+				newClient:        newStubGCPClientFactory(tc.client),
+				protectionLevel:  kmspb.ProtectionLevel_SOFTWARE,
+				waitBackoffLimit: 1,
+			}
 
-	// Create KEK
+			err := client.CreateKEK(context.Background(), "test-key", tc.importKey)
-	assert.NoError(client.CreateKEK(ctx, kekName, nil))
+			if tc.errExpected {
-
+				assert.Error(err)
-	// Encrypt and save new DEK
+			} else {
-	client.clientOpts = []option.ClientOption{getConnection(lis.Addr().String(), require)}
+				assert.NoError(err)
-	err = client.putDEK(ctx, kekName, dekName, plainDEK)
+				if len(tc.importKey) != 0 {
-	assert.NoError(err)
+					assert.True(tc.client.updateCryptoKeyPrimaryVersionCalled)
-	savedDEK, err := store.Get(ctx, dekName)
+				} else {
-	require.NoError(err)
+					assert.True(tc.client.createCryptoKeyCalled)
-	assert.NotEqual(plainDEK, savedDEK)
+				}
-
+			}
-	// Decrypt DEK
+		})
-	client.clientOpts = []option.ClientOption{getConnection(lis.Addr().String(), require)}
+	}
-	res, err := client.GetDEK(ctx, kekName, dekName, config.SymmetricKeyLength)
-	assert.NoError(err)
-	assert.Equal(plainDEK, res)
-
-	// Import a key
-	client.clientOpts = []option.ClientOption{getConnection(lis.Addr().String(), require)}
-	assert.NoError(client.CreateKEK(ctx, kekName, testKey))
 }
 
-func TestGetNewDEK(t *testing.T) {
+func TestGetDEK(t *testing.T) {
-	assert := assert.New(t)
+	someErr := errors.New("error")
-	require := require.New(t)
+	testKey := []byte("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA")
 
-	serv := grpc.NewServer()
+	testCases := map[string]struct {
-	defer serv.GracefulStop()
+		client      *stubGCPClient
-	var mockKeyManagement mockKeyManagementServer
+		storage     kmsInterface.Storage
-	kmspb.RegisterKeyManagementServiceServer(serv, &mockKeyManagement)
+		errExpected bool
-	lis, err := net.Listen("tcp", "localhost:0")
+	}{
-	require.NoError(err)
+		"GetDEK successful for new key": {
-	go serv.Serve(lis)
+			client:  &stubGCPClient{},
-
+			storage: &stubStorage{getErr: storage.ErrDEKUnset},
-	project := "test-project"
-	location := "global"
-	keyRing := "test-key-ring"
-	kekName := "test-kek"
-	dekName := "test-dek"
-	largeDEKName := "test-dek-large"
-
-	store := storage.NewMemMapStorage()
-	client := New(project, location, keyRing, store, kmspb.ProtectionLevel_SOFTWARE)
-
-	mockKeyManagement.resps = []proto.Message{
-		&kmspb.EncryptResponse{
-			Name: dekName,
 		},
-		&kmspb.DecryptResponse{},
+		"GetDEK successful for existing key": {
-		&kmspb.EncryptResponse{
+			client:  &stubGCPClient{decryptResponse: testKey},
-			Name: largeDEKName,
+			storage: &stubStorage{key: testKey},
+		},
+		"Get from storage fails": {
+			client:      &stubGCPClient{},
+			storage:     &stubStorage{getErr: someErr},
+			errExpected: true,
+		},
+		"Encrypt fails": {
+			client:      &stubGCPClient{encryptErr: someErr},
+			storage:     &stubStorage{getErr: storage.ErrDEKUnset},
+			errExpected: true,
+		},
+		"Encrypt fails with notfound error": {
+			client:      &stubGCPClient{encryptErr: status.Error(codes.NotFound, "error")},
+			storage:     &stubStorage{getErr: storage.ErrDEKUnset},
+			errExpected: true,
+		},
+		"Put to storage fails": {
+			client: &stubGCPClient{},
+			storage: &stubStorage{
+				getErr: storage.ErrDEKUnset,
+				putErr: someErr,
+			},
+			errExpected: true,
+		},
+		"Decrypt fails": {
+			client:      &stubGCPClient{decryptErr: someErr},
+			storage:     &stubStorage{key: testKey},
+			errExpected: true,
+		},
+		"Decrypt fails with notfound error": {
+			client:      &stubGCPClient{decryptErr: status.Error(codes.NotFound, "error")},
+			storage:     &stubStorage{key: testKey},
+			errExpected: true,
+		},
+		"creating client fails": {
+			client:      &stubGCPClient{createErr: someErr},
+			storage:     &stubStorage{getErr: storage.ErrDEKUnset},
+			errExpected: true,
 		},
 	}
-	ctx := context.Background()
 
-	// Requesting an unset DEK should generate a new one, which we can then fetch in a second request
+	for name, tc := range testCases {
-	client.clientOpts = []option.ClientOption{getConnection(lis.Addr().String(), require)}
+		t.Run(name, func(t *testing.T) {
-	res1, err := client.GetDEK(ctx, kekName, dekName, config.SymmetricKeyLength)
+			assert := assert.New(t)
-	assert.NoError(err)
-	client.clientOpts = []option.ClientOption{getConnection(lis.Addr().String(), require)}
-	res2, err := client.GetDEK(ctx, kekName, dekName, config.SymmetricKeyLength)
-	assert.NoError(err)
-	assert.Equal(res1, res2)
 
-	// Requesting larger key sizes should be possible
+			client := &KMSClient{
-	client.clientOpts = []option.ClientOption{getConnection(lis.Addr().String(), require)}
+				projectID:        "test-project",
-	res3, err := client.GetDEK(ctx, kekName, largeDEKName, 96)
+				locationID:       "global",
-	assert.NoError(err)
+				keyRingID:        "test-ring",
-	assert.Len(res3, 96)
+				newClient:        newStubGCPClientFactory(tc.client),
+				protectionLevel:  kmspb.ProtectionLevel_SOFTWARE,
+				waitBackoffLimit: 1,
+				storage:          tc.storage,
+			}
+
+			dek, err := client.GetDEK(context.Background(), "test-key", "volume-01", 32)
+			if tc.errExpected {
+				assert.Error(err)
+			} else {
+				assert.NoError(err)
+				assert.Len(dek, 32)
+			}
+		})
+	}
 }
 
-func TestUnknownKEK(t *testing.T) {
+func TestConnection(t *testing.T) {
-	assert := assert.New(t)
+	someErr := errors.New("error")
-	require := require.New(t)
+	testCases := map[string]struct {
+		client      *stubGCPClient
+		errExpected bool
+	}{
+		"success": {
+			client: &stubGCPClient{},
+		},
+		"newClient fails": {
+			client:      &stubGCPClient{createErr: someErr},
+			errExpected: true,
+		},
+		"GetKeyRing fails": {
+			client:      &stubGCPClient{getKeyRingErr: someErr},
+			errExpected: true,
+		},
+	}
 
-	serv := grpc.NewServer()
+	for name, tc := range testCases {
-	defer serv.GracefulStop()
+		t.Run(name, func(t *testing.T) {
-	var mockKeyManagement mockKeyManagementServer
+			assert := assert.New(t)
-	kmspb.RegisterKeyManagementServiceServer(serv, &mockKeyManagement)
-	lis, err := net.Listen("tcp", "localhost:0")
-	require.NoError(err)
-	go serv.Serve(lis)
 
-	mockKeyManagement.err = errors.New("rpc error: code = NotFound")
+			client := &KMSClient{
+				projectID:        "test-project",
+				locationID:       "global",
+				keyRingID:        "test-ring",
+				newClient:        newStubGCPClientFactory(tc.client),
+				protectionLevel:  kmspb.ProtectionLevel_SOFTWARE,
+				waitBackoffLimit: 1,
+			}
 
-	store := storage.NewMemMapStorage()
+			err := client.testConnection(context.Background())
-	client := New("test-project", "global", "test-key-ring", store, kmspb.ProtectionLevel_SOFTWARE)
+			if tc.errExpected {
-	ctx := context.Background()
+				assert.Error(err)
-
+			} else {
-	client.clientOpts = []option.ClientOption{getConnection(lis.Addr().String(), require)}
+				assert.NoError(err)
-	err = client.putDEK(ctx, "invalid-kek", "test-dek", []byte("dek"))
+			}
-	assert.Error(err)
+		})
-	assert.ErrorIs(err, kmsInterface.ErrKEKUnknown)
+	}
-
-	require.NoError(store.Put(ctx, "test-dek", []byte("Test Key")))
-	client.clientOpts = []option.ClientOption{getConnection(lis.Addr().String(), require)}
-	_, err = client.GetDEK(ctx, "invalid-kek", "test-dek", config.SymmetricKeyLength)
-	assert.Error(err)
-	assert.ErrorIs(err, kmsInterface.ErrKEKUnknown)
 }
 
-func getConnection(lisAddr string, r *require.Assertions) option.ClientOption {
+func TestWrapCryptoKey(t *testing.T) {
-	conn, err := grpc.Dial(lisAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
-	r.NoError(err)
-	return option.WithGRPCConn(conn)
-}
-
-func TestWrapKeyRSA(t *testing.T) {
 	assert := assert.New(t)
 
 	rsaPub, err := util.ParsePEMtoPublicKeyRSA([]byte(testKeyRSA))
 	assert.NoError(err)
 
-	res, err := wrapCryptoKey(testKey, rsaPub)
+	res, err := wrapCryptoKey([]byte("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"), rsaPub)
 	assert.NoError(err)
 	assert.Equal(552, len(res))
+
+	_, err = wrapCryptoKey([]byte{0x1}, rsaPub)
+	assert.Error(err)
 }

kms/test/gcp_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/edgelesssys/constellation/kms/kms/gcp"
 	"github.com/edgelesssys/constellation/kms/storage"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	kmspb "google.golang.org/genproto/googleapis/cloud/kms/v1"
 )
 
@@ -27,15 +28,18 @@ func TestCreateGcpKEK(t *testing.T) {
 		t.Skip("Skipping Google KMS key creation test")
 	}
 	assert := assert.New(t)
+	require := require.New(t)
 	store := storage.NewMemMapStorage()
 
 	kekName := addSuffix("test-kek")
 	dekName := "test-dek"
 
-	kmsClient := gcp.New(gcpProjectID, gcpLocation, gcpKeyRing, store, kmspb.ProtectionLevel_SOFTWARE)
-
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
 	defer cancel()
+
+	kmsClient, err := gcp.New(ctx, gcpProjectID, gcpLocation, gcpKeyRing, store, kmspb.ProtectionLevel_SOFTWARE)
+	require.NoError(err)
+
 	// Key name is random, but there is a chance we try to create a key that already exists, in that case the test fails
 	assert.NoError(kmsClient.CreateKEK(ctx, kekName, nil))
 
@@ -57,16 +61,19 @@ func TestImportGcpKEK(t *testing.T) {
 		t.Skip("Skipping Google KMS key import test")
 	}
 	assert := assert.New(t)
+	require := require.New(t)
 	store := storage.NewMemMapStorage()
 
 	kekName := addSuffix("test-kek")
 	kekData := []byte{0x52, 0xFD, 0xFC, 0x07, 0x21, 0x82, 0x65, 0x4F, 0x16, 0x3F, 0x5F, 0x0F, 0x9A, 0x62, 0x1D, 0x72, 0x95, 0x66, 0xC7, 0x4D, 0x10, 0x03, 0x7C, 0x4D, 0x7B, 0xBB, 0x04, 0x07, 0xD1, 0xE2, 0xC6, 0x49}
 	dekName := "test-dek"
 
-	kmsClient := gcp.New(gcpProjectID, gcpLocation, gcpKeyRing, store, kmspb.ProtectionLevel_SOFTWARE)
-
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
 	defer cancel()
+
+	kmsClient, err := gcp.New(ctx, gcpProjectID, gcpLocation, gcpKeyRing, store, kmspb.ProtectionLevel_SOFTWARE)
+	require.NoError(err)
+
 	assert.NoError(kmsClient.CreateKEK(ctx, kekName, kekData))
 
 	res, err := kmsClient.GetDEK(ctx, kekName, dekName, config.SymmetricKeyLength)