implement MultiHostReverseProxy

pkg/command/server.go

@@ -2,6 +2,7 @@ package command
 
 import (
 	"context"
+	gohttp "net/http"
 	"net/http/httputil"
 	"net/url"
 	"os"
@@ -26,6 +27,12 @@ import (
 	"go.opencensus.io/trace"
 )
 
+// ReverseProxy extends httputil to support multiple hosts with diffent policies
+type ReverseProxy struct {
+	httputil.ReverseProxy
+	Directors map[string]map[string]func(req *gohttp.Request)
+}
+
 // Server is the entrypoint for the server command.
 func Server(cfg *config.Config) *cli.Command {
 	return &cli.Command{
@@ -135,8 +142,11 @@ func Server(cfg *config.Config) *cli.Command {
 
 			defer cancel()
 
+			rp := NewMultiHostReverseProxy(cfg)
+
 			{
 				server, err := http.Server(
+					http.Handler(rp),
 					http.Logger(logger),
 					http.Namespace(httpNamespace),
 					http.Context(ctx),
@@ -146,17 +156,6 @@ func Server(cfg *config.Config) *cli.Command {
 					http.Flags(flagset.ServerWithConfig(cfg)),
 				)
 
-				for _, ep := range cfg.Routes {
-					uri, err := url.Parse(ep.Backend)
-					if err != nil {
-						logger.Info().
-							Str("server", "http").
-							Msg("error while parsing URL")
-					}
-
-					server.Handle(ep.Endpoint, httputil.NewSingleHostReverseProxy(uri))
-				}
-
 				if err != nil {
 					logger.Error().
 						Err(err).
@@ -233,3 +232,72 @@ func Server(cfg *config.Config) *cli.Command {
 		},
 	}
 }
+
+// NewMultiHostReverseProxy undocummented
+func NewMultiHostReverseProxy(conf *config.Config) *ReverseProxy {
+	reverseProxy := &ReverseProxy{Directors: make(map[string]map[string]func(req *gohttp.Request))}
+
+	for _, target := range conf.Routes {
+		uri, err := url.Parse(target.Backend)
+		if err != nil { /* do something with err */
+		}
+		reverseProxy.AddHost(target.Policy, uri, target.Endpoint)
+	}
+
+	return reverseProxy
+}
+
+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
+}
+
+// AddHost undocumented
+func (p *ReverseProxy) AddHost(policy string, target *url.URL, endpoint string) {
+	targetQuery := target.RawQuery
+	if p.Directors[policy] == nil {
+		p.Directors[policy] = make(map[string]func(req *gohttp.Request))
+	}
+	p.Directors[policy][endpoint] = func(req *gohttp.Request) {
+		req.URL.Scheme = target.Scheme
+		req.URL.Host = target.Host
+		req.URL.Path = singleJoiningSlash(target.Path, req.URL.Path)
+		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", "")
+		}
+	}
+}
+
+func (p *ReverseProxy) ServeHTTP(rw gohttp.ResponseWriter, req *gohttp.Request) {
+	// TODO need to fetch from the accounts service
+	policy := "reva"
+	var hit bool
+
+	for k := range p.Directors[policy] {
+		if strings.HasPrefix(req.URL.Path, k) && k != "/" {
+			p.Director = p.Directors[policy][k]
+			hit = true
+		}
+	}
+
+	// override default director with root. If any
+	if !hit && p.Directors[policy]["/"] != nil {
+		p.Director = p.Directors[policy]["/"]
+	}
+
+	// Call upstream ServeHTTP
+	p.ReverseProxy.ServeHTTP(rw, req)
+}