From d4f2e84f07032aa89e2b3962df37209b1eca98f7 Mon Sep 17 00:00:00 2001
From: Frank Schroeder <frschroeder@ebay.com>
Date: Thu, 30 Jun 2016 19:28:31 +0200
Subject: [PATCH] Issue #122: Use local port for X-Forwarded-Port

---
 proxy/util.go      | 19 +++++++++++++++++--
 proxy/util_test.go | 47 ++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 62 insertions(+), 4 deletions(-)

diff --git a/proxy/util.go b/proxy/util.go
index 20dbdf28c..9e2e0e0f3 100644
--- a/proxy/util.go
+++ b/proxy/util.go
@@ -5,6 +5,7 @@ import (
 	"log"
 	"net"
 	"net/http"
+	"strings"
 
 	"github.com/eBay/fabio/config"
 	"github.com/eBay/fabio/route"
@@ -19,7 +20,7 @@ import (
 // * TLS connection: Set header with name from `cfg.TLSHeader` to `cfg.TLSHeaderValue`
 //
 func addHeaders(r *http.Request, cfg config.Proxy) error {
-	remoteIP, remotePort, err := net.SplitHostPort(r.RemoteAddr)
+	remoteIP, _, err := net.SplitHostPort(r.RemoteAddr)
 	if err != nil {
 		return errors.New("cannot parse " + r.RemoteAddr)
 	}
@@ -59,7 +60,7 @@ func addHeaders(r *http.Request, cfg config.Proxy) error {
 	}
 
 	if r.Header.Get("X-Forwarded-Port") == "" {
-		r.Header.Set("X-Forwarded-Port", remotePort)
+		r.Header.Set("X-Forwarded-Port", localPort(r))
 	}
 
 	fwd := r.Header.Get("Forwarded")
@@ -96,3 +97,17 @@ func target(r *http.Request) *route.Target {
 	}
 	return t
 }
+
+func localPort(r *http.Request) string {
+	if r == nil {
+		return ""
+	}
+	n := strings.Index(r.Host, ":")
+	if n > 0 && n < len(r.Host)-1 {
+		return r.Host[n+1:]
+	}
+	if r.TLS != nil {
+		return "443"
+	}
+	return "80"
+}
diff --git a/proxy/util_test.go b/proxy/util_test.go
index dd9e5889f..336a1489c 100644
--- a/proxy/util_test.go
+++ b/proxy/util_test.go
@@ -163,10 +163,31 @@ func TestAddHeaders(t *testing.T) {
 			"",
 		},
 
-		{"set X-Forwarded-Port",
+		{"set X-Forwarded-Port for HTTP",
 			&http.Request{RemoteAddr: "1.2.3.4:5555"},
 			config.Proxy{},
-			http.Header{"X-Forwarded-Port": {"5555"}},
+			http.Header{"X-Forwarded-Port": {"80"}},
+			"",
+		},
+
+		{"set X-Forwarded-Port for HTTPS",
+			&http.Request{RemoteAddr: "1.2.3.4:5555", TLS: &tls.ConnectionState{}},
+			config.Proxy{},
+			http.Header{"X-Forwarded-Port": {"443"}},
+			"",
+		},
+
+		{"set X-Forwarded-Port from Host",
+			&http.Request{RemoteAddr: "1.2.3.4:5555", Host: "5.6.7.8:1234"},
+			config.Proxy{},
+			http.Header{"X-Forwarded-Port": {"1234"}},
+			"",
+		},
+
+		{"set X-Forwarded-Port from Host for HTTPS",
+			&http.Request{RemoteAddr: "1.2.3.4:5555", Host: "5.6.7.8:1234", TLS: &tls.ConnectionState{}},
+			config.Proxy{},
+			http.Header{"X-Forwarded-Port": {"1234"}},
 			"",
 		},
 
@@ -217,3 +238,25 @@ func TestAddHeaders(t *testing.T) {
 		}
 	}
 }
+
+func TestLocalPort(t *testing.T) {
+	tests := []struct {
+		r    *http.Request
+		port string
+	}{
+		{nil, ""},
+		{&http.Request{Host: ""}, "80"},
+		{&http.Request{Host: ":"}, "80"},
+		{&http.Request{Host: "1.2.3.4:5678"}, "5678"},
+		{&http.Request{Host: "1.2.3.4"}, "80"},
+		{&http.Request{Host: "1.2.3.4", TLS: &tls.ConnectionState{}}, "443"},
+		{&http.Request{Host: "1.2.3.4:"}, "80"},
+		{&http.Request{Host: "1.2.3.4:", TLS: &tls.ConnectionState{}}, "443"},
+	}
+
+	for i, tt := range tests {
+		if got, want := localPort(tt.r), tt.port; got != want {
+			t.Errorf("%d: got %q want %q", i, got, want)
+		}
+	}
+}