From 065f82a4ed5118ac0b055daa5617be4ff787ae3e Mon Sep 17 00:00:00 2001
From: Daniel Hoelbling-Inzko <daniel.hoelbling-inzko@bitmovin.com>
Date: Fri, 22 Apr 2022 15:12:44 +0200
Subject: [PATCH] #342 - Implemented Label Value Sanitizer

Signed-off-by: Daniel Hoelbling-Inzko <daniel.hoelbling-inzko@bitmovin.com>
---
 .../client/LabelValueSanitizer.java           |  5 +++
 .../io/prometheus/client/SimpleCollector.java | 45 +++++++++++++++++++
 ...SimpleCollectorWithLabelSanitizerTest.java | 35 +++++++++++++++
 3 files changed, 85 insertions(+)
 create mode 100644 simpleclient/src/main/java/io/prometheus/client/LabelValueSanitizer.java
 create mode 100644 simpleclient/src/test/java/io/prometheus/client/SimpleCollectorWithLabelSanitizerTest.java

diff --git a/simpleclient/src/main/java/io/prometheus/client/LabelValueSanitizer.java b/simpleclient/src/main/java/io/prometheus/client/LabelValueSanitizer.java
new file mode 100644
index 000000000..6d6b34d6d
--- /dev/null
+++ b/simpleclient/src/main/java/io/prometheus/client/LabelValueSanitizer.java
@@ -0,0 +1,5 @@
+package io.prometheus.client;
+
+public interface LabelValueSanitizer {
+    String[] sanitize(String... labelValue);
+}
diff --git a/simpleclient/src/main/java/io/prometheus/client/SimpleCollector.java b/simpleclient/src/main/java/io/prometheus/client/SimpleCollector.java
index 5c5bf7c37..eb0a71434 100644
--- a/simpleclient/src/main/java/io/prometheus/client/SimpleCollector.java
+++ b/simpleclient/src/main/java/io/prometheus/client/SimpleCollector.java
@@ -51,6 +51,7 @@ public abstract class SimpleCollector<Child> extends Collector {
   protected final String help;
   protected final String unit;
   protected final List<String> labelNames;
+  protected final LabelValueSanitizer labelValueSanitizer;
 
   protected final ConcurrentMap<List<String>, Child> children = new ConcurrentHashMap<List<String>, Child>();
   protected Child noLabelsChild;
@@ -64,6 +65,7 @@ public Child labels(String... labelValues) {
     if (labelValues.length != labelNames.size()) {
       throw new IllegalArgumentException("Incorrect number of labels.");
     }
+    labelValues = labelValueSanitizer.sanitize(labelValues);
     for (String label: labelValues) {
       if (label == null) {
         throw new IllegalArgumentException("Label cannot be null.");
@@ -174,12 +176,42 @@ protected SimpleCollector(Builder b) {
     for (String n: labelNames) {
       checkMetricLabelName(n);
     }
+    labelValueSanitizer = b.labelValueSanitizer;
 
     if (!b.dontInitializeNoLabelsChild) {
       initializeNoLabelsChild();
     }
   }
 
+  /**
+   * Default Sanitizer - Will not perform any manipulation on labels
+   */
+  static LabelValueSanitizer NoOpLabelValueSanitizer() {
+    return new LabelValueSanitizer() {
+      @Override
+      public String[] sanitize(String... labelValue) {
+        return labelValue;
+      }
+    };
+  }
+
+  /**
+   * Transforms null labels to an empty string to guard against IllegalArgument runtime exceptions
+   */
+  static LabelValueSanitizer TransformNullLabelsToEmptyString() {
+    return new LabelValueSanitizer() {
+      @Override
+      public String[] sanitize(String... labelValue) {
+        for (int i = 0; i < labelValue.length; i++) {
+          if (labelValue[i] == null) {
+            labelValue[i] = "";
+          }
+        }
+        return labelValue;
+      }
+    };
+  }
+
   /**
    * Builders let you configure and then create collectors.
    */
@@ -191,6 +223,7 @@ public abstract static class Builder<B extends Builder<B, C>, C extends SimpleCo
     String unit = "";
     String help = "";
     String[] labelNames = new String[]{};
+    LabelValueSanitizer labelValueSanitizer = NoOpLabelValueSanitizer();
     // Some metrics require additional setup before the initialization can be done.
     boolean dontInitializeNoLabelsChild;
 
@@ -239,6 +272,18 @@ public B labelNames(String... labelNames) {
       return (B)this;
     }
 
+    /**
+     * Sanitize labels. Optional, defaults to no manipulation of labels.
+     * Useful to scrub credentials from labels
+     * or avoid Null values causing runtime exceptions
+     * @param handler
+     * @return new array of labels to use
+     */
+    public B labelValueSanitizer(LabelValueSanitizer handler) {
+      this.labelValueSanitizer = handler;
+      return (B)this;
+    }
+
     /**
      * Return the constructed collector.
      * <p>
diff --git a/simpleclient/src/test/java/io/prometheus/client/SimpleCollectorWithLabelSanitizerTest.java b/simpleclient/src/test/java/io/prometheus/client/SimpleCollectorWithLabelSanitizerTest.java
new file mode 100644
index 000000000..f964fe51c
--- /dev/null
+++ b/simpleclient/src/test/java/io/prometheus/client/SimpleCollectorWithLabelSanitizerTest.java
@@ -0,0 +1,35 @@
+package io.prometheus.client;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+
+import static org.junit.Assert.*;
+import static org.junit.rules.ExpectedException.none;
+
+
+public class SimpleCollectorWithLabelSanitizerTest {
+
+  CollectorRegistry registry;
+  Gauge metric;
+
+  @Rule
+  public final ExpectedException thrown = none();
+
+  @Before
+  public void setUp() {
+    registry = new CollectorRegistry();
+    metric = Gauge.build().name("nonulllabels").help("help").labelNames("l").labelValueSanitizer(Gauge.TransformNullLabelsToEmptyString()).register(registry);
+  }
+
+  private Double getValue(String labelValue) {
+    return registry.getSampleValue("nonulllabels", new String[]{"l"}, new String[]{labelValue});
+  }
+
+  @Test
+  public void testNullLabelDoesntThrowWithLabelSanitizer() {
+    metric.labels(new String[]{null}).inc();
+    assertEquals(1.0, getValue("").doubleValue(), .001);
+  }
+}