github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/tools/querytee/proxy_backend.go (about) 1 package querytee 2 3 import ( 4 "context" 5 "io" 6 "io/ioutil" 7 "net" 8 "net/http" 9 "net/url" 10 "path" 11 "time" 12 13 "github.com/pkg/errors" 14 ) 15 16 // ProxyBackend holds the information of a single backend. 17 type ProxyBackend struct { 18 name string 19 endpoint *url.URL 20 client *http.Client 21 timeout time.Duration 22 23 // Whether this is the preferred backend from which picking up 24 // the response and sending it back to the client. 25 preferred bool 26 } 27 28 // NewProxyBackend makes a new ProxyBackend 29 func NewProxyBackend(name string, endpoint *url.URL, timeout time.Duration, preferred bool) *ProxyBackend { 30 return &ProxyBackend{ 31 name: name, 32 endpoint: endpoint, 33 timeout: timeout, 34 preferred: preferred, 35 client: &http.Client{ 36 CheckRedirect: func(_ *http.Request, _ []*http.Request) error { 37 return errors.New("the query-tee proxy does not follow redirects") 38 }, 39 Transport: &http.Transport{ 40 Proxy: http.ProxyFromEnvironment, 41 DialContext: (&net.Dialer{ 42 Timeout: 30 * time.Second, 43 KeepAlive: 30 * time.Second, 44 }).DialContext, 45 MaxIdleConns: 100, 46 MaxIdleConnsPerHost: 100, // see https://github.com/golang/go/issues/13801 47 IdleConnTimeout: 90 * time.Second, 48 DisableCompression: true, 49 }, 50 }, 51 } 52 } 53 54 func (b *ProxyBackend) ForwardRequest(orig *http.Request, body io.ReadCloser) (int, []byte, error) { 55 req := b.createBackendRequest(orig, body) 56 return b.doBackendRequest(req) 57 } 58 59 func (b *ProxyBackend) createBackendRequest(orig *http.Request, body io.ReadCloser) *http.Request { 60 req := orig.Clone(context.Background()) 61 req.Body = body 62 // RequestURI can't be set on a cloned request. It's only for handlers. 63 req.RequestURI = "" 64 // Replace the endpoint with the backend one. 65 req.URL.Scheme = b.endpoint.Scheme 66 req.URL.Host = b.endpoint.Host 67 68 // Prepend the endpoint path to the request path. 69 req.URL.Path = path.Join(b.endpoint.Path, req.URL.Path) 70 71 // Set the correct host header for the backend 72 req.Header.Set("Host", b.endpoint.Host) 73 74 // Replace the auth: 75 // - If the endpoint has user and password, use it. 76 // - If the endpoint has user only, keep it and use the request password (if any). 77 // - If the endpoint has no user and no password, use the request auth (if any). 78 clientUser, clientPass, clientAuth := orig.BasicAuth() 79 endpointUser := b.endpoint.User.Username() 80 endpointPass, _ := b.endpoint.User.Password() 81 82 req.Header.Del("Authorization") 83 if endpointUser != "" && endpointPass != "" { 84 req.SetBasicAuth(endpointUser, endpointPass) 85 } else if endpointUser != "" { 86 req.SetBasicAuth(endpointUser, clientPass) 87 } else if clientAuth { 88 req.SetBasicAuth(clientUser, clientPass) 89 } 90 91 // Remove Accept-Encoding header to avoid sending compressed responses 92 req.Header.Del("Accept-Encoding") 93 94 return req 95 } 96 97 func (b *ProxyBackend) doBackendRequest(req *http.Request) (int, []byte, error) { 98 // Honor the read timeout. 99 ctx, cancel := context.WithTimeout(context.Background(), b.timeout) 100 defer cancel() 101 102 // Execute the request. 103 res, err := b.client.Do(req.WithContext(ctx)) 104 if err != nil { 105 return 0, nil, errors.Wrap(err, "executing backend request") 106 } 107 108 // Read the entire response body. 109 defer res.Body.Close() 110 body, err := ioutil.ReadAll(res.Body) 111 if err != nil { 112 return 0, nil, errors.Wrap(err, "reading backend response") 113 } 114 115 return res.StatusCode, body, nil 116 }