gohttp/handler/proxy.go

package handler

import (
	"net/http"
	"net/http/httputil"
	"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
@param upstream upstream server url
@param path http path config
@param index the proxy index in upstreams

@return httputil.ReverseProxy instance

what is the directive? its stands for update of request, like HostSchemas, Path, RemoveCookie, etc.
eg: HostSchemas $target
it stande for replace req url host and schema according to $target url. $target == upstream
*/
func makeProxy(upstream string, path *model.HttpPath, index int) *httputil.ReverseProxy {
	p := &httputil.ReverseProxy{}
	directiveHandlers := []func(r *http.Request){}
	if len(path.Directives) > 0 {
		for _, directive := range path.Directives {
			d := strings.Replace(string(directive), "$target", upstream, 1)
			dh := GetUpdaterFn(d)
			if dh != nil {
				directiveHandlers = append(directiveHandlers, dh)
			}
		}
	}
	p.Director = func(req *http.Request) {
		for _, handler := range directiveHandlers {
			handler(req)
		}
	}

	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
}