diff --git a/spring-cloud-vault-config/src/main/java/org/springframework/cloud/vault/config/VaultConfigDataLoader.java b/spring-cloud-vault-config/src/main/java/org/springframework/cloud/vault/config/VaultConfigDataLoader.java index 39d191a5..7e4cb5e8 100644 --- a/spring-cloud-vault-config/src/main/java/org/springframework/cloud/vault/config/VaultConfigDataLoader.java +++ b/spring-cloud-vault-config/src/main/java/org/springframework/cloud/vault/config/VaultConfigDataLoader.java @@ -19,6 +19,7 @@ import java.lang.reflect.Field; import java.lang.reflect.Modifier; import java.util.Arrays; +import java.util.Collection; import java.util.Collections; import java.util.List; import java.util.concurrent.atomic.AtomicReference; @@ -27,8 +28,11 @@ import java.util.function.Supplier; import org.apache.commons.logging.Log; + import org.springframework.beans.factory.BeanFactory; +import org.springframework.beans.factory.DisposableBean; import org.springframework.beans.factory.config.ConfigurableListableBeanFactory; +import org.springframework.beans.factory.support.DefaultSingletonBeanRegistry; import org.springframework.boot.BootstrapContext; import org.springframework.boot.BootstrapRegistry; import org.springframework.boot.BootstrapRegistryInitializer; @@ -266,7 +270,8 @@ private void registerSecretLeaseContainer(ConfigurableBootstrapContext bootstrap container.start(); return container; - }, ConfigurableApplicationContext::registerShutdownHook); + }, ConfigurableApplicationContext::registerShutdownHook, + List.of("vaultTaskScheduler", "vaultSessionManager", "reactiveVaultSessionManager")); } private PropertySource createVaultPropertySource(VaultConfigOperations configOperations, boolean failFast, @@ -343,23 +348,24 @@ private RequestedSecret getRequestedSecret(SecretBackendMetadata accessor) { static void registerIfAbsent(ConfigurableBootstrapContext bootstrap, String beanName, Class instanceType, Supplier instanceSupplier) { registerIfAbsent(bootstrap, beanName, instanceType, ctx -> instanceSupplier.get(), ctx -> { - }); + }, Collections.emptyList()); } static void registerIfAbsent(ConfigurableBootstrapContext bootstrap, String beanName, Class instanceType, Supplier instanceSupplier, Consumer contextCustomizer) { - registerIfAbsent(bootstrap, beanName, instanceType, ctx -> instanceSupplier.get(), contextCustomizer); + registerIfAbsent(bootstrap, beanName, instanceType, ctx -> instanceSupplier.get(), contextCustomizer, + Collections.emptyList()); } static void registerIfAbsent(ConfigurableBootstrapContext bootstrap, String beanName, Class instanceType, Function instanceSupplier) { registerIfAbsent(bootstrap, beanName, instanceType, instanceSupplier, ctx -> { - }); + }, Collections.emptyList()); } static void registerIfAbsent(ConfigurableBootstrapContext bootstrap, String beanName, Class instanceType, - Function instanceSupplier, - Consumer contextCustomizer) { + Function instanceSupplier, Consumer contextCustomizer, + Collection dependsOn) { bootstrap.registerIfAbsent(instanceType, instanceSupplier::apply); @@ -377,6 +383,17 @@ static void registerIfAbsent(ConfigurableBootstrapContext bootstrap, String T instance = event.getBootstrapContext().get(instanceType); factory.registerSingleton(beanName, instance); + + if (instance instanceof DisposableBean db) { + + if (factory instanceof DefaultSingletonBeanRegistry dsbr) { + dsbr.registerDisposableBean(beanName, db); + + for (String dependencyName : dependsOn) { + dsbr.registerDependentBean(dependencyName, beanName); + } + } + } }); } @@ -504,7 +521,8 @@ void registerVaultSessionManager() { ctx.get(RestTemplateFactory.class)); reconfigureLogger(sessionManager, this.logFactory); return sessionManager; - }); + }, ctx -> { + }, List.of("clientHttpRequestFactoryWrapper")); } } @@ -575,7 +593,9 @@ void registerReactiveSessionManager() { registerIfAbsent(this.bootstrap, "reactiveVaultSessionManager", ReactiveSessionManager.class, ctx -> this.configuration.createReactiveSessionManager(ctx.get(VaultTokenSupplier.class), () -> ctx.get(TaskSchedulerWrapper.class).getTaskScheduler(), - ctx.get(WebClientFactory.class))); + ctx.get(WebClientFactory.class)), + ctx -> { + }, List.of("clientHttpConnectorWrapper")); } void registerSessionManager() { diff --git a/spring-cloud-vault-config/src/test/java/org/springframework/cloud/vault/config/ConfigDataShutdownTests.java b/spring-cloud-vault-config/src/test/java/org/springframework/cloud/vault/config/ConfigDataShutdownTests.java new file mode 100644 index 00000000..640d9c39 --- /dev/null +++ b/spring-cloud-vault-config/src/test/java/org/springframework/cloud/vault/config/ConfigDataShutdownTests.java @@ -0,0 +1,52 @@ +/* + * Copyright 2016-2021 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.vault.config; + +import org.junit.Test; + +import org.springframework.boot.SpringApplication; +import org.springframework.boot.autoconfigure.SpringBootApplication; +import org.springframework.boot.builder.SpringApplicationBuilder; +import org.springframework.cloud.vault.util.IntegrationTestSupport; +import org.springframework.context.ConfigurableApplicationContext; +import org.springframework.test.util.ReflectionTestUtils; +import org.springframework.vault.core.lease.SecretLeaseContainer; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Tests verifying shutdown behavior. + * + * @author Mark Paluch + */ +@SpringBootApplication +public class ConfigDataShutdownTests extends IntegrationTestSupport { + + @Test + public void contextShutdownDestroysSecretLeaseContainer() { + ConfigurableApplicationContext context = new SpringApplicationBuilder().sources(ConfigDataShutdownTests.class) + .run("--server.port=0", "--spring.cloud.bootstrap.enabled=false", "--spring.cloud.vault.failFast=true", + "--spring.cloud.vault.config.lifecycle.enabled=true", "--spring.config.import=vault://"); + + SecretLeaseContainer container = context.getBean(SecretLeaseContainer.class); + + assertThat((Integer) ReflectionTestUtils.getField(container, "status")).isEqualTo(1); // started + SpringApplication.exit(context); + assertThat((Integer) ReflectionTestUtils.getField(container, "status")).isEqualTo(2); // destroyed + } + +}