Improve routing table association management (#1636)

* Improve routing table association management

* Simplify routing table association validation logic

* Refactor routing table association

* Improve function name

Author: hebelsan

pkg/aws/client/client.go

@@ -2085,6 +2085,38 @@ func (c *Client) DeleteRouteTableAssociation(ctx context.Context, associationId
 	return ignoreNotFound(err)
 }
 
+// GetRouteTableAssociationIDs gets all route table associations for the given vpc and subnet IDs.
+// If no subnet IDs are given, all association IDs for the VPC are returned.
+func (c *Client) GetRouteTableAssociationIDs(ctx context.Context, vpc string, subnetIDs []string) ([]string, error) {
+	var associationIDs []string
+	input := &ec2.DescribeRouteTablesInput{
+		Filters: []ec2types.Filter{
+			{
+				Name:   aws.String("vpc-id"),
+				Values: []string{vpc},
+			},
+			{
+				Name:   aws.String("association.subnet-id"),
+				Values: subnetIDs,
+			},
+		},
+	}
+
+	routeTablesOutput, err := c.EC2.DescribeRouteTables(ctx, input)
+	if ignoreNotFound(err) != nil {
+		return nil, err
+	}
+
+	for _, routeTable := range routeTablesOutput.RouteTables {
+		for _, assoc := range routeTable.Associations {
+			if assoc.RouteTableAssociationId != nil {
+				associationIDs = append(associationIDs, aws.ToString(assoc.RouteTableAssociationId))
+			}
+		}
+	}
+	return associationIDs, nil
+}
+
 // CreateIAMRole creates an IAM role resource.
 func (c *Client) CreateIAMRole(ctx context.Context, role *IAMRole) (*IAMRole, error) {
 	input := &iam.CreateRoleInput{

pkg/aws/client/mock/mocks.go

@@ -146,6 +146,7 @@ type Interface interface {
 	// Route table associations
 	CreateRouteTableAssociation(ctx context.Context, routeTableId, subnetId string) (associationId *string, err error)
 	DeleteRouteTableAssociation(ctx context.Context, associationId string) error
+	GetRouteTableAssociationIDs(ctx context.Context, vpc string, subnetIDs []string) ([]string, error)
 
 	// Elastic IP
 	CreateElasticIP(ctx context.Context, eip *ElasticIP) (*ElasticIP, error)

pkg/aws/client/types.go

@@ -1479,18 +1479,70 @@ func (c *FlowContext) deletePrivateRoutingTable(zoneName string) flow.TaskFn {
 	}
 }
 
+func (c *FlowContext) routingAssociationSpecs() []routeTableAssociationSpec {
+	return []routeTableAssociationSpec{
+		{IdentifierZoneSubnetPublic, IdentifierZoneSubnetPublicRouteTableAssoc, false},
+		{IdentifierZoneSubnetPrivate, IdentifierZoneSubnetPrivateRouteTableAssoc, true},
+		{IdentifierZoneSubnetWorkers, IdentifierZoneSubnetWorkersRouteTableAssoc, true},
+	}
+}
+
+// validateAndPruneRoutingTableAssocState checks whether the routing table associations stored in the state
+// still exist in AWS. If not, it removes them from the state.
+func (c *FlowContext) validateAndPruneRoutingTableAssocState(ctx context.Context, zoneName string, specs []routeTableAssociationSpec) error {
+	child := c.getSubnetZoneChild(zoneName)
+	log := LogFromContext(ctx)
+
+	// should validate only if at least one association ID is present in state
+	if !hasRouteTableAssociationInState(child.Get, specs) {
+		return nil
+	}
+
+	subnetIDs := make([]string, 0, len(specs))
+	for _, spec := range specs {
+		id := child.Get(spec.subnetKey)
+		if id == nil {
+			return fmt.Errorf("missing subnet id for key %s", spec.subnetKey)
+		}
+		subnetIDs = append(subnetIDs, *id)
+	}
+
+	vpc := c.state.Get(IdentifierVPC)
+	if vpc == nil {
+		return fmt.Errorf("VPC ID not found in state")
+	}
+
+	routeTableAssociations, err := c.client.GetRouteTableAssociationIDs(ctx, *vpc, subnetIDs)
+	if err != nil {
+		return err
+	}
+
+	for _, spec := range specs {
+		if assocID := child.Get(spec.assocKey); assocID != nil && !slices.Contains(routeTableAssociations, *assocID) {
+			log.Info("route table association not found in AWS, removing from state",
+				"AssociationID", *assocID)
+			child.Delete(spec.assocKey)
+		}
+	}
+	return nil
+}
+
 func (c *FlowContext) ensureRoutingTableAssociations(zoneName string) flow.TaskFn {
 	return func(ctx context.Context) error {
-		if err := c.ensureZoneRoutingTableAssociation(ctx, zoneName, false,
-			IdentifierZoneSubnetPublic, IdentifierZoneSubnetPublicRouteTableAssoc); err != nil {
+		specs := c.routingAssociationSpecs()
+
+		err := c.validateAndPruneRoutingTableAssocState(ctx, zoneName, specs)
+		if err != nil {
 			return err
 		}
-		if err := c.ensureZoneRoutingTableAssociation(ctx, zoneName, true,
-			IdentifierZoneSubnetPrivate, IdentifierZoneSubnetPrivateRouteTableAssoc); err != nil {
-			return err
+
+		for _, spec := range specs {
+			err := c.ensureZoneRoutingTableAssociation(ctx, zoneName, spec.zoneRouteTable, spec.subnetKey, spec.assocKey)
+			if err != nil {
+				return err
+			}
 		}
-		return c.ensureZoneRoutingTableAssociation(ctx, zoneName, true,
-			IdentifierZoneSubnetWorkers, IdentifierZoneSubnetWorkersRouteTableAssoc)
+		return nil
 	}
 }
 
@@ -1570,16 +1622,16 @@ func (c *FlowContext) ensureVPCEndpointZoneRoutingTableAssociation(ctx context.C
 
 func (c *FlowContext) deleteRoutingTableAssociations(zoneName string) flow.TaskFn {
 	return func(ctx context.Context) error {
-		if err := c.deleteZoneRoutingTableAssociation(ctx, zoneName, false,
-			IdentifierZoneSubnetPublic, IdentifierZoneSubnetPublicRouteTableAssoc); err != nil {
-			return err
-		}
-		if err := c.deleteZoneRoutingTableAssociation(ctx, zoneName, true,
-			IdentifierZoneSubnetPrivate, IdentifierZoneSubnetPrivateRouteTableAssoc); err != nil {
-			return err
+		specs := c.routingAssociationSpecs()
+
+		for _, spec := range specs {
+			err := c.deleteZoneRoutingTableAssociation(ctx, zoneName, spec.zoneRouteTable, spec.subnetKey, spec.assocKey)
+			if err != nil {
+				return err
+			}
 		}
-		return c.deleteZoneRoutingTableAssociation(ctx, zoneName, true,
-			IdentifierZoneSubnetWorkers, IdentifierZoneSubnetWorkersRouteTableAssoc)
+
+		return nil
 	}
 }
 

pkg/controller/infrastructure/infraflow/reconcile.go

@@ -329,3 +329,20 @@ func BuildInfrastructureStatus(
 
 	return status
 }
+
+// routeTableAssociationSpec contains the specification to associate a route table with a subnet.
+type routeTableAssociationSpec struct {
+	subnetKey      string
+	assocKey       string
+	zoneRouteTable bool
+}
+
+// hasRouteTableAssociationInState checks if any of the route table associations specified in specs exist in the state.
+func hasRouteTableAssociationInState(getAssociationID func(string) *string, specs []routeTableAssociationSpec) bool {
+	for _, spec := range specs {
+		if assocID := getAssociationID(spec.assocKey); assocID != nil {
+			return true
+		}
+	}
+	return false
+}