Replace DescribeNetworkInterfaces with paginated version

pkg/awsutils/awsutils.go

@@ -66,6 +66,9 @@ const (
 
 	// Http client timeout env for sessions
 	httpTimeoutEnv = "HTTP_TIMEOUT"
+
+	// the default page size when paginating the DescribeNetworkInterfaces call
+	describeENIPageSize = 1000
 )
 
 var (
@@ -355,7 +358,7 @@ func New(useCustomNetworking bool) (*EC2InstanceMetadataCache, error) {
 
 	sess, err := session.NewSession(
 		&aws.Config{
-			Region: aws.String(cache.region), 
+			Region: aws.String(cache.region),
 			MaxRetries: aws.Int(15),
 			HTTPClient: &http.Client{
 				Timeout: httpTimeoutValue,
@@ -1425,14 +1428,16 @@ func (cache *EC2InstanceMetadataCache) getFilteredListOfNetworkInterfaces() ([]*
 
 	input := &ec2.DescribeNetworkInterfacesInput{
 		Filters: []*ec2.Filter{tagFilter, statusFilter},
+		MaxResults: aws.Int64(describeENIPageSize),
 	}
-	result, err := cache.ec2SVC.DescribeNetworkInterfacesWithContext(context.Background(), input, userAgent)
+
+	outputENIs, err := cache.getENIsFromPaginatedDescribeNetworkInterfaces(input)
 	if err != nil {
 		return nil, errors.Wrap(err, "awsutils: unable to obtain filtered list of network interfaces")
 	}
 
 	networkInterfaces := make([]*ec2.NetworkInterface, 0)
-	for _, networkInterface := range result.NetworkInterfaces {
+	for _, networkInterface := range outputENIs {
 		// Verify the description starts with "aws-K8S-"
 		if !strings.HasPrefix(aws.StringValue(networkInterface.Description), eniDescriptionPrefix) {
 			continue
@@ -1515,3 +1520,22 @@ func (cache *EC2InstanceMetadataCache) IsUnmanagedENI(eniID string) bool {
 	}
 	return false
 }
+
+func (cache *EC2InstanceMetadataCache) getENIsFromPaginatedDescribeNetworkInterfaces(
+	input *ec2.DescribeNetworkInterfacesInput) ([]*ec2.NetworkInterface, error) {
+	outputENIs := make([]*ec2.NetworkInterface, 0)
+	pageNum := 0
+	log.Debugf("Paginating describe ENI has page size: %d", *input.MaxResults)
+	pageFn := func(output *ec2.DescribeNetworkInterfacesOutput, lastPage bool) (nextPage bool) {
+		pageNum++
+		log.Debugf("EC2 DescribeNetworkInterfaces succeeded with %d results on page %d",
+			len(output.NetworkInterfaces), pageNum)
+
+		outputENIs = append(outputENIs, output.NetworkInterfaces...)
+		// Loop is guided by nextToken, the func expect a false to exit.
+		return output.NextToken != nil
+	}
+
+	err := cache.ec2SVC.DescribeNetworkInterfacesPagesWithContext(context.Background(), input, pageFn, userAgent)
+	return outputENIs, err
+}

pkg/awsutils/awsutils_test.go

@@ -24,6 +24,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
@@ -711,10 +712,8 @@ func TestEC2InstanceMetadataCache_getFilteredListOfNetworkInterfaces_OneResult(t
 	attachment := &ec2.NetworkInterfaceAttachment{AttachmentId: &attachmentID}
 	cureniID := eniID
 
-	result := &ec2.DescribeNetworkInterfacesOutput{
-		NetworkInterfaces: []*ec2.NetworkInterface{{Attachment: attachment, Status: &status, TagSet: tag, Description: &description, NetworkInterfaceId: &cureniID}}}
-	mockEC2.EXPECT().DescribeNetworkInterfacesWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(result, nil)
-
+	interfaces := []*ec2.NetworkInterface{{Attachment: attachment, Status: &status, TagSet: tag, Description: &description, NetworkInterfaceId: &cureniID}}
+	setupDescribeNetworkInterfacesPagesWithContextMock(t, mockEC2, interfaces, nil, 1)
 	ins := &EC2InstanceMetadataCache{ec2SVC: mockEC2}
 	got, err := ins.getFilteredListOfNetworkInterfaces()
 	assert.NotNil(t, got)
@@ -725,10 +724,7 @@ func TestEC2InstanceMetadataCache_getFilteredListOfNetworkInterfaces_NoResult(t
 	ctrl, mockEC2 := setup(t)
 	defer ctrl.Finish()
 
-	result := &ec2.DescribeNetworkInterfacesOutput{
-		NetworkInterfaces: []*ec2.NetworkInterface{}}
-	mockEC2.EXPECT().DescribeNetworkInterfacesWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(result, nil)
-
+	setupDescribeNetworkInterfacesPagesWithContextMock(t, mockEC2, []*ec2.NetworkInterface{}, nil, 1)
 	ins := &EC2InstanceMetadataCache{ec2SVC: mockEC2}
 	got, err := ins.getFilteredListOfNetworkInterfaces()
 	assert.Nil(t, got)
@@ -739,7 +735,12 @@ func TestEC2InstanceMetadataCache_getFilteredListOfNetworkInterfaces_Error(t *te
 	ctrl, mockEC2 := setup(t)
 	defer ctrl.Finish()
 
-	mockEC2.EXPECT().DescribeNetworkInterfacesWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, errors.New("dummy error"))
+	interfaces := []*ec2.NetworkInterface{{
+		TagSet: []*ec2.Tag{
+			{Key: aws.String("foo"), Value: aws.String("foo-value")},
+		},
+	}}
+	setupDescribeNetworkInterfacesPagesWithContextMock(t, mockEC2, interfaces, errors.New("dummy error"), 1)
 
 	ins := &EC2InstanceMetadataCache{ec2SVC: mockEC2}
 	got, err := ins.getFilteredListOfNetworkInterfaces()
@@ -857,19 +858,30 @@ func TestEC2InstanceMetadataCache_cleanUpLeakedENIsInternal(t *testing.T) {
 	defer ctrl.Finish()
 
 	description := eniDescriptionPrefix + "test"
-	result := &ec2.DescribeNetworkInterfacesOutput{
-		NetworkInterfaces: []*ec2.NetworkInterface{{
-			Description: &description,
-			TagSet: []*ec2.Tag{
-				{Key: aws.String(eniNodeTagKey), Value: aws.String("test-value")},
-			},
-		}},
-	}
+	interfaces := []*ec2.NetworkInterface{{
+		Description: &description,
+		TagSet: []*ec2.Tag{
+			{Key: aws.String(eniNodeTagKey), Value: aws.String("test-value")},
+		},
+	}}
 
-	mockEC2.EXPECT().DescribeNetworkInterfacesWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(result, nil)
+	setupDescribeNetworkInterfacesPagesWithContextMock(t, mockEC2, interfaces, nil, 1)
 	mockEC2.EXPECT().CreateTagsWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil)
 
 	ins := &EC2InstanceMetadataCache{ec2SVC: mockEC2}
 	// Test checks that both mocks gets called.
 	ins.cleanUpLeakedENIsInternal(time.Millisecond)
 }
+
+func setupDescribeNetworkInterfacesPagesWithContextMock(
+	t *testing.T, mockEC2 *mock_ec2wrapper.MockEC2, interfaces []*ec2.NetworkInterface, err error, times int) {
+	mockEC2.EXPECT().
+		DescribeNetworkInterfacesPagesWithContext(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Times(times).
+		DoAndReturn(func(_ context.Context, _ *ec2.DescribeNetworkInterfacesInput,
+			fn func(*ec2.DescribeNetworkInterfacesOutput, bool) bool, userAgent request.Option) error {
+			assert.Equal(t, false, fn(&ec2.DescribeNetworkInterfacesOutput{
+				NetworkInterfaces: interfaces,
+			}, true))
+			return err
+		})
+}

pkg/ec2wrapper/client.go

@@ -33,6 +33,7 @@ type EC2 interface {
 	DescribeNetworkInterfacesWithContext(ctx aws.Context, input *ec2svc.DescribeNetworkInterfacesInput, opts ...request.Option) (*ec2svc.DescribeNetworkInterfacesOutput, error)
 	ModifyNetworkInterfaceAttributeWithContext(ctx aws.Context, input *ec2svc.ModifyNetworkInterfaceAttributeInput, opts ...request.Option) (*ec2svc.ModifyNetworkInterfaceAttributeOutput, error)
 	CreateTagsWithContext(ctx aws.Context, input *ec2svc.CreateTagsInput, opts ...request.Option) (*ec2svc.CreateTagsOutput, error)
+	DescribeNetworkInterfacesPagesWithContext(ctx aws.Context, input *ec2svc.DescribeNetworkInterfacesInput, fn func(*ec2svc.DescribeNetworkInterfacesOutput, bool) bool, opts ...request.Option) error
 }
 
 // New creates a new EC2 wrapper