Include customized Host header in curl and HTTPie request snippets

Previously, the curl and HTTPie request snippets always excluded the
Host header, relying on the Host header that’s auto-generated by curl
and HTTPie instead. This worked well for the most part, but meant that
incorrect snippets were generated when the request was being sent with
a custom host header that did not match the header that would be
auto-generated.

This commit updates CliOperationRequest to only filter out the Host
header if it’s the same as the header that would be auto-generated
by curl or HTTPie.

Closes gh-258

Author: wilkinsona

spring-restdocs-core/src/main/java/org/springframework/restdocs/cli/CliOperationRequest.java

@@ -17,8 +17,8 @@
 package org.springframework.restdocs.cli;
 
 import java.net.URI;
+import java.util.Arrays;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -41,20 +41,15 @@
  */
 final class CliOperationRequest implements OperationRequest {
 
-	private static final Set<HeaderFilter> HEADER_FILTERS;
-
-	static {
-		Set<HeaderFilter> headerFilters = new HashSet<>();
-		headerFilters.add(new NamedHeaderFilter(HttpHeaders.HOST));
-		headerFilters.add(new NamedHeaderFilter(HttpHeaders.CONTENT_LENGTH));
-		headerFilters.add(new BasicAuthHeaderFilter());
-		HEADER_FILTERS = Collections.unmodifiableSet(headerFilters);
-	}
+	private final Set<HeaderFilter> headerFilters;
 
 	private final OperationRequest delegate;
 
 	CliOperationRequest(OperationRequest delegate) {
 		this.delegate = delegate;
+		this.headerFilters = new HashSet<>(Arrays.asList(
+				new NamedHeaderFilter(HttpHeaders.CONTENT_LENGTH),
+				new BasicAuthHeaderFilter(), new HostHeaderFilter(delegate.getUri())));
 	}
 
 	Parameters getUniqueParameters() {
@@ -121,11 +116,20 @@ public HttpHeaders getHeaders() {
 	}
 
 	private boolean allowedHeader(Map.Entry<String, List<String>> header) {
-		for (HeaderFilter headerFilter : HEADER_FILTERS) {
+		for (HeaderFilter headerFilter : this.headerFilters) {
 			if (!headerFilter.allow(header.getKey(), header.getValue())) {
 				return false;
 			}
 		}
+		if (HttpHeaders.HOST.equalsIgnoreCase(header.getKey())) {
+			if (!header.getValue().isEmpty()) {
+				String value = header.getValue().get(0);
+				if (value.equals(this.delegate.getUri().getHost() + ":"
+						+ this.delegate.getUri().getPort())) {
+					return false;
+				}
+			}
+		}
 		return true;
 	}
 
@@ -187,4 +191,27 @@ public boolean allow(String name, List<String> value) {
 
 	}
 
+	private static final class HostHeaderFilter implements HeaderFilter {
+
+		private final URI uri;
+
+		private HostHeaderFilter(URI uri) {
+			this.uri = uri;
+		}
+
+		@Override
+		public boolean allow(String name, List<String> value) {
+			if (value.isEmpty() || this.getImplicitHostHeader().equals(value.get(0))) {
+				return false;
+			}
+			return true;
+		}
+
+		private String getImplicitHostHeader() {
+			return this.uri.getHost()
+					+ ((this.uri.getPort() == -1) ? "" : ":" + this.uri.getPort());
+		}
+
+	}
+
 }

spring-restdocs-core/src/test/java/org/springframework/restdocs/cli/CurlRequestSnippetTests.java

@@ -330,4 +330,18 @@ public void customAttributes() throws IOException {
 												.request("http://localhost/foo").build());
 	}
 
+	@Test
+	public void customHostHeaderIsIncluded() throws IOException {
+		this.snippet.expectCurlRequest("custom-host-header")
+				.withContents(codeBlock("bash").content(
+						"$ curl 'http://localhost/foo' -i" + " -H 'Host: api.example.com'"
+								+ " -H 'Content-Type: application/json' -H 'a: alpha'"));
+		new CurlRequestSnippet().document(
+				operationBuilder("custom-host-header").request("http://localhost/foo")
+						.header(HttpHeaders.HOST, "api.example.com")
+						.header(HttpHeaders.CONTENT_TYPE,
+								MediaType.APPLICATION_JSON_VALUE)
+						.header("a", "alpha").build());
+	}
+
 }

spring-restdocs-core/src/test/java/org/springframework/restdocs/cli/HttpieRequestSnippetTests.java

@@ -332,4 +332,18 @@ public void customAttributes() throws IOException {
 												.request("http://localhost/foo").build());
 	}
 
+	@Test
+	public void customHostHeaderIsIncluded() throws IOException {
+		this.snippet.expectHttpieRequest("custom-host-header")
+				.withContents(codeBlock("bash").content(
+						"$ http GET 'http://localhost/foo' 'Host:api.example.com'"
+								+ " 'Content-Type:application/json' 'a:alpha'"));
+		new HttpieRequestSnippet().document(
+				operationBuilder("custom-host-header").request("http://localhost/foo")
+						.header(HttpHeaders.HOST, "api.example.com")
+						.header(HttpHeaders.CONTENT_TYPE,
+								MediaType.APPLICATION_JSON_VALUE)
+						.header("a", "alpha").build());
+	}
+
 }