Fix proxy options additionalRequestHeaders bug (#1274)

* Add test for proxyBaseUrl with additional request header * Fix proxy options additionalRequestHeaders bug
wiremock · Mar 5, 2020 · 2d40d90 · 2d40d90
1 parent 2010bb1
commit 2d40d90
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 4 deletions.
diff --git a/...github/tomakehurst/wiremock/extension/responsetemplating/ResponseTemplateTransformer.java b/...github/tomakehurst/wiremock/extension/responsetemplating/ResponseTemplateTransformer.java
@@ -181,10 +181,31 @@ public HttpHeader apply(final HttpHeader header) {
         if (responseDefinition.getProxyBaseUrl() != null) {
             HandlebarsOptimizedTemplate proxyBaseUrlTemplate = getTemplate(TemplateCacheKey.forProxyUrl(responseDefinition), responseDefinition.getProxyBaseUrl());
             String newProxyBaseUrl = uncheckedApplyTemplate(proxyBaseUrlTemplate, model);
-            newResponseDefBuilder.proxiedFrom(newProxyBaseUrl);
-        }
 
-        return newResponseDefBuilder.build();
+            ResponseDefinitionBuilder.ProxyResponseDefinitionBuilder newProxyResponseDefBuilder = newResponseDefBuilder.proxiedFrom(newProxyBaseUrl);
+
+            if (responseDefinition.getAdditionalProxyRequestHeaders() != null) {
+                Iterable<HttpHeader> newResponseHeaders = Iterables.transform(responseDefinition.getAdditionalProxyRequestHeaders().all(), new Function<HttpHeader, HttpHeader>() {
+                    @Override
+                    public HttpHeader apply(final HttpHeader header) {
+                        ImmutableList.Builder<String> valueListBuilder = ImmutableList.builder();
+                        int index = 0;
+                        for (String headerValue: header.values()) {
+                            HandlebarsOptimizedTemplate template = getTemplate(TemplateCacheKey.forHeader(responseDefinition, header.key(), index++), headerValue);
+                            valueListBuilder.add(uncheckedApplyTemplate(template, model));
+                        }
+                        return new HttpHeader(header.key(), valueListBuilder.build());
+                    }
+                });
+                HttpHeaders proxyHttpHeaders = new HttpHeaders(newResponseHeaders);
+                for (String key: proxyHttpHeaders.keys()) {
+                    newProxyResponseDefBuilder.withAdditionalRequestHeader(key, proxyHttpHeaders.getHeader(key).firstValue());
+                }
+            }
+            return newProxyResponseDefBuilder.build();
+        } else {
+            return newResponseDefBuilder.build();
+        }
     }
 
     /**

diff --git a/...ub/tomakehurst/wiremock/extension/responsetemplating/ResponseTemplateTransformerTest.java b/...ub/tomakehurst/wiremock/extension/responsetemplating/ResponseTemplateTransformerTest.java
@@ -35,6 +35,7 @@
 import static com.github.tomakehurst.wiremock.testsupport.NoFileSource.noFileSource;
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.is;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThat;
 
 public class ResponseTemplateTransformerTest {
@@ -307,16 +308,24 @@ public void areConditionalHelpersLoaded() {
 
 
     @Test
-    public void proxyBaseUrl() {
+    public void proxyBaseUrlWithAdditionalRequestHeader() {
         ResponseDefinition transformedResponseDef = transform(mockRequest()
                 .url("/things")
                 .header("X-WM-Uri", "http://localhost:8000"),
             aResponse().proxiedFrom("{{request.headers.X-WM-Uri}}")
+                .withAdditionalRequestHeader("X-Origin-Url", "{{request.url}}")
         );
 
         assertThat(transformedResponseDef.getProxyBaseUrl(), is(
             "http://localhost:8000"
         ));
+        assertNotNull(transformedResponseDef.getAdditionalProxyRequestHeaders());
+        assertThat(transformedResponseDef
+                        .getAdditionalProxyRequestHeaders()
+                        .getHeader("X-Origin-Url")
+                        .firstValue(),
+                is("/things")
+        );
     }
 
     @Test