Ver Fonte

Merge pull request #540 from kubecost/AjayTripathy-prom-basicauth

add prometheus basic auth
Ajay Tripathy há 5 anos atrás
pai
commit
1fa0516f15
3 ficheiros alterados com 47 adições e 15 exclusões
  1. 6 2
      pkg/costmodel/router.go
  2. 22 0
      pkg/env/costmodelenv.go
  3. 19 13
      pkg/prom/prom.go

+ 6 - 2
pkg/costmodel/router.go

@@ -46,8 +46,12 @@ const (
 var (
 	// gitCommit is set by the build system
 	gitCommit                       string
+	dbBasicAuthUsername             string = env.GetDBBasicAuthUsername()
+	dbBasicAuthPW                   string = env.GetDBBasicAuthUserPassword()
+	dbBearerToken                   string = env.GetDBBearerToken()
 	multiclusterDBBasicAuthUsername string = env.GetMultiClusterBasicAuthUsername()
 	multiclusterDBBasicAuthPW       string = env.GetMultiClusterBasicAuthPassword()
+	multiClusterBearerToken         string = env.GetMultiClusterBearerToken()
 )
 
 var Router = httprouter.New()
@@ -739,7 +743,7 @@ func Initialize(additionalConfigWatchers ...ConfigWatchers) {
 		Address:      address,
 		RoundTripper: LongTimeoutRoundTripper,
 	}
-	promCli, _ := prom.NewRateLimitedClient(pc, queryConcurrency, "", "")
+	promCli, _ := prom.NewRateLimitedClient(pc, queryConcurrency, dbBasicAuthUsername, dbBasicAuthPW, dbBearerToken)
 
 	m, err := ValidatePrometheus(promCli, false)
 	if err != nil || m.Running == false {
@@ -973,7 +977,7 @@ func Initialize(additionalConfigWatchers ...ConfigWatchers) {
 				RoundTripper: thanosRT,
 			}
 
-			thanosCli, _ := prom.NewRateLimitedClient(thanosConfig, queryConcurrency, multiclusterDBBasicAuthUsername, multiclusterDBBasicAuthPW)
+			thanosCli, _ := prom.NewRateLimitedClient(thanosConfig, queryConcurrency, multiclusterDBBasicAuthUsername, multiclusterDBBasicAuthPW, multiClusterBearerToken)
 
 			_, err = ValidatePrometheus(thanosCli, true)
 			if err != nil {

+ 22 - 0
pkg/env/costmodelenv.go

@@ -28,8 +28,13 @@ const (
 	ErrorReportingEnabledEnvVar   = "ERROR_REPORTING_ENABLED"
 	ValuesReportingEnabledEnvVar  = "VALUES_REPORTING_ENABLED"
 
+	DBBasicAuthUsername = "DB_BASIC_AUTH_USERNAME"
+	DBBasicAuthPassword = "DB_BASIC_AUTH_PW"
+	DBBearerToken       = "DB_BEARER_TOKEN"
+
 	MultiClusterBasicAuthUsername = "MC_BASIC_AUTH_USERNAME"
 	MultiClusterBasicAuthPassword = "MC_BASIC_AUTH_PW"
+	MultiClusterBearerToken       = "MC_BEARER_TOKEN"
 
 	InsecureSkipVerify = "INSECURE_SKIP_VERIFY"
 )
@@ -178,6 +183,19 @@ func GetMaxQueryConcurrency() int {
 	return GetInt(MaxQueryConcurrencyEnvVar, 5)
 }
 
+func GetDBBasicAuthUsername() string {
+	return Get(DBBasicAuthUsername, "")
+}
+
+func GetDBBasicAuthUserPassword() string {
+	return Get(DBBasicAuthPassword, "")
+
+}
+
+func GetDBBearerToken() string {
+	return Get(DBBearerToken, "")
+}
+
 // GetMultiClusterBasicAuthUsername returns the environemnt variable value for MultiClusterBasicAuthUsername
 func GetMultiClusterBasicAuthUsername() string {
 	return Get(MultiClusterBasicAuthUsername, "")
@@ -187,3 +205,7 @@ func GetMultiClusterBasicAuthUsername() string {
 func GetMultiClusterBasicAuthPassword() string {
 	return Get(MultiClusterBasicAuthPassword, "")
 }
+
+func GetMultiClusterBearerToken() string {
+	return Get(MultiClusterBearerToken, "")
+}

+ 19 - 13
pkg/prom/prom.go

@@ -13,12 +13,13 @@ import (
 // Creates a new prometheus client which limits the total number of concurrent outbound requests
 // allowed at a given moment.
 type RateLimitedPrometheusClient struct {
-	client   prometheus.Client
-	limiter  *util.Semaphore
-	requests *util.AtomicInt32
-	outbound *util.AtomicInt32
-	username string
-	password string
+	client      prometheus.Client
+	limiter     *util.Semaphore
+	requests    *util.AtomicInt32
+	outbound    *util.AtomicInt32
+	username    string
+	password    string
+	bearerToken string
 }
 
 // requestCounter is used to determine if the prometheus client keeps track of
@@ -30,7 +31,7 @@ type requestCounter interface {
 
 // NewRateLimitedClient creates a prometheus client which limits the number of concurrent outbound
 // prometheus requests.
-func NewRateLimitedClient(config prometheus.Config, maxConcurrency int, username, password string) (prometheus.Client, error) {
+func NewRateLimitedClient(config prometheus.Config, maxConcurrency int, username, password, bearerToken string) (prometheus.Client, error) {
 	c, err := prometheus.NewClient(config)
 	if err != nil {
 		return nil, err
@@ -41,12 +42,13 @@ func NewRateLimitedClient(config prometheus.Config, maxConcurrency int, username
 	outbound := util.NewAtomicInt32(0)
 
 	return &RateLimitedPrometheusClient{
-		client:   c,
-		limiter:  limiter,
-		requests: requests,
-		outbound: outbound,
-		username: username,
-		password: password,
+		client:      c,
+		limiter:     limiter,
+		requests:    requests,
+		outbound:    outbound,
+		username:    username,
+		password:    password,
+		bearerToken: bearerToken,
 	}, nil
 }
 
@@ -84,6 +86,10 @@ func (rlpc *RateLimitedPrometheusClient) Do(ctx context.Context, req *http.Reque
 	if rlpc.username != "" {
 		req.SetBasicAuth(rlpc.username, rlpc.password)
 	}
+	if rlpc.bearerToken != "" {
+		token := "Bearer " + rlpc.bearerToken
+		req.Header.Add("Authorization", token)
+	}
 	// Increment the total request counter first
 	rlpc.requests.Increment()
 	defer rlpc.requests.Decrement()