From 0b6469b742742b447cd20967a8192b5c6f327510 Mon Sep 17 00:00:00 2001
From: Michael Quigley <michael@quigley.com>
Date: Tue, 26 Jul 2022 13:30:19 -0400
Subject: [PATCH] adjust reverse proxy infrastructure to selectively proxy
 based on host header (#7)

---
 proxy/proxy.go | 11 ++++++-
 util/proxy.go  | 79 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/proxy/proxy.go b/proxy/proxy.go
index 7141d16c..278e9e7e 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -5,6 +5,7 @@ import (
 	"github.com/openziti/sdk-golang/ziti"
 	"github.com/openziti/sdk-golang/ziti/config"
 	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
 	"net/http"
 )
 
@@ -18,10 +19,18 @@ func Run(cfg *Config) error {
 	zTransport := http.DefaultTransport.(*http.Transport).Clone()
 	zTransport.DialContext = zDialCtx.Dial
 
-	proxy, err := util.NewProxy("http://zrok")
+	proxy, err := util.NewServiceProxy(&resolver{})
 	if err != nil {
 		return err
 	}
 	proxy.Transport = zTransport
 	return http.ListenAndServe(cfg.Address, util.NewProxyHandler(proxy))
 }
+
+type resolver struct {
+}
+
+func (r *resolver) Service(host string) string {
+	logrus.Infof("host = '%v'", host)
+	return "zrok"
+}
diff --git a/util/proxy.go b/util/proxy.go
index 874c48fe..2aca1566 100644
--- a/util/proxy.go
+++ b/util/proxy.go
@@ -1,12 +1,18 @@
 package util
 
 import (
+	"fmt"
 	"github.com/sirupsen/logrus"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
+	"strings"
 )
 
+type ProxyServiceResolver interface {
+	Service(host string) string
+}
+
 func NewProxy(target string) (*httputil.ReverseProxy, error) {
 	targetURL, err := url.Parse(target)
 	if err != nil {
@@ -30,6 +36,24 @@ func NewProxy(target string) (*httputil.ReverseProxy, error) {
 	return proxy, nil
 }
 
+func NewServiceProxy(p ProxyServiceResolver) (*httputil.ReverseProxy, error) {
+	proxy := hostTargetReverseProxy(p)
+	director := proxy.Director
+	proxy.Director = func(req *http.Request) {
+		director(req)
+		logrus.Infof("-> %v", req.URL.String())
+		req.Header.Set("X-Proxy", "zrok")
+	}
+	proxy.ModifyResponse = func(resp *http.Response) error {
+		return nil
+	}
+	proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
+		logrus.Errorf("error proxying: %v", err)
+	}
+
+	return proxy, nil
+}
+
 type proxyHandler struct {
 	proxy *httputil.ReverseProxy
 }
@@ -41,3 +65,58 @@ func NewProxyHandler(proxy *httputil.ReverseProxy) *proxyHandler {
 func (self *proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	self.proxy.ServeHTTP(w, r)
 }
+
+func hostTargetReverseProxy(r ProxyServiceResolver) *httputil.ReverseProxy {
+	director := func(req *http.Request) {
+		targetSvc := r.Service(req.Host)
+		if target, err := url.Parse(fmt.Sprintf("http://%v", targetSvc)); err == nil {
+			targetQuery := target.RawQuery
+			req.URL.Scheme = target.Scheme
+			req.URL.Host = target.Host
+			req.URL.Path, req.URL.RawPath = joinURLPath(target, req.URL)
+			if targetQuery == "" || req.URL.RawQuery == "" {
+				req.URL.RawQuery = targetQuery + req.URL.RawQuery
+			} else {
+				req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
+			}
+			if _, ok := req.Header["User-Agent"]; !ok {
+				// explicitly disable User-Agent so it's not set to default value
+				req.Header.Set("User-Agent", "")
+			}
+		}
+	}
+	return &httputil.ReverseProxy{Director: director}
+}
+
+func joinURLPath(a, b *url.URL) (path, rawpath string) {
+	if a.RawPath == "" && b.RawPath == "" {
+		return singleJoiningSlash(a.Path, b.Path), ""
+	}
+	// Same as singleJoiningSlash, but uses EscapedPath to determine
+	// whether a slash should be added
+	apath := a.EscapedPath()
+	bpath := b.EscapedPath()
+
+	aslash := strings.HasSuffix(apath, "/")
+	bslash := strings.HasPrefix(bpath, "/")
+
+	switch {
+	case aslash && bslash:
+		return a.Path + b.Path[1:], apath + bpath[1:]
+	case !aslash && !bslash:
+		return a.Path + "/" + b.Path, apath + "/" + bpath
+	}
+	return a.Path + b.Path, apath + bpath
+}
+
+func singleJoiningSlash(a, b string) string {
+	aslash := strings.HasSuffix(a, "/")
+	bslash := strings.HasPrefix(b, "/")
+	switch {
+	case aslash && bslash:
+		return a + b[1:]
+	case !aslash && !bslash:
+		return a + "/" + b
+	}
+	return a + b
+}