Use X-Forwarded-Proto to determine scheme

2 files changed, 24 insertions, 0 deletions

diff --git a/api/api.go b/api/api.go
index 7db80fb..3c76565 100644
--- a/api/api.go
+++ b/api/api.go
@@ -36,6 +36,10 @@ func urlPrefix(r *http.Request) string {
 	if r.TLS != nil {
 		scheme = "https"
 	}
+	forwardedProto := r.Header["X-Forwarded-Proto"]
+	if len(forwardedProto) > 0 {
+		scheme = forwardedProto[0]
+	}
 	url := url.URL{Scheme: scheme, Host: host}
 	return url.String()
 }
diff --git a/api/api_test.go b/api/api_test.go
index 9e389b0..5824878 100644
--- a/api/api_test.go
+++ b/api/api_test.go
@@ -1,6 +1,7 @@
 package api
 
 import (
+	"crypto/tls"
 	"fmt"
 	"io/ioutil"
 	"log"
@@ -101,6 +102,25 @@ func TestAPI(t *testing.T) {
 	}
 }
 
+func TestURLPrefix(t *testing.T) {
+	var tests = []struct {
+		in  *http.Request
+		out string
+	}{
+		{&http.Request{Host: "foo"}, "http://foo"},
+		{&http.Request{Host: "", RemoteAddr: "127.0.0.1"}, "http://127.0.0.1"},
+		{&http.Request{Host: "bar", TLS: &tls.ConnectionState{}}, "https://bar"},
+		{&http.Request{Host: "baz", Header: map[string][]string{"X-Forwarded-Proto": []string{"https"}}}, "https://baz"},
+		{&http.Request{Host: "qux", Header: map[string][]string{"X-Forwarded-Proto": []string{}}}, "http://qux"},
+	}
+	for _, tt := range tests {
+		prefix := urlPrefix(tt.in)
+		if prefix != tt.out {
+			t.Errorf("want %s, got %s", tt.out, prefix)
+		}
+	}
+}
+
 func TestGetBusStops(t *testing.T) {
 	server := atbTestServer()
 	defer server.Close()