From 9fc8a1845c69ab14950be1c4a97b21ee1b3d98d7 Mon Sep 17 00:00:00 2001
From: Josh Powers <powersj@fastmail.com>
Date: Thu, 2 May 2024 08:08:27 -0600
Subject: [PATCH] fix(outputs.loki): Option to sanitize label names

fixes: #15274
---
 plugins/outputs/loki/README.md    |  5 ++++
 plugins/outputs/loki/loki.go      | 42 +++++++++++++++++++++----------
 plugins/outputs/loki/loki_test.go | 35 ++++++++++++++++++++++++++
 plugins/outputs/loki/sample.conf  |  5 ++++
 4 files changed, 74 insertions(+), 13 deletions(-)

diff --git a/plugins/outputs/loki/README.md b/plugins/outputs/loki/README.md
index 5c0f8ad0782cc..428e4676b5179 100644
--- a/plugins/outputs/loki/README.md
+++ b/plugins/outputs/loki/README.md
@@ -53,6 +53,11 @@ to use them.
   # tls_cert = "/etc/telegraf/cert.pem"
   # tls_key = "/etc/telegraf/key.pem"
 
+  ## Sanitize Tag Names
+  ## If true, all tag names will have invalid characters replaced with
+  ## underscores that do not match the regex: ^[a-zA-Z_:][a-zA-Z0-9_:]*.
+  # sanitize_label_names = false
+
   ## Metric Name Label
   ## Label to use for the metric name to when sending metrics. If set to an
   ## empty string, this will not add the label. This is NOT suggested as there
diff --git a/plugins/outputs/loki/loki.go b/plugins/outputs/loki/loki.go
index e2f226ef84208..ca256b390f60a 100644
--- a/plugins/outputs/loki/loki.go
+++ b/plugins/outputs/loki/loki.go
@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"regexp"
 	"sort"
 	"strconv"
 	"strings"
@@ -34,18 +35,19 @@ const (
 )
 
 type Loki struct {
-	Domain          string            `toml:"domain"`
-	Endpoint        string            `toml:"endpoint"`
-	Timeout         config.Duration   `toml:"timeout"`
-	Username        config.Secret     `toml:"username"`
-	Password        config.Secret     `toml:"password"`
-	Headers         map[string]string `toml:"http_headers"`
-	ClientID        string            `toml:"client_id"`
-	ClientSecret    string            `toml:"client_secret"`
-	TokenURL        string            `toml:"token_url"`
-	Scopes          []string          `toml:"scopes"`
-	GZipRequest     bool              `toml:"gzip_request"`
-	MetricNameLabel string            `toml:"metric_name_label"`
+	Domain             string            `toml:"domain"`
+	Endpoint           string            `toml:"endpoint"`
+	Timeout            config.Duration   `toml:"timeout"`
+	Username           config.Secret     `toml:"username"`
+	Password           config.Secret     `toml:"password"`
+	Headers            map[string]string `toml:"http_headers"`
+	ClientID           string            `toml:"client_id"`
+	ClientSecret       string            `toml:"client_secret"`
+	TokenURL           string            `toml:"token_url"`
+	Scopes             []string          `toml:"scopes"`
+	GZipRequest        bool              `toml:"gzip_request"`
+	MetricNameLabel    string            `toml:"metric_name_label"`
+	SanitizeLabelNames bool              `toml:"sanitize_label_names"`
 
 	url    string
 	client *http.Client
@@ -127,8 +129,13 @@ func (l *Loki) Write(metrics []telegraf.Metric) error {
 		}
 
 		tags := m.TagList()
-		var line string
+		if l.SanitizeLabelNames {
+			for _, t := range tags {
+				t.Key = sanitizeLabelName(t.Key)
+			}
+		}
 
+		var line string
 		for _, f := range m.FieldList() {
 			line += fmt.Sprintf("%s=\"%v\" ", f.Key, f.Value)
 		}
@@ -200,6 +207,15 @@ func (l *Loki) writeMetrics(s Streams) error {
 	return nil
 }
 
+// Verify the label name matches the regex [a-zA-Z_:][a-zA-Z0-9_:]*
+func sanitizeLabelName(name string) string {
+	re := regexp.MustCompile(`^[^a-zA-Z_:]`)
+	result := re.ReplaceAllString(name, "_")
+
+	re = regexp.MustCompile(`[^a-zA-Z0-9_:]`)
+	return re.ReplaceAllString(result, "_")
+}
+
 func init() {
 	outputs.Add("loki", func() telegraf.Output {
 		return &Loki{
diff --git a/plugins/outputs/loki/loki_test.go b/plugins/outputs/loki/loki_test.go
index fa7dc2bac07eb..08c2acd7d8016 100644
--- a/plugins/outputs/loki/loki_test.go
+++ b/plugins/outputs/loki/loki_test.go
@@ -469,3 +469,38 @@ func TestMetricSorting(t *testing.T) {
 		require.NoError(t, err)
 	})
 }
+
+func TestSanitizeLabelName(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    string
+		expected string
+	}{
+		{
+			name:     "no change",
+			input:    "foobar",
+			expected: "foobar",
+		},
+		{
+			name:     "replace invalid first charachter",
+			input:    "3foobar",
+			expected: "_foobar",
+		},
+		{
+			name:     "replace invalid later charachter",
+			input:    "foobar.foobar",
+			expected: "foobar_foobar",
+		},
+		{
+			name:     "numbers allowed later",
+			input:    "foobar.foobar.2002",
+			expected: "foobar_foobar_2002",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require.Equal(t, tt.expected, sanitizeLabelName(tt.input))
+		})
+	}
+}
diff --git a/plugins/outputs/loki/sample.conf b/plugins/outputs/loki/sample.conf
index 93695764851be..e899f5bf91306 100644
--- a/plugins/outputs/loki/sample.conf
+++ b/plugins/outputs/loki/sample.conf
@@ -24,6 +24,11 @@
   # tls_cert = "/etc/telegraf/cert.pem"
   # tls_key = "/etc/telegraf/key.pem"
 
+  ## Sanitize Tag Names
+  ## If true, all tag names will have invalid characters replaced with
+  ## underscores that do not match the regex: ^[a-zA-Z_:][a-zA-Z0-9_:]*.
+  # sanitize_label_names = false
+
   ## Metric Name Label
   ## Label to use for the metric name to when sending metrics. If set to an
   ## empty string, this will not add the label. This is NOT suggested as there