From 2c6fa4d1b331295d891c327c0aa12e74208465a7 Mon Sep 17 00:00:00 2001
From: Vladimir Lyubitelev <vladimir.lyubitelev@optile.de>
Date: Tue, 15 Apr 2014 12:12:36 +0200
Subject: [PATCH] #170 - LinkBuilder takes X-Forwarded-Port header into account
 if present.

ControllerLinkBuilder now inspects the X-Forwarded-Port and uses that if present. If an X-Forwarded-Host header doesn't contain a :, we now also reset the port to prevent the local one from leaking into links generated.
---
 .../hateoas/mvc/ControllerLinkBuilder.java    | 13 +++++++--
 .../mvc/ControllerLinkBuilderUnitTest.java    | 29 +++++++++++++++++++
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/src/main/java/org/springframework/hateoas/mvc/ControllerLinkBuilder.java b/src/main/java/org/springframework/hateoas/mvc/ControllerLinkBuilder.java
index 4f335228..9e74e3d7 100755
--- a/src/main/java/org/springframework/hateoas/mvc/ControllerLinkBuilder.java
+++ b/src/main/java/org/springframework/hateoas/mvc/ControllerLinkBuilder.java
@@ -174,13 +174,13 @@ public class ControllerLinkBuilder extends LinkBuilderSupport<ControllerLinkBuil
 			builder.scheme("https");
 		}
 
-		String header = request.getHeader("X-Forwarded-Host");
+		String host = request.getHeader("X-Forwarded-Host");
 
-		if (!StringUtils.hasText(header)) {
+		if (!StringUtils.hasText(host)) {
 			return builder;
 		}
 
-		String[] hosts = StringUtils.commaDelimitedListToStringArray(header);
+		String[] hosts = StringUtils.commaDelimitedListToStringArray(host);
 		String hostToUse = hosts[0];
 
 		if (hostToUse.contains(":")) {
@@ -192,6 +192,13 @@ public class ControllerLinkBuilder extends LinkBuilderSupport<ControllerLinkBuil
 
 		} else {
 			builder.host(hostToUse);
+			builder.port(-1); // reset port if it was forwarded from default port
+		}
+
+		String port = request.getHeader("X-Forwarded-Port");
+
+		if (StringUtils.hasText(port)) {
+			builder.port(Integer.parseInt(port));
 		}
 
 		return builder;
diff --git a/src/test/java/org/springframework/hateoas/mvc/ControllerLinkBuilderUnitTest.java b/src/test/java/org/springframework/hateoas/mvc/ControllerLinkBuilderUnitTest.java
index 40bcb411..06754413 100644
--- a/src/test/java/org/springframework/hateoas/mvc/ControllerLinkBuilderUnitTest.java
+++ b/src/test/java/org/springframework/hateoas/mvc/ControllerLinkBuilderUnitTest.java
@@ -29,6 +29,7 @@ import org.springframework.hateoas.Identifiable;
 import org.springframework.hateoas.Link;
 import org.springframework.hateoas.TestUtils;
 import org.springframework.http.HttpEntity;
+import org.springframework.mock.web.MockHttpServletRequest;
 import org.springframework.util.MultiValueMap;
 import org.springframework.web.bind.annotation.PathVariable;
 import org.springframework.web.bind.annotation.RequestBody;
@@ -295,6 +296,34 @@ public class ControllerLinkBuilderUnitTest extends TestUtils {
 		linkTo(methodOn(ControllerWithMethods.class).methodWithRequestParam(null));
 	}
 
+	/**
+	 * @see #170
+	 */
+	@Test
+	public void usesForwardedPortFromHeader() {
+
+		request.addHeader("X-Forwarded-Host", "foobarhost");
+		request.addHeader("X-Forwarded-Port", "9090");
+		request.setServerPort(8080);
+
+		Link link = linkTo(PersonControllerImpl.class).withSelfRel();
+
+		assertThat(link.getHref(), startsWith("http://foobarhost:9090/"));
+	}
+
+	/**
+	 * @see #170
+	 */
+	@Test
+	public void usesForwardedHostFromHeaderWithDefaultPort() {
+
+		request.addHeader("X-Forwarded-Host", "foobarhost");
+		request.setServerPort(8080);
+
+		Link link = linkTo(PersonControllerImpl.class).withSelfRel();
+		assertThat(link.getHref(), startsWith("http://foobarhost/"));
+	}
+
 	private static UriComponents toComponents(Link link) {
 		return UriComponentsBuilder.fromUriString(link.getHref()).build();
 	}