servers/prom_remote_write/
validation.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Prometheus remote write validation utilities.
16
17use prost::DecodeError;
18use serde::{Deserialize, Serialize};
19
20const IS_VALID_LABEL_REST: [bool; 256] = {
21    let mut table = [false; 256];
22    let mut i = 0;
23    while i < 256 {
24        let b = i as u8;
25        table[i] = b.is_ascii_alphanumeric() || b == b'_';
26        i += 1;
27    }
28    table
29};
30
31/// Validates that the given bytes form a valid Prometheus label name.
32#[inline]
33pub fn validate_label_name(name: &[u8]) -> bool {
34    if name.is_empty() {
35        return false;
36    }
37    let first = name[0];
38    if !(first.is_ascii_alphabetic() || first == b'_') {
39        return false;
40    }
41
42    let mut rest = &name[1..];
43    while rest.len() >= 8 {
44        let res = IS_VALID_LABEL_REST[rest[0] as usize] as u8
45            & IS_VALID_LABEL_REST[rest[1] as usize] as u8
46            & IS_VALID_LABEL_REST[rest[2] as usize] as u8
47            & IS_VALID_LABEL_REST[rest[3] as usize] as u8
48            & IS_VALID_LABEL_REST[rest[4] as usize] as u8
49            & IS_VALID_LABEL_REST[rest[5] as usize] as u8
50            & IS_VALID_LABEL_REST[rest[6] as usize] as u8
51            & IS_VALID_LABEL_REST[rest[7] as usize] as u8;
52
53        if res == 0 {
54            return false;
55        }
56        rest = &rest[8..];
57    }
58
59    for &b in rest {
60        if !IS_VALID_LABEL_REST[b as usize] {
61            return false;
62        }
63    }
64
65    true
66}
67
68/// Validation mode for decoding Prometheus remote write requests.
69#[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
70#[serde(rename_all = "snake_case")]
71#[derive(Default)]
72pub enum PromValidationMode {
73    #[default]
74    Strict,
75    Lossy,
76    Unchecked,
77}
78
79impl PromValidationMode {
80    pub fn decode_string(&self, bytes: &[u8]) -> Result<String, DecodeError> {
81        let result = match self {
82            PromValidationMode::Strict => match String::from_utf8(bytes.to_vec()) {
83                Ok(s) => s,
84                Err(e) => {
85                    common_telemetry::debug!(
86                        "Invalid UTF-8 string value: {:?}, error: {:?}",
87                        bytes,
88                        e
89                    );
90                    return Err(DecodeError::new("invalid utf-8"));
91                }
92            },
93            PromValidationMode::Lossy => String::from_utf8_lossy(bytes).to_string(),
94            PromValidationMode::Unchecked => unsafe { String::from_utf8_unchecked(bytes.to_vec()) },
95        };
96        Ok(result)
97    }
98
99    pub fn decode_label_name<'a>(&self, bytes: &'a [u8]) -> Result<&'a str, DecodeError> {
100        if !validate_label_name(bytes) {
101            common_telemetry::debug!(
102                "Invalid Prometheus label name: {:?}, must match [a-zA-Z_][a-zA-Z0-9_]*",
103                bytes
104            );
105            return Err(DecodeError::new(format!(
106                "invalid prometheus label name: '{}', must match [a-zA-Z_][a-zA-Z0-9_]*",
107                String::from_utf8_lossy(bytes)
108            )));
109        }
110        Ok(unsafe { std::str::from_utf8_unchecked(bytes) })
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[test]
119    fn test_validate_label_name() {
120        assert!(validate_label_name(b"__name__"));
121        assert!(validate_label_name(b"job"));
122        assert!(validate_label_name(b"_"));
123        assert!(validate_label_name(b"A"));
124        assert!(validate_label_name(b"abc123"));
125        assert!(!validate_label_name(b""));
126        assert!(!validate_label_name(b"0starts_with_digit"));
127        assert!(!validate_label_name(b"has-dash"));
128    }
129}
servers/prom_remote_write/validation.rs

servers/prom_remote_write/
validation.rs
