Update ReverseProxy code from official golang repo (#2051)

Fix #1192
This commit is contained in:
yuyulei 2020-11-05 16:34:17 +08:00 committed by GitHub
parent 3c2e2bcea5
commit 644a0cfdb6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -13,6 +13,7 @@ import (
"log"
"net"
"net/http"
"net/textproto"
"net/url"
"strings"
"sync"
@ -24,6 +25,19 @@ import (
// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
//
// ReverseProxy by default sets the client IP as the value of the
// X-Forwarded-For header.
//
// If an X-Forwarded-For header already exists, the client IP is
// appended to the existing values. As a special case, if the header
// exists in the Request.Header map but has a nil value (such as when
// set by the Director func), the X-Forwarded-For header is
// not modified.
//
// To prevent IP spoofing, be sure to delete any pre-existing
// X-Forwarded-For header coming from the client or
// an untrusted proxy.
type ReverseProxy struct {
// Director must be a function which modifies
// the request into a new request to be sent
@ -44,9 +58,9 @@ type ReverseProxy struct {
// A negative value means to flush immediately
// after each write to the client.
// The FlushInterval is ignored when ReverseProxy
// recognizes a response as a streaming response;
// for such responses, writes are flushed to the client
// immediately.
// recognizes a response as a streaming response, or
// if its ContentLength is -1; for such responses, writes
// are flushed to the client immediately.
FlushInterval time.Duration
// ErrorLog specifies an optional logger for errors
@ -97,6 +111,27 @@ func singleJoiningSlash(a, b string) string {
return a + b
}
func joinURLPath(a, b *url.URL) (path, rawpath string) {
if a.RawPath == "" && b.RawPath == "" {
return singleJoiningSlash(a.Path, b.Path), ""
}
// Same as singleJoiningSlash, but uses EscapedPath to determine
// whether a slash should be added
apath := a.EscapedPath()
bpath := b.EscapedPath()
aslash := strings.HasSuffix(apath, "/")
bslash := strings.HasPrefix(bpath, "/")
switch {
case aslash && bslash:
return a.Path + b.Path[1:], apath + bpath[1:]
case !aslash && !bslash:
return a.Path + "/" + b.Path, apath + "/" + bpath
}
return a.Path + b.Path, apath + bpath
}
// NewSingleHostReverseProxy returns a new ReverseProxy that routes
// URLs to the scheme, host, and base path provided in target. If the
// target's path is "/base" and the incoming request was for "/dir",
@ -109,7 +144,7 @@ func NewSingleHostReverseProxy(target *url.URL) *ReverseProxy {
director := func(req *http.Request) {
req.URL.Scheme = target.Scheme
req.URL.Host = target.Host
req.URL.Path = singleJoiningSlash(target.Path, req.URL.Path)
req.URL.Path, req.URL.RawPath = joinURLPath(target, req.URL)
if targetQuery == "" || req.URL.RawQuery == "" {
req.URL.RawQuery = targetQuery + req.URL.RawQuery
} else {
@ -195,16 +230,19 @@ func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
}()
}
outreq := req.WithContext(ctx)
outreq := req.Clone(ctx)
if req.ContentLength == 0 {
outreq.Body = nil // Issue 16036: nil Body for http.Transport retries
}
if outreq.Header == nil {
outreq.Header = make(http.Header) // Issue 33142: historical behavior was to always allocate
}
// =============================
// Modified for frp
outreq = outreq.WithContext(context.WithValue(outreq.Context(), RouteInfoURL, req.URL.Path))
outreq = outreq.WithContext(context.WithValue(outreq.Context(), RouteInfoHost, req.Host))
outreq = outreq.WithContext(context.WithValue(outreq.Context(), RouteInfoRemote, req.RemoteAddr))
outreq = outreq.Clone(context.WithValue(outreq.Context(), RouteInfoURL, req.URL.Path))
outreq = outreq.Clone(context.WithValue(outreq.Context(), RouteInfoHost, req.Host))
outreq = outreq.Clone(context.WithValue(outreq.Context(), RouteInfoRemote, req.RemoteAddr))
// =============================
p.Director(outreq)
@ -244,10 +282,14 @@ func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
// If we aren't the first proxy retain prior
// X-Forwarded-For information as a comma+space
// separated list and fold multiple headers into one.
if prior, ok := outreq.Header["X-Forwarded-For"]; ok {
prior, ok := outreq.Header["X-Forwarded-For"]
omit := ok && prior == nil // Issue 38079: nil now means don't populate the header
if len(prior) > 0 {
clientIP = strings.Join(prior, ", ") + ", " + clientIP
}
outreq.Header.Set("X-Forwarded-For", clientIP)
if !omit {
outreq.Header.Set("X-Forwarded-For", clientIP)
}
}
res, err := transport.RoundTrip(outreq)
@ -290,7 +332,7 @@ func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
rw.WriteHeader(res.StatusCode)
err = p.copyResponse(rw, res.Body, p.flushInterval(req, res))
err = p.copyResponse(rw, res.Body, p.flushInterval(res))
if err != nil {
defer res.Body.Close()
// Since we're streaming the response, if we run into an error all we can do
@ -353,7 +395,7 @@ func shouldPanicOnCopyError(req *http.Request) bool {
func removeConnectionHeaders(h http.Header) {
for _, f := range h["Connection"] {
for _, sf := range strings.Split(f, ",") {
if sf = strings.TrimSpace(sf); sf != "" {
if sf = textproto.TrimString(sf); sf != "" {
h.Del(sf)
}
}
@ -362,7 +404,7 @@ func removeConnectionHeaders(h http.Header) {
// flushInterval returns the p.FlushInterval value, conditionally
// overriding its value for a specific request/response.
func (p *ReverseProxy) flushInterval(req *http.Request, res *http.Response) time.Duration {
func (p *ReverseProxy) flushInterval(res *http.Response) time.Duration {
resCT := res.Header.Get("Content-Type")
// For Server-Sent Events responses, flush immediately.
@ -371,7 +413,11 @@ func (p *ReverseProxy) flushInterval(req *http.Request, res *http.Response) time
return -1 // negative means immediately
}
// TODO: more specific cases? e.g. res.ContentLength == -1?
// We might have the case of streaming for which Content-Length might be unset.
if res.ContentLength == -1 {
return -1
}
return p.FlushInterval
}
@ -510,8 +556,6 @@ func (p *ReverseProxy) handleUpgradeResponse(rw http.ResponseWriter, req *http.R
return
}
copyHeader(res.Header, rw.Header())
hj, ok := rw.(http.Hijacker)
if !ok {
p.getErrorHandler()(rw, req, fmt.Errorf("can't switch protocols using non-Hijacker ResponseWriter type %T", rw))
@ -522,13 +566,30 @@ func (p *ReverseProxy) handleUpgradeResponse(rw http.ResponseWriter, req *http.R
p.getErrorHandler()(rw, req, fmt.Errorf("internal error: 101 switching protocols response with non-writable body"))
return
}
defer backConn.Close()
backConnCloseCh := make(chan bool)
go func() {
// Ensure that the cancelation of a request closes the backend.
// See issue https://golang.org/issue/35559.
select {
case <-req.Context().Done():
case <-backConnCloseCh:
}
backConn.Close()
}()
defer close(backConnCloseCh)
conn, brw, err := hj.Hijack()
if err != nil {
p.getErrorHandler()(rw, req, fmt.Errorf("Hijack failed on protocol switch: %v", err))
return
}
defer conn.Close()
copyHeader(rw.Header(), res.Header)
res.Header = rw.Header()
res.Body = nil // so res.Write only writes the headers; we have res.Body in backConn above
if err := res.Write(brw); err != nil {
p.getErrorHandler()(rw, req, fmt.Errorf("response write: %v", err))