gohttp/handler/proxy.go

105 lines
2.4 KiB
Go

package handler
import (
"fmt"
"net/http"
"net/http/httputil"
"strconv"
"strings"
"git.pyer.club/kingecg/gohttpd/model"
"git.pyer.club/kingecg/gologger"
)
type ProxyHandler struct {
proxy []*httputil.ReverseProxy
Upstreams []string
count int
}
func (p *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
l := gologger.GetLogger("Proxy")
originalUrl := r.Host + r.URL.String()
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)
}
l.Info(fmt.Sprintf("proxy %s to %s", originalUrl, p.Upstreams[proxyIndex]))
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{
Upstreams: p.Upstreams,
}
ph.proxy = make([]*httputil.ReverseProxy, upstreamCount)
for index, upstream := range p.Upstreams {
ph.proxy[index] = makeProxy(upstream, p, index)
}
return ph
}