From 07456d125979bde6c69b806c5daa09064f5153d8 Mon Sep 17 00:00:00 2001
From: yanli <38203792+tianyanli@users.noreply.github.com>
Date: Thu, 27 Jun 2024 22:26:55 +0800
Subject: [PATCH] =?UTF-8?q?=E6=94=AF=E6=8C=81=E5=BC=80=E5=90=AFhttp?=
 =?UTF-8?q?=E6=9C=8D=E5=8A=A1=E5=90=8Epushtateway=E6=8E=A5=E5=8F=A3?=
 =?UTF-8?q?=E5=90=8E=E9=80=9A=E8=BF=87=E9=85=8D=E7=BD=AE=E6=96=87=E4=BB=B6?=
 =?UTF-8?q?=E4=B8=ADagent=5Fhost=5Ftag=E4=BB=8E=E6=8E=A8=E9=80=81=E7=9A=84?=
 =?UTF-8?q?=E6=95=B0=E6=8D=AE=E4=B8=AD=E6=8C=87=E5=AE=9Aagent=5Fhostname?=
 =?UTF-8?q?=E7=9A=84=E5=80=BC=20(#984)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* iptables插件

* 添加注册

* pushgateway接口支持用户指定值置换agent_hostname

* pushgateway接口支持用户指定值置换agent_hostname

* 修改命名
---
 api/router_func.go        | 24 ++++++++++++------------
 api/router_pushgateway.go | 22 ++++++++++++++++++----
 conf/config.toml          |  1 +
 config/config.go          | 12 +++++++-----
 4 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/api/router_func.go b/api/router_func.go
index 6c1794ec..fa4ae546 100644
--- a/api/router_func.go
+++ b/api/router_func.go
@@ -4,7 +4,6 @@ import (
 	"compress/gzip"
 	"errors"
 	"io"
-	"io/ioutil"
 	"net/http"
 
 	"github.com/gogo/protobuf/proto"
@@ -15,29 +14,30 @@ import (
 const agentHostnameLabelKey = "agent_hostname"
 
 func readerGzipBody(contentEncoding string, request *http.Request) (bytes []byte, err error) {
-	if contentEncoding == "gzip" {
-		var (
-			r *gzip.Reader
-		)
+	switch contentEncoding {
+	case "gzip":
+		var r *gzip.Reader
 		r, err = gzip.NewReader(request.Body)
 		if err != nil {
 			return nil, err
 		}
-
 		defer r.Close()
-		bytes, err = ioutil.ReadAll(r)
-	} else if contentEncoding == "snappy" {
+
+		bytes, err = io.ReadAll(r)
+	case "snappy":
 		defer request.Body.Close()
-		compressed, err := ioutil.ReadAll(request.Body)
+		var compressed []byte
+		compressed, err = io.ReadAll(request.Body)
 		if err != nil {
 			return nil, err
 		}
 
 		bytes, err = snappy.Decode(nil, compressed)
-	} else {
+	default:
 		defer request.Body.Close()
-		bytes, err = ioutil.ReadAll(request.Body)
+		bytes, err = io.ReadAll(request.Body)
 	}
+
 	if err != nil || len(bytes) == 0 {
 		return nil, errors.New("request parameter error")
 	}
@@ -48,7 +48,7 @@ func readerGzipBody(contentEncoding string, request *http.Request) (bytes []byte
 // DecodeWriteRequest from an io.Reader into a prompb.WriteRequest, handling
 // snappy decompression.
 func DecodeWriteRequest(r io.Reader) (*prompb.WriteRequest, error) {
-	compressed, err := ioutil.ReadAll(r)
+	compressed, err := io.ReadAll(r)
 	if err != nil {
 		return nil, err
 	}
diff --git a/api/router_pushgateway.go b/api/router_pushgateway.go
index 23fd880b..15ed1061 100644
--- a/api/router_pushgateway.go
+++ b/api/router_pushgateway.go
@@ -70,6 +70,11 @@ func pushgateway(c *gin.Context) {
 
 	ignoreHostname := config.Config.HTTP.IgnoreHostname || c.GetBool("ignore_hostname")
 	ignoreGlobalLabels := config.Config.HTTP.IgnoreGlobalLabels || c.GetBool("ignore_global_labels")
+	// 获取 AgentHostTag 的值
+	agentHostTag := config.Config.HTTP.AgentHostTag
+	if agentHostTag == "" {
+		agentHostTag = c.GetString("agent_host_tag")
+	}
 
 	now := time.Now()
 
@@ -92,16 +97,25 @@ func pushgateway(c *gin.Context) {
 		// add url labels
 		for k, v := range labels {
 			samples[i].Labels[k] = v
+
 		}
 
 		// add label: agent_hostname
-		if _, has := samples[i].Labels[agentHostnameLabelKey]; !has && !ignoreHostname {
-			samples[i].Labels[agentHostnameLabelKey] = config.Config.GetHostname()
+		if !ignoreHostname {
+			if agentHostTag == "" {
+				if _, has := samples[i].Labels[agentHostnameLabelKey]; !has {
+					samples[i].Labels[agentHostnameLabelKey] = config.Config.GetHostname()
+				}
+			} else {
+				// 从当前现有的 Labels 中找到 key 等于 config.Config.HTTP.AgentHostTag 的值
+				if value, exists := samples[i].Labels[agentHostTag]; exists {
+					samples[i].Labels[agentHostnameLabelKey] = value
+				}
+			}
 		}
-
 	}
 	writer.WriteSamples(samples)
-	c.String(200, "forwarding...")
+	c.String(http.StatusOK, "forwarding...")
 }
 
 // fork prometheus/pushgateway handler/push.go
diff --git a/conf/config.toml b/conf/config.toml
index 13bf40e6..39221909 100644
--- a/conf/config.toml
+++ b/conf/config.toml
@@ -84,6 +84,7 @@ address = ":9100"
 print_access = false
 run_mode = "release"
 ignore_hostname = false
+agent_host_tag = ""
 ignore_global_labels = false
 
 [ibex]
diff --git a/config/config.go b/config/config.go
index c2543f2d..54c28acf 100644
--- a/config/config.go
+++ b/config/config.go
@@ -66,11 +66,13 @@ type WriterOption struct {
 }
 
 type HTTP struct {
-	Enable             bool   `toml:"enable"`
-	Address            string `toml:"address"`
-	PrintAccess        bool   `toml:"print_access"`
-	RunMode            string `toml:"run_mode"`
-	IgnoreHostname     bool   `toml:"ignore_hostname"`
+	Enable         bool   `toml:"enable"`
+	Address        string `toml:"address"`
+	PrintAccess    bool   `toml:"print_access"`
+	RunMode        string `toml:"run_mode"`
+	IgnoreHostname bool   `toml:"ignore_hostname"`
+	// The tag used to name the agent host
+	AgentHostTag       string `toml:"agent_host_tag"`
 	IgnoreGlobalLabels bool   `toml:"ignore_global_labels"`
 	CertFile           string `toml:"cert_file"`
 	KeyFile            string `toml:"key_file"`