feat(kms): waiter for wrapping keys

Author: bahkauv70

services/kms/wait/wait.go

@@ -25,9 +25,17 @@ const (
 	StatusKeyVersionDestroyed           = "destroyed"
 )
 
+const (
+	StatusWrappingKeyActive              = "active"
+	StatusWrappingKeyKeyMaterialNotReady = "key_material_not_ready"
+	StatusWrappingKeyExpired             = "expired"
+	StatusWrappingKeyDeleting            = "deleting"
+)
+
 type ApiKmsClient interface {
 	GetKeyExecute(ctx context.Context, projectId string, regionId string, keyRingId string, keyId string) (*kms.Key, error)
 	GetVersionExecute(ctx context.Context, projectId string, regionId string, keyRingId string, keyId string, versionNumber int64) (*kms.Version, error)
+	GetWrappingKeyExecute(ctx context.Context, projectId string, regionId string, keyRingId string, wrappingKeyId string) (*kms.WrappingKey, error)
 }
 
 func CreateOrUpdateKeyWaitHandler(ctx context.Context, client ApiKmsClient, projectId, region, keyRingId, keyId string) *wait.AsyncActionHandler[kms.Key] {
@@ -115,3 +123,25 @@ func DisableKeyVersionWaitHandler(ctx context.Context, client ApiKmsClient, proj
 	handler.SetTimeout(10 * time.Minute)
 	return handler
 }
+
+func CreateWrappingKeyWaitHandler(ctx context.Context, client ApiKmsClient, projectId, region, keyRingId, wrappingKeyId string) *wait.AsyncActionHandler[kms.WrappingKey] {
+	handler := wait.New(func() (bool, *kms.WrappingKey, error) {
+		response, err := client.GetWrappingKeyExecute(ctx, projectId, region, keyRingId, wrappingKeyId)
+		if err != nil {
+			return false, nil, err
+		}
+
+		if state := response.State; state != nil {
+			switch *state {
+			case StatusWrappingKeyKeyMaterialNotReady:
+				return false, nil, nil
+			default:
+				return true, response, nil
+			}
+		}
+
+		return false, nil, nil
+	})
+	handler.SetTimeout(10 * time.Minute)
+	return handler
+}

services/kms/wait/wait_test.go

@@ -35,11 +35,26 @@ type versionResponse struct {
 	err     error
 }
 
+type wrappingKeyResponse struct {
+	wrappingKey *kms.WrappingKey
+	err         error
+}
+
 type apiKmsMocked struct {
-	idxKeyResponse     int
-	idxVersionResponse int
-	keyResponses       []keyResponse
-	versionResponses   []versionResponse
+	idxKeyResponse         int
+	idxVersionResponse     int
+	idxWrappingKeyResponse int
+	keyResponses           []keyResponse
+	versionResponses       []versionResponse
+	wrappingKeyResponses   []wrappingKeyResponse
+}
+
+// GetWrappingKeyExecute implements ApiKmsClient.
+func (a *apiKmsMocked) GetWrappingKeyExecute(_ context.Context, _, _, _, _ string) (*kms.WrappingKey, error) {
+	resp := a.wrappingKeyResponses[a.idxWrappingKeyResponse]
+	a.idxWrappingKeyResponse++
+	a.idxWrappingKeyResponse %= len(a.wrappingKeyResponses)
+	return resp.wrappingKey, resp.err
 }
 
 // GetVersionExecute implements ApiKmsClient.
@@ -74,6 +89,22 @@ func fixtureKey(state string) *kms.Key {
 	}
 }
 
+func fixtureWrappingKey(state string) *kms.WrappingKey {
+	return &kms.WrappingKey{
+		Algorithm:   kms.WRAPPINGALGORITHM__2048_OAEP_SHA256.Ptr(),
+		Backend:     kms.BACKEND_SOFTWARE.Ptr(),
+		CreatedAt:   &testDate,
+		Description: utils.Ptr("test-description"),
+		DisplayName: utils.Ptr("test-displayname"),
+		Id:          &testKeyId,
+		KeyRingId:   &testKeyRingId,
+		Purpose:     kms.WRAPPINGPURPOSE_SYMMETRIC_KEY.Ptr(),
+		State:       &state,
+		ExpiresAt:   &testDate,
+		PublicKey:   &testPublicKey,
+	}
+}
+
 func fixtureVersion(version int, disabled bool, state string) *kms.Version {
 	return &kms.Version{
 		CreatedAt:   &testDate,
@@ -414,3 +445,81 @@ func TestDisableKeyVersionWaitHandler(t *testing.T) {
 		})
 	}
 }
+
+func TestCreateWrappingWaitHandler(t *testing.T) {
+	tests := []struct {
+		name      string
+		responses []wrappingKeyResponse
+		want      *kms.WrappingKey
+		wantErr   bool
+	}{
+		{
+			"create succeeded immediately",
+			[]wrappingKeyResponse{
+				{fixtureWrappingKey(StatusWrappingKeyActive), nil},
+			},
+			fixtureWrappingKey(StatusWrappingKeyActive),
+			false,
+		},
+		{
+			"create succeeded delayed",
+			[]wrappingKeyResponse{
+				{fixtureWrappingKey(StatusWrappingKeyKeyMaterialNotReady), nil},
+				{fixtureWrappingKey(StatusWrappingKeyKeyMaterialNotReady), nil},
+				{fixtureWrappingKey(StatusWrappingKeyKeyMaterialNotReady), nil},
+				{fixtureWrappingKey(StatusWrappingKeyActive), nil},
+			},
+			fixtureWrappingKey(StatusWrappingKeyActive),
+			false,
+		},
+		{
+			"create failed delayed",
+			[]wrappingKeyResponse{
+				{fixtureWrappingKey(StatusWrappingKeyKeyMaterialNotReady), nil},
+				{fixtureWrappingKey(StatusWrappingKeyKeyMaterialNotReady), nil},
+				{fixtureWrappingKey(StatusWrappingKeyKeyMaterialNotReady), nil},
+				{fixtureWrappingKey(StatusWrappingKeyDeleting), nil},
+			},
+			fixtureWrappingKey(StatusWrappingKeyDeleting),
+			false,
+		},
+		{
+			"timeout",
+			[]wrappingKeyResponse{
+				{fixtureWrappingKey(StatusWrappingKeyKeyMaterialNotReady), nil},
+			},
+			nil,
+			true,
+		},
+		{
+			"broken state",
+			[]wrappingKeyResponse{
+				{fixtureWrappingKey("bogus"), nil},
+			},
+			fixtureWrappingKey("bogus"),
+			false,
+		},
+		// no special update tests needed as the states are the same
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx := context.Background()
+			client := &apiKmsMocked{
+				wrappingKeyResponses: tt.responses,
+			}
+
+			handler := CreateWrappingKeyWaitHandler(ctx, client, testProject, testRegion, testKeyRingId, testKeyId)
+			got, err := handler.SetTimeout(1 * time.Second).
+				SetThrottle(250 * time.Millisecond).
+				WaitWithContext(ctx)
+
+			if (err != nil) != tt.wantErr {
+				t.Fatalf("unexpected error response. want %v but got %v ", tt.wantErr, err)
+			}
+
+			if diff := cmp.Diff(tt.want, got); diff != "" {
+				t.Errorf("differing key %s", diff)
+			}
+		})
+	}
+}