Use built-in pagination instead of loops. (#121)

* Use built-in pagination instead of loops. * Mock paginated method for tests.
segmentio · Sep 19, 2018 · 7b5f2b8 · 7b5f2b8
1 parent 1ac83b0
commit 7b5f2b8
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 70 deletions.
diff --git a/store/ssmstore.go b/store/ssmstore.go
@@ -247,40 +247,30 @@ func (s *SSMStore) readLatest(id SecretId) (Secret, error) {
 func (s *SSMStore) List(service string, includeValues bool) ([]Secret, error) {
 	secrets := map[string]Secret{}
 
-	var nextToken *string
-	for {
-		var describeParametersInput *ssm.DescribeParametersInput
-
-		if s.usePaths {
-			describeParametersInput = &ssm.DescribeParametersInput{
-				ParameterFilters: []*ssm.ParameterStringFilter{
-					{
-						Key:    aws.String("Path"),
-						Option: aws.String("OneLevel"),
-						Values: []*string{aws.String("/" + service)},
-					},
-				},
-				MaxResults: aws.Int64(50),
-				NextToken:  nextToken,
-			}
-		} else {
-			describeParametersInput = &ssm.DescribeParametersInput{
-				Filters: []*ssm.ParametersFilter{
-					{
-						Key:    aws.String("Name"),
-						Values: []*string{aws.String(service + ".")},
-					},
+	var describeParametersInput *ssm.DescribeParametersInput
+
+	if s.usePaths {
+		describeParametersInput = &ssm.DescribeParametersInput{
+			ParameterFilters: []*ssm.ParameterStringFilter{
+				{
+					Key:    aws.String("Path"),
+					Option: aws.String("OneLevel"),
+					Values: []*string{aws.String("/" + service)},
 				},
-				MaxResults: aws.Int64(50),
-				NextToken:  nextToken,
-			}
+			},
 		}
-
-		resp, err := s.svc.DescribeParameters(describeParametersInput)
-		if err != nil {
-			return nil, err
+	} else {
+		describeParametersInput = &ssm.DescribeParametersInput{
+			Filters: []*ssm.ParametersFilter{
+				{
+					Key:    aws.String("Name"),
+					Values: []*string{aws.String(service + ".")},
+				},
+			},
 		}
+	}
 
+	err := s.svc.DescribeParametersPages(describeParametersInput, func(resp *ssm.DescribeParametersOutput, lastPage bool) bool {
 		for _, meta := range resp.Parameters {
 			if !s.validateName(*meta.Name) {
 				continue
@@ -291,12 +281,10 @@ func (s *SSMStore) List(service string, includeValues bool) ([]Secret, error) {
 				Meta:  secretMeta,
 			}
 		}
-
-		if resp.NextToken == nil {
-			break
-		}
-
-		nextToken = resp.NextToken
+		return true
+	})
+	if err != nil {
+		return nil, err
 	}
 
 	if includeValues {
@@ -325,6 +313,7 @@ func (s *SSMStore) List(service string, includeValues bool) ([]Secret, error) {
 			}
 		}
 	}
+
 	return values(secrets), nil
 }
 
@@ -334,37 +323,12 @@ func (s *SSMStore) List(service string, includeValues bool) ([]Secret, error) {
 func (s *SSMStore) ListRaw(service string) ([]RawSecret, error) {
 	if s.usePaths {
 		secrets := map[string]RawSecret{}
-		var nextToken *string
-		for {
-			getParametersByPathInput := &ssm.GetParametersByPathInput{
-				MaxResults:     aws.Int64(10),
-				NextToken:      nextToken,
-				Path:           aws.String("/" + service + "/"),
-				WithDecryption: aws.Bool(true),
-			}
-
-			resp, err := s.svc.GetParametersByPath(getParametersByPathInput)
-			if err != nil {
-				// If the error is an access-denied exception
-				awsErr, isAwserr := err.(awserr.Error)
-				if isAwserr {
-					if awsErr.Code() == "AccessDeniedException" && strings.Contains(awsErr.Message(), "is not authorized to perform: ssm:GetParametersByPath on resource") {
-						// Fall-back to using the old list method in case some users haven't updated their IAM permissions yet, but warn about it and
-						// tell them to fix their permissions
-						fmt.Fprintf(
-							os.Stderr,
-							"Warning: %s\nFalling-back to using ssm:DescribeParameters. This may cause delays or failures due to AWS rate-limiting.\n"+
-								"This is behavior deprecated and will be removed in a future version of chamber. Please update your IAM permissions to grant ssm:GetParametersByPath.\n\n",
-							awsErr)
-
-						// Delegate to List
-						return s.listRawViaList(service)
-					}
-				}
-
-				return nil, err
-			}
+		getParametersByPathInput := &ssm.GetParametersByPathInput{
+			Path:           aws.String("/" + service + "/"),
+			WithDecryption: aws.Bool(true),
+		}
 
+		err := s.svc.GetParametersByPathPages(getParametersByPathInput, func(resp *ssm.GetParametersByPathOutput, lastPage bool) bool {
 			for _, param := range resp.Parameters {
 				if !s.validateName(*param.Name) {
 					continue
@@ -375,12 +339,28 @@ func (s *SSMStore) ListRaw(service string) ([]RawSecret, error) {
 					Key:   *param.Name,
 				}
 			}
+			return true
+		})
 
-			if resp.NextToken == nil {
-				break
+		if err != nil {
+			// If the error is an access-denied exception
+			awsErr, isAwserr := err.(awserr.Error)
+			if isAwserr {
+				if awsErr.Code() == "AccessDeniedException" && strings.Contains(awsErr.Message(), "is not authorized to perform: ssm:GetParametersByPath on resource") {
+					// Fall-back to using the old list method in case some users haven't updated their IAM permissions yet, but warn about it and
+					// tell them to fix their permissions
+					fmt.Fprintf(
+						os.Stderr,
+						"Warning: %s\nFalling-back to using ssm:DescribeParameters. This may cause delays or failures due to AWS rate-limiting.\n"+
+							"This is behavior deprecated and will be removed in a future version of chamber. Please update your IAM permissions to grant ssm:GetParametersByPath.\n\n",
+						awsErr)
+
+					// Delegate to List
+					return s.listRawViaList(service)
+				}
 			}
 
-			nextToken = resp.NextToken
+			return nil, err
 		}
 
 		rawSecrets := make([]RawSecret, len(secrets))
@@ -390,7 +370,6 @@ func (s *SSMStore) ListRaw(service string) ([]RawSecret, error) {
 			i += 1
 		}
 		return rawSecrets, nil
-
 	}
 
 	// Delete to List (which uses the DescribeParameters API)

diff --git a/store/ssmstore_test.go b/store/ssmstore_test.go
@@ -173,6 +173,15 @@ func (m *mockSSMClient) GetParametersByPath(i *ssm.GetParametersByPathInput) (*s
 	}, nil
 }
 
+func (m *mockSSMClient) GetParametersByPathPages(i *ssm.GetParametersByPathInput, fn func(*ssm.GetParametersByPathOutput, bool) bool) error {
+	o, err := m.GetParametersByPath(i)
+	if err != nil {
+		return err
+	}
+	fn(o, true)
+	return nil
+}
+
 func (m *mockSSMClient) DescribeParametersPages(i *ssm.DescribeParametersInput, fn func(*ssm.DescribeParametersOutput, bool) bool) error {
 	o, err := m.DescribeParameters(i)
 	if err != nil {