fix: keep all values of multivalued headers instead of keeping only t…

…he last one https://gravitee.atlassian.net/browse/APIM-1906 gravitee-io/issues#9085
gravitee-io · Jun 27, 2023 · b33a3a6 · b33a3a6
1 parent fb637e9
commit b33a3a6
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 4 deletions.
diff --git a/src/main/java/io/gravitee/policy/urlrewriting/URLRewritingPolicy.java b/src/main/java/io/gravitee/policy/urlrewriting/URLRewritingPolicy.java
@@ -29,9 +29,7 @@
 import io.gravitee.policy.api.annotations.OnResponse;
 import io.gravitee.policy.api.annotations.OnResponseContent;
 import io.gravitee.policy.urlrewriting.configuration.URLRewritingPolicyConfiguration;
-import java.util.Map;
-import java.util.Set;
-import java.util.TreeSet;
+import java.util.*;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
@@ -83,7 +81,15 @@ public void onResponse(Request request, Response response, ExecutionContext exec
     private void rewriteHeaders(HttpHeaders headers, ExecutionContext executionContext) {
         LOGGER.debug("Rewrite HTTP response headers");
 
-        headers.names().forEach(header -> headers.set(header, rewrite(headers.get(header), executionContext)));
+        Set<String> names = new HashSet<>(headers.names());
+        names.forEach(headerName -> {
+            List<String> headerValues = headers.getAll(headerName);
+            headers.remove(headerName);
+            headerValues.forEach(headerValue -> {
+                String rewrittenHeaderValue = rewrite(headerValue, executionContext);
+                headers.add(headerName, rewrittenHeaderValue);
+            });
+        });
     }
 
     @OnResponseContent

diff --git a/src/test/java/io/gravitee/policy/urlrewriting/URLRewritingPolicyTest.java b/src/test/java/io/gravitee/policy/urlrewriting/URLRewritingPolicyTest.java
@@ -28,6 +28,7 @@
 import io.gravitee.gateway.api.stream.ReadWriteStream;
 import io.gravitee.policy.api.PolicyChain;
 import io.gravitee.policy.urlrewriting.configuration.URLRewritingPolicyConfiguration;
+import java.util.List;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
@@ -99,6 +100,34 @@ public void test_rewriteHeaders() {
         verify(policyChain).doNext(any(Request.class), any(Response.class));
     }
 
+    @Test
+    public void test_rewriteHeadersWithMultipleValues() {
+        // Prepare
+        final HttpHeaders headers = HttpHeaders
+            .create()
+            .add(HttpHeaderNames.SET_COOKIE, "SID=ABAN12398123NJHJZEHDK123012039301U93274923U4KADNZKN; Path=/test")
+            .add(HttpHeaderNames.SET_COOKIE, "JSESSIONID=123456789; Path=/test");
+
+        when(response.headers()).thenReturn(headers);
+
+        when(configuration.isRewriteResponseHeaders()).thenReturn(true);
+        when(configuration.getFromRegex()).thenReturn("Path=/test");
+        when(configuration.getToReplacement()).thenReturn("Path=/updated-path");
+
+        when(executionContext.getTemplateEngine()).thenReturn(TemplateEngine.templateEngine());
+
+        urlRewritingPolicy.onResponse(request, response, executionContext, policyChain);
+
+        Assert.assertEquals(
+            List.of(
+                "SID=ABAN12398123NJHJZEHDK123012039301U93274923U4KADNZKN; Path=/updated-path",
+                "JSESSIONID=123456789; Path=/updated-path"
+            ),
+            response.headers().getAll(HttpHeaderNames.SET_COOKIE)
+        );
+        verify(policyChain).doNext(any(Request.class), any(Response.class));
+    }
+
     @Test
     public void test_rewriteResponse_disabled() {
         // Prepare