fix: allow to set custom header providers in reactive context (#2142)

* fix: allow to set custom header providers in reactive context

* chore: update copyright year information

* chore: rearrange code and add tests

* Replace lambda with Method reference

Co-authored-by: ulrichschulte <[email protected]>

Author: SteKoe

spring-boot-admin-server/src/main/java/de/codecentric/boot/admin/server/config/AdminServerInstanceWebClientConfiguration.java

@@ -43,6 +43,8 @@
 import de.codecentric.boot.admin.server.web.client.cookies.CookieStoreCleanupTrigger;
 import de.codecentric.boot.admin.server.web.client.cookies.JdkPerInstanceCookieStore;
 import de.codecentric.boot.admin.server.web.client.cookies.PerInstanceCookieStore;
+import de.codecentric.boot.admin.server.web.client.reactive.CompositeReactiveHttpHeadersProvider;
+import de.codecentric.boot.admin.server.web.client.reactive.ReactiveHttpHeadersProvider;
 
 @Configuration(proxyBeanMethods = false)
 @Lazy(false)
@@ -86,6 +88,16 @@ public InstanceExchangeFilterFunction addHeadersInstanceExchangeFilter(
 				return InstanceExchangeFilterFunctions.addHeaders(new CompositeHttpHeadersProvider(headersProviders));
 			}
 
+			@Bean
+			@Order(0)
+			@ConditionalOnBean(ReactiveHttpHeadersProvider.class)
+			@ConditionalOnMissingBean(name = "addReactiveHeadersInstanceExchangeFilter")
+			public InstanceExchangeFilterFunction addReactiveHeadersInstanceExchangeFilter(
+					List<ReactiveHttpHeadersProvider> reactiveHeadersProviders) {
+				return InstanceExchangeFilterFunctions
+						.addHeadersReactive(new CompositeReactiveHttpHeadersProvider(reactiveHeadersProviders));
+			}
+
 			@Bean
 			@Order(10)
 			@ConditionalOnMissingBean(name = "rewriteEndpointUrlInstanceExchangeFilter")

spring-boot-admin-server/src/main/java/de/codecentric/boot/admin/server/web/client/InstanceExchangeFilterFunctions.java

@@ -42,6 +42,7 @@
 import de.codecentric.boot.admin.server.domain.values.InstanceId;
 import de.codecentric.boot.admin.server.web.client.cookies.PerInstanceCookieStore;
 import de.codecentric.boot.admin.server.web.client.exception.ResolveEndpointException;
+import de.codecentric.boot.admin.server.web.client.reactive.ReactiveHttpHeadersProvider;
 
 import static java.util.Arrays.asList;
 import static java.util.Collections.singletonList;
@@ -71,6 +72,15 @@ public static InstanceExchangeFilterFunction addHeaders(HttpHeadersProvider http
 		};
 	}
 
