github.com/grafana/pyroscope@v1.18.0/cmd/profilecli/client.go (about)

     1  package main
     2  
     3  import (
     4  	"fmt"
     5  	"net/http"
     6  	"strings"
     7  
     8  	"connectrpc.com/connect"
     9  	"github.com/prometheus/common/version"
    10  	"gopkg.in/alecthomas/kingpin.v2"
    11  )
    12  
    13  const (
    14  	envPrefix = "PROFILECLI_"
    15  
    16  	protocolTypeConnect = "connect"
    17  	protocolTypeGRPC    = "grpc"
    18  	protocolTypeGRPCWeb = "grpc-web"
    19  
    20  	acceptHeaderMimeType = "*/*"
    21  )
    22  
    23  var acceptHeaderClientCapabilities = []string{
    24  	"allow-utf8-labelnames=true",
    25  }
    26  
    27  var userAgentHeader = fmt.Sprintf("pyroscope/%s", version.Version)
    28  
    29  func addClientCapabilitiesHeader(r *http.Request, mime string, clientCapabilities []string) {
    30  	missingClientCapabilities := make([]string, 0, len(clientCapabilities))
    31  	for _, capability := range clientCapabilities {
    32  		found := false
    33  		// Check if any header value already contains this capability
    34  		for _, value := range r.Header.Values("Accept") {
    35  			if strings.Contains(value, capability) {
    36  				found = true
    37  				break
    38  			}
    39  		}
    40  
    41  		if !found {
    42  			missingClientCapabilities = append(missingClientCapabilities, capability)
    43  		}
    44  	}
    45  
    46  	if len(missingClientCapabilities) > 0 {
    47  		acceptHeader := mime
    48  		acceptHeader += ";" + strings.Join(missingClientCapabilities, ";")
    49  		r.Header.Add("Accept", acceptHeader)
    50  	}
    51  }
    52  
    53  type phlareClient struct {
    54  	TenantID    string
    55  	URL         string
    56  	BearerToken string
    57  	BasicAuth   struct {
    58  		Username string
    59  		Password string
    60  	}
    61  	defaultTransport http.RoundTripper
    62  	client           *http.Client
    63  	protocol         string
    64  }
    65  
    66  type authRoundTripper struct {
    67  	client *phlareClient
    68  	next   http.RoundTripper
    69  }
    70  
    71  func (a *authRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
    72  	if c := a.client; c != nil {
    73  		if c.TenantID != "" {
    74  			req.Header.Set("X-Scope-OrgID", c.TenantID)
    75  		}
    76  		if c.BasicAuth.Username != "" || c.BasicAuth.Password != "" {
    77  			req.SetBasicAuth(c.BasicAuth.Username, c.BasicAuth.Password)
    78  		} else if c.BearerToken != "" {
    79  			req.Header.Set("Authorization", "Bearer "+c.BearerToken)
    80  		}
    81  	}
    82  
    83  	addClientCapabilitiesHeader(req, acceptHeaderMimeType, acceptHeaderClientCapabilities)
    84  	req.Header.Set("User-Agent", userAgentHeader)
    85  
    86  	return a.next.RoundTrip(req)
    87  }
    88  
    89  func (c *phlareClient) httpClient() *http.Client {
    90  	if c.client == nil {
    91  		if c.defaultTransport == nil {
    92  			c.defaultTransport = http.DefaultTransport
    93  		}
    94  		c.client = &http.Client{Transport: &authRoundTripper{
    95  			client: c,
    96  			next:   c.defaultTransport,
    97  		}}
    98  	}
    99  	return c.client
   100  }
   101  
   102  func (c *phlareClient) protocolOption() connect.ClientOption {
   103  	switch c.protocol {
   104  	case protocolTypeGRPC:
   105  		return connect.WithGRPC()
   106  	case protocolTypeGRPCWeb:
   107  		return connect.WithGRPCWeb()
   108  	case protocolTypeConnect:
   109  		return connect.WithClientOptions()
   110  	default:
   111  		return connect.WithClientOptions()
   112  	}
   113  }
   114  
   115  type commander interface {
   116  	Flag(name, help string) *kingpin.FlagClause
   117  	Arg(name, help string) *kingpin.ArgClause
   118  }
   119  
   120  func addPhlareClient(cmd commander) *phlareClient {
   121  	client := &phlareClient{}
   122  
   123  	cmd.Flag("url", "URL of the Pyroscope Endpoint. (Examples: https://profiles-prod-001.grafana.net for a Grafana Cloud endpoint, https://grafana.example.net/api/datasources/proxy/uid/<uid> when using the Grafana data source proxy)").Default("http://localhost:4040").Envar(envPrefix + "URL").StringVar(&client.URL)
   124  	cmd.Flag("tenant-id", "The tenant ID to be used for the X-Scope-OrgID header.").Default("").Envar(envPrefix + "TENANT_ID").StringVar(&client.TenantID)
   125  	cmd.Flag("token", "The bearer token to be used for communication with the server. Particularly useful when connecting to Grafana data source URLs (bearer token should be a Grafana Service Account token of the form 'glsa_[...]')").Default("").Envar(envPrefix + "TOKEN").StringVar(&client.BearerToken)
   126  	cmd.Flag("username", "The username to be used for basic auth.").Default("").Envar(envPrefix + "USERNAME").StringVar(&client.BasicAuth.Username)
   127  	cmd.Flag("password", "The password to be used for basic auth.").Default("").Envar(envPrefix + "PASSWORD").StringVar(&client.BasicAuth.Password)
   128  	cmd.Flag("protocol", "The protocol to be used for communicating with the server.").Default(protocolTypeConnect).EnumVar(&client.protocol,
   129  		protocolTypeConnect, protocolTypeGRPC, protocolTypeGRPCWeb)
   130  	return client
   131  }