Skip to main content

netget/server/prometheus_remote_write/
codec.rs

1//! Published Remote Write 1.0 wire facts; native bounded protobuf and Snappy block.
2use anyhow::{bail, ensure, Context, Result};
3use serde::{Deserialize, Serialize};
4use std::collections::{BTreeMap, BTreeSet};
5
6pub const DEFAULT_LLM_FALLBACK: bool = false;
7pub const DEFAULT_RETRY_429: bool = false;
8pub const DEFAULT_PATH: &str = "/api/v1/write";
9pub const MAX_BODY_BYTES: usize = 256 * 1024;
10pub const MAX_SERIES: usize = 128;
11pub const MAX_SAMPLES: usize = 2048;
12pub const MAX_LABELS: usize = 32;
13pub const MAX_NAME_BYTES: usize = 128;
14pub const MAX_VALUE_BYTES: usize = 2048;
15pub const MAX_PROTO_FIELDS: usize = 32768;
16pub const STALE_BITS: u64 = 0x7ff0000000000002;
17
18#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq)]
19#[serde(untagged)]
20pub enum SampleValue {
21    Number(f64),
22    Special(SpecialValue),
23}
24#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq)]
25pub enum SpecialValue {
26    #[serde(rename = "nan")]
27    Nan,
28    #[serde(rename = "+inf")]
29    PositiveInfinity,
30    #[serde(rename = "-inf")]
31    NegativeInfinity,
32    #[serde(rename = "stale")]
33    Stale,
34}
35impl SampleValue {
36    pub fn bits(self) -> Result<u64> {
37        Ok(match self {
38            Self::Number(n) => {
39                ensure!(n.is_finite(), "use typed special sample values");
40                n.to_bits()
41            }
42            Self::Special(SpecialValue::Nan) => 0x7ff8000000000000,
43            Self::Special(SpecialValue::PositiveInfinity) => f64::INFINITY.to_bits(),
44            Self::Special(SpecialValue::NegativeInfinity) => f64::NEG_INFINITY.to_bits(),
45            Self::Special(SpecialValue::Stale) => STALE_BITS,
46        })
47    }
48    pub fn from_bits(bits: u64) -> Self {
49        let n = f64::from_bits(bits);
50        if bits == STALE_BITS {
51            Self::Special(SpecialValue::Stale)
52        } else if n.is_nan() {
53            Self::Special(SpecialValue::Nan)
54        } else if n == f64::INFINITY {
55            Self::Special(SpecialValue::PositiveInfinity)
56        } else if n == f64::NEG_INFINITY {
57            Self::Special(SpecialValue::NegativeInfinity)
58        } else {
59            Self::Number(n)
60        }
61    }
62}
63#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
64#[serde(deny_unknown_fields)]
65pub struct Sample {
66    pub timestamp_ms: i64,
67    pub value: SampleValue,
68}
69#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
70#[serde(deny_unknown_fields)]
71pub struct Series {
72    pub labels: BTreeMap<String, String>,
73    pub samples: Vec<Sample>,
74}
75#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
76#[serde(deny_unknown_fields)]
77pub struct WriteBatch {
78    pub series: Vec<Series>,
79}
80#[derive(Clone, Debug, Serialize)]
81pub struct DecodedBatch {
82    pub series: Vec<Series>,
83    pub ignored_fields: usize,
84}
85
86pub fn validate_token(token: &str) -> Result<()> {
87    ensure!(
88        !token.is_empty()
89            && token.len() <= 1024
90            && token.bytes().all(|b| (0x21..=0x7e).contains(&b)),
91        "token printable ASCII byte limit"
92    );
93    Ok(())
94}
95pub fn validate_path(path: &str) -> Result<()> {
96    ensure!(
97        path.starts_with('/')
98            && !path.starts_with("//")
99            && path.len() <= 1024
100            && path.bytes().all(|b| (0x21..=0x7e).contains(&b))
101            && !path.contains(['?', '#', '\\']),
102        "absolute HTTP path byte/character limit"
103    );
104    let uri: hyper::Uri = path.parse().context("invalid write path")?;
105    ensure!(
106        uri.scheme().is_none() && uri.authority().is_none() && uri.query().is_none(),
107        "write path only"
108    );
109    Ok(())
110}
111fn identifier(name: &str, metric: bool) -> Result<()> {
112    let mut bytes = name.bytes();
113    ensure!(
114        !name.is_empty()
115            && name.len() <= MAX_NAME_BYTES
116            && bytes
117                .next()
118                .is_some_and(|b| b.is_ascii_alphabetic() || b == b'_' || (metric && b == b':'))
119            && bytes.all(|b| b.is_ascii_alphanumeric() || b == b'_' || (metric && b == b':')),
120        "legacy v1 metric/label name character/byte limit"
121    );
122    Ok(())
123}
124pub fn validate_batch(batch: &WriteBatch) -> Result<()> {
125    ensure!(batch.series.len() <= MAX_SERIES, "series count limit");
126    let mut samples = 0usize;
127    let mut seen = BTreeSet::new();
128    for series in &batch.series {
129        ensure!(
130            !series.labels.is_empty() && series.labels.len() <= MAX_LABELS,
131            "label count limit"
132        );
133        ensure!(seen.insert(&series.labels), "duplicate series label set");
134        for (name, value) in &series.labels {
135            identifier(name, false)?;
136            ensure!(
137                !value.is_empty() && value.len() <= MAX_VALUE_BYTES,
138                "label value empty/byte limit"
139            );
140            if name == "__name__" {
141                identifier(value, true)?;
142            }
143        }
144        ensure!(!series.samples.is_empty(), "series requires float samples");
145        samples = samples
146            .checked_add(series.samples.len())
147            .context("sample count overflow")?;
148        ensure!(samples <= MAX_SAMPLES, "sample count limit");
149        let mut previous = None;
150        for sample in &series.samples {
151            sample.value.bits()?;
152            ensure!(
153                previous.is_none_or(|p| p <= sample.timestamp_ms),
154                "samples must be timestamp ordered within a series"
155            );
156            previous = Some(sample.timestamp_ms);
157        }
158    }
159    Ok(())
160}
161fn varint(out: &mut Vec<u8>, mut n: u64) {
162    while n >= 128 {
163        out.push(n as u8 | 128);
164        n >>= 7;
165    }
166    out.push(n as u8);
167}
168fn blob(out: &mut Vec<u8>, field: u64, value: &[u8]) {
169    varint(out, field << 3 | 2);
170    varint(out, value.len() as u64);
171    out.extend_from_slice(value);
172}
173pub fn encode_proto(batch: &WriteBatch) -> Result<Vec<u8>> {
174    validate_batch(batch)?;
175    let mut out = Vec::new();
176    for series in &batch.series {
177        let mut ts = Vec::new();
178        for (name, value) in &series.labels {
179            let mut label = Vec::new();
180            blob(&mut label, 1, name.as_bytes());
181            blob(&mut label, 2, value.as_bytes());
182            blob(&mut ts, 1, &label);
183        }
184        for sample in &series.samples {
185            let mut scalar = vec![9];
186            scalar.extend_from_slice(&sample.value.bits()?.to_le_bytes());
187            // int64 (not sint64): negative timestamps use a ten-byte varint.
188            scalar.push(16);
189            varint(&mut scalar, sample.timestamp_ms as u64);
190            blob(&mut ts, 2, &scalar);
191        }
192        blob(&mut out, 1, &ts);
193        ensure!(out.len() <= MAX_BODY_BYTES, "uncompressed body byte limit");
194    }
195    Ok(out)
196}
197pub fn encode_batch(batch: &WriteBatch) -> Result<Vec<u8>> {
198    let body = snappy_encode(&encode_proto(batch)?);
199    ensure!(body.len() <= MAX_BODY_BYTES, "wire body byte limit");
200    Ok(body)
201}
202struct Reader<'a> {
203    bytes: &'a [u8],
204    at: usize,
205}
206impl<'a> Reader<'a> {
207    fn new(bytes: &'a [u8]) -> Self {
208        Self { bytes, at: 0 }
209    }
210    fn take(&mut self, n: usize) -> Result<&'a [u8]> {
211        let end = self.at.checked_add(n).context("length overflow")?;
212        let value = self.bytes.get(self.at..end).context("truncated field")?;
213        self.at = end;
214        Ok(value)
215    }
216    fn number(&mut self) -> Result<u64> {
217        let mut n = 0;
218        for shift in (0..70).step_by(7) {
219            let b = self.take(1)?[0];
220            ensure!(shift != 63 || b <= 1, "varint overflow");
221            n |= u64::from(b & 127) << shift;
222            if b & 128 == 0 {
223                return Ok(n);
224            }
225        }
226        bail!("varint overflow")
227    }
228    fn field(&mut self, budget: &mut usize) -> Result<Option<(u64, u8)>> {
229        if self.at == self.bytes.len() {
230            return Ok(None);
231        }
232        *budget = budget
233            .checked_sub(1)
234            .context("protobuf field count limit")?;
235        let key = self.number()?;
236        ensure!(
237            key >> 3 > 0 && key >> 3 <= 0x1fffffff,
238            "invalid protobuf field number"
239        );
240        Ok(Some((key >> 3, (key & 7) as u8)))
241    }
242    fn blob(&mut self, wire: u8) -> Result<&'a [u8]> {
243        ensure!(wire == 2, "expected length-delimited field");
244        let n = usize::try_from(self.number()?)?;
245        self.take(n)
246    }
247    fn scalar(&mut self, wire: u8) -> Result<u64> {
248        ensure!(wire == 0, "expected integer field");
249        self.number()
250    }
251    fn skip(&mut self, wire: u8) -> Result<()> {
252        match wire {
253            0 => {
254                self.number()?;
255            }
256            1 => {
257                self.take(8)?;
258            }
259            2 => {
260                self.blob(wire)?;
261            }
262            5 => {
263                self.take(4)?;
264            }
265            _ => bail!("protobuf groups/invalid wire types unsupported"),
266        }
267        Ok(())
268    }
269}
270fn label(bytes: &[u8], budget: &mut usize, ignored: &mut usize) -> Result<(String, String)> {
271    let mut r = Reader::new(bytes);
272    let (mut name, mut value) = (None, None);
273    while let Some((field, wire)) = r.field(budget)? {
274        match field {
275            1 => {
276                ensure!(name.is_none(), "duplicate label name field");
277                name = Some(std::str::from_utf8(r.blob(wire)?)?.to_owned());
278            }
279            2 => {
280                ensure!(value.is_none(), "duplicate label value field");
281                value = Some(std::str::from_utf8(r.blob(wire)?)?.to_owned());
282            }
283            _ => {
284                r.skip(wire)?;
285                *ignored += 1;
286            }
287        }
288    }
289    Ok((name.unwrap_or_default(), value.unwrap_or_default()))
290}
291fn sample(bytes: &[u8], budget: &mut usize, ignored: &mut usize) -> Result<Sample> {
292    let mut r = Reader::new(bytes);
293    let (mut value, mut timestamp) = (None, None);
294    while let Some((field, wire)) = r.field(budget)? {
295        match field {
296            1 => {
297                ensure!(
298                    value.is_none() && wire == 1,
299                    "duplicate/invalid sample double"
300                );
301                value = Some(u64::from_le_bytes(r.take(8)?.try_into()?));
302            }
303            2 => {
304                ensure!(timestamp.is_none(), "duplicate sample timestamp");
305                timestamp = Some(r.scalar(wire)? as i64);
306            }
307            _ => {
308                r.skip(wire)?;
309                *ignored += 1;
310            }
311        }
312    }
313    Ok(Sample {
314        timestamp_ms: timestamp.unwrap_or(0),
315        value: SampleValue::from_bits(value.unwrap_or(0)),
316    })
317}
318fn series(
319    bytes: &[u8],
320    budget: &mut usize,
321    ignored: &mut usize,
322    total_samples: &mut usize,
323) -> Result<Series> {
324    let mut r = Reader::new(bytes);
325    let mut labels = BTreeMap::new();
326    let mut samples = vec![];
327    let mut previous = None;
328    while let Some((field, wire)) = r.field(budget)? {
329        match field {
330            1 => {
331                ensure!(labels.len() < MAX_LABELS, "label count limit");
332                let (name, value) = label(r.blob(wire)?, budget, ignored)?;
333                ensure!(
334                    previous.as_ref().is_none_or(|p| p < &name),
335                    "labels must be sorted and unique"
336                );
337                previous = Some(name.clone());
338                labels.insert(name, value);
339            }
340            2 => {
341                ensure!(*total_samples < MAX_SAMPLES, "sample count limit");
342                samples.push(sample(r.blob(wire)?, budget, ignored)?);
343                *total_samples += 1;
344            }
345            3 | 4 => bail!("v1 exemplars/native histograms outside selected float-sample scope"),
346            _ => {
347                r.skip(wire)?;
348                *ignored += 1;
349            }
350        }
351    }
352    Ok(Series { labels, samples })
353}
354pub fn decode_proto(bytes: &[u8]) -> Result<DecodedBatch> {
355    ensure!(
356        bytes.len() <= MAX_BODY_BYTES,
357        "uncompressed body byte limit"
358    );
359    let mut r = Reader::new(bytes);
360    let mut budget = MAX_PROTO_FIELDS;
361    let (mut ignored, mut total_samples) = (0, 0);
362    let mut output = vec![];
363    while let Some((field, wire)) = r.field(&mut budget)? {
364        if field == 1 {
365            ensure!(output.len() < MAX_SERIES, "series count limit");
366            output.push(series(
367                r.blob(wire)?,
368                &mut budget,
369                &mut ignored,
370                &mut total_samples,
371            )?);
372        } else {
373            // Reserved v1 source/metadata and future optional fields are discarded,
374            // not exposed as bytes or interpreted as supported extension data.
375            r.skip(wire)?;
376            ignored += 1;
377        }
378    }
379    validate_batch(&WriteBatch {
380        series: output.clone(),
381    })?;
382    Ok(DecodedBatch {
383        series: output,
384        ignored_fields: ignored,
385    })
386}
387pub fn decode_batch(bytes: &[u8]) -> Result<DecodedBatch> {
388    ensure!(bytes.len() <= MAX_BODY_BYTES, "wire body byte limit");
389    decode_proto(&snappy_decode(bytes)?)
390}
391
392// Literal-only encoding is valid Snappy; the decoder accepts all three copy forms.
393pub fn snappy_encode(bytes: &[u8]) -> Vec<u8> {
394    let mut out = Vec::new();
395    varint(&mut out, bytes.len() as u64);
396    if !bytes.is_empty() {
397        let n = bytes.len() - 1;
398        if n < 60 {
399            out.push((n as u8) << 2);
400        } else {
401            let width = ((usize::BITS - n.leading_zeros()) as usize).div_ceil(8);
402            out.push(((59 + width) as u8) << 2);
403            out.extend_from_slice(&(n as u32).to_le_bytes()[..width]);
404        }
405        out.extend_from_slice(bytes);
406    }
407    out
408}
409pub fn snappy_decode(bytes: &[u8]) -> Result<Vec<u8>> {
410    let mut r = Reader::new(bytes);
411    let n = usize::try_from(r.number()?)?;
412    ensure!(r.at <= 5, "Snappy length preamble overflow");
413    ensure!(n <= MAX_BODY_BYTES, "Snappy decompressed byte limit");
414    let mut out = Vec::with_capacity(n);
415    while r.at < bytes.len() {
416        let tag = r.take(1)?[0];
417        let (kind, len, offset) = match tag & 3 {
418            0 => {
419                let v = usize::from(tag >> 2);
420                let len = if v < 60 {
421                    v + 1
422                } else {
423                    let width = v - 59;
424                    let mut raw = [0; 4];
425                    raw[..width].copy_from_slice(r.take(width)?);
426                    usize::try_from(u32::from_le_bytes(raw))?
427                        .checked_add(1)
428                        .context("literal length overflow")?
429                };
430                (0, len, 0)
431            }
432            1 => (
433                1,
434                usize::from((tag >> 2) & 7) + 4,
435                (usize::from(tag & 0xe0) << 3) | usize::from(r.take(1)?[0]),
436            ),
437            2 => {
438                let b = r.take(2)?;
439                (
440                    2,
441                    usize::from(tag >> 2) + 1,
442                    usize::from(u16::from_le_bytes([b[0], b[1]])),
443                )
444            }
445            _ => {
446                let b = r.take(4)?;
447                (
448                    3,
449                    usize::from(tag >> 2) + 1,
450                    usize::try_from(u32::from_le_bytes([b[0], b[1], b[2], b[3]]))?,
451                )
452            }
453        };
454        ensure!(
455            len <= n.saturating_sub(out.len()),
456            "Snappy output length mismatch/limit"
457        );
458        if kind == 0 {
459            out.extend_from_slice(r.take(len)?);
460        } else {
461            ensure!(
462                offset > 0 && offset <= out.len(),
463                "invalid Snappy backreference"
464            );
465            for _ in 0..len {
466                out.push(out[out.len() - offset]);
467            }
468        }
469    }
470    ensure!(out.len() == n, "Snappy truncated output");
471    Ok(out)
472}