gohttp/hander/proxy.go

package handler

import (
	"net/http"
	"net/http/httputil"
	"net/url"
	"strconv"
	"strings"

	"git.pyer.club/kingecg/gohttpd/model"
)

type ProxyHandler struct {
	proxy []*httputil.ReverseProxy
	count int
}

func (p *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	s, err := r.Cookie("s")
	var proxyIndex int
	if err != nil {
		proxyIndex = p.count
		p.count++
		if p.count >= len(p.proxy) {
			p.count = 0
		}
	} else {
		proxyIndex, _ = strconv.Atoi(s.Value)
	}
	p.proxy[proxyIndex].ServeHTTP(w, r)
}

// init httputil.ReverseProxy instance and add path rewrite and add session-sticky cookie to response
func makeProxy(upstream string, path *model.HttpPath, index int) *httputil.ReverseProxy {
	p := &httputil.ReverseProxy{}
	p.Director = func(req *http.Request) {
		turl, _ := url.Parse(upstream)
		req.URL.Host = turl.Host
		req.URL.Scheme = turl.Scheme
		pw := path.Rewrite
		if pw.Replace != "" {
			req.URL.Path = strings.TrimPrefix(req.URL.Path, pw.Replace)
			if req.URL.RawPath != "" {
				req.URL.RawPath = strings.TrimPrefix(req.URL.RawPath, pw.Replace)
			}
		}
		if pw.With != "" {
			req.URL.Path = strings.TrimSuffix(pw.With, "/") + "/" + req.URL.Path
			if req.URL.RawPath != "" {
				req.URL.RawPath = strings.TrimSuffix(pw.With, "/") + "/" + req.URL.RawPath
			}
		}

		if len(path.Headers) > 0 {
			for _, header := range path.Headers {
				//req.Header.Add(header.Name, header.Value)
				value := ""
				switch header.Value {
				case string(model.ProxyHost):
					value = turl.Host
				default:
					value = header.Value
				}
				req.Header.Set(header.Name, value)
			}
		}
	}

	p.ModifyResponse = func(resp *http.Response) error {
		hasSticky := false
		for _, cookie := range resp.Cookies() {
			if cookie.Name == "s" {
				hasSticky = true
				break
			}
		}
		if !hasSticky {
			c := http.Cookie{
				Name:  "s",
				Value: strconv.Itoa(index),
			}
			resp.Header.Add("Set-Cookie", c.String())
		}
		return nil
	}
	return p
}

func NewProxyHandler(p *model.HttpPath) *ProxyHandler {
	upstreamCount := len(p.Upstreams)
	if upstreamCount == 0 {
		panic("no upstream defined")
	}
	ph := &ProxyHandler{}
	ph.proxy = make([]*httputil.ReverseProxy, upstreamCount)

	for index, upstream := range p.Upstreams {
		ph.proxy[index] = makeProxy(upstream, p, index)
	}
	return ph
}