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 }