From 704ca63cc7d1134696ea6df270960dee647e7ced Mon Sep 17 00:00:00 2001 From: Joe Gresock Date: Fri, 14 Oct 2022 07:33:53 -0400 Subject: [PATCH] NIFI-10643: Paging secrets in AwsSecretsManagerParameterProvider This closes #6532 Signed-off-by: David Handermann --- .../AwsSecretsManagerParameterProvider.java | 17 +++++-- ...estAwsSecretsManagerParameterProvider.java | 47 +++++++++++++++---- 2 files changed, 51 insertions(+), 13 deletions(-) diff --git a/nifi-nar-bundles/nifi-aws-bundle/nifi-aws-parameter-providers/src/main/java/org/apache/nifi/parameter/aws/AwsSecretsManagerParameterProvider.java b/nifi-nar-bundles/nifi-aws-bundle/nifi-aws-parameter-providers/src/main/java/org/apache/nifi/parameter/aws/AwsSecretsManagerParameterProvider.java index 4cab6fa9a2..7f9e4239f4 100644 --- a/nifi-nar-bundles/nifi-aws-bundle/nifi-aws-parameter-providers/src/main/java/org/apache/nifi/parameter/aws/AwsSecretsManagerParameterProvider.java +++ b/nifi-nar-bundles/nifi-aws-bundle/nifi-aws-parameter-providers/src/main/java/org/apache/nifi/parameter/aws/AwsSecretsManagerParameterProvider.java @@ -139,10 +139,19 @@ public class AwsSecretsManagerParameterProvider extends AbstractParameterProvide AWSSecretsManager secretsManager = this.configureClient(context); final List groups = new ArrayList<>(); - final ListSecretsRequest listSecretsRequest = new ListSecretsRequest(); - final ListSecretsResult listSecretsResult = secretsManager.listSecrets(listSecretsRequest); - for (final SecretListEntry entry : listSecretsResult.getSecretList()) { - groups.addAll(fetchSecret(secretsManager, context, entry.getName())); + ListSecretsRequest listSecretsRequest = new ListSecretsRequest(); + ListSecretsResult listSecretsResult = secretsManager.listSecrets(listSecretsRequest); + while(!listSecretsResult.getSecretList().isEmpty()) { + for (final SecretListEntry entry : listSecretsResult.getSecretList()) { + groups.addAll(fetchSecret(secretsManager, context, entry.getName())); + } + final String nextToken = listSecretsResult.getNextToken(); + if (nextToken == null) { + break; + } + + listSecretsRequest.setNextToken(nextToken); + listSecretsResult = secretsManager.listSecrets(listSecretsRequest); } return groups; diff --git a/nifi-nar-bundles/nifi-aws-bundle/nifi-aws-parameter-providers/src/test/java/org/apache/nifi/parameter/aws/TestAwsSecretsManagerParameterProvider.java b/nifi-nar-bundles/nifi-aws-bundle/nifi-aws-parameter-providers/src/test/java/org/apache/nifi/parameter/aws/TestAwsSecretsManagerParameterProvider.java index 8b96a6e7f0..572053c039 100644 --- a/nifi-nar-bundles/nifi-aws-bundle/nifi-aws-parameter-providers/src/test/java/org/apache/nifi/parameter/aws/TestAwsSecretsManagerParameterProvider.java +++ b/nifi-nar-bundles/nifi-aws-bundle/nifi-aws-parameter-providers/src/test/java/org/apache/nifi/parameter/aws/TestAwsSecretsManagerParameterProvider.java @@ -20,6 +20,7 @@ import com.amazonaws.services.secretsmanager.AWSSecretsManager; import com.amazonaws.services.secretsmanager.model.AWSSecretsManagerException; import com.amazonaws.services.secretsmanager.model.GetSecretValueRequest; import com.amazonaws.services.secretsmanager.model.GetSecretValueResult; +import com.amazonaws.services.secretsmanager.model.ListSecretsRequest; import com.amazonaws.services.secretsmanager.model.ListSecretsResult; import com.amazonaws.services.secretsmanager.model.SecretListEntry; import com.fasterxml.jackson.core.JsonProcessingException; @@ -46,6 +47,7 @@ import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.stream.Collectors; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -64,7 +66,7 @@ public class TestAwsSecretsManagerParameterProvider { private AWSSecretsManager defaultSecretsManager; @Mock - private ListSecretsResult listSecretsResult; + private ListSecretsResult emptyListSecretsResult; final ObjectMapper objectMapper = new ObjectMapper(); @@ -105,8 +107,9 @@ public class TestAwsSecretsManagerParameterProvider { @Test public void testFetchParametersGetSecretFailure() throws InitializationException { final List secretList = Collections.singletonList(new SecretListEntry().withName("MySecret")); + final ListSecretsResult listSecretsResult = mock(ListSecretsResult.class); when(listSecretsResult.getSecretList()).thenReturn(secretList); - when(defaultSecretsManager.listSecrets(any())).thenReturn(listSecretsResult); + when(defaultSecretsManager.listSecrets(argThat(ListSecretsRequestMatcher.hasToken(null)))).thenReturn(listSecretsResult); when(defaultSecretsManager.getSecretValue(argThat(matchesGetSecretValueRequest("MySecret")))).thenThrow(new AWSSecretsManagerException("Fake exception")); runProviderTest(defaultSecretsManager, 0, ConfigVerificationResult.Outcome.FAILED); } @@ -115,16 +118,24 @@ public class TestAwsSecretsManagerParameterProvider { return spy(new AwsSecretsManagerParameterProvider()); } - private AWSSecretsManager mockSecretsManager(final List mockGroup) { + private AWSSecretsManager mockSecretsManager(final List mockParameterGroups) { final AWSSecretsManager secretsManager = mock(AWSSecretsManager.class); + when(emptyListSecretsResult.getSecretList()).thenReturn(Collections.emptyList()); - final List secretList = mockGroup.stream() - .map(group -> new SecretListEntry().withName(group.getGroupName())) - .collect(Collectors.toList()); - when(listSecretsResult.getSecretList()).thenReturn(secretList); - when(secretsManager.listSecrets(any())).thenReturn(listSecretsResult); + String currentToken = null; + for (int i = 0; i < mockParameterGroups.size(); i++) { + final ParameterGroup group = mockParameterGroups.get(i); + final List secretList = Collections.singletonList(new SecretListEntry().withName(group.getGroupName())); + final ListSecretsResult listSecretsResult = mock(ListSecretsResult.class); + when(listSecretsResult.getSecretList()).thenReturn(secretList); + when(secretsManager.listSecrets(argThat(ListSecretsRequestMatcher.hasToken(currentToken)))).thenReturn(listSecretsResult); - mockGroup.forEach(group -> { + currentToken = "token-" + i; + when(listSecretsResult.getNextToken()).thenReturn(currentToken); + } + when(secretsManager.listSecrets(argThat(ListSecretsRequestMatcher.hasToken(currentToken)))).thenReturn(emptyListSecretsResult); + + mockParameterGroups.forEach(group -> { final String groupName = group.getGroupName(); final Map keyValues = group.getParameters().stream().collect(Collectors.toMap( param -> param.getDescriptor().getName(), @@ -196,4 +207,22 @@ public class TestAwsSecretsManagerParameterProvider { return argument != null && argument.getSecretId().equals(secretId); } } + + private static class ListSecretsRequestMatcher implements ArgumentMatcher { + + private static ListSecretsRequestMatcher hasToken(final String token) { + return new ListSecretsRequestMatcher(token); + } + + private final String token; + + private ListSecretsRequestMatcher(final String token) { + this.token = token; + } + + @Override + public boolean matches(final ListSecretsRequest argument) { + return argument != null && Objects.equals(argument.getNextToken(), token); + } + } }