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  }