87 lines
1.9 KiB
Go
87 lines
1.9 KiB
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, pw model.PathRewrite, 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
|
||
|
|
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
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
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.Rewrite, index)
|
||
|
|
}
|
||
|
|
return ph
|
||
|
|
}
|