Rewrite GetClusterSubnets() using EC2 specific API

internal/alb/lb/loadbalancer.go

@@ -380,7 +380,6 @@ func (controller *defaultController) resolveSubnets(ctx context.Context, scheme
 }
 
 func (controller *defaultController) clusterSubnets(ctx context.Context, scheme string) ([]string, error) {
-	var subnetIds []string
 	var useableSubnets []*ec2.Subnet
 	var out []string
 	var key string
@@ -393,27 +392,12 @@ func (controller *defaultController) clusterSubnets(ctx context.Context, scheme
 		return nil, fmt.Errorf("invalid scheme [%s]", scheme)
 	}
 
-	clusterSubnets, err := controller.cloud.GetClusterSubnets()
+	clusterSubnets, err := controller.cloud.GetClusterSubnets(key)
 	if err != nil {
-		return nil, fmt.Errorf("failed to get AWS tags. Error: %s", err.Error())
+		return nil, fmt.Errorf("unable to fetch subnets. Error: %s", err.Error())
 	}
 
-	for arn, subnetTags := range clusterSubnets {
-		for _, tag := range subnetTags {
-			if aws.StringValue(tag.Key) == key {
-				p := strings.Split(arn, "/")
-				subnetID := p[len(p)-1]
-				subnetIds = append(subnetIds, subnetID)
-			}
-		}
-	}
-
-	o, err := controller.cloud.GetSubnetsByNameOrID(ctx, subnetIds)
-	if err != nil {
-		return nil, fmt.Errorf("unable to fetch subnets due to %v", err)
-	}
-
-	for _, subnet := range o {
+	for _, subnet := range clusterSubnets {
 		if subnetIsUsable(subnet, useableSubnets) {
 			useableSubnets = append(useableSubnets, subnet)
 			out = append(out, aws.StringValue(subnet.SubnetId))

internal/aws/ec2.go

@@ -44,6 +44,9 @@ type EC2API interface {
 	// GetSecurityGroupsByName retrieves securityGroups by securityGroupName(SecurityGroup names within vpc are unique)
 	GetSecurityGroupsByName(context.Context, []string) ([]*ec2.SecurityGroup, error)
 
+	// GetClusterSubnets retrieves the subnets associated with the cluster, by matching tags
+	GetClusterSubnets(string) ([]*ec2.Subnet, error)
+
 	// DeleteSecurityGroupByID delete securityGroup by securityGroupID
 	DeleteSecurityGroupByID(context.Context, string) error
 
@@ -164,6 +167,26 @@ func (c *Cloud) GetSubnetsByNameOrID(ctx context.Context, nameOrIDs []string) (s
 	return
 }
 
+func (c *Cloud) GetClusterSubnets(tagSubnetType string) ([]*ec2.Subnet, error) {
+	in := &ec2.DescribeSubnetsInput{Filters: []*ec2.Filter{
+		{
+			Name:   aws.String("tag:kubernetes.io/cluster/" + c.clusterName),
+			Values: aws.StringSlice([]string{"owned", "shared"}),
+		},
+		{
+			Name:   aws.String("tag:" + tagSubnetType),
+			Values: aws.StringSlice([]string{"", "1"}),
+		},
+	}}
+
+	result, err := c.describeSubnetsHelper(in)
+	if err != nil {
+		return nil, err
+	}
+
+	return result, nil
+}
+
 func (c *Cloud) GetSecurityGroupsByName(ctx context.Context, names []string) (groups []*ec2.SecurityGroup, err error) {
 	in := &ec2.DescribeSecurityGroupsInput{Filters: []*ec2.Filter{
 		{
@@ -272,6 +295,16 @@ func (c *Cloud) describeSecurityGroupsHelper(params *ec2.DescribeSecurityGroupsI
 	return results, err
 }
 
+// describeSubnetsHelper is a helper to handle pagination for DescribeSubnets API call
+func (c *Cloud) describeSubnetsHelper(params *ec2.DescribeSubnetsInput) (result []*ec2.Subnet, err error) {
+	err = c.ec2.DescribeSubnetsPages(params, func(output *ec2.DescribeSubnetsOutput, _ bool) bool {
+		result = append(result, output.Subnets...)
+		return true
+	})
+
+	return result, err
+}
+
 func (c *Cloud) describeInstancesHelper(params *ec2.DescribeInstancesInput) (result []*ec2.Reservation, err error) {
 	err = c.ec2.DescribeInstancesPages(params, func(output *ec2.DescribeInstancesOutput, _ bool) bool {
 		result = append(result, output.Reservations...)

internal/aws/ec2_test.go

@@ -5,6 +5,8 @@ import (
 	"errors"
 	"testing"
 
+	"github.com/stretchr/testify/mock"
+
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/kubernetes-sigs/aws-alb-ingress-controller/mocks"
@@ -147,3 +149,122 @@ func TestCloud_RevokeSecurityGroupIngressWithContext(t *testing.T) {
 		svc.AssertExpectations(t)
 	})
 }
+
+func TestCloud_GetClusterSubnets(t *testing.T) {
+	clusterName := "clusterName"
+	internalSubnet1 := &ec2.Subnet{
+		SubnetId: aws.String("arn:aws:ec2:region:account-id:subnet/subnet-id1"),
+		Tags: []*ec2.Tag{
+			{
+				Key:   aws.String("kubernetes.io/cluster/" + clusterName),
+				Value: aws.String("owned"),
+			},
+			{
+				Key:   aws.String("kubernetes.io/role/internal-elb"),
+				Value: aws.String("1"),
+			},
+		},
+	}
+	internalSubnet2 := &ec2.Subnet{
+		SubnetId: aws.String("arn:aws:ec2:region:account-id:subnet/subnet-id2"),
+		Tags: []*ec2.Tag{
+			{
+				Key:   aws.String("kubernetes.io/cluster/" + clusterName),
+				Value: aws.String("owned"),
+			},
+			{
+				Key:   aws.String("kubernetes.io/role/internal-elb"),
+				Value: aws.String(""),
+			},
+		},
+	}
+	publicSubnet := &ec2.Subnet{
+		SubnetId: aws.String("arn:aws:ec2:region:account-id:subnet/subnet-id3"),
+		Tags: []*ec2.Tag{
+			{
+				Key:   aws.String("kubernetes.io/cluster/" + clusterName),
+				Value: aws.String("shared"),
+			},
+			{
+				Key:   aws.String("kubernetes.io/role/elb"),
+				Value: aws.String("1"),
+			},
+		},
+	}
+
+	for _, tc := range []struct {
+		Name                  string
+		DescribeSubnetsOutput *ec2.DescribeSubnetsOutput
+		DescribeSubnetsError  error
+		TagSubnetType         string
+		ExpectedResult        []*ec2.Subnet
+		ExpectedError         error
+	}{
+		{
+			Name:          "No subnets returned",
+			TagSubnetType: TagNameSubnetInternalELB,
+			DescribeSubnetsOutput: &ec2.DescribeSubnetsOutput{
+				NextToken: nil,
+				Subnets:   []*ec2.Subnet{},
+			},
+		},
+		{
+			Name:          "Two internal subnets returned",
+			TagSubnetType: TagNameSubnetInternalELB,
+			DescribeSubnetsOutput: &ec2.DescribeSubnetsOutput{
+				NextToken: nil,
+				Subnets:   []*ec2.Subnet{internalSubnet1, internalSubnet2},
+			},
+			ExpectedResult: []*ec2.Subnet{internalSubnet1, internalSubnet2},
+		},
+		{
+			Name:          "One public subnet returned",
+			TagSubnetType: TagNameSubnetPublicELB,
+			DescribeSubnetsOutput: &ec2.DescribeSubnetsOutput{
+				NextToken: nil,
+				Subnets:   []*ec2.Subnet{publicSubnet},
+			},
+			ExpectedResult: []*ec2.Subnet{publicSubnet},
+		},
+		{
+			Name:          "Error from API call",
+			TagSubnetType: TagNameSubnetPublicELB,
+			DescribeSubnetsOutput: &ec2.DescribeSubnetsOutput{
+				NextToken: nil,
+				Subnets:   []*ec2.Subnet{},
+			},
+			DescribeSubnetsError: errors.New("Some API error"),
+			ExpectedError:        errors.New("Some API error"),
+		},
+	} {
+		t.Run(tc.Name, func(t *testing.T) {
+			svc := &mocks.EC2API{}
+
+			svc.On("DescribeSubnetsPages",
+				&ec2.DescribeSubnetsInput{Filters: []*ec2.Filter{
+					{
+						Name:   aws.String("tag:kubernetes.io/cluster/" + clusterName),
+						Values: aws.StringSlice([]string{"owned", "shared"}),
+					},
+					{
+						Name:   aws.String("tag:" + tc.TagSubnetType),
+						Values: aws.StringSlice([]string{"", "1"}),
+					}},
+				},
+				mock.AnythingOfType("func(*ec2.DescribeSubnetsOutput, bool) bool"),
+			).Return(tc.DescribeSubnetsError).Run(func(args mock.Arguments) {
+				arg := args.Get(1).(func(*ec2.DescribeSubnetsOutput, bool) bool)
+				arg(tc.DescribeSubnetsOutput, false)
+			})
+
+			cloud := &Cloud{
+				clusterName: clusterName,
+				ec2:         svc,
+			}
+			subnets, err := cloud.GetClusterSubnets(tc.TagSubnetType)
+			assert.Equal(t, tc.ExpectedResult, subnets)
+			assert.Equal(t, tc.ExpectedError, err)
+			svc.AssertExpectations(t)
+		})
+	}
+}

internal/aws/rgt.go

@@ -2,15 +2,9 @@ package aws
 
 import (
 	"context"
-	"strings"
-
-	"github.com/aws/aws-sdk-go/service/ec2"
 
 	"github.com/aws/aws-sdk-go/aws"
-
 	"github.com/aws/aws-sdk-go/service/resourcegroupstaggingapi"
-
-	util "github.com/kubernetes-sigs/aws-alb-ingress-controller/pkg/util/types"
 )
 
 const (
@@ -20,8 +14,6 @@ const (
 )
 
 type ResourceGroupsTaggingAPIAPI interface {
-	GetClusterSubnets() (map[string]util.EC2Tags, error)
-
 	// GetResourcesByFilters fetches resources ARNs by tagFilters and 0 or more resourceTypesFilters
 	GetResourcesByFilters(tagFilters map[string][]string, resourceTypeFilters ...string) ([]string, error)
 
@@ -36,76 +28,6 @@ func (c *Cloud) UntagResourcesWithContext(ctx context.Context, i *resourcegroups
 	return c.rgt.UntagResourcesWithContext(ctx, i)
 }
 
-// GetClusterSubnets looks up all subnets in AWS that are tagged for the cluster.
-func (c *Cloud) GetClusterSubnets() (map[string]util.EC2Tags, error) {
-	subnets := make(map[string]util.EC2Tags)
-
-	paramSets := []*resourcegroupstaggingapi.GetResourcesInput{
-		{
-			ResourcesPerPage: aws.Int64(50),
-			ResourceTypeFilters: []*string{
-				aws.String("ec2"),
-			},
-			TagFilters: []*resourcegroupstaggingapi.TagFilter{
-				{
-					Key:    aws.String("kubernetes.io/role/internal-elb"),
-					Values: []*string{aws.String(""), aws.String("1")},
-				},
-				{
-					Key:    aws.String("kubernetes.io/cluster/" + c.clusterName),
-					Values: []*string{aws.String("owned"), aws.String("shared")},
-				},
-			},
-		},
-		{
-			ResourcesPerPage: aws.Int64(50),
-			ResourceTypeFilters: []*string{
-				aws.String("ec2"),
-			},
-			TagFilters: []*resourcegroupstaggingapi.TagFilter{
-				{
-					Key:    aws.String("kubernetes.io/role/elb"),
-					Values: []*string{aws.String(""), aws.String("1")},
-				},
-				{
-					Key:    aws.String("kubernetes.io/cluster/" + c.clusterName),
-					Values: []*string{aws.String("owned"), aws.String("shared")},
-				},
-			},
-		},
-	}
-
-	for _, paramSet := range paramSets {
-		err := c.rgt.GetResourcesPages(paramSet, func(page *resourcegroupstaggingapi.GetResourcesOutput, lastPage bool) bool {
-			if page == nil {
-				return false
-			}
-			for _, rtm := range page.ResourceTagMappingList {
-				switch {
-				case strings.Contains(*rtm.ResourceARN, ":subnet/"):
-					subnets[*rtm.ResourceARN] = rgtTagAsEC2Tag(rtm.Tags)
-				}
-			}
-			return true
-		})
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	return subnets, nil
-}
-
-func rgtTagAsEC2Tag(in []*resourcegroupstaggingapi.Tag) (tags util.EC2Tags) {
-	for _, t := range in {
-		tags = append(tags, &ec2.Tag{
-			Key:   t.Key,
-			Value: t.Value,
-		})
-	}
-	return tags
-}
-
 func (c *Cloud) GetResourcesByFilters(tagFilters map[string][]string, resourceTypeFilters ...string) ([]string, error) {
 	var awsTagFilters []*resourcegroupstaggingapi.TagFilter
 	for k, v := range tagFilters {