github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/tools/querytee/proxy_backend.go (about)

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