Fix restoration of infra with migrated network layout

pkg/apis/azure/helper/scheme.go

@@ -5,6 +5,7 @@
 package helper
 
 import (
+	"encoding/json"
 	"fmt"
 
 	"github.com/gardener/gardener/extensions/pkg/controller"
@@ -130,11 +131,35 @@ func DNSRecordConfigFromDNSRecord(dnsRecord *extensionsv1alpha1.DNSRecord) (api.
 	return dnsRecordConfig, nil
 }
 
+// HasFlowState returns true if the group version of the State field in the provided
+// `extensionsv1alpha1.InfrastructureStatus` is azure.provider.extensions.gardener.cloud/v1alpha1.
+func HasFlowState(status extensionsv1alpha1.InfrastructureStatus) (bool, error) {
+	if status.State == nil {
+		return false, nil
+	}
+
+	flowState := runtime.TypeMeta{}
+	stateJson, err := status.State.MarshalJSON()
+	if err != nil {
+		return false, err
+	}
+
+	if err := json.Unmarshal(stateJson, &flowState); err != nil {
+		return false, err
+	}
+
+	if flowState.GroupVersionKind().GroupVersion() == apiv1alpha1.SchemeGroupVersion {
+		return true, nil
+	}
+
+	return false, nil
+}
+
 // InfrastructureStateFromRaw extracts the state from the Infrastructure. If no state was available, it returns a "zero" value InfrastructureState object.
 func InfrastructureStateFromRaw(raw *runtime.RawExtension) (*api.InfrastructureState, error) {
 	state := &api.InfrastructureState{}
 	if raw != nil && raw.Raw != nil {
-		if _, _, err := lenientDecoder.Decode(raw.Raw, nil, state); err != nil {
+		if _, _, err := decoder.Decode(raw.Raw, nil, state); err != nil {
 			return nil, err
 		}
 	}

pkg/apis/azure/helper/scheme_test.go

@@ -0,0 +1,31 @@
+// SPDX-FileCopyrightText: 2024 SAP SE or an SAP affiliate company and Gardener contributors
+//
+// SPDX-License-Identifier: Apache-2.0
+
+package helper_test
+
+import (
+	. "github.com/onsi/ginkgo/v2"
+	"k8s.io/apimachinery/pkg/runtime"
+
+	. "github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure/helper"
+	extensionsv1alpha1 "github.com/gardener/gardener/pkg/apis/extensions/v1alpha1"
+)
+
+var _ = Describe("Scheme", func() {
+	DescribeTable("#HasFlowState",
+		func(state *runtime.RawExtension, expectedHasFlowState, expectedErr bool) {
+			infraStatus := extensionsv1alpha1.InfrastructureStatus{
+				DefaultStatus: extensionsv1alpha1.DefaultStatus{
+					State: state,
+				},
+			}
+			hasFlowState, err := HasFlowState(infraStatus)
+			expectResults(hasFlowState, expectedHasFlowState, err, expectedErr)
+		},
+		Entry("when state is nil", nil, false, false),
+		Entry("when state is invalid json", &runtime.RawExtension{Raw: []byte(`foo`)}, false, true),
+		Entry("when state is not in 'azure.provider.extensions.gardener.cloud/v1alpha1' group version", &runtime.RawExtension{Raw: []byte(`{"apiVersion":"foo.bar/v1alpha1","kind":"InfrastructureState"}`)}, false, false),
+		Entry("when state is in 'azure.provider.extensions.gardener.cloud/v1alpha1' group version", &runtime.RawExtension{Raw: []byte(`{"apiVersion":"azure.provider.extensions.gardener.cloud/v1alpha1","kind":"InfrastructureState"}`)}, true, false),
+	)
+})

pkg/controller/infrastructure/actuator_helper.go

@@ -6,15 +6,11 @@ package infrastructure
 
 import (
 	"context"
-	"encoding/json"
 	"strconv"
 
 	"github.com/gardener/gardener/extensions/pkg/terraformer"
-	extensionsv1alpha1 "github.com/gardener/gardener/pkg/apis/extensions/v1alpha1"
 	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"k8s.io/apimachinery/pkg/runtime"
 
-	"github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure/v1alpha1"
 	azuretypes "github.com/gardener/gardener-extension-provider-azure/pkg/azure"
 )
 
@@ -29,28 +25,6 @@ func CleanupTerraformerResources(ctx context.Context, tf terraformer.Terraformer
 	return tf.RemoveTerraformerFinalizerFromConfig(ctx)
 }
 
-func hasFlowState(status extensionsv1alpha1.InfrastructureStatus) (bool, error) {
-	if status.State == nil {
-		return false, nil
-	}
-
-	flowState := runtime.TypeMeta{}
-	stateJson, err := status.State.MarshalJSON()
-	if err != nil {
-		return false, err
-	}
-
-	if err := json.Unmarshal(stateJson, &flowState); err != nil {
-		return false, err
-	}
-
-	if flowState.GroupVersionKind().GroupVersion() == v1alpha1.SchemeGroupVersion {
-		return true, nil
-	}
-
-	return false, nil
-}
-
 // GetFlowAnnotationValue returns the boolean value of the expected flow annotation. Returns false if the annotation was not found, if it couldn't be converted to bool,
 // or had a "false" value.
 func GetFlowAnnotationValue(o v1.Object) bool {

pkg/controller/infrastructure/flow_reconciler.go

@@ -46,7 +46,7 @@ func (f *FlowReconciler) Reconcile(ctx context.Context, infra *extensionsv1alpha
 		infraState *azure.InfrastructureState
 		err        error
 	)
-	fsOk, err := hasFlowState(infra.Status)
+	fsOk, err := helper.HasFlowState(infra.Status)
 	if err != nil {
 		return err
 	}

pkg/controller/infrastructure/infraflow/ensurer.go

@@ -19,6 +19,7 @@ import (
 
 	"github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure/helper"
 	"github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure/v1alpha1"
+	azuretypes "github.com/gardener/gardener-extension-provider-azure/pkg/azure"
 	"github.com/gardener/gardener-extension-provider-azure/pkg/azure/client"
 	"github.com/gardener/gardener-extension-provider-azure/pkg/controller/infrastructure/infraflow/shared"
 	"github.com/gardener/gardener-extension-provider-azure/pkg/internal/infrastructure"
@@ -728,20 +729,16 @@ func (fctx *FlowContext) GetInfrastructureState() *runtime.RawExtension {
 		ManagedItems: fctx.inventory.ToList(),
 	}
 
-	return &runtime.RawExtension{
-		Object: state,
+	if migratedZone, ok := fctx.infra.Annotations[azuretypes.NetworkLayoutZoneMigrationAnnotation]; ok {
+		if state.Data == nil {
+			state.Data = make(map[string]string)
+		}
+		state.Data[azuretypes.NetworkLayoutZoneMigrationAnnotation] = migratedZone
 	}
-}
 
-func (fctx *FlowContext) enrichStatusWithIdentity(_ context.Context, status *v1alpha1.InfrastructureStatus) error {
-	if identity := fctx.cfg.Identity; identity != nil {
-		status.Identity = &v1alpha1.IdentityStatus{
-			ID:        *fctx.whiteboard.Get(KeyManagedIdentityId),
-			ClientID:  *fctx.whiteboard.Get(KeyManagedIdentityClientId),
-			ACRAccess: identity.ACRAccess != nil && *identity.ACRAccess,
-		}
+	return &runtime.RawExtension{
+		Object: state,
 	}
-	return nil
 }
 
 // DeleteResourceGroup deletes the shoot's resource group.

pkg/controller/infrastructure/strategyselector.go

@@ -8,6 +8,7 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure/helper"
 	"github.com/gardener/gardener/extensions/pkg/controller"
 	extensionsv1alpha1 "github.com/gardener/gardener/pkg/apis/extensions/v1alpha1"
 	"github.com/gardener/gardener/pkg/extensions"
@@ -60,7 +61,7 @@ type SelectorFunc func(*extensionsv1alpha1.Infrastructure, *extensions.Cluster)
 
 // OnReconcile returns true if the operation should use the Flow for the given cluster.
 func OnReconcile(infra *extensionsv1alpha1.Infrastructure, _ *extensions.Cluster) (bool, error) {
-	hasState, err := hasFlowState(infra.Status)
+	hasState, err := helper.HasFlowState(infra.Status)
 	if err != nil {
 		return false, err
 	}
@@ -69,7 +70,7 @@ func OnReconcile(infra *extensionsv1alpha1.Infrastructure, _ *extensions.Cluster
 
 // OnDelete returns true if the operation should use the Flow deletion for the given cluster.
 func OnDelete(infra *extensionsv1alpha1.Infrastructure, _ *extensions.Cluster) (bool, error) {
-	return hasFlowState(infra.Status)
+	return helper.HasFlowState(infra.Status)
 }
 
 // OnRestore decides the reconciler used on migration.

pkg/webhook/infrastructure/layout.go

@@ -6,16 +6,19 @@ package infrastructure
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 
 	extensionswebhook "github.com/gardener/gardener/extensions/pkg/webhook"
+	v1beta1constants "github.com/gardener/gardener/pkg/apis/core/v1beta1/constants"
 	extensionsv1alpha1 "github.com/gardener/gardener/pkg/apis/extensions/v1alpha1"
 	"github.com/go-logr/logr"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	"github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure"
 	"github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure/helper"
 	azuretypes "github.com/gardener/gardener-extension-provider-azure/pkg/azure"
+	"github.com/gardener/gardener-extension-provider-azure/pkg/internal/infrastructure"
 )
 
 // MutateFunc is a function that can perform a mutation on Infrastructure objects.
@@ -109,6 +112,11 @@ func NetworkLayoutMigrationMutate(_ context.Context, logger logr.Logger, newInfr
 		return nil
 	}
 
+	// If the current operation is restore, then mutation might be necessary as the zone migration annotation is not preserved during control plane migration.
+	if operation, ok := newInfra.Annotations[v1beta1constants.GardenerOperation]; ok && operation == v1beta1constants.GardenerOperationRestore {
+		return addMigratedZoneAnnotationDuringRestore(newInfra)
+	}
+
 	// if the old configuration is not using zones or if it is already using a multi-subnet layout, no mutation is necessary.
 	if !oldProviderCfg.Zoned || len(oldProviderCfg.Networks.Zones) > 0 {
 		return nil
@@ -142,3 +150,45 @@ func NetworkLayoutMigrationMutate(_ context.Context, logger logr.Logger, newInfr
 
 	return nil
 }
+
+func addMigratedZoneAnnotationDuringRestore(infra *extensionsv1alpha1.Infrastructure) error {
+	if infra.Status.State == nil || infra.Status.State.Raw == nil {
+		return nil
+	}
+
+	fsOk, err := helper.HasFlowState(infra.Status)
+	if err != nil {
+		return err
+	}
+
+	if fsOk {
+		infraState, err := helper.InfrastructureStateFromRaw(infra.Status.State)
+		if err != nil {
+			return err
+		}
+
+		if migratedZone, ok := infraState.Data[azuretypes.NetworkLayoutZoneMigrationAnnotation]; ok {
+			infra.Annotations[azuretypes.NetworkLayoutZoneMigrationAnnotation] = migratedZone
+		}
+
+		return nil
+	}
+
+	infraState := &infrastructure.InfrastructureState{}
+	if err := json.Unmarshal(infra.Status.State.Raw, infraState); err != nil {
+		return err
+	}
+
+	infrastructureStatus, err := helper.InfrastructureStatusFromRaw(infraState.SavedProviderStatus)
+	if err != nil {
+		return err
+	}
+
+	for _, subnet := range infrastructureStatus.Networks.Subnets {
+		if subnet.Migrated && subnet.Zone != nil {
+			infra.Annotations[azuretypes.NetworkLayoutZoneMigrationAnnotation] = *subnet.Zone
+			break
+		}
+	}
+	return nil
+}

pkg/webhook/infrastructure/layout_test.go

@@ -18,8 +18,10 @@ import (
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/utils/ptr"
 
+	"github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure/helper"
 	azurev1alpha1 "github.com/gardener/gardener-extension-provider-azure/pkg/apis/azure/v1alpha1"
 	"github.com/gardener/gardener-extension-provider-azure/pkg/azure"
+	"github.com/gardener/gardener-extension-provider-azure/pkg/internal/infrastructure"
 )
 
 const (
@@ -117,6 +119,72 @@ var _ = Describe("Mutate", func() {
 				_, ok := getLayoutMigrationAnnotation(newInfra)
 				Expect(ok).To(BeFalse())
 			})
+			It("should mutate the resource if the current 'gardener.cloud/operation' annotation is 'restore' and has flow state", func() {
+				newInfra := generateInfrastructureWithProviderConfig(zonesConfig, nil)
+				newInfra.Annotations = map[string]string{
+					"gardener.cloud/operation": "restore",
+				}
+
+				state := &azurev1alpha1.InfrastructureState{
+					TypeMeta: helper.InfrastructureStateTypeMeta,
+					Data: map[string]string{
+						azure.NetworkLayoutZoneMigrationAnnotation: "2",
+					},
+				}
+				marshalled, err := json.Marshal(state)
+				Expect(err).To(BeNil())
+				newInfra.Status.State = &runtime.RawExtension{Raw: marshalled}
+
+				err = mutator.Mutate(context.TODO(), newInfra, newInfra)
+				Expect(err).To(BeNil())
+				v, ok := getLayoutMigrationAnnotation(newInfra)
+				Expect(ok).To(BeTrue())
+				Expect(v).To(Equal("2"))
+			})
+			It("should mutate the resource if the current 'gardener.cloud/operation' annotation is restore and has terraform state", func() {
+				newInfra := generateInfrastructureWithProviderConfig(zonesConfig, nil)
+				newInfra.Annotations = map[string]string{
+					"gardener.cloud/operation": "restore",
+				}
+
+				status := &azurev1alpha1.InfrastructureStatus{
+					TypeMeta: metav1.TypeMeta{
+						APIVersion: azurev1alpha1.SchemeGroupVersion.String(),
+						Kind:       "InfrastructureStatus",
+					},
+					Networks: azurev1alpha1.NetworkStatus{
+						Subnets: []azurev1alpha1.Subnet{
+							{
+								Name:     "subnet-zone1",
+								Zone:     ptr.To("1"),
+								Migrated: false,
+							},
+							{
+								Name:     "subnet",
+								Zone:     ptr.To("2"),
+								Migrated: true,
+							},
+						},
+					},
+				}
+				marshalled, err := json.Marshal(status)
+				Expect(err).To(BeNil())
+
+				state := &infrastructure.InfrastructureState{
+					SavedProviderStatus: &runtime.RawExtension{
+						Raw: marshalled,
+					},
+				}
+				marshalledState, err := json.Marshal(state)
+				Expect(err).To(BeNil())
+				newInfra.Status.State = &runtime.RawExtension{Raw: marshalledState}
+
+				err = mutator.Mutate(context.TODO(), newInfra, newInfra)
+				Expect(err).To(BeNil())
+				v, ok := getLayoutMigrationAnnotation(newInfra)
+				Expect(ok).To(BeTrue())
+				Expect(v).To(Equal("2"))
+			})
 		})
 
 		Context("remove migration annotation", func() {