From 1885e570918e6ed34ad88f3c4c789cd8ee520919 Mon Sep 17 00:00:00 2001
From: Enny Frick <enny.frick@conductorone.com>
Date: Wed, 15 Nov 2023 16:22:12 -0800
Subject: [PATCH] use AWS SDK paginator funcs for pagination loops

---
 pkg/connector/account.go   | 28 +++++++++++++++-------------
 pkg/connector/connector.go |  9 +++------
 2 files changed, 18 insertions(+), 19 deletions(-)

diff --git a/pkg/connector/account.go b/pkg/connector/account.go
index 3c98fa9e..3d338cd1 100644
--- a/pkg/connector/account.go
+++ b/pkg/connector/account.go
@@ -21,7 +21,7 @@ import (
 )
 
 const (
-	accountMemeberEntitlement = "member"
+	accountMemberEntitlement = "member"
 )
 
 type accountResourceType struct {
@@ -115,7 +115,7 @@ func (o *accountResourceType) Entitlements(ctx context.Context, resource *v2.Res
 		annos.Update(&v2.V1Identifier{
 			Id: b.String(),
 		})
-		member := sdk.NewAssignmentEntitlement(resource, accountMemeberEntitlement, resourceTypeAccount)
+		member := sdk.NewAssignmentEntitlement(resource, accountMemberEntitlement, resourceTypeAccount)
 		member.Description = awsSdk.ToString(ps.Description)
 		member.Annotations = annos
 		member.Id = b.String()
@@ -132,8 +132,9 @@ func (o *accountResourceType) Grants(ctx context.Context, resource *v2.Resource,
 		InstanceArn: o.identityInstance.InstanceArn,
 	}
 
+	psBindingPaginator := awsSsoAdmin.NewListPermissionSetsProvisionedToAccountPaginator(o.ssoAdminClient, psBindingInput)
 	for {
-		psBindingsResp, err := o.ssoAdminClient.ListPermissionSetsProvisionedToAccount(ctx, psBindingInput)
+		psBindingsResp, err := psBindingPaginator.NextPage(ctx)
 		if err != nil {
 			return nil, "", nil, fmt.Errorf("aws-connector: ssoadmin.ListPermissionSetsProvisionedToAccount failed: %w", err)
 		}
@@ -168,8 +169,9 @@ func (o *accountResourceType) Grants(ctx context.Context, resource *v2.Resource,
 				PermissionSetArn: ps.PermissionSetArn,
 			}
 
+			assignmentsPaginator := awsSsoAdmin.NewListAccountAssignmentsPaginator(o.ssoAdminClient, assignmentsInput)
 			for {
-				assignmentsResp, err := o.ssoAdminClient.ListAccountAssignments(ctx, assignmentsInput)
+				assignmentsResp, err := assignmentsPaginator.NextPage(ctx)
 				if err != nil {
 					return nil, "", nil, fmt.Errorf("aws-connector: ssoadmin.ListAccountAssignments failed: %w", err)
 				}
@@ -227,13 +229,13 @@ func (o *accountResourceType) Grants(ctx context.Context, resource *v2.Resource,
 					}
 				}
 				assignmentsInput.NextToken = assignmentsResp.NextToken
-				if assignmentsResp.NextToken == nil {
+				if !assignmentsPaginator.HasMorePages() {
 					break
 				}
 			} // end pagination loop for assignments
 		} // end range ange psBindingsResp.PermissionSets
-		psBindingInput.NextToken = psBindingsResp.NextToken
-		if psBindingsResp.NextToken == nil {
+
+		if !psBindingPaginator.HasMorePages() {
 			break
 		}
 	} // end pagination loop for permission set to account binding
@@ -300,8 +302,9 @@ func (o *accountResourceType) getGroupMembers(ctx context.Context, groupId strin
 		GroupId:         awsSdk.String(groupId),
 	}
 	userIds := make([]string, 0, 16)
+	paginator := awsIdentityStore.NewListGroupMembershipsPaginator(o.identityClient, input)
 	for {
-		resp, err := o.identityClient.ListGroupMemberships(ctx, input)
+		resp, err := paginator.NextPage(ctx)
 		if err != nil {
 			return nil, err
 		}
@@ -312,10 +315,9 @@ func (o *accountResourceType) getGroupMembers(ctx context.Context, groupId strin
 			}
 			userIds = append(userIds, member.Value)
 		}
-		if resp.NextToken == nil {
+		if !paginator.HasMorePages() {
 			break
 		}
-		input.NextToken = resp.NextToken
 	}
 	o._groupMembersCache.Store(groupId, userIds)
 	return userIds, nil
@@ -332,16 +334,16 @@ func (o *accountResourceType) getPermissionSets(ctx context.Context) ([]*awsSsoA
 	input := &awsSsoAdmin.ListPermissionSetsInput{
 		InstanceArn: o.identityInstance.InstanceArn,
 	}
+	paginator := awsSsoAdmin.NewListPermissionSetsPaginator(o.ssoAdminClient, input)
 	for {
-		resp, err := o.ssoAdminClient.ListPermissionSets(ctx, input)
+		resp, err := paginator.NextPage(ctx)
 		if err != nil {
 			return nil, err
 		}
 		permissionSetIDs = append(permissionSetIDs, resp.PermissionSets...)
-		if resp.NextToken == nil {
+		if !paginator.HasMorePages() {
 			break
 		}
-		input.NextToken = resp.NextToken
 	}
 	for _, psId := range permissionSetIDs {
 		ps, err := o.getPermissionSet(ctx, psId)
diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go
index a89ab813..35eb6d0f 100644
--- a/pkg/connector/connector.go
+++ b/pkg/connector/connector.go
@@ -361,11 +361,9 @@ func (c *AWS) getIdentityInstance(ctx context.Context) (*awsSsoAdminTypes.Instan
 		return nil, err
 	}
 
-	nextToken := awsSdk.String("")
+	paginator := awsSsoAdmin.NewListInstancesPaginator(ssoClient, &awsSsoAdmin.ListInstancesInput{})
 	for {
-		resp, err := ssoClient.ListInstances(ctx, &awsSsoAdmin.ListInstancesInput{
-			NextToken: nextToken,
-		})
+		resp, err := paginator.NextPage(ctx)
 		if err != nil {
 			c._identityInstancesCacheErr = err
 			return nil, err
@@ -375,8 +373,7 @@ func (c *AWS) getIdentityInstance(ctx context.Context) (*awsSsoAdminTypes.Instan
 				func(i awsSsoAdminTypes.InstanceMetadata) *awsSsoAdminTypes.InstanceMetadata { return &i },
 			)...,
 		)
-		nextToken = resp.NextToken
-		if nextToken == nil {
+		if !paginator.HasMorePages() {
 			break
 		}
 	}