+	public static InstanceExchangeFilterFunction addHeadersReactive(ReactiveHttpHeadersProvider httpHeadersProvider) {
+		return (instance, request, next) -> httpHeadersProvider.getHeaders(instance).flatMap((httpHeaders) -> {
+			ClientRequest requestWithAdditionalHeaders = ClientRequest.from(request)
+					.headers((headers) -> headers.addAll(httpHeaders)).build();
+
+			return next.exchange(requestWithAdditionalHeaders);
+		}).switchIfEmpty(Mono.defer(() -> next.exchange(request)));
+	}
+
 	public static InstanceExchangeFilterFunction rewriteEndpointUrl() {
 		return (instance, request, next) -> {
 			if (request.url().isAbsolute()) {

spring-boot-admin-server/src/main/java/de/codecentric/boot/admin/server/web/client/reactive/CompositeReactiveHttpHeadersProvider.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2014-2022 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 de.codecentric.boot.admin.server.web.client.reactive;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.List;
+
+import org.springframework.http.HttpHeaders;
+import reactor.core.publisher.Mono;
+
+import de.codecentric.boot.admin.server.domain.entities.Instance;
+
+import static java.util.stream.Collectors.toList;
+
+public class CompositeReactiveHttpHeadersProvider implements ReactiveHttpHeadersProvider {
+
+	private final Collection<ReactiveHttpHeadersProvider> delegates;
+
+	public CompositeReactiveHttpHeadersProvider(Collection<ReactiveHttpHeadersProvider> delegates) {
+		this.delegates = delegates;
+	}
+
+	@Override
+	public Mono<HttpHeaders> getHeaders(Instance instance) {
+		List<Mono<HttpHeaders>> headers = delegates.stream()
+				.map((reactiveHttpHeadersProvider) -> reactiveHttpHeadersProvider.getHeaders(instance))
+				.collect(toList());
+
+		return Mono.zip(headers, this::mergeMonosToHeaders);
+	}
+
+	private HttpHeaders mergeMonosToHeaders(Object[] e) {
+		return Arrays.stream(e).map(HttpHeaders.class::cast).reduce(new HttpHeaders(), (h1, h2) -> {
+			h1.addAll(h2);
+			return h1;
+		});
+	}
+
+}

spring-boot-admin-server/src/main/java/de/codecentric/boot/admin/server/web/client/reactive/ReactiveHttpHeadersProvider.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2014-2022 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 de.codecentric.boot.admin.server.web.client.reactive;
+
+import org.springframework.http.HttpHeaders;
+import reactor.core.publisher.Mono;
+
+import de.codecentric.boot.admin.server.domain.entities.Instance;
+
+/**
+ * Is responsible to provide the {@link HttpHeaders} used to interact with the given
+ * {@link Instance}.
+ */
+public interface ReactiveHttpHeadersProvider {
+
+	Mono<HttpHeaders> getHeaders(Instance instance);
+
+}

spring-boot-admin-server/src/test/java/de/codecentric/boot/admin/server/web/client/InstanceExchangeFilterFunctionsTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2014-2021 the original author or authors.
+ * Copyright 2014-2022 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.
@@ -287,6 +287,46 @@ void should_add_headers_from_provider() {
 
 	}
 
+	@Nested
+	class AddHeadersReactive {
+
+		@Test
+		void should_add_headers_from_provider() {
+			InstanceExchangeFilterFunction filter = InstanceExchangeFilterFunctions.addHeadersReactive((i) -> {
+				HttpHeaders headers = new HttpHeaders();
+				headers.add("X-INSTANCE-ID", i.getId().getValue());
+				return Mono.just(headers);
+			});
+
+			ClientRequest request = ClientRequest.create(HttpMethod.GET, URI.create("/test"))
+					.attribute(ATTRIBUTE_INSTANCE, INSTANCE).build();
+
+			Mono<ClientResponse> response = filter.filter(INSTANCE, request, (req) -> {
+				assertThat(req.headers().get("X-INSTANCE-ID")).containsExactly(INSTANCE.getId().getValue());
+				return Mono.just(ClientResponse.create(HttpStatus.OK).build());
+			});
+
+			StepVerifier.create(response).expectNextCount(1).verifyComplete();
+		}
+
+		@Test
+		void should_pass_on_mono_empty() {
+			InstanceExchangeFilterFunction filter = InstanceExchangeFilterFunctions
+					.addHeadersReactive((i) -> Mono.empty());
+
+			ClientRequest request = ClientRequest.create(HttpMethod.GET, URI.create("/test"))
+					.attribute(ATTRIBUTE_INSTANCE, INSTANCE).build();
+
+			Mono<ClientResponse> response = filter.filter(INSTANCE, request, (req) -> {
+				assertThat(req.headers().size()).isEqualTo(0);
+				return Mono.just(ClientResponse.create(HttpStatus.OK).build());
+			});
+
+			StepVerifier.create(response).expectNextCount(1).verifyComplete();
+		}
+
+	}
+
 	@Nested
 	class AddDefaultHeaders {
 

spring-boot-admin-server/src/test/java/de/codecentric/boot/admin/server/web/client/reactive/CompositeReactiveHttpHeadersProviderTest.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2014-2018 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 de.codecentric.boot.admin.server.web.client.reactive;
+
+import org.junit.jupiter.api.Test;
+import org.springframework.http.HttpHeaders;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.emptyList;
+import static java.util.Collections.singletonList;
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class CompositeReactiveHttpHeadersProviderTest {
+
+	@Test
+	public void should_return_all_headers() {
+		ReactiveHttpHeadersProvider provider = new CompositeReactiveHttpHeadersProvider(asList((i) -> {
+			HttpHeaders headers = new HttpHeaders();
+			headers.set("a", "1");
+			headers.set("b", "2-a");
+			return Mono.just(headers);
+		}, (i) -> {
+			HttpHeaders headers = new HttpHeaders();
+			headers.set("b", "2-b");
+			headers.set("c", "3");
+			return Mono.just(headers);
+		}));
+
+		StepVerifier.create(provider.getHeaders(null)).thenConsumeWhile((headers) -> {
+			assertThat(headers.get("a")).isEqualTo(singletonList("1"));
+			assertThat(headers.get("b")).isEqualTo(asList("2-a", "2-b"));
+			assertThat(headers.get("c")).isEqualTo(singletonList("3"));
+			return true;
+		}).verifyComplete();
+	}
+
+	@Test
+	public void should_return_empty_headers() {
+		CompositeReactiveHttpHeadersProvider provider = new CompositeReactiveHttpHeadersProvider(emptyList());
+
+		StepVerifier.create(provider.getHeaders(null)).thenConsumeWhile((headers) -> {
+			assertThat(headers.size()).isEqualTo(0);
+			return true;
+		}).verifyComplete();
+	}
+
+}