GoDoxy/internal/net/http/middleware/vars.go
2024-12-03 11:09:30 +08:00

114 lines
3.3 KiB
Go

package middleware
import (
"net"
"net/http"
"regexp"
"strconv"
"strings"
gphttp "github.com/yusing/go-proxy/internal/net/http"
)
type varReplaceFunc func(string) string
var (
reArg = regexp.MustCompile(`\$arg\([\w-_]+\)`)
reHeader = regexp.MustCompile(`\$header\([\w-]+\)`)
reRespHeader = regexp.MustCompile(`\$resp_header\([\w-]+\)`)
reStatic = regexp.MustCompile(`\$[\w_]+`)
)
func varSubsMap(req *Request, resp *Response) map[string]func() string {
reqHost, reqPort, err := net.SplitHostPort(req.Host)
if err != nil {
reqHost = req.Host
}
reqAddr := reqHost
if reqPort != "" {
reqAddr += ":" + reqPort
}
pairs := map[string]func() string{
"$req_method": func() string { return req.Method },
"$req_scheme": func() string { return req.URL.Scheme },
"$req_host": func() string { return reqHost },
"$req_port": func() string { return reqPort },
"$req_addr": func() string { return reqAddr },
"$req_path": func() string { return req.URL.Path },
"$req_query": func() string { return req.URL.RawQuery },
"$req_url": func() string { return req.URL.String() },
"$req_uri": req.URL.RequestURI,
"$req_content_type": func() string { return req.Header.Get("Content-Type") },
"$req_content_length": func() string { return strconv.FormatInt(req.ContentLength, 10) },
"$remote_addr": func() string { return req.RemoteAddr },
}
if resp != nil {
pairs["$resp_content_type"] = func() string { return resp.Header.Get("Content-Type") }
pairs["$resp_content_length"] = func() string { return resp.Header.Get("Content-Length") }
pairs["$status_code"] = func() string { return strconv.Itoa(resp.StatusCode) }
}
upScheme := req.Header.Get(gphttp.HeaderUpstreamScheme)
if upScheme == "" {
return pairs
}
upHost := req.Header.Get(gphttp.HeaderUpstreamHost)
upPort := req.Header.Get(gphttp.HeaderUpstreamPort)
upAddr := upHost
if upPort != "" {
upAddr += ":" + upPort
}
upURL := upScheme + "://" + upAddr
pairs["$upstream_scheme"] = func() string { return upScheme }
pairs["$upstream_host"] = func() string { return upHost }
pairs["$upstream_port"] = func() string { return upPort }
pairs["$upstream_addr"] = func() string { return upAddr }
pairs["$upstream_url"] = func() string { return upURL }
return pairs
}
func varReplacer(req *Request, resp *Response) varReplaceFunc {
pairs := varSubsMap(req, resp)
return func(s string) string {
// Replace query parameters
s = reArg.ReplaceAllStringFunc(s, func(match string) string {
name := match[5 : len(match)-1]
for k, v := range req.URL.Query() {
if strings.EqualFold(k, name) {
return v[0]
}
}
return ""
})
// Replace headers
s = reHeader.ReplaceAllStringFunc(s, func(match string) string {
header := http.CanonicalHeaderKey(match[8 : len(match)-1])
return req.Header.Get(header)
})
if resp != nil {
s = reRespHeader.ReplaceAllStringFunc(s, func(match string) string {
header := http.CanonicalHeaderKey(match[14 : len(match)-1])
return resp.Header.Get(header)
})
}
// Replace static variables
return reStatic.ReplaceAllStringFunc(s, func(match string) string {
if fn, ok := pairs[match]; ok {
return fn()
}
return match
})
}
}
func varReplacerDummy(s string) string {
return s
}