From b893483d26ccce68c381b717f3551d5da96ed622 Mon Sep 17 00:00:00 2001
From: Yening Qin <710leo@gmail.com>
Date: Tue, 18 Oct 2022 20:44:38 +0800
Subject: [PATCH] prom client support add header (#1203)

* prom client support add header
---
 src/webapi/prom/prom.go                | 34 ++++++++++++++++++++++++--
 src/webapi/router/router_prometheus.go |  2 +-
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/src/webapi/prom/prom.go b/src/webapi/prom/prom.go
index b06da7b6..d67574fe 100644
--- a/src/webapi/prom/prom.go
+++ b/src/webapi/prom/prom.go
@@ -7,6 +7,7 @@ import (
 	"math/rand"
 	"net"
 	"net/http"
+	"sort"
 	"strings"
 	"sync"
 	"time"
@@ -120,7 +121,8 @@ type DSReply struct {
 					PrometheusUser string `json:"prometheus.user"`
 					PrometheusPass string `json:"prometheus.password"`
 				} `json:"prometheus.basic"`
-				PrometheusTimeout int64 `json:"prometheus.timeout"`
+				Headers           map[string]string `json:"prometheus.headers"`
+				PrometheusTimeout int64             `json:"prometheus.timeout"`
 			} `json:"settings,omitempty"`
 		} `json:"items"`
 	} `json:"data"`
@@ -192,7 +194,8 @@ func loadClustersFromAPI() {
 			old.Opts.BasicAuthUser != item.Settings.PrometheusBasic.PrometheusUser ||
 			old.Opts.BasicAuthPass != item.Settings.PrometheusBasic.PrometheusPass ||
 			old.Opts.Timeout != item.Settings.PrometheusTimeout ||
-			old.Opts.Prom != item.Settings.PrometheusAddr {
+			old.Opts.Prom != item.Settings.PrometheusAddr ||
+			!equalHeader(old.Opts.Headers, transformHeader(item.Settings.Headers)) {
 			opt := config.ClusterOptions{
 				Name:                item.Name,
 				Prom:                item.Settings.PrometheusAddr,
@@ -201,6 +204,7 @@ func loadClustersFromAPI() {
 				Timeout:             item.Settings.PrometheusTimeout,
 				DialTimeout:         5000,
 				MaxIdleConnsPerHost: 32,
+				Headers:             transformHeader(item.Settings.Headers),
 			}
 
 			if strings.HasPrefix(opt.Prom, "https") {
@@ -260,3 +264,29 @@ func newClusterByOption(opt config.ClusterOptions) *ClusterType {
 
 	return cluster
 }
+
+func equalHeader(a, b []string) bool {
+	sort.Strings(a)
+	sort.Strings(b)
+
+	if len(a) != len(b) {
+		return false
+	}
+
+	for i := range a {
+		if a[i] != b[i] {
+			return false
+		}
+	}
+
+	return true
+}
+
+func transformHeader(header map[string]string) []string {
+	var headers []string
+	for k, v := range header {
+		headers = append(headers, k)
+		headers = append(headers, v)
+	}
+	return headers
+}
diff --git a/src/webapi/router/router_prometheus.go b/src/webapi/router/router_prometheus.go
index 5bb2b87b..9f3af605 100644
--- a/src/webapi/router/router_prometheus.go
+++ b/src/webapi/router/router_prometheus.go
@@ -112,7 +112,7 @@ func prometheusProxy(c *gin.Context) {
 		headerCount := len(cluster.Opts.Headers)
 		if headerCount > 0 && headerCount%2 == 0 {
 			for i := 0; i < len(cluster.Opts.Headers); i += 2 {
-				req.Header.Add(cluster.Opts.Headers[i], cluster.Opts.Headers[i+1])
+				req.Header.Set(cluster.Opts.Headers[i], cluster.Opts.Headers[i+1])
 				if cluster.Opts.Headers[i] == "Host" {
 					req.Host = cluster.Opts.Headers[i+1]
 				}
-- 
GitLab