support multiple backends

1 files changed, 32 insertions, 32 deletions

diff --git a/src/daq/s5proxy/src/s5proxy/proxy.go b/src/daq/s5proxy/src/s5proxy/proxy.go
index 50075c6..0e4fe21 100644
--- a/src/daq/s5proxy/src/s5proxy/proxy.go
+++ b/src/daq/s5proxy/src/s5proxy/proxy.go
@@ -50,12 +50,12 @@ func generateTime(d time.Duration) string {
 }
 
 type s5proxyResponseWriter struct {
-	wrapped   http.ResponseWriter
-	conf      *Config
-	StatsCh   chan<- *Client
-	IP        string
-	Port      int
-	UserAgent string
+	wrapped    http.ResponseWriter
+	respHeader []HeaderOperation
+	StatsCh    chan<- *Client
+	IP         string
+	Port       int
+	UserAgent  string
 }
 
 func (r s5proxyResponseWriter) Header() http.Header {
@@ -75,7 +75,7 @@ func (r s5proxyResponseWriter) Write(data []byte) (int, error) {
 }
 
 func (r s5proxyResponseWriter) WriteHeader(status int) {
-	for _, h := range r.conf.ResponseHeader {
+	for _, h := range r.respHeader {
 		switch h.Operation {
 		case OpAdd:
 			r.wrapped.Header().Add(h.Header, h.Value)
@@ -90,9 +90,9 @@ func (r s5proxyResponseWriter) WriteHeader(status int) {
 	r.wrapped.WriteHeader(status)
 }
 
-func proxyHandler(p *Proxy, w http.ResponseWriter, r *http.Request) {
-	pw := s5proxyResponseWriter{wrapped: w, conf: p.conf}
-	if p.stats != nil {
+func proxyHandler(p *httputil.ReverseProxy, s *Stats, respHeader []HeaderOperation, w http.ResponseWriter, r *http.Request) {
+	pw := s5proxyResponseWriter{wrapped: w, respHeader: respHeader}
+	if s != nil {
 		if ip, port, err := net.SplitHostPort(r.RemoteAddr); err != nil {
 			s5l.Printf("PROXY: client '%s' error: invalid address/port info: %v, no statistics will be gathered", r.RemoteAddr, err)
 		} else {
@@ -101,38 +101,32 @@ func proxyHandler(p *Proxy, w http.ResponseWriter, r *http.Request) {
 				s5l.Printf("PROXY: client '%s' warning: can't parse port or service name: %v", r.RemoteAddr, err)
 			}
 			pw.UserAgent = r.UserAgent()
-			pw.StatsCh = p.stats.GetUpdateChannel(r.URL.String())
+			pw.StatsCh = s.GetUpdateChannel(r.URL.String())
 		}
 	}
 
-	p.proxy.ServeHTTP(pw, r)
+	p.ServeHTTP(pw, r)
 }
 
 type Proxy struct {
 	conf  *Config
 	stats *Stats
-	proxy *httputil.ReverseProxy
 	mux   *http.ServeMux
 	srv   *http.Server
 }
 
-func NewProxy(conf *Config, stats *Stats) (p *Proxy, err error) {
-
-	p = &Proxy{conf: conf, stats: stats}
-
-	var remote *url.URL
-	remote, err = url.Parse(conf.ConnectAddr)
+func (p *Proxy) installProxyHandler(location, connectAddr string, reqHeader, respHeader []HeaderOperation) error {
+	remote, err := url.Parse(connectAddr)
 	if err != nil {
-		return
+		return err
 	}
-	s5l.Printf("PROXY: forwarding traffic to '%s'", remote.String())
-
-	p.proxy = httputil.NewSingleHostReverseProxy(remote)
-	origDir := p.proxy.Director
+	s5l.Printf("PROXY: forwarding traffic for %s to '%s'", location, remote.String())
 
-	p.proxy.Director = func(req *http.Request) {
+	proxy := httputil.NewSingleHostReverseProxy(remote)
+	origDir := proxy.Director
+	proxy.Director = func(req *http.Request) {
 		origDir(req)
-		for _, h := range conf.RequestHeader {
+		for _, h := range reqHeader {
 			switch h.Operation {
 			case OpAdd:
 				req.Header.Add(h.Header, h.Value)
@@ -146,15 +140,21 @@ func NewProxy(conf *Config, stats *Stats) (p *Proxy, err error) {
 		}
 	}
 
-	p.mux = http.NewServeMux()
-	// TODO: add additional handler for conf.Locations
-	p.mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
-		proxyHandler(p, w, r)
+	p.mux.HandleFunc(location, func(w http.ResponseWriter, r *http.Request) {
+		proxyHandler(proxy, p.stats, respHeader, w, r)
 	})
+	return nil
+}
+
+func NewProxy(conf *Config, stats *Stats) (p *Proxy, err error) {
+	p = &Proxy{conf: conf, stats: stats}
+	p.mux = http.NewServeMux()
 
-	p.srv = &http.Server{
-		Handler: p.mux,
+	for location, l := range conf.Locations {
+		p.installProxyHandler(location, l.ConnectAddr, l.RequestHeader, l.ResponseHeader)
 	}
+	p.installProxyHandler("/", conf.ConnectAddr, conf.RequestHeader, conf.ResponseHeader)
+	p.srv = &http.Server{Handler: p.mux}
 	return
 }