Add AWSAuthFactory.createCloseable to release the assume-role STS client

buildStsCredentialsProvider hands a fresh StsClient to
StsAssumeRoleCredentialsProvider, which never closes it:
StsCredentialsProvider.close() only clears the session cache. A caller
that builds a provider per operation therefore leaks the STS client's
Apache connection pool on every call.

Add createCloseable(...), returning a CredentialsProviderHandle that
bundles the provider with its STS client so a caller can close both.
The existing create(...) overloads delegate to it and discard the
handle, so their behaviour is unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: patrickmann

graylog2-server/src/main/java/org/graylog/integrations/aws/AWSAuthFactory.java

@@ -30,8 +30,11 @@
 import software.amazon.awssdk.services.sts.StsClient;
 import software.amazon.awssdk.services.sts.StsClientBuilder;
 import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
+import software.amazon.awssdk.services.sts.auth.StsCredentialsProvider;
 import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
 import software.amazon.awssdk.services.sts.model.GetCallerIdentityRequest;
+import software.amazon.awssdk.utils.IoUtils;
+import software.amazon.awssdk.utils.SdkAutoCloseable;
 
 import javax.annotation.Nullable;
 import java.util.Locale;
@@ -96,6 +99,29 @@ public AwsCredentialsProvider create(boolean requireKeySecret,
                                          @Nullable String assumeRoleArn,
                                          @Nullable String externalId,
                                          @Nullable ApacheHttpClient.Builder stsHttpClientBuilder) {
+        // Discards the STS client handle. Callers on this overload cannot close the STS client backing an
+        // assume-role provider, so its connection pool leaks for the process lifetime; see createCloseable.
+        return createCloseable(requireKeySecret, stsRegion, accessKey, secretKey, assumeRoleArn, externalId, stsHttpClientBuilder)
+                .provider();
+    }
+
+    /**
+     * Same resolution as {@link #create(boolean, String, String, String, String, String, ApacheHttpClient.Builder)},
+     * but returns a {@link CredentialsProviderHandle} so the caller can close the STS client backing an assume-role
+     * provider.
+     *
+     * <p>For an assume-role configuration the STS client is created here and handed to
+     * {@link StsAssumeRoleCredentialsProvider}, which never closes it: {@link StsCredentialsProvider#close()} only
+     * clears the session cache. Callers that build a fresh provider per operation must therefore close the handle,
+     * or the STS client's Apache connection pool leaks on every call.
+     */
+    public CredentialsProviderHandle createCloseable(boolean requireKeySecret,
+                                                     @Nullable String stsRegion,
+                                                     @Nullable String accessKey,
+                                                     @Nullable String secretKey,
+                                                     @Nullable String assumeRoleArn,
+                                                     @Nullable String externalId,
+                                                     @Nullable ApacheHttpClient.Builder stsHttpClientBuilder) {
         AwsCredentialsProvider awsCredentials = requireKeySecret ? getKeySecretCredentialsProvider(accessKey, secretKey) :
                 getAwsCredentialsProvider(accessKey, secretKey);
 
@@ -109,7 +135,7 @@ public AwsCredentialsProvider create(boolean requireKeySecret,
             return buildStsCredentialsProvider(awsCredentials, stsRegion, assumeRoleArn, accessKey, externalId, stsHttpClientBuilder);
         }
 
-        return awsCredentials;
+        return new CredentialsProviderHandle(awsCredentials, null);
     }
 
     private static AwsCredentialsProvider getAwsCredentialsProvider(String accessKey, String secretKey) {
@@ -135,10 +161,10 @@ private static AwsCredentialsProvider getKeySecretCredentialsProvider(String acc
      * Note: In order to assume a role, a role must be provided to the AWS STS client a role that has the "sts:AssumeRole"
      * permission, which provides authorization for a role to be assumed.
      */
-    private static AwsCredentialsProvider buildStsCredentialsProvider(AwsCredentialsProvider awsCredentials, String stsRegion,
-                                                                       String assumeRoleArn, @Nullable String accessKey,
-                                                                       @Nullable String externalId,
-                                                                       @Nullable ApacheHttpClient.Builder stsHttpClientBuilder) {
+    private static CredentialsProviderHandle buildStsCredentialsProvider(AwsCredentialsProvider awsCredentials, String stsRegion,
+                                                                         String assumeRoleArn, @Nullable String accessKey,
+                                                                         @Nullable String externalId,
+                                                                         @Nullable ApacheHttpClient.Builder stsHttpClientBuilder) {
 
         final StsClientBuilder stsClientBuilder = StsClient.builder()
                 .region(Region.of(stsRegion))
@@ -166,9 +192,41 @@ private static AwsCredentialsProvider buildStsCredentialsProvider(AwsCredentials
         if (!isNullOrEmpty(externalId)) {
             assumeRoleRequestBuilder.externalId(externalId);
         }
-        return StsAssumeRoleCredentialsProvider.builder()
+        final StsAssumeRoleCredentialsProvider provider = StsAssumeRoleCredentialsProvider.builder()
                 .refreshRequest(assumeRoleRequestBuilder.build())
                 .stsClient(stsClient)
                 .build();
+        return new CredentialsProviderHandle(provider, stsClient);
+    }
+
+    /**
+     * A resolved {@link AwsCredentialsProvider} paired with the STS client backing an assume-role provider, if any.
+     *
+     * <p>{@link StsCredentialsProvider#close()} only clears the provider's session cache; the {@link StsClient} it
+     * was handed is caller-supplied and {@code final}, so it is never closed and its Apache connection pool leaks.
+     * This handle owns both resources so a caller can release them together. {@link #close()} closes the credentials
+     * provider first -- stopping its background session refresh -- and then the STS client.
+     */
+    public static final class CredentialsProviderHandle implements AutoCloseable {
+        private final AwsCredentialsProvider provider;
+        @Nullable
+        private final SdkAutoCloseable stsClient;
+
+        CredentialsProviderHandle(AwsCredentialsProvider provider, @Nullable SdkAutoCloseable stsClient) {
+            this.provider = provider;
+            this.stsClient = stsClient;
+        }
+
+        public AwsCredentialsProvider provider() {
+            return provider;
+        }
+
+        @Override
+        public void close() {
+            if (provider instanceof SdkAutoCloseable closeableProvider) {
+                IoUtils.closeQuietly(closeableProvider, LOG);
+            }
+            IoUtils.closeQuietly(stsClient, LOG);
+        }
     }
 }

graylog2-server/src/test/java/org/graylog/integrations/aws/AWSAuthFactoryTest.java

@@ -34,9 +34,11 @@
 
 import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
 import static org.assertj.core.api.AssertionsForClassTypes.assertThatExceptionOfType;
+import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 public class AWSAuthFactoryTest {
@@ -214,4 +216,48 @@ public void testExternalIdIgnoredWhenNoAssumeRole() {
         assertThat(withoutExternalId).isExactlyInstanceOf(StaticCredentialsProvider.class);
         assertThat(withExternalId).isExactlyInstanceOf(StaticCredentialsProvider.class);
     }
+
+    // --- createCloseable: the returned handle must close the STS client that the credentials provider never owns.
+    // StsCredentialsProvider.close() only clears the session cache, leaving the StsClient's Apache connection pool
+    // to leak on every call; the handle is what lets a caller release it. ---
+
+    @Test
+    public void createCloseable_assumeRole_closeClosesStsClientAndProvider() {
+        final StsClient mockStsClient = buildMockStsClient();
+        final StsClientBuilder mockStsClientBuilder = mock(StsClientBuilder.class);
+        when(mockStsClientBuilder.region(any())).thenReturn(mockStsClientBuilder);
+        when(mockStsClientBuilder.credentialsProvider(any())).thenReturn(mockStsClientBuilder);
+        when(mockStsClientBuilder.build()).thenReturn(mockStsClient);
+
+        final StsAssumeRoleCredentialsProvider mockProvider = mock(StsAssumeRoleCredentialsProvider.class);
+        final StsAssumeRoleCredentialsProvider.Builder mockProviderBuilder = mock(StsAssumeRoleCredentialsProvider.Builder.class);
+        when(mockProviderBuilder.refreshRequest(any(AssumeRoleRequest.class))).thenReturn(mockProviderBuilder);
+        when(mockProviderBuilder.stsClient(any())).thenReturn(mockProviderBuilder);
+        when(mockProviderBuilder.build()).thenReturn(mockProvider);
+
+        try (MockedStatic<StsClient> mockedStsClient = mockStatic(StsClient.class);
+             MockedStatic<StsAssumeRoleCredentialsProvider> mockedProvider = mockStatic(StsAssumeRoleCredentialsProvider.class)) {
+            mockedStsClient.when(StsClient::builder).thenReturn(mockStsClientBuilder);
+            mockedProvider.when(StsAssumeRoleCredentialsProvider::builder).thenReturn(mockProviderBuilder);
+
+            final AWSAuthFactory.CredentialsProviderHandle handle = awsAuthFactory.createCloseable(
+                    false, "us-east-1", "key", "secret", "arn:aws:iam::123456789012:role/TestRole", null, (ApacheHttpClient.Builder) null);
+
+            assertThat(handle.provider()).isSameAs(mockProvider);
+
+            handle.close();
+
+            verify(mockStsClient).close();
+            verify(mockProvider).close();
+        }
+    }
+
+    @Test
+    public void createCloseable_staticCredentials_hasNoStsClientAndCloseIsSafe() {
+        final AWSAuthFactory.CredentialsProviderHandle handle = awsAuthFactory.createCloseable(
+                false, null, "key", "secret", null, null, (ApacheHttpClient.Builder) null);
+
+        assertThat(handle.provider()).isExactlyInstanceOf(StaticCredentialsProvider.class);
+        assertThatCode(handle::close).doesNotThrowAnyException();
+    }
 }