From b3dd9d34cb982ee03c2a2c01d8e7956d78eee23d Mon Sep 17 00:00:00 2001
From: trsteel <trsteel@github.com>
Date: Wed, 12 Jul 2023 12:27:40 +0800
Subject: [PATCH] Support Basic/Bearer Authorization for prometheus metrics
 query

Signed-off-by: trsteel <trsteel@github.com>
---
 .../source/metrics_client_prometheus.go       | 13 ++++-
 .../source/metrics_client_prometheus_test.go  | 48 +++++++++++++++++++
 2 files changed, 60 insertions(+), 1 deletion(-)
 create mode 100644 pkg/scheduler/metrics/source/metrics_client_prometheus_test.go

diff --git a/pkg/scheduler/metrics/source/metrics_client_prometheus.go b/pkg/scheduler/metrics/source/metrics_client_prometheus.go
index 4b870628c1b..e155c2bd574 100644
--- a/pkg/scheduler/metrics/source/metrics_client_prometheus.go
+++ b/pkg/scheduler/metrics/source/metrics_client_prometheus.go
@@ -28,6 +28,7 @@ import (
 	"github.com/prometheus/client_golang/api"
 	prometheusv1 "github.com/prometheus/client_golang/api/prometheus/v1"
 	pmodel "github.com/prometheus/common/model"
+	"k8s.io/client-go/transport"
 	"k8s.io/klog/v2"
 )
 
@@ -52,11 +53,21 @@ func (p *PrometheusMetricsClient) NodeMetricsAvg(ctx context.Context, nodeName s
 	var client api.Client
 	var err error
 	insecureSkipVerify := p.conf["tls.insecureSkipVerify"] == "true"
-	tr := &http.Transport{
+	var tr http.RoundTripper = &http.Transport{
 		TLSClientConfig: &tls.Config{
 			InsecureSkipVerify: insecureSkipVerify,
 		},
 	}
+	tConf := &transport.Config{
+		Username:        p.conf["username"],
+		Password:        p.conf["password"],
+		BearerToken:     p.conf["bearertoken"],
+		BearerTokenFile: p.conf["bearertokenfile"],
+	}
+	if tr, err = transport.HTTPWrappersForConfig(tConf, tr); err != nil {
+		klog.Errorf("Error while wrap http round tripper with error: %v\n", err)
+		return nil, err
+	}
 	client, err = api.NewClient(api.Config{
 		Address:      p.address,
 		RoundTripper: tr,
diff --git a/pkg/scheduler/metrics/source/metrics_client_prometheus_test.go b/pkg/scheduler/metrics/source/metrics_client_prometheus_test.go
new file mode 100644
index 00000000000..c01e4b84ca2
--- /dev/null
+++ b/pkg/scheduler/metrics/source/metrics_client_prometheus_test.go
@@ -0,0 +1,48 @@
+package source
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+)
+
+const (
+	basicAuthUser   = "username"
+	basicAuthPwd    = "password"
+	bearerAuthToken = "bearertoken"
+)
+
+func TestPrometheusMetricsClient_NodeMetricsAvg(t *testing.T) {
+	testServer := httptest.NewServer(http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
+		if strings.HasPrefix(req.URL.Path, "/basic-auth") {
+			userName, password, ok := req.BasicAuth()
+			if !ok || userName != basicAuthUser || password != basicAuthPwd {
+				t.Errorf("basic auth missmatch, ok: %v, userName: %v, password: %v", ok, userName, password)
+				res.WriteHeader(http.StatusUnauthorized)
+			}
+
+		} else if strings.HasPrefix(req.URL.Path, "/token-auth") {
+			if auth := req.Header.Get("Authorization"); auth != "Bearer "+bearerAuthToken {
+				t.Errorf("bearer token missmatch, token: %s", auth)
+				res.WriteHeader(http.StatusUnauthorized)
+			}
+		}
+		res.Write([]byte(`{"status":"success","data":{"resultType":"vector","result":[]}}`))
+		res.WriteHeader(http.StatusOK)
+	}))
+	defer testServer.Close()
+	// test basic auth
+	conf := map[string]string{"username": basicAuthUser, "password": basicAuthPwd}
+	metricsClient, _ := NewPrometheusMetricsClient(testServer.URL+"/basic-auth", conf)
+	if _, err := metricsClient.NodeMetricsAvg(context.TODO(), "node-name", "5m"); err != nil {
+		t.Errorf("Get Node Metric Avg with basic auth err %v", err)
+	}
+	// test token auth
+	conf = map[string]string{"bearertoken": bearerAuthToken}
+	metricsClient, _ = NewPrometheusMetricsClient(testServer.URL+"/token-auth", conf)
+	if _, err := metricsClient.NodeMetricsAvg(context.TODO(), "node-name", "5m"); err != nil {
+		t.Errorf("Get Node Metric Avg with token auth err %v", err)
+	}
+}