Enhance Trust Forwarded Headers

vendor/github.com/vulcand/oxy/forward/rewrite.go

@@ -15,30 +15,36 @@ type HeaderRewriter struct {
 }
 
 func (rw *HeaderRewriter) Rewrite(req *http.Request) {
+	if !rw.TrustForwardHeader {
+		utils.RemoveHeaders(req.Header, XHeaders...)
+	}
+
 	if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
-		if rw.TrustForwardHeader {
-			if prior, ok := req.Header[XForwardedFor]; ok {
-				clientIP = strings.Join(prior, ", ") + ", " + clientIP
-			}
+		if prior, ok := req.Header[XForwardedFor]; ok {
+			req.Header.Set(XForwardedFor, strings.Join(prior, ", ")+", "+clientIP)
+		} else {
+			req.Header.Set(XForwardedFor, clientIP)
+		}
+
+		if req.Header.Get(XRealIp) == "" {
+			req.Header.Set(XRealIp, clientIP)
 		}
-		req.Header.Set(XForwardedFor, clientIP)
 	}
 
-	if xfp := req.Header.Get(XForwardedProto); xfp != "" && rw.TrustForwardHeader {
-		req.Header.Set(XForwardedProto, xfp)
-	} else if req.TLS != nil {
-		req.Header.Set(XForwardedProto, "https")
-	} else {
-		req.Header.Set(XForwardedProto, "http")
+	xfProto := req.Header.Get(XForwardedProto)
+	if xfProto == "" {
+		if req.TLS != nil {
+			req.Header.Set(XForwardedProto, "https")
+		} else {
+			req.Header.Set(XForwardedProto, "http")
+		}
 	}
 
-	if xfp := req.Header.Get(XForwardedPort); xfp != "" && rw.TrustForwardHeader {
-		req.Header.Set(XForwardedPort, xfp)
+	if xfp := req.Header.Get(XForwardedPort); xfp == "" {
+		req.Header.Set(XForwardedPort, forwardedPort(req))
 	}
 
-	if xfh := req.Header.Get(XForwardedHost); xfh != "" && rw.TrustForwardHeader {
-		req.Header.Set(XForwardedHost, xfh)
-	} else if req.Host != "" {
+	if xfHost := req.Header.Get(XForwardedHost); xfHost == "" && req.Host != "" {
 		req.Header.Set(XForwardedHost, req.Host)
 	}
 
@@ -50,3 +56,19 @@ func (rw *HeaderRewriter) Rewrite(req *http.Request) {
 	// connection, regardless of what the client sent to us.
 	utils.RemoveHeaders(req.Header, HopHeaders...)
 }
+
+func forwardedPort(req *http.Request) string {
+	if req == nil {
+		return ""
+	}
+
+	if _, port, err := net.SplitHostPort(req.Host); err == nil && port != "" {
+		return port
+	}
+
+	if req.TLS != nil {
+		return "443"
+	}
+
+	return "80"
+}