Retry CloudKMS GetPublicKey.

On HSM keys are not generated instantly and the GetPublicKey fails
with a FailedPrecondition error. This change will retry GetPublicKey
if this happens.
This commit is contained in:
Mariano Cano 2020-02-21 12:36:29 -08:00
parent 806abb6232
commit 1138cc4207
2 changed files with 65 additions and 6 deletions

View file

@ -175,6 +175,7 @@ func TestCloudKMS_CreateKey(t *testing.T) {
t.Fatal(err)
}
var retries int
type fields struct {
client KeyManagementClient
}
@ -236,6 +237,24 @@ func TestCloudKMS_CreateKey(t *testing.T) {
}},
args{&apiv1.CreateKeyRequest{Name: keyName, ProtectionLevel: apiv1.HSM, SignatureAlgorithm: apiv1.ECDSAWithSHA256}},
&apiv1.CreateKeyResponse{Name: keyName + "/cryptoKeyVersions/2", PublicKey: pk, CreateSignerRequest: apiv1.CreateSignerRequest{SigningKey: keyName + "/cryptoKeyVersions/2"}}, false},
{"ok with retries", fields{
&MockClient{
getKeyRing: func(_ context.Context, _ *kmspb.GetKeyRingRequest, _ ...gax.CallOption) (*kmspb.KeyRing, error) {
return &kmspb.KeyRing{}, nil
},
createCryptoKey: func(_ context.Context, _ *kmspb.CreateCryptoKeyRequest, _ ...gax.CallOption) (*kmspb.CryptoKey, error) {
return &kmspb.CryptoKey{Name: keyName}, nil
},
getPublicKey: func(_ context.Context, _ *kmspb.GetPublicKeyRequest, _ ...gax.CallOption) (*kmspb.PublicKey, error) {
if retries != 2 {
retries++
return nil, status.Error(codes.FailedPrecondition, "key is not enabled, current state is: PENDING_GENERATION")
}
return &kmspb.PublicKey{Pem: string(pemBytes)}, nil
},
}},
args{&apiv1.CreateKeyRequest{Name: keyName, ProtectionLevel: apiv1.HSM, SignatureAlgorithm: apiv1.ECDSAWithSHA256}},
&apiv1.CreateKeyResponse{Name: keyName + "/cryptoKeyVersions/1", PublicKey: pk, CreateSignerRequest: apiv1.CreateSignerRequest{SigningKey: keyName + "/cryptoKeyVersions/1"}}, false},
{"fail name", fields{&MockClient{}}, args{&apiv1.CreateKeyRequest{}}, nil, true},
{"fail protection level", fields{&MockClient{}}, args{&apiv1.CreateKeyRequest{Name: keyName, ProtectionLevel: apiv1.ProtectionLevel(100)}}, nil, true},
{"fail signature algorithm", fields{&MockClient{}}, args{&apiv1.CreateKeyRequest{Name: keyName, ProtectionLevel: apiv1.Software, SignatureAlgorithm: apiv1.SignatureAlgorithm(100)}}, nil, true},
@ -322,6 +341,7 @@ func TestCloudKMS_GetPublicKey(t *testing.T) {
t.Fatal(err)
}
var retries int
type fields struct {
client KeyManagementClient
}
@ -342,6 +362,17 @@ func TestCloudKMS_GetPublicKey(t *testing.T) {
},
}},
args{&apiv1.GetPublicKeyRequest{Name: keyName}}, pk, false},
{"ok with retries", fields{
&MockClient{
getPublicKey: func(_ context.Context, _ *kmspb.GetPublicKeyRequest, _ ...gax.CallOption) (*kmspb.PublicKey, error) {
if retries != 2 {
retries++
return nil, status.Error(codes.FailedPrecondition, "key is not enabled, current state is: PENDING_GENERATION")
}
return &kmspb.PublicKey{Pem: string(pemBytes)}, nil
},
}},
args{&apiv1.GetPublicKeyRequest{Name: keyName}}, pk, false},
{"fail name", fields{&MockClient{}}, args{&apiv1.GetPublicKeyRequest{}}, nil, true},
{"fail get public key", fields{
&MockClient{