Core: Fix Connection Header Issue (#114)

prebid · Mar 4, 2024 · 7e056a0 · 7e056a0
1 parent 39a8b17
commit 7e056a0
Showing 2 changed files with 52 additions and 33 deletions.
diff --git a/src/main/java/org/prebid/cache/builders/PrebidServerResponseBuilder.java b/src/main/java/org/prebid/cache/builders/PrebidServerResponseBuilder.java
@@ -18,6 +18,7 @@
 import java.time.ZonedDateTime;
 import java.time.format.DateTimeFormatter;
 import java.util.Date;
+import java.util.List;
 import java.util.function.Predicate;
 
 import static org.springframework.web.reactive.function.BodyInserters.fromValue;
@@ -26,8 +27,10 @@
 @Component
 @Slf4j
 public class PrebidServerResponseBuilder {
+
     private static final String HEADER_CONNECTION_KEEPALIVE = "keep-alive";
     private static final String HEADER_CONNECTION_CLOSE = "close";
+
     private final ApiConfig apiConfig;
 
     @Autowired
@@ -48,13 +51,13 @@ public Mono<ServerResponse> createResponseMono(final ServerRequest request,
     }
 
     private ServerResponse.BodyBuilder ok(final ServerRequest request, final MediaType mediaType) {
-        ServerResponse.BodyBuilder builder =
-                ServerResponse.ok()
-                        .contentType(mediaType)
-                        .header(HttpHeaders.DATE, ZonedDateTime.now().format(DateTimeFormatter.RFC_1123_DATE_TIME))
-                        .varyBy(HttpHeaders.ACCEPT_ENCODING)
-                        .cacheControl(CacheControl.noCache());
-        builder = applyHeaders(builder, request);
+        final String now = ZonedDateTime.now().format(DateTimeFormatter.RFC_1123_DATE_TIME);
+        ServerResponse.BodyBuilder builder = ServerResponse.ok()
+                                                     .contentType(mediaType)
+                                                     .header(HttpHeaders.DATE, now)
+                                                     .varyBy(HttpHeaders.ACCEPT_ENCODING)
+                                                     .cacheControl(CacheControl.noCache());
+        applyHeaders(builder, request);
         return builder;
     }
 
@@ -87,26 +90,21 @@ private static ServerResponse.BodyBuilder addHeaders(final ServerResponse.BodyBu
 
     private static ServerResponse.BodyBuilder applyHeaders(final ServerResponse.BodyBuilder builder,
                                                            final ServerRequest request) {
-        if (isConnectionKeepAlive(request))
+
+        final List<String> connectionHeaders = request.headers().header(HttpHeaders.CONNECTION);
+        if (hasConnectionValue(connectionHeaders, HEADER_CONNECTION_KEEPALIVE)) {
             builder.header(HttpHeaders.CONNECTION, HEADER_CONNECTION_KEEPALIVE);
-        if (isConnectionClose(request))
+        }
+        if (hasConnectionValue(connectionHeaders, HEADER_CONNECTION_CLOSE)) {
             builder.header(HttpHeaders.CONNECTION, HEADER_CONNECTION_CLOSE);
+        }
         return builder;
     }
 
-    private static boolean isConnectionKeepAlive(final ServerRequest request) {
-        return request.headers()
-                .header(HttpHeaders.CONNECTION)
-                .stream()
-                .map(String::toLowerCase)
-                .allMatch(Predicate.isEqual(PrebidServerResponseBuilder.HEADER_CONNECTION_KEEPALIVE));
+    private static boolean hasConnectionValue(List<String> connectionHeaders, String value) {
+        return !connectionHeaders.isEmpty() && connectionHeaders.stream()
+                                                       .map(String::toLowerCase)
+                                                       .allMatch(Predicate.isEqual(value));
     }
 
-    private static boolean isConnectionClose(final ServerRequest request) {
-        return request.headers()
-                .header(HttpHeaders.CONNECTION)
-                .stream()
-                .map(String::toLowerCase)
-                .allMatch(Predicate.isEqual(PrebidServerResponseBuilder.HEADER_CONNECTION_CLOSE));
-    }
 }
diff --git a/src/test/java/org/prebid/cache/builders/PrebidServerResponseBuilderTests.java b/src/test/java/org/prebid/cache/builders/PrebidServerResponseBuilderTests.java
@@ -11,6 +11,7 @@
 import org.prebid.cache.routers.ApiConfig;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.mock.web.reactive.function.server.MockServerRequest;
@@ -22,7 +23,6 @@
 import reactor.core.publisher.Signal;
 import reactor.test.StepVerifier;
 
-import java.util.Collections;
 import java.util.function.Consumer;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -56,20 +56,21 @@ private void subscribeAndVerify(final Mono<ServerResponse> mono,
 
         mono.doOnEach(consumer).subscribe();
         StepVerifier.create(mono)
-                    .expectSubscription()
-                    .expectNextMatches(t -> true)
-                    .expectComplete()
-                    .verify();
+                .expectSubscription()
+                .expectNextMatches(t -> true)
+                .expectComplete()
+                .verify();
     }
 
-    private void verifyServerResponse(MediaType mediaType) {
-        final var request = MockServerRequest.builder().build();
+    private void verifyServerResponse(MediaType mediaType, HttpHeaders requestHeaders, HttpHeaders expectedHeaders) {
+        final var request = MockServerRequest.builder().headers(requestHeaders).build();
         final Consumer<Signal<ServerResponse>> consumer = signal -> {
             assertTrue(signal.isOnComplete());
 
             final ServerResponse response = signal.get();
             assertEquals(200, response.statusCode().value());
-            assertTrue(response.headers().containsValue(Collections.singletonList(mediaType.toString())));
+            assertEquals(response.headers().getContentType(), expectedHeaders.getContentType());
+            assertEquals(response.headers().getConnection(), expectedHeaders.getConnection());
         };
 
         subscribeAndVerify(createResponseMono(request, mediaType), consumer);
@@ -99,16 +100,36 @@ private Mono<ServerResponse> createErrorMono(final ServerRequest request, final
 
     @Test
     void verifyXmlServerResponse() {
-        verifyServerResponse(APPLICATION_XML);
+        final HttpHeaders requestHeaders = new HttpHeaders();
+        requestHeaders.add(HttpHeaders.CONNECTION, "keep-alive");
+
+        final HttpHeaders responseHeaders = new HttpHeaders();
+        responseHeaders.add(HttpHeaders.CONNECTION, "keep-alive");
+        responseHeaders.add(HttpHeaders.CONTENT_TYPE, APPLICATION_XML.toString());
+
+        verifyServerResponse(APPLICATION_XML, requestHeaders, responseHeaders);
     }
 
     @Test
     void verifyJsonServerResponse() {
-        verifyServerResponse(APPLICATION_JSON_UTF8);
+        final HttpHeaders requestHeaders = new HttpHeaders();
+        requestHeaders.add(HttpHeaders.CONNECTION, "close");
+
+        final HttpHeaders responseHeaders = new HttpHeaders();
+        responseHeaders.add(HttpHeaders.CONNECTION, "close");
+        responseHeaders.add(HttpHeaders.CONTENT_TYPE, APPLICATION_JSON_UTF8.toString());
+
+        verifyServerResponse(APPLICATION_JSON_UTF8, requestHeaders, responseHeaders);
     }
 
     @Test
-    void verifyJsonUTF8ServerResponse() { verifyServerResponse(APPLICATION_JSON); }
+    void verifyJsonUTF8ServerResponse() {
+        final HttpHeaders requestHeaders = new HttpHeaders();
+
+        final HttpHeaders responseHeaders = new HttpHeaders();
+        responseHeaders.add(HttpHeaders.CONTENT_TYPE, APPLICATION_JSON.toString());
+        verifyServerResponse(APPLICATION_JSON, requestHeaders, responseHeaders);
+    }
 
     @Test
     void verifyNotFound() { verifyErrorResponse(HttpStatus.NOT_FOUND); }