From d7207f7ef8591c1a0bf2ec3ccdabece8cc7cd063 Mon Sep 17 00:00:00 2001
From: Daniel Worthington-Bodart <dan@bodar.com>
Date: Mon, 11 Sep 2017 14:38:52 +0100
Subject: [PATCH] Added support for X-Forwarded-Host in reverse proxy scenarios

---
 src/com/googlecode/utterlyidle/BaseUri.java   | 17 ++++------
 .../googlecode/utterlyidle/HttpHeaders.java   |  1 +
 .../googlecode/utterlyidle/BaseUriTest.java   | 34 +++++++++++++++++++
 3 files changed, 41 insertions(+), 11 deletions(-)
 create mode 100644 test/com/googlecode/utterlyidle/BaseUriTest.java

diff --git a/src/com/googlecode/utterlyidle/BaseUri.java b/src/com/googlecode/utterlyidle/BaseUri.java
index 1ff3ce5c..449705fa 100644
--- a/src/com/googlecode/utterlyidle/BaseUri.java
+++ b/src/com/googlecode/utterlyidle/BaseUri.java
@@ -4,14 +4,14 @@
 import com.googlecode.totallylazy.io.Uri;
 
 import static com.googlecode.totallylazy.io.Uri.uri;
+import static com.googlecode.utterlyidle.HttpHeaders.HOST;
+import static com.googlecode.utterlyidle.HttpHeaders.X_FORWARDED_HOST;
 import static com.googlecode.utterlyidle.HttpHeaders.X_FORWARDED_PROTO;
 import static java.lang.String.format;
 
-public class BaseUri implements Value<Uri> {
-    private final Uri uri;
-
+public class BaseUri extends Value.Type<Uri> implements Value<Uri> {
     public BaseUri(Uri uri) {
-        this.uri = uri;
+        super(uri);
     }
 
     public static BaseUri baseUri(Uri uri){
@@ -23,7 +23,7 @@ public static BaseUri baseUri(String uri){
     }
 
     public static BaseUri baseUri(Request request, BasePath basePath) {
-        String host = request.headers().getValue(HttpHeaders.HOST);
+        String host = request.headers().valueOption(X_FORWARDED_HOST).getOrElse(request.headers().getValue(HOST));
         if (host == null) {
             return new BaseUri(uri(basePath.toString()));
         }
@@ -33,13 +33,8 @@ public static BaseUri baseUri(Request request, BasePath basePath) {
         return new BaseUri(uri(format("%s://%s%s", scheme, host, basePath)));
     }
 
-    @Override
-    public Uri value() {
-        return uri;
-    }
-
     @Override
     public String toString() {
-        return uri.toString();
+        return value().toString();
     }
 }
diff --git a/src/com/googlecode/utterlyidle/HttpHeaders.java b/src/com/googlecode/utterlyidle/HttpHeaders.java
index 059032e5..4ad561b3 100644
--- a/src/com/googlecode/utterlyidle/HttpHeaders.java
+++ b/src/com/googlecode/utterlyidle/HttpHeaders.java
@@ -30,6 +30,7 @@ public class HttpHeaders {
     public static final String SET_COOKIE = "Set-Cookie";
     public static final String X_FORWARDED_FOR = "X-Forwarded-For";
     public static final String X_FORWARDED_PROTO = "X-Forwarded-Proto";
+    public static final String X_FORWARDED_HOST = "X-Forwarded-Host";
     public static final String X_FRAME_OPTIONS = "X-Frame-Options";
     public static final String X_CORRELATION_ID = "X-CorrelationID";
     public static final String TRANSFER_ENCODING = "Transfer-Encoding";
diff --git a/test/com/googlecode/utterlyidle/BaseUriTest.java b/test/com/googlecode/utterlyidle/BaseUriTest.java
new file mode 100644
index 00000000..69bd7187
--- /dev/null
+++ b/test/com/googlecode/utterlyidle/BaseUriTest.java
@@ -0,0 +1,34 @@
+package com.googlecode.utterlyidle;
+
+import org.junit.Test;
+
+import static com.googlecode.totallylazy.Assert.assertThat;
+import static com.googlecode.totallylazy.predicates.Predicates.is;
+import static com.googlecode.utterlyidle.BasePath.basePath;
+import static com.googlecode.utterlyidle.BaseUri.baseUri;
+import static com.googlecode.utterlyidle.HttpHeaders.HOST;
+import static com.googlecode.utterlyidle.HttpHeaders.X_FORWARDED_HOST;
+import static com.googlecode.utterlyidle.HttpHeaders.X_FORWARDED_PROTO;
+import static com.googlecode.utterlyidle.Request.get;
+
+public class BaseUriTest {
+    @Test
+    public void whenNoHostDefaultToBasePath() throws Exception {
+        assertThat(baseUri(get("foo"), basePath("/root/")), is(baseUri("/root/")));
+    }
+
+    @Test
+    public void whenHostIsSetUseThat() throws Exception {
+        assertThat(baseUri(get("foo").header(HOST, "server"), basePath("/root/")), is(baseUri("http://server/root/")));
+    }
+
+    @Test
+    public void whenXForwardedHostIsSetUseThatOverHost() throws Exception {
+        assertThat(baseUri(get("foo").header(HOST, "internal").header(X_FORWARDED_HOST, "external"), basePath("/root/")), is(baseUri("http://external/root/")));
+    }
+
+    @Test
+    public void whenXForwardedProtoIsSetUseThatOverDefaultHttp() throws Exception {
+        assertThat(baseUri(get("foo").header(HOST, "server").header(X_FORWARDED_PROTO, "https"), basePath("/root/")), is(baseUri("https://server/root/")));
+    }
+}