Improve proxy

cmd/proxy/proxy.go

@@ -1,10 +1,13 @@
+// This is honestly terrible
 package main
 
 import (
 	"io"
 	"log"
 	"net/http"
+	"net/url"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/miekg/dns"
@@ -12,12 +15,14 @@ import (
 
 var endpoint string
 var endpointSecure bool
+var behindProxy bool
 var client *http.Client
 
 func main() {
 	// Read configuration
 	endpoint = os.Getenv("MINIO_ENDPOINT")
 	endpointSecure = os.Getenv("MINIO_ENDPOINT_SECURE") == "true"
+	behindProxy = os.Getenv("BEHIND_PROXY") == "true"
 
 	// Setup HTTP client
 	transport := &http.Transport{
@@ -44,37 +49,34 @@ func main() {
 }
 
 func handle(w http.ResponseWriter, r *http.Request) {
-	if len(r.Host) == 0 {
-		w.WriteHeader(http.StatusNotFound)
-		return
-	}
+	// Validate host
 	if _, ok := dns.IsDomainName(r.Host); !ok {
 		w.WriteHeader(http.StatusBadRequest)
 		return
 	}
 
-	cloneRequest := r.Clone(r.Context())
-
-	cloneRequest.URL.Scheme = "http"
+	// Rewrite URL
+	path := mapPath(r.URL.Path, r.Host)
+	url := &url.URL{
+		Scheme: "http",
+		Host:   endpoint,
+		Path:   path,
+	}
 	if endpointSecure {
-		cloneRequest.URL.Scheme = "https"
+		url.Scheme = "https"
 	}
 
-	cloneRequest.URL.Host = endpoint
+	// Update Forwarded headers
+	header := safeCloneHeader(r.Header)
 
-	cloneRequest.URL.Path = "/" + r.Host + r.URL.Path
-	if r.URL.Path == "/" {
-		cloneRequest.URL.Path = "/" + r.Host + "/index.html"
-	}
-	if r.URL.Path == "/manage" {
-		cloneRequest.URL.Path = "/" + r.Host + "/manage/index.html"
+	// Create forwarded request
+	out := &http.Request{
+		URL:    url,
+		Header: header,
 	}
 
-	cloneRequest.Host = ""
-	cloneRequest.RequestURI = ""
-
-	log.Println(cloneRequest.URL)
-	resp, err := client.Do(cloneRequest)
+	log.Println(out.URL)
+	resp, err := client.Do(out)
 	if err != nil {
 		w.WriteHeader(http.StatusInternalServerError)
 		log.Println(err)
@@ -82,11 +84,46 @@ func handle(w http.ResponseWriter, r *http.Request) {
 	}
 	defer resp.Body.Close()
 
-	for key, value := range resp.Header {
-		for _, v := range value {
-			w.Header().Add(key, v)
-		}
-	}
 	w.WriteHeader(resp.StatusCode)
 	io.Copy(w, resp.Body)
 }
+
+func mapPath(path, host string) string {
+	if strings.HasSuffix(path, "/") {
+		path += "index.html"
+	}
+	return "/" + host + path
+}
+
+// Further validation is required for this function
+func safeCloneHeader(h http.Header) http.Header {
+	clone := make(http.Header, 0)
+	for key, value := range h {
+		if !safeHeaderName(key) {
+			continue
+		}
+		for _, v := range value {
+			clone.Add(key, v)
+		}
+	}
+	return clone
+}
+
+// Comparisons in this function MUST be in canonical form
+func safeHeaderName(name string) bool {
+	if behindProxy {
+		if name == "X-Forwarded-For" {
+			return true
+		}
+		if name == "X-Forwarded-Host" {
+			return true
+		}
+		if name == "X-Forwarded-Proto" {
+			return true
+		}
+		if name == "Forwarded" {
+			return true
+		}
+	}
+	return false
+}