From bfcd4bbfabe11aae4b035424ee9724582e288501 Mon Sep 17 00:00:00 2001 From: Fabian Fett Date: Fri, 29 Mar 2024 08:02:23 +0100 Subject: [PATCH] Merge pull request from GHSA-x768-cvr2-345r --- .../PrometheusCollectorRegistry.swift | 120 +++++++++++++++++- Tests/PrometheusTests/ValidNamesTests.swift | 95 ++++++++++++++ 2 files changed, 211 insertions(+), 4 deletions(-) create mode 100644 Tests/PrometheusTests/ValidNamesTests.swift diff --git a/Sources/Prometheus/PrometheusCollectorRegistry.swift b/Sources/Prometheus/PrometheusCollectorRegistry.swift index caaef10..1284112 100644 --- a/Sources/Prometheus/PrometheusCollectorRegistry.swift +++ b/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 + } + } + } +} diff --git a/Tests/PrometheusTests/ValidNamesTests.swift b/Tests/PrometheusTests/ValidNamesTests.swift new file mode 100644 index 0000000..8b1e172 --- /dev/null +++ b/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 + """ + ) + } +}