Include request URI’s port, if any, in default Host header

Previously, OperationRequestFactory would add a Host header if one did
not already exist in the request’s headers, however it did not include
the port. This meant that when the request was being made to a
non-standard port (a port other than 80 for an HTTP request and 443 for
an HTTPS request) the Host header was incorrect.

This commit updates OperationRequestFactory to check the URI for a port
and, if it has one, include it in the Host header.
UriModifyingOperationPreprocessor has also been updated to correctly
include the port in the Host header.

Closes gh-269

Author: wilkinsona

spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/OperationRequestFactory.java

@@ -100,10 +100,17 @@ public OperationRequest createFrom(OperationRequest original,
 	private HttpHeaders augmentHeaders(HttpHeaders originalHeaders, URI uri,
 			byte[] content) {
 		return new HttpHeadersHelper(originalHeaders)
-				.addIfAbsent(HttpHeaders.HOST, uri.getHost())
+				.addIfAbsent(HttpHeaders.HOST, createHostHeader(uri))
 				.setContentLengthHeader(content).getHeaders();
 	}
 
+	private String createHostHeader(URI uri) {
+		if (uri.getPort() == -1) {
+			return uri.getHost();
+		}
+		return uri.getHost() + ":" + uri.getPort();
+	}
+
 	private HttpHeaders getUpdatedHeaders(HttpHeaders originalHeaders,
 			byte[] updatedContent) {
 		return new HttpHeadersHelper(originalHeaders)

spring-restdocs-core/src/test/java/org/springframework/restdocs/http/HttpRequestSnippetTests.java

@@ -59,6 +59,16 @@ public void getRequest() throws IOException {
 				.request("http://localhost/foo").header("Alpha", "a").build());
 	}
 
+	@Test
+	public void getRequestWithPort() throws IOException {
+		this.snippet.expectHttpRequest("get-request")
+				.withContents(httpRequest(RequestMethod.GET, "/foo").header("Alpha", "a")
+						.header(HttpHeaders.HOST, "localhost:8080"));
+
+		new HttpRequestSnippet().document(operationBuilder("get-request")
+				.request("http://localhost:8080/foo").header("Alpha", "a").build());
+	}
+
 	@Test
 	public void getRequestWithQueryString() throws IOException {
 		this.snippet.expectHttpRequest("get-request-with-query-string")

spring-restdocs-mockmvc/src/test/java/org/springframework/restdocs/mockmvc/MockMvcRestDocumentationIntegrationTests.java

@@ -375,7 +375,7 @@ public void preprocessedRequest() throws Exception {
 										.header("Content-Type", "application/json")
 										.header("Accept",
 												MediaType.APPLICATION_JSON_VALUE)
-										.header("Host", "localhost")
+										.header("Host", "localhost:8080")
 										.header("Content-Length", "13")
 										.content("{\"a\":\"alpha\"}"))));
 		String prettyPrinted = String.format("{%n  \"a\" : \"<<beta>>\"%n}");

spring-restdocs-restassured/src/main/java/org/springframework/restdocs/restassured/operation/preprocess/UriModifyingOperationPreprocessor.java

@@ -16,6 +16,7 @@
 
 package org.springframework.restdocs.restassured.operation.preprocess;
 
+import java.net.URI;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -125,8 +126,10 @@ public OperationRequest preprocess(OperationRequest request) {
 		if (this.scheme != null) {
 			uriBuilder.scheme(this.scheme);
 		}
+		HttpHeaders modifiedHeaders = modify(request.getHeaders());
 		if (this.host != null) {
 			uriBuilder.host(this.host);
+			modifiedHeaders.set(HttpHeaders.HOST, this.host);
 		}
 		if (this.port != null) {
 			if (StringUtils.hasText(this.port)) {
@@ -136,10 +139,9 @@ public OperationRequest preprocess(OperationRequest request) {
 				uriBuilder.port(null);
 			}
 		}
-		HttpHeaders modifiedHeaders = modify(request.getHeaders());
-		if (this.host != null) {
-			modifiedHeaders.set(HttpHeaders.HOST, this.host);
-		}
+		URI modifiedUri = uriBuilder.build(true).toUri();
+		modifiedHeaders.set(HttpHeaders.HOST, modifiedUri.getHost()
+				+ (modifiedUri.getPort() == -1 ? "" : ":" + modifiedUri.getPort()));
 		return this.contentModifyingDelegate.preprocess(
 				new OperationRequestFactory().create(uriBuilder.build(true).toUri(),
 						request.getMethod(), request.getContent(), modifiedHeaders,

spring-restdocs-restassured/src/test/java/org/springframework/restdocs/restassured/RestAssuredRequestConverterTests.java

@@ -142,7 +142,7 @@ public void headers() {
 		assertThat(request.getHeaders().get("Foo"), is(equalTo(Arrays.asList("bar"))));
 		assertThat(request.getHeaders().get("Accept"), is(equalTo(Arrays.asList("*/*"))));
 		assertThat(request.getHeaders().get("Host"),
-				is(equalTo(Arrays.asList("localhost"))));
+				is(equalTo(Arrays.asList("localhost:" + this.port))));
 	}
 
 	@Test

spring-restdocs-restassured/src/test/java/org/springframework/restdocs/restassured/RestAssuredRestDocumentationIntegrationTests.java

@@ -268,9 +268,9 @@ public void preprocessedRequest() throws Exception {
 				.filter(document("original-request"))
 				.filter(document("preprocessed-request",
 						preprocessRequest(prettyPrint(),
-								removeHeaders("a", HttpHeaders.HOST,
-										HttpHeaders.CONTENT_LENGTH),
-								replacePattern(pattern, "\"<<beta>>\""))))
+								replacePattern(pattern, "\"<<beta>>\""),
+								modifyUris().removePort(),
+								removeHeaders("a", HttpHeaders.CONTENT_LENGTH))))
 				.get("/").then().statusCode(200);
 		assertThat(
 				new File("build/generated-snippets/original-request/http-request.adoc"),
@@ -279,7 +279,7 @@ public void preprocessedRequest() throws Exception {
 								.header("a", "alpha").header("b", "bravo")
 								.header("Accept", MediaType.APPLICATION_JSON_VALUE)
 								.header("Content-Type", "application/json; charset=UTF-8")
-								.header("Host", "localhost")
+								.header("Host", "localhost:" + this.port)
 								.header("Content-Length", "13")
 								.content("{\"a\":\"alpha\"}"))));
 		String prettyPrinted = String.format("{%n  \"a\" : \"<<beta>>\"%n}");
@@ -291,7 +291,7 @@ public void preprocessedRequest() throws Exception {
 								.header("b", "bravo")
 								.header("Accept", MediaType.APPLICATION_JSON_VALUE)
 								.header("Content-Type", "application/json; charset=UTF-8")
-								.content(prettyPrinted))));
+								.header("Host", "localhost").content(prettyPrinted))));
 	}
 
 	@Test

spring-restdocs-restassured/src/test/java/org/springframework/restdocs/restassured/operation/preprocess/UriModifyingOperationPreprocessorTests.java

@@ -63,9 +63,11 @@ public void requestUriSchemeCanBeModified() {
 	public void requestUriHostCanBeModified() {
 		this.preprocessor.host("api.example.com");
 		OperationRequest processed = this.preprocessor
-				.preprocess(createRequestWithUri("http://api.example.com:12345"));
+				.preprocess(createRequestWithUri("http://api.foo.com:12345"));
 		assertThat(processed.getUri(),
 				is(equalTo(URI.create("http://api.example.com:12345"))));
+		assertThat(processed.getHeaders().getFirst(HttpHeaders.HOST),
+				is(equalTo("api.example.com:12345")));
 	}
 
 	@Test
@@ -75,6 +77,8 @@ public void requestUriPortCanBeModified() {
 				.preprocess(createRequestWithUri("http://api.example.com:12345"));
 		assertThat(processed.getUri(),
 				is(equalTo(URI.create("http://api.example.com:23456"))));
+		assertThat(processed.getHeaders().getFirst(HttpHeaders.HOST),
+				is(equalTo("api.example.com:23456")));
 	}
 
 	@Test
@@ -83,6 +87,8 @@ public void requestUriPortCanBeRemoved() {
 		OperationRequest processed = this.preprocessor
 				.preprocess(createRequestWithUri("http://api.example.com:12345"));
 		assertThat(processed.getUri(), is(equalTo(URI.create("http://api.example.com"))));
+		assertThat(processed.getHeaders().getFirst(HttpHeaders.HOST),
+				is(equalTo("api.example.com")));
 	}
 
 	@Test