Merge pull request from GHSA-x768-cvr2-345r

Author: fabianfett

Sources/Prometheus/PrometheusCollectorRegistry.swift

@@ -73,7 +73,8 @@ public final class PrometheusCollectorRegistry: Sendable {
     /// - Parameter name: A name to identify ``Counter``'s value.
     /// - Returns: A ``Counter`` that is registered with this ``PrometheusCollectorRegistry``
     public func makeCounter(name: String) -> Counter {
-        self.box.withLockedValue { store -> Counter in
+        let name = name.ensureValidMetricName()
+        return self.box.withLockedValue { store -> Counter in
             guard let value = store[name] else {
                 let counter = Counter(name: name, labels: [])
                 store[name] = .counter(counter)
@@ -106,6 +107,9 @@ public final class PrometheusCollectorRegistry: Sendable {
             return self.makeCounter(name: name)
         }
 
+        let name = name.ensureValidMetricName()
+        let labels = labels.ensureValidLabelNames()
+
         return self.box.withLockedValue { store -> Counter in
             guard let value = store[name] else {
                 let labelNames = labels.allLabelNames
@@ -154,7 +158,8 @@ public final class PrometheusCollectorRegistry: Sendable {
     /// - Parameter name: A name to identify ``Gauge``'s value.
     /// - Returns: A ``Gauge`` that is registered with this ``PrometheusCollectorRegistry``
     public func makeGauge(name: String) -> Gauge {
-        self.box.withLockedValue { store -> Gauge in
+        let name = name.ensureValidMetricName()
+        return self.box.withLockedValue { store -> Gauge in
             guard let value = store[name] else {
                 let gauge = Gauge(name: name, labels: [])
                 store[name] = .gauge(gauge)
@@ -187,6 +192,9 @@ public final class PrometheusCollectorRegistry: Sendable {
             return self.makeGauge(name: name)
         }
 
+        let name = name.ensureValidMetricName()
+        let labels = labels.ensureValidLabelNames()
+
         return self.box.withLockedValue { store -> Gauge in
             guard let value = store[name] else {
                 let labelNames = labels.allLabelNames
@@ -236,7 +244,8 @@ public final class PrometheusCollectorRegistry: Sendable {
     /// - Parameter buckets: Define the buckets that shall be used within the ``DurationHistogram``
     /// - Returns: A ``DurationHistogram`` that is registered with this ``PrometheusCollectorRegistry``
     public func makeDurationHistogram(name: String, buckets: [Duration]) -> DurationHistogram {
-        self.box.withLockedValue { store -> DurationHistogram in
+        let name = name.ensureValidMetricName()
+        return self.box.withLockedValue { store -> DurationHistogram in
             guard let value = store[name] else {
                 let gauge = DurationHistogram(name: name, labels: [], buckets: buckets)
                 store[name] = .durationHistogram(gauge)
@@ -274,6 +283,9 @@ public final class PrometheusCollectorRegistry: Sendable {
             return self.makeDurationHistogram(name: name, buckets: buckets)
         }
 
+        let name = name.ensureValidMetricName()
+        let labels = labels.ensureValidLabelNames()
+
         return self.box.withLockedValue { store -> DurationHistogram in
             guard let value = store[name] else {
                 let labelNames = labels.allLabelNames
@@ -335,7 +347,8 @@ public final class PrometheusCollectorRegistry: Sendable {
     /// - Parameter buckets: Define the buckets that shall be used within the ``ValueHistogram``
     /// - Returns: A ``ValueHistogram`` that is registered with this ``PrometheusCollectorRegistry``
     public func makeValueHistogram(name: String, buckets: [Double]) -> ValueHistogram {
-        self.box.withLockedValue { store -> ValueHistogram in
+        let name = name.ensureValidMetricName()
+        return self.box.withLockedValue { store -> ValueHistogram in
             guard let value = store[name] else {
                 let gauge = ValueHistogram(name: name, labels: [], buckets: buckets)
                 store[name] = .valueHistogram(gauge)
@@ -364,6 +377,9 @@ public final class PrometheusCollectorRegistry: Sendable {
             return self.makeValueHistogram(name: name, buckets: buckets)
         }
 
+        let name = name.ensureValidMetricName()
+        let labels = labels.ensureValidLabelNames()
+
         return self.box.withLockedValue { store -> ValueHistogram in
             guard let value = store[name] else {
                 let labelNames = labels.allLabelNames
@@ -560,6 +576,14 @@ extension [(String, String)] {
         result = result.sorted()
         return result
     }
+
+    fileprivate func ensureValidLabelNames() -> [(String, String)] {
+        if self.allSatisfy({ $0.0.isValidLabelName() }) {
+            return self
+        } else {
+            return self.map { ($0.ensureValidLabelName(), $1) }
+        }
+    }
 }
 
 extension [UInt8] {
@@ -595,3 +619,91 @@ extension PrometheusMetric {
         return prerendered
     }
 }
+
+extension String {
+    fileprivate func isValidMetricName() -> Bool {
+        var isFirstCharacter = true
+        for ascii in self.utf8 {
+            defer { isFirstCharacter = false }
+            switch ascii {
+            case UInt8(ascii: "A")...UInt8(ascii: "Z"),
+                 UInt8(ascii: "a")...UInt8(ascii: "z"),
+                 UInt8(ascii: "_"), UInt8(ascii: ":"):
+                continue
+            case UInt8(ascii: "0"), UInt8(ascii: "9"):
+                if isFirstCharacter {
+                    return false
+                }
+                continue
+            default:
+                return false
+            }
+        }
+        return true
+    }
+
+    fileprivate func isValidLabelName() -> Bool {
+        var isFirstCharacter = true
+        for ascii in self.utf8 {
+            defer { isFirstCharacter = false }
+            switch ascii {
+            case UInt8(ascii: "A")...UInt8(ascii: "Z"),
+                 UInt8(ascii: "a")...UInt8(ascii: "z"),
+                 UInt8(ascii: "_"):
+                continue
+            case UInt8(ascii: "0"), UInt8(ascii: "9"):
+                if isFirstCharacter {
+                    return false
+                }
+                continue
+            default:
+                return false
+            }
+        }
+        return true
+    }
+
+    fileprivate func ensureValidMetricName() -> String {
+        if self.isValidMetricName() {
+            return self
+        } else {
+            var new = self
+            new.fixPrometheusName(allowColon: true)
+            return new
+        }
+    }
+
+    fileprivate func ensureValidLabelName() -> String {
+        if self.isValidLabelName() {
+            return self
+        } else {
+            var new = self
+            new.fixPrometheusName(allowColon: false)
+            return new
+        }
+    }
+
+    fileprivate mutating func fixPrometheusName(allowColon: Bool) {
+        var startIndex = self.startIndex
+        var isFirstCharacter = true
+        while let fixIndex = self[startIndex...].firstIndex(where: { character in
+            defer { isFirstCharacter = false }
+            switch character {
+            case "A"..."Z", "a"..."z", "_":
+                return false
+            case ":":
+                return !allowColon
+            case "0"..."9":
+                return isFirstCharacter
+            default:
+                return true
+            }
+        }) {
+            self.replaceSubrange(fixIndex...fixIndex, with: CollectionOfOne("_"))
+            startIndex = fixIndex
+            if startIndex == self.endIndex {
+                break
+            }
+        }
+    }
+}

Tests/PrometheusTests/ValidNamesTests.swift

@@ -0,0 +1,95 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftPrometheus open source project
+//
+// Copyright (c) 2024 SwiftPrometheus project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftPrometheus project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import Prometheus
+import XCTest
+
+final class ValidNamesTests: XCTestCase {
+    func testCounterWithEmoji() {
+        let client = PrometheusCollectorRegistry()
+        let counter = client.makeCounter(name: "coffee☕️", labels: [])
+        counter.increment()
+
+        var buffer = [UInt8]()
+        client.emit(into: &buffer)
+        XCTAssertEqual(
+            String(decoding: buffer, as: Unicode.UTF8.self),
+            """
+            # TYPE coffee_ counter
+            coffee_ 1
+
+            """
+        )
+    }
+
+    func testIllegalMetricNames() async throws {
+        let registry = PrometheusCollectorRegistry()
+
+        /// Notably, newlines must not allow creating whole new metric root
+        let tests = [
+            "name",
+            """
+            name{bad="haha"} 121212121
+            bad_bad 12321323
+            """
+        ]
+
+        for test in tests {
+            registry.makeCounter(
+                name: test,
+                labels: []
+            ).increment()
+        }
+
+        var buffer = [UInt8]()
+        registry.emit(into: &buffer)
+        XCTAssertEqual(
+            String(decoding: buffer, as: Unicode.UTF8.self).split(separator: "\n").sorted().joined(separator: "\n"),
+            """
+            # TYPE name counter
+            # TYPE name_bad__haha___121212121_bad_bad_12321323 counter
+            name 1
+            name_bad__haha___121212121_bad_bad_12321323 1
+            """
+        )
+    }
+
+    func testIllegalLabelNames() async throws {
+        let registry = PrometheusCollectorRegistry()
+
+        let tests = [
+            """
+            name{bad="haha"} 121212121
+            bad_bad 12321323
+            """
+        ]
+
+        for test in tests {
+            registry.makeCounter(
+                name: "metric",
+                labels: [(test, "value")]
+            ).increment()
+        }
+
+        var buffer = [UInt8]()
+        registry.emit(into: &buffer)
+        XCTAssertEqual(
+            String(decoding: buffer, as: Unicode.UTF8.self).split(separator: "\n").sorted().joined(separator: "\n"),
+            """
+            # TYPE metric counter
+            metric{name_bad__haha___121212121_bad_bad_12321323="value"} 1
+            """
+        )
+    }
+}