From 921ad0771f2c8637dc70499a8f84e2a9ca43bfb8 Mon Sep 17 00:00:00 2001
From: Greg Chambers <gregory.w.chambers@gmail.com>
Date: Sun, 2 Jun 2024 11:54:28 -0700
Subject: [PATCH] Added aggregate type checking with test validation

Signed-off-by: Greg Chambers <gregory.w.chambers@gmail.com>
---
 pkg/line/line.go      | 35 ++++++++++++++++++++++++++---------
 pkg/line/line_test.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/pkg/line/line.go b/pkg/line/line.go
index e591b30..33d9106 100644
--- a/pkg/line/line.go
+++ b/pkg/line/line.go
@@ -229,17 +229,34 @@ func (p *Parser) LineToEvents(line string, sampleErrors prometheus.CounterVec, s
 			return events
 		}
 
-		// handle DogStatsD extended aggregation
-		lineParts := strings.SplitN(elements[1], "|", 2)
+		lineParts := strings.SplitN(elements[1], "|", 3)
 		if strings.Contains(lineParts[0], ":") {
-			aggValues := strings.Split(lineParts[0], ":")
-			aggLines := make([]string, len(aggValues))
-			tags := lineParts[1]
-
-			for i, aggValue := range aggValues {
-				aggLines[i] = strings.Join([]string{aggValue, tags}, "|")
+			// handle DogStatsD extended aggregation
+			var isValidAggType bool
+			switch lineParts[1] {
+			case
+				"ms", // timer
+				"h",  // histogram
+				"d":  // distribution
+				isValidAggType = true
+			default:
+				isValidAggType = false
+			}
+
+			if isValidAggType {
+				aggValues := strings.Split(lineParts[0], ":")
+				aggLines := make([]string, len(aggValues))
+
+				for i, aggValue := range aggValues {
+					aggLines[i] = strings.Join([]string{aggValue, lineParts[1], lineParts[2]}, "|")
+				}
+
+				samples = aggLines
+			} else {
+				sampleErrors.WithLabelValues("invalid_extended_aggregate_type").Inc()
+				level.Debug(logger).Log("msg", "Bad line (invalid extended aggregate type) from StatsD", "line", line)
+				return events
 			}
-			samples = aggLines
 		} else {
 			// disable multi-metrics
 			samples = elements[1:]
diff --git a/pkg/line/line_test.go b/pkg/line/line_test.go
index c084b22..56032dd 100644
--- a/pkg/line/line_test.go
+++ b/pkg/line/line_test.go
@@ -414,6 +414,34 @@ func TestLineToEvents(t *testing.T) {
 				&event.ObserverEvent{OMetricName: "foo_timing", OValue: 0.00001, OLabels: map[string]string{"tag1": "bar", "tag2": "baz"}},
 			},
 		},
+		"datadog histogram with extended aggregation values": {
+			in: "foo_histogram:0.5:120:3000:10:20000:0.01|h|#tag1:bar,#tag2:baz",
+			out: event.Events{
+				&event.ObserverEvent{OMetricName: "foo_histogram", OValue: 0.5, OLabels: map[string]string{"tag1": "bar", "tag2": "baz"}},
+				&event.ObserverEvent{OMetricName: "foo_histogram", OValue: 120, OLabels: map[string]string{"tag1": "bar", "tag2": "baz"}},
+				&event.ObserverEvent{OMetricName: "foo_histogram", OValue: 3000, OLabels: map[string]string{"tag1": "bar", "tag2": "baz"}},
+				&event.ObserverEvent{OMetricName: "foo_histogram", OValue: 10, OLabels: map[string]string{"tag1": "bar", "tag2": "baz"}},
+				&event.ObserverEvent{OMetricName: "foo_histogram", OValue: 20000, OLabels: map[string]string{"tag1": "bar", "tag2": "baz"}},
+				&event.ObserverEvent{OMetricName: "foo_histogram", OValue: 0.01, OLabels: map[string]string{"tag1": "bar", "tag2": "baz"}},
+			},
+		},
+		"datadog distribution with extended aggregation values": {
+			in: "foo_distribution:0.5:120:3000:10:20000:0.01|d|#tag1:bar,#tag2:baz",
+			out: event.Events{
+				&event.ObserverEvent{OMetricName: "foo_distribution", OValue: 0.5, OLabels: map[string]string{"tag1": "bar", "tag2": "baz"}},
+				&event.ObserverEvent{OMetricName: "foo_distribution", OValue: 120, OLabels: map[string]string{"tag1": "bar", "tag2": "baz"}},
+				&event.ObserverEvent{OMetricName: "foo_distribution", OValue: 3000, OLabels: map[string]string{"tag1": "bar", "tag2": "baz"}},
+				&event.ObserverEvent{OMetricName: "foo_distribution", OValue: 10, OLabels: map[string]string{"tag1": "bar", "tag2": "baz"}},
+				&event.ObserverEvent{OMetricName: "foo_distribution", OValue: 20000, OLabels: map[string]string{"tag1": "bar", "tag2": "baz"}},
+				&event.ObserverEvent{OMetricName: "foo_distribution", OValue: 0.01, OLabels: map[string]string{"tag1": "bar", "tag2": "baz"}},
+			},
+		},
+		"datadog counter with invalid extended aggregation values": {
+			in: "foo_counter:0.5:120:3000:10:20000:0.01|c|#tag1:bar,#tag2:baz",
+		},
+		"datadog gauge with invalid extended aggregation values": {
+			in: "foo_gauge:0.5:120:3000:10:20000:0.01|g|#tag1:bar,#tag2:baz",
+		},
 		"timings with sampling factor": {
 			in: "foo.timing:0.5|ms|@0.1",
 			out: event.Events{