Skip to content

Commit d2335e9

Browse files
astreterVladimir Kuznichenkov
astreter
authored and
Vladimir Kuznichenkov
committed
Implement tests for YC KMS provider and credentials forward
gRPC calls `Encrypt` and `Decrypt` are mocked with dummy responses using base64 instead of actual encryption. Since YC KMS response with binary data we are storing it as base64, together with mocked server where we encode it one more time instead of actual encryption we are using double encoding and double decoding in tests cases. Signed-off-by: Vladimir Kuznichenkov <kuzaxak.tech@gmail.com>
1 parent 195e070 commit d2335e9

File tree

2 files changed

+236
-19
lines changed

2 files changed

+236
-19
lines changed

‎yckms/keysource.go‎

Lines changed: 59 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
ycsdk "github.com/yandex-cloud/go-sdk"
1111
"github.com/yandex-cloud/go-sdk/gen/kmscrypto"
1212
"github.com/yandex-cloud/go-sdk/iamkey"
13+
"google.golang.org/grpc"
1314
"os"
1415
"strings"
1516
"time"
@@ -40,17 +41,24 @@ type MasterKey struct {
4041
// CreationDate is the creation timestamp of the MasterKey. Used
4142
// for NeedsRotation.
4243
CreationDate time.Time
44+
45+
credentials ycsdk.Credentials
46+
47+
// grpcConn can be used to inject a custom YC KMS client connection.
48+
// Mostly useful for testing at present, to wire the client to a mock
49+
// server.
50+
grpcConn *grpc.ClientConn
4351
}
4452

4553
func (key *MasterKey) TypeToIdentifier() string {
4654
return KeyTypeIdentifier
4755
}
4856

4957
func NewMasterKeyFromKeyID(keyID string) *MasterKey {
50-
k := &MasterKey{}
51-
k.KeyID = keyID
52-
k.CreationDate = time.Now().UTC()
53-
return k
58+
return &MasterKey{
59+
KeyID: keyID,
60+
CreationDate: time.Now().UTC(),
61+
}
5462
}
5563

5664
func NewMasterKeyFromKeyIDString(keyID string) []*MasterKey {
@@ -59,18 +67,34 @@ func NewMasterKeyFromKeyIDString(keyID string) []*MasterKey {
5967
return keys
6068
}
6169
for _, s := range strings.Split(keyID, ",") {
62-
keys = append(keys, NewMasterKeyFromKeyID(s))
70+
keys = append(keys, NewMasterKeyFromKeyID(strings.TrimSpace(s)))
6371
}
6472
return keys
6573
}
6674

67-
func (key *MasterKey) Encrypt(dataKey []byte) error {
68-
client, ctx, err := key.newKMSClient()
75+
// YCCredentials is a ycsdk.Credentials used for authenticating towards YC KMS
76+
type YCCredentials struct {
77+
credentials ycsdk.Credentials
78+
}
79+
80+
// NewYCCredentials creates a new YCCredentials with the provided ycsdk.Credentials.
81+
func NewYCCredentials(credentials ycsdk.Credentials) *YCCredentials {
82+
return &YCCredentials{credentials: credentials}
83+
}
84+
85+
// ApplyToMasterKey configures the TokenCredential on the provided key.
86+
func (c YCCredentials) ApplyToMasterKey(key *MasterKey) {
87+
key.credentials = c.credentials
88+
}
89+
90+
func (key *MasterKey) Encrypt(dataKey []byte) (err error) {
91+
client, err := key.newKMSClient()
6992
if err != nil {
7093
log.WithError(err).WithField("keyID", key.KeyID).Error("Encryption failed")
7194
return fmt.Errorf("cannot create YC KMS service: %w", err)
7295
}
73-
ciphertextResponse, err := client.Encrypt(ctx, &yckms.SymmetricEncryptRequest{
96+
97+
ciphertextResponse, err := client.Encrypt(context.Background(), &yckms.SymmetricEncryptRequest{
7498
KeyId: key.KeyID,
7599
Plaintext: dataKey,
76100
})
@@ -105,13 +129,14 @@ func (key *MasterKey) EncryptIfNeeded(dataKey []byte) error {
105129
// Decrypt decrypts the EncryptedKey field with YC KMS and returns
106130
// the result.
107131
func (key *MasterKey) Decrypt() ([]byte, error) {
108-
client, ctx, err := key.newKMSClient()
132+
client, err := key.newKMSClient()
109133
if err != nil {
110134
log.WithError(err).WithField("keyID", key.KeyID).Error("Decryption failed")
111135
return nil, fmt.Errorf("cannot create YC KMS service: %w", err)
112136
}
137+
113138
decodedCipher, err := base64.StdEncoding.DecodeString(string(key.EncryptedDataKey()))
114-
plaintextResponse, err := client.Decrypt(ctx, &yckms.SymmetricDecryptRequest{
139+
plaintextResponse, err := client.Decrypt(context.Background(), &yckms.SymmetricDecryptRequest{
115140
KeyId: key.KeyID,
116141
Ciphertext: decodedCipher,
117142
})
@@ -134,7 +159,7 @@ func (key *MasterKey) ToString() string {
134159
}
135160

136161
// ToMap converts the MasterKey to a map for serialization purposes.
137-
func (key MasterKey) ToMap() map[string]interface{} {
162+
func (key *MasterKey) ToMap() map[string]interface{} {
138163
out := make(map[string]interface{})
139164
out["key_id"] = key.KeyID
140165
out["created_at"] = key.CreationDate.UTC().Format(time.RFC3339)
@@ -146,21 +171,36 @@ func (key MasterKey) ToMap() map[string]interface{} {
146171
// and/or grpcConn, falling back to environmental defaults.
147172
// It returns an error if the ResourceID is invalid, or if the setup of the
148173
// client fails.
149-
func (key *MasterKey) newKMSClient() (*kms.SymmetricCryptoServiceClient, context.Context, error) {
150-
ctx := context.Background()
151-
cred, err := getYandexCloudCredentials()
152-
if err != nil {
153-
return nil, nil, err
174+
func (key *MasterKey) newKMSClient() (*kms.SymmetricCryptoServiceClient, error) {
175+
var (
176+
cred ycsdk.Credentials
177+
err error
178+
)
179+
180+
switch {
181+
case key.credentials != nil:
182+
cred = key.credentials
183+
default:
184+
cred, err = getYandexCloudCredentials()
185+
if err != nil {
186+
return nil, err
187+
}
154188
}
155189

156-
client, err := ycsdk.Build(ctx, ycsdk.Config{
190+
client, err := ycsdk.Build(context.Background(), ycsdk.Config{
157191
Credentials: cred,
158192
})
159193
if err != nil {
160-
return nil, nil, err
194+
return nil, err
195+
}
196+
197+
if key.grpcConn != nil {
198+
return kms.NewKMSCrypto(func(ctx context.Context) (*grpc.ClientConn, error) {
199+
return key.grpcConn, nil
200+
}).SymmetricCrypto(), nil
161201
}
162202

163-
return client.KMSCrypto().SymmetricCrypto(), ctx, nil
203+
return client.KMSCrypto().SymmetricCrypto(), nil
164204
}
165205

166206
// getYandexCloudCredentials trying to locate credentials in the following order

‎yckms/keysource_test.go‎

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
1+
package yckms
2+
3+
import (
4+
"context"
5+
"encoding/base64"
6+
yckms "github.com/yandex-cloud/go-genproto/yandex/cloud/kms/v1"
7+
"google.golang.org/grpc"
8+
"google.golang.org/grpc/codes"
9+
"google.golang.org/grpc/credentials/insecure"
10+
"google.golang.org/grpc/status"
11+
"google.golang.org/grpc/test/bufconn"
12+
"net"
13+
"testing"
14+
"time"
15+
16+
"github.com/stretchr/testify/assert"
17+
)
18+
19+
const (
20+
dummyKey = "920aff2e-c5f1-4040-943a-047fa387b27e"
21+
anotherDummyKey = "920aff2e-c5f1-4040-943a-047fa587b27e"
22+
dummyKeys = dummyKey + ", " + anotherDummyKey
23+
decodedKey = "I want to be a DJ"
24+
)
25+
26+
const bufSize = 1024 * 1024
27+
28+
var lis *bufconn.Listener
29+
30+
func init() {
31+
lis = bufconn.Listen(bufSize)
32+
s := grpc.NewServer()
33+
yckms.RegisterSymmetricCryptoServiceServer(s, mockSymmetricCryptoServiceServer{})
34+
go func() {
35+
if err := s.Serve(lis); err != nil {
36+
log.Fatalf("Server exited with error: %v", err)
37+
}
38+
}()
39+
}
40+
41+
func bufDialer(context.Context, string) (net.Conn, error) {
42+
return lis.Dial()
43+
}
44+
45+
type mockSymmetricCryptoServiceServer struct {
46+
}
47+
48+
func (mockSymmetricCryptoServiceServer) Encrypt(ctx context.Context, req *yckms.SymmetricEncryptRequest) (*yckms.SymmetricEncryptResponse, error) {
49+
return &yckms.SymmetricEncryptResponse{
50+
Ciphertext: []byte(base64.StdEncoding.EncodeToString(req.Plaintext)),
51+
}, nil
52+
}
53+
54+
func (mockSymmetricCryptoServiceServer) Decrypt(ctx context.Context, req *yckms.SymmetricDecryptRequest) (*yckms.SymmetricDecryptResponse, error) {
55+
plain, err := base64.StdEncoding.DecodeString(string(req.Ciphertext))
56+
if err != nil {
57+
return nil, err
58+
}
59+
return &yckms.SymmetricDecryptResponse{
60+
Plaintext: plain,
61+
}, nil
62+
}
63+
func (mockSymmetricCryptoServiceServer) ReEncrypt(context.Context, *yckms.SymmetricReEncryptRequest) (*yckms.SymmetricReEncryptResponse, error) {
64+
return nil, status.Errorf(codes.Unimplemented, "method ReEncrypt not implemented")
65+
}
66+
func (mockSymmetricCryptoServiceServer) GenerateDataKey(context.Context, *yckms.GenerateDataKeyRequest) (*yckms.GenerateDataKeyResponse, error) {
67+
return nil, status.Errorf(codes.Unimplemented, "method GenerateDataKey not implemented")
68+
}
69+
70+
func TestNewMasterKeyFromKeyID(t *testing.T) {
71+
key := NewMasterKeyFromKeyID(dummyKey)
72+
assert.Equal(t, dummyKey, key.KeyID)
73+
assert.NotNil(t, key.CreationDate)
74+
}
75+
76+
func TestNewMasterKeyFromKeyIDString(t *testing.T) {
77+
keys := NewMasterKeyFromKeyIDString(dummyKeys)
78+
assert.Len(t, keys, 2)
79+
80+
k1 := keys[0]
81+
k2 := keys[1]
82+
83+
assert.Equal(t, dummyKey, k1.KeyID)
84+
assert.Equal(t, anotherDummyKey, k2.KeyID)
85+
}
86+
87+
func TestMasterKey_Encrypt(t *testing.T) {
88+
t.Run("encrypt", func(t *testing.T) {
89+
grpcConn, err := createMockGRPCClient()
90+
assert.NoError(t, err)
91+
92+
key := &MasterKey{
93+
grpcConn: grpcConn,
94+
}
95+
96+
dataKey := []byte(decodedKey)
97+
err = key.Encrypt(dataKey)
98+
assert.NoError(t, err)
99+
100+
// Double base64 is used because encrypted data stored as base64
101+
// and our mock uses base64 instead of actual encryption
102+
assert.EqualValues(t, base64.StdEncoding.EncodeToString([]byte(base64.StdEncoding.EncodeToString([]byte(decodedKey)))), key.EncryptedDataKey())
103+
})
104+
}
105+
106+
func TestMasterKey_Decrypt(t *testing.T) {
107+
t.Run("decrypt", func(t *testing.T) {
108+
grpcConn, err := createMockGRPCClient()
109+
assert.NoError(t, err)
110+
111+
// Double base64 is used because encrypted data stored as base64
112+
// and our mock uses base64 instead of actual encryption
113+
key := &MasterKey{
114+
EncryptedKey: base64.StdEncoding.EncodeToString([]byte(base64.StdEncoding.EncodeToString([]byte(decodedKey)))),
115+
grpcConn: grpcConn,
116+
}
117+
118+
got, err := key.Decrypt()
119+
assert.NoError(t, err)
120+
assert.Equal(t, []byte(decodedKey), got)
121+
})
122+
}
123+
124+
func TestMasterKey_EncryptedDataKey(t *testing.T) {
125+
key := &MasterKey{EncryptedKey: "some key"}
126+
assert.EqualValues(t, key.EncryptedKey, key.EncryptedDataKey())
127+
}
128+
129+
func TestMasterKey_SetEncryptedDataKey(t *testing.T) {
130+
key := &MasterKey{}
131+
data := []byte("some data")
132+
key.SetEncryptedDataKey(data)
133+
assert.EqualValues(t, data, key.EncryptedKey)
134+
}
135+
136+
func TestMasterKey_NeedsRotation(t *testing.T) {
137+
t.Run("false", func(t *testing.T) {
138+
k := &MasterKey{}
139+
k.CreationDate = time.Now().UTC()
140+
141+
assert.False(t, k.NeedsRotation())
142+
})
143+
144+
t.Run("true", func(t *testing.T) {
145+
k := &MasterKey{}
146+
k.CreationDate = time.Now().UTC().Add(-kmsTTL - 1)
147+
148+
assert.True(t, k.NeedsRotation())
149+
})
150+
}
151+
152+
func TestMasterKey_ToString(t *testing.T) {
153+
key := NewMasterKeyFromKeyID(dummyKey)
154+
155+
assert.Equal(t, dummyKey, key.ToString())
156+
}
157+
158+
func TestMasterKey_ToMap(t *testing.T) {
159+
key := NewMasterKeyFromKeyID(dummyKey)
160+
161+
data := []byte("some data")
162+
key.SetEncryptedDataKey(data)
163+
164+
res := key.ToMap()
165+
assert.Equal(t, dummyKey, res["key_id"])
166+
assert.Equal(t, key.CreationDate.UTC().Format(time.RFC3339), res["created_at"])
167+
assert.Equal(t, "some data", res["enc"])
168+
}
169+
170+
func createMockGRPCClient() (*grpc.ClientConn, error) {
171+
return grpc.DialContext(
172+
context.Background(),
173+
"bufnet",
174+
grpc.WithContextDialer(bufDialer),
175+
grpc.WithTransportCredentials(insecure.NewCredentials()),
176+
)
177+
}

0 commit comments

Comments
 (0)