Tests for YC KMS provider

- Tests for MasterKey public methods. gRPC calls `Encrypt` and `Decrypt` are mocked with dummy responses.
- Minor adjustments for MasterKey constructors: the constructors return an error

Author: astreter

cmd/sops/main.go

@@ -432,7 +432,12 @@ func main() {
 							group = append(group, gcpkms.NewMasterKeyFromResourceID(kms))
 						}
 						for _, kms := range ycKmses {
-							group = append(group, yckms.NewMasterKeyFromKeyID(kms))
+							k, err := yckms.NewMasterKeyFromKeyID(kms)
+							if err != nil {
+								log.WithError(err).Error("Failed to add key")
+								continue
+							}
+							group = append(group, k)
 						}
 						for _, uri := range vaultURIs {
 							k, err := hcvault.NewMasterKeyFromURI(uri)
@@ -1105,7 +1110,11 @@ func keyGroups(c *cli.Context, file string) ([]sops.KeyGroup, error) {
 		}
 	}
 	if c.String("yc-kms") != "" {
-		for _, k := range yckms.NewMasterKeyFromKeyIDString(c.String("yc-kms")) {
+		ycKeys, err := yckms.NewMasterKeyFromKeyIDString(c.String("yc-kms"))
+		if err != nil {
+			return nil, err
+		}
+		for _, k := range ycKeys {
 			ycKmsKeys = append(ycKmsKeys, k)
 		}
 	}

config/config.go

@@ -178,7 +178,11 @@ func getKeyGroupsFromCreationRule(cRule *creationRule, kmsEncryptionContext map[
 				keyGroup = append(keyGroup, gcpkms.NewMasterKeyFromResourceID(k.ResourceID))
 			}
 			for _, k := range group.YCKMS {
-				keyGroup = append(keyGroup, yckms.NewMasterKeyFromKeyID(k.KeyID))
+				if masterKey, err := yckms.NewMasterKeyFromKeyID(k.KeyID); err == nil {
+					keyGroup = append(keyGroup, masterKey)
+				} else {
+					return nil, err
+				}
 			}
 			for _, k := range group.AzureKV {
 				keyGroup = append(keyGroup, azkv.NewMasterKey(k.VaultURL, k.Key, k.Version))
@@ -213,7 +217,11 @@ func getKeyGroupsFromCreationRule(cRule *creationRule, kmsEncryptionContext map[
 		for _, k := range gcpkms.MasterKeysFromResourceIDString(cRule.GCPKMS) {
 			keyGroup = append(keyGroup, k)
 		}
-		for _, k := range yckms.NewMasterKeyFromKeyIDString(cRule.YCKMS) {
+		ycKeys, err := yckms.NewMasterKeyFromKeyIDString(cRule.YCKMS)
+		if err != nil {
+			return nil, err
+		}
+		for _, k := range ycKeys {
 			keyGroup = append(keyGroup, k)
 		}
 		azureKeys, err := azkv.MasterKeysFromURLs(cRule.AzureKeyVault)

yckms/keysource.go

@@ -39,33 +39,42 @@ type MasterKey struct {
 	// CreationDate is the creation timestamp of the MasterKey. Used
 	// for NeedsRotation.
 	CreationDate time.Time
+
+	client *kms.SymmetricCryptoServiceClient
 }
 
-func NewMasterKeyFromKeyID(keyID string) *MasterKey {
-	k := &MasterKey{}
-	k.KeyID = keyID
-	k.CreationDate = time.Now().UTC()
-	return k
+func NewMasterKeyFromKeyID(keyID string) (k *MasterKey, err error) {
+	k = &MasterKey{
+		KeyID:        keyID,
+		CreationDate: time.Now().UTC(),
+	}
+
+	k.client, err = NewKMSClient()
+	if err != nil {
+		log.WithError(err).WithField("keyID", keyID).Error("Encryption failed")
+		return nil, fmt.Errorf("cannot create YC KMS service: %w", err)
+	}
+
+	return k, nil
 }
 
-func NewMasterKeyFromKeyIDString(keyID string) []*MasterKey {
+func NewMasterKeyFromKeyIDString(keyID string) ([]*MasterKey, error) {
 	var keys []*MasterKey
 	if keyID == "" {
-		return keys
+		return keys, nil
 	}
 	for _, s := range strings.Split(keyID, ",") {
-		keys = append(keys, NewMasterKeyFromKeyID(s))
+		key, err := NewMasterKeyFromKeyID(strings.TrimSpace(s))
+		if err != nil {
+			return nil, err
+		}
+		keys = append(keys, key)
 	}
-	return keys
+	return keys, nil
 }
 
 func (key *MasterKey) Encrypt(dataKey []byte) error {
-	client, ctx, err := key.newKMSClient()
-	if err != nil {
-		log.WithError(err).WithField("keyID", key.KeyID).Error("Encryption failed")
-		return fmt.Errorf("cannot create YC KMS service: %w", err)
-	}
-	ciphertextResponse, err := client.Encrypt(ctx, &yckms.SymmetricEncryptRequest{
+	ciphertextResponse, err := key.client.Encrypt(context.Background(), &yckms.SymmetricEncryptRequest{
 		KeyId:     key.KeyID,
 		Plaintext: dataKey,
 	})
@@ -100,13 +109,8 @@ func (key *MasterKey) EncryptIfNeeded(dataKey []byte) error {
 // Decrypt decrypts the EncryptedKey field with YC KMS and returns
 // the result.
 func (key *MasterKey) Decrypt() ([]byte, error) {
-	client, ctx, err := key.newKMSClient()
-	if err != nil {
-		log.WithError(err).WithField("keyID", key.KeyID).Error("Decryption failed")
-		return nil, fmt.Errorf("cannot create YC KMS service: %w", err)
-	}
 	decodedCipher, err := base64.StdEncoding.DecodeString(string(key.EncryptedDataKey()))
-	plaintextResponse, err := client.Decrypt(ctx, &yckms.SymmetricDecryptRequest{
+	plaintextResponse, err := key.client.Decrypt(context.Background(), &yckms.SymmetricDecryptRequest{
 		KeyId:      key.KeyID,
 		Ciphertext: decodedCipher,
 	})
@@ -129,33 +133,32 @@ func (key *MasterKey) ToString() string {
 }
 
 // ToMap converts the MasterKey to a map for serialization purposes.
-func (key MasterKey) ToMap() map[string]interface{} {
+func (key *MasterKey) ToMap() map[string]interface{} {
 	out := make(map[string]interface{})
 	out["key_id"] = key.KeyID
 	out["created_at"] = key.CreationDate.UTC().Format(time.RFC3339)
 	out["enc"] = key.EncryptedKey
 	return out
 }
 
-// newKMSClient returns a YC KMS client configured with the credentialsStore
+// NewKMSClient returns a YC KMS client configured with the credentialsStore
 // and/or grpcConn, falling back to environmental defaults.
 // It returns an error if the ResourceID is invalid, or if the setup of the
 // client fails.
-func (key *MasterKey) newKMSClient() (*kms.SymmetricCryptoServiceClient, context.Context, error) {
-	ctx := context.Background()
+func NewKMSClient() (*kms.SymmetricCryptoServiceClient, error) {
 	cred, err := getYandexCloudCredentials()
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 
-	client, err := ycsdk.Build(ctx, ycsdk.Config{
+	client, err := ycsdk.Build(context.Background(), ycsdk.Config{
 		Credentials: cred,
 	})
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 
-	return client.KMSCrypto().SymmetricCrypto(), ctx, nil
+	return client.KMSCrypto().SymmetricCrypto(), nil
 }
 
 // getYandexCloudCredentials trying to locate credentials in the following order

yckms/keysource_test.go

@@ -0,0 +1,192 @@
+package yckms
+
+import (
+	"context"
+	"encoding/base64"
+	yckms "github.com/yandex-cloud/go-genproto/yandex/cloud/kms/v1"
+	kms "github.com/yandex-cloud/go-sdk/gen/kmscrypto"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/status"
+	"google.golang.org/grpc/test/bufconn"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+const (
+	dummyKey        = "my:dummy:ru-east-1:12345678:key/920aff2e-c5f1-4040-943a-047fa387b27e+my:dummy:iam::927034868273:role/sops-dev"
+	anotherDummyKey = "my:dummy:ru-east-1:12345678:key/920aff2e-c5f1-4040-943a-047fa387b27e+my:dummy:iam::927034868273"
+	dummyKeys       = dummyKey + ", " + anotherDummyKey
+	decodedKey      = "I want to be a DJ"
+)
+
+const bufSize = 1024 * 1024
+
+var lis *bufconn.Listener
+
+func init() {
+	lis = bufconn.Listen(bufSize)
+	s := grpc.NewServer()
+	yckms.RegisterSymmetricCryptoServiceServer(s, mockSymmetricCryptoServiceServer{})
+	go func() {
+		if err := s.Serve(lis); err != nil {
+			log.Fatalf("Server exited with error: %v", err)
+		}
+	}()
+}
+
+func bufDialer(context.Context, string) (net.Conn, error) {
+	return lis.Dial()
+}
+
+type mockSymmetricCryptoServiceServer struct {
+}
+
+func (mockSymmetricCryptoServiceServer) Encrypt(context.Context, *yckms.SymmetricEncryptRequest) (*yckms.SymmetricEncryptResponse, error) {
+	key, _ := base64.StdEncoding.DecodeString("test")
+	return &yckms.SymmetricEncryptResponse{
+		Ciphertext: key,
+	}, nil
+}
+func (mockSymmetricCryptoServiceServer) Decrypt(context.Context, *yckms.SymmetricDecryptRequest) (*yckms.SymmetricDecryptResponse, error) {
+	return &yckms.SymmetricDecryptResponse{
+		Plaintext: []byte(decodedKey),
+	}, nil
+}
+func (mockSymmetricCryptoServiceServer) ReEncrypt(context.Context, *yckms.SymmetricReEncryptRequest) (*yckms.SymmetricReEncryptResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method ReEncrypt not implemented")
+}
+func (mockSymmetricCryptoServiceServer) GenerateDataKey(context.Context, *yckms.GenerateDataKeyRequest) (*yckms.GenerateDataKeyResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method GenerateDataKey not implemented")
+}
+
+func TestNewMasterKeyFromKeyID(t *testing.T) {
+	key, err := NewMasterKeyFromKeyID(dummyKey)
+	assert.NoError(t, err)
+	assert.Equal(t, dummyKey, key.KeyID)
+	assert.NotNil(t, key.CreationDate)
+}
+
+func TestNewMasterKeyFromKeyIDString(t *testing.T) {
+	keys, err := NewMasterKeyFromKeyIDString(dummyKeys)
+	assert.NoError(t, err)
+	assert.Len(t, keys, 2)
+
+	k1 := keys[0]
+	k2 := keys[1]
+
+	assert.Equal(t, dummyKey, k1.KeyID)
+	assert.Equal(t, anotherDummyKey, k2.KeyID)
+}
+
+func TestMasterKey_Encrypt(t *testing.T) {
+	t.Run("encrypt", func(t *testing.T) {
+		kmsClient, err := createTestKMSClient()
+		assert.NoError(t, err)
+
+		key := &MasterKey{
+			client: kmsClient,
+		}
+
+		dataKey := []byte(decodedKey)
+		assert.NoError(t, key.Encrypt(dataKey))
+
+		decodedCipher, err := base64.StdEncoding.DecodeString(key.EncryptedKey)
+		assert.NoError(t, err)
+
+		decrypted, err := kmsClient.Decrypt(context.TODO(), &yckms.SymmetricDecryptRequest{
+			KeyId:      key.KeyID,
+			Ciphertext: decodedCipher,
+		})
+		if assert.NoError(t, err) {
+			assert.NotNil(t, decrypted)
+			assert.Equal(t, dataKey, decrypted.Plaintext)
+		}
+	})
+}
+
+func TestMasterKey_Decrypt(t *testing.T) {
+	t.Run("decrypt", func(t *testing.T) {
+		kmsClient, err := createTestKMSClient()
+		assert.NoError(t, err)
+
+		key := &MasterKey{
+			client: kmsClient,
+		}
+
+		dataKey := []byte(decodedKey)
+		out, err := kmsClient.Encrypt(context.TODO(), &yckms.SymmetricEncryptRequest{
+			KeyId:     dummyKey,
+			Plaintext: dataKey,
+		})
+		assert.NoError(t, err)
+
+		key.EncryptedKey = base64.StdEncoding.EncodeToString(out.Ciphertext)
+		got, err := key.Decrypt()
+		assert.NoError(t, err)
+		assert.Equal(t, dataKey, got)
+	})
+}
+
+func TestMasterKey_EncryptedDataKey(t *testing.T) {
+	key := &MasterKey{EncryptedKey: "some key"}
+	assert.EqualValues(t, key.EncryptedKey, key.EncryptedDataKey())
+}
+
+func TestMasterKey_SetEncryptedDataKey(t *testing.T) {
+	key := &MasterKey{}
+	data := []byte("some data")
+	key.SetEncryptedDataKey(data)
+	assert.EqualValues(t, data, key.EncryptedKey)
+}
+
+func TestMasterKey_NeedsRotation(t *testing.T) {
+	t.Run("false", func(t *testing.T) {
+		k := &MasterKey{}
+		k.CreationDate = time.Now().UTC()
+
+		assert.False(t, k.NeedsRotation())
+	})
+
+	t.Run("true", func(t *testing.T) {
+		k := &MasterKey{}
+		k.CreationDate = time.Now().UTC().Add(-kmsTTL - 1)
+
+		assert.True(t, k.NeedsRotation())
+	})
+}
+
+func TestMasterKey_ToString(t *testing.T) {
+	key, err := NewMasterKeyFromKeyID(dummyKey)
+	assert.NoError(t, err)
+
+	assert.Equal(t, dummyKey, key.ToString())
+}
+
+func TestMasterKey_ToMap(t *testing.T) {
+	key, err := NewMasterKeyFromKeyID(dummyKey)
+	assert.NoError(t, err)
+
+	data := []byte("some data")
+	key.SetEncryptedDataKey(data)
+
+	res := key.ToMap()
+	assert.Equal(t, dummyKey, res["key_id"])
+	assert.Equal(t, key.CreationDate.UTC().Format(time.RFC3339), res["created_at"])
+	assert.Equal(t, "some data", res["enc"])
+}
+
+func createTestKMSClient() (*kms.SymmetricCryptoServiceClient, error) {
+	return kms.NewKMSCrypto(func(ctx context.Context) (*grpc.ClientConn, error) {
+		return grpc.DialContext(
+			ctx,
+			"bufnet",
+			grpc.WithContextDialer(bufDialer),
+			grpc.WithTransportCredentials(insecure.NewCredentials()),
+		)
+	}).SymmetricCrypto(), nil
+}