netget/server/prometheus_remote_write/
mod.rs

1//! Published Remote Write 1.0 float-sample HTTP collector with owned tasks.
2pub mod actions;
3pub mod codec;
4use crate::protocol::{Event, SpawnContext};
5use crate::server::{
6    accept_bounded::{accept_bounded, ConnectionLimiter, DEFAULT_MAX_CONNECTIONS},
7    connection::ConnectionId,
8};
9use crate::state::{
10    server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo},
11    AccessLogOwner,
12};
13use crate::{console_error, console_info};
14use anyhow::Result;
15use bytes::Bytes;
16use http_body_util::{BodyExt, Full, Limited};
17use hyper::{
18    body::Incoming,
19    header::{HeaderValue, ALLOW, AUTHORIZATION, CONTENT_TYPE, RETRY_AFTER},
20    server::conn::http1,
21    service::service_fn,
22    Method, Request, Response, StatusCode,
23};
24use hyper_util::rt::{TokioIo, TokioTimer};
25use serde_json::json;
26use std::{convert::Infallible, net::SocketAddr, sync::Arc, time::Duration};
27use tokio::sync::Notify;
28
29pub const HEADER_TIMEOUT: Duration = Duration::from_secs(30);
30pub const BODY_TIMEOUT: Duration = Duration::from_secs(30);
31pub const WRITE_TIMEOUT: Duration = Duration::from_secs(10);
32pub const MAX_HEADER_BYTES: usize = 32 * 1024;
33pub const MAX_HEADERS: usize = 64;
34const CAP_REFUSAL:&[u8]=b"HTTP/1.1 503 Service Unavailable\r\nContent-Type: text/plain; charset=utf-8\r\nContent-Length: 19\r\nConnection: close\r\nRetry-After: 5\r\n\r\nconnection capacity";
35#[derive(Clone)]
36struct Config {
37    token: Option<String>,
38    llm_fallback: bool,
39    path: String,
40}
41pub struct PrometheusRemoteWriteServer;
42impl PrometheusRemoteWriteServer {
43    pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
44        let token = ctx
45            .startup_params
46            .as_ref()
47            .map(|p| p.get_optional_string("auth_token"))
48            .transpose()?
49            .flatten();
50        if let Some(token) = &token {
51            codec::validate_token(token)?;
52        }
53        let llm_fallback = ctx
54            .startup_params
55            .as_ref()
56            .map(|p| p.get_optional_bool("llm_fallback"))
57            .transpose()?
58            .flatten()
59            .unwrap_or(codec::DEFAULT_LLM_FALLBACK);
60        let path = ctx
61            .startup_params
62            .as_ref()
63            .map(|p| p.get_optional_string("path"))
64            .transpose()?
65            .flatten()
66            .unwrap_or_else(|| codec::DEFAULT_PATH.into());
67        codec::validate_path(&path)?;
68        let config = Arc::new(Config {
69            token,
70            llm_fallback,
71            path,
72        });
73        let listener = tokio::net::TcpListener::bind(ctx.legacy_listen_addr()).await?;
74        let local = listener.local_addr()?;
75        console_info!(
76            ctx.status_tx,
77            "PrometheusRemoteWrite push listening on {} (auth_required={},llm_fallback={})",
78            local,
79            config.token.is_some(),
80            llm_fallback
81        );
82        let owner = ctx.state.clone();
83        let server_id = ctx.server_id;
84        owner.spawn_server_task(server_id, async move {
85            let limiter = ConnectionLimiter::new(DEFAULT_MAX_CONNECTIONS);
86            loop {
87                let (stream, peer, permit) = match accept_bounded(&listener, &limiter, CAP_REFUSAL, "PrometheusRemoteWrite", Some(&ctx.status_tx)).await {
88                    Ok(value) => value,
89                    Err(error) => {
90                        console_error!(ctx.status_tx, "PrometheusRemoteWrite decision=fail_closed_accept_error error={}", error);
91                        break;
92                    }
93                };
94                let id = ConnectionId::new(ctx.state.get_next_unified_id().await);
95                let now = crate::utils::clock::Instant::now();
96                ctx.state.add_connection_to_server(server_id, ConnectionState {
97                    id, remote_addr: peer, local_addr: local, bytes_sent: 0,
98                    bytes_received: 0, packets_sent: 0, packets_received: 0,
99                    last_activity: now, status: ConnectionStatus::Active,
100                    status_changed_at: now, protocol_info: ProtocolConnectionInfo::empty(),
101                }).await;
102                let child_ctx = ctx.clone();
103                let config = config.clone();
104                ctx.state.spawn_server_task(server_id, async move {
105                    let _permit = permit;
106                    let write_ready = Arc::new(Notify::new());
107                    let service_ready = write_ready.clone();
108                    let request_ctx = child_ctx.clone();
109                    let service = service_fn(move |request| {
110                        let ctx = request_ctx.clone();
111                        let config = config.clone();
112                        let ready = service_ready.clone();
113                        async move {
114                            let response = handle_request(request, peer, id, &ctx, &config).await;
115                            ready.notify_one();
116                            Ok::<_, Infallible>(response)
117                        }
118                    });
119                    let mut builder = http1::Builder::new();
120                    builder.keep_alive(false).timer(TokioTimer::new())
121                        .header_read_timeout(HEADER_TIMEOUT).max_headers(MAX_HEADERS)
122                        .max_buf_size(MAX_HEADER_BYTES);
123                    let connection = builder.serve_connection(TokioIo::new(stream), service);
124                    tokio::pin!(connection);
125                    tokio::select! {
126                        result = &mut connection => {
127                            if let Err(error) = result {
128                                console_error!(child_ctx.status_tx,
129                                    "PrometheusRemoteWrite decision=fail_closed_http_framing error={}", error);
130                            }
131                        }
132                        _ = async {
133                            write_ready.notified().await;
134                            tokio::time::sleep(WRITE_TIMEOUT).await;
135                        } => {
136                            console_error!(child_ctx.status_tx,
137                                "PrometheusRemoteWrite decision=fail_closed_response_write_timeout");
138                        }
139                    }
140                    child_ctx.state.remove_connection_from_server(server_id, id).await;
141                    let _ = child_ctx.status_tx.send("__UPDATE_UI__".into());
142                }).await;
143            }
144        }).await;
145        Ok(local)
146    }
147}
148fn response(status: u16, message: &str, retry: Option<u16>) -> Response<Full<Bytes>> {
149    let bytes = if status == 204 {
150        Vec::new()
151    } else {
152        message.as_bytes().to_vec()
153    };
154    let mut r = Response::new(Full::new(Bytes::from(bytes)));
155    *r.status_mut() = StatusCode::from_u16(status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
156    r.headers_mut().insert(
157        CONTENT_TYPE,
158        HeaderValue::from_static("text/plain; charset=utf-8"),
159    );
160    if let Some(n) = retry {
161        if let Ok(v) = HeaderValue::from_str(&n.to_string()) {
162            r.headers_mut().insert(RETRY_AFTER, v);
163        }
164    }
165    r.headers_mut().insert(
166        "X-Prometheus-Remote-Write-Version",
167        HeaderValue::from_static("0.1.0"),
168    );
169    r
170}
171fn error(status: u16, message: &str) -> Response<Full<Bytes>> {
172    response(status, message, None)
173}
174fn authorized(headers: &hyper::HeaderMap, expected: Option<&str>) -> bool {
175    let Some(expected) = expected else {
176        return true;
177    };
178    let values = headers.get_all(AUTHORIZATION).iter().collect::<Vec<_>>();
179    if values.len() != 1 {
180        return false;
181    }
182    let Some((scheme, token)) = values[0].to_str().ok().and_then(|v| v.split_once(' ')) else {
183        return false;
184    };
185    let token = token.trim_start_matches(' ');
186    if !scheme.eq_ignore_ascii_case("Bearer") || codec::validate_token(token).is_err() {
187        return false;
188    }
189    let mut difference = expected.len() ^ token.len();
190    for at in 0..expected.len().max(token.len()) {
191        difference |= usize::from(
192            expected.as_bytes().get(at).copied().unwrap_or(0)
193                ^ token.as_bytes().get(at).copied().unwrap_or(0),
194        );
195    }
196    difference == 0
197}
198// MIME tokens and parameter names are case-insensitive. The protobuf message
199// identifier is case-sensitive; v2 must not be decoded as an empty v1 request.
200fn request_encoding(content_type: &str) -> bool {
201    let mut parts = content_type.split(';');
202    if !parts
203        .next()
204        .is_some_and(|v| v.trim().eq_ignore_ascii_case("application/x-protobuf"))
205    {
206        return false;
207    }
208    let mut proto_seen = false;
209    for parameter in parts {
210        let Some((name, value)) = parameter.trim().split_once('=') else {
211            return false;
212        };
213        if proto_seen || !name.trim().eq_ignore_ascii_case("proto") {
214            return false;
215        }
216        let value = value.trim();
217        let value = if value.starts_with('"') {
218            match value.strip_prefix('"').and_then(|v| v.strip_suffix('"')) {
219                Some(v) => v,
220                None => return false,
221            }
222        } else {
223            value
224        };
225        if value != "prometheus.WriteRequest" {
226            return false;
227        }
228        proto_seen = true;
229    }
230    true
231}
232async fn handle_request(
233    request: Request<Incoming>,
234    peer: SocketAddr,
235    id: ConnectionId,
236    ctx: &SpawnContext,
237    config: &Config,
238) -> Response<Full<Bytes>> {
239    let response = process_request(request, peer, id, ctx, config).await;
240    if response.status().is_success() {
241        console_info!(
242            ctx.status_tx,
243            "PrometheusRemoteWrite decision=accept status={}",
244            response.status().as_u16()
245        );
246    } else {
247        console_error!(
248            ctx.status_tx,
249            "PrometheusRemoteWrite decision=fail_closed_http_reject status={}",
250            response.status().as_u16()
251        );
252    }
253    response
254}
255async fn process_request(
256    request: Request<Incoming>,
257    peer: SocketAddr,
258    id: ConnectionId,
259    ctx: &SpawnContext,
260    config: &Config,
261) -> Response<Full<Bytes>> {
262    let (parts, body) = request.into_parts();
263    if parts.headers.len() > MAX_HEADERS
264        || parts
265            .headers
266            .iter()
267            .map(|(k, v)| k.as_str().len() + v.as_bytes().len() + 4)
268            .sum::<usize>()
269            > MAX_HEADER_BYTES
270    {
271        return error(413, "Request header byte/count limit");
272    }
273    let bytes = match tokio::time::timeout(
274        BODY_TIMEOUT,
275        Limited::new(body, codec::MAX_BODY_BYTES).collect(),
276    )
277    .await
278    {
279        Err(_) => return error(408, "Request body deadline exceeded"),
280        Ok(Err(error)) => {
281            return if error
282                .downcast_ref::<http_body_util::LengthLimitError>()
283                .is_some()
284            {
285                self::error(413, "Request body limit exceeded")
286            } else {
287                self::error(400, "Incomplete HTTP body")
288            };
289        }
290        Ok(Ok(body)) => body.to_bytes(),
291    };
292    ctx.state
293        .update_connection_stats(
294            ctx.server_id,
295            id,
296            Some(bytes.len() as u64),
297            None,
298            Some(1),
299            None,
300        )
301        .await;
302    if parts.uri.path() != config.path {
303        return error(404, "Only the configured write endpoint is implemented");
304    }
305    if parts.method != Method::POST {
306        let mut r = error(405, "Only POST is supported");
307        r.headers_mut()
308            .insert(ALLOW, HeaderValue::from_static("POST"));
309        return r;
310    }
311    if parts.uri.query().is_some() {
312        return error(400, "Write query parameters unsupported");
313    }
314    if !authorized(&parts.headers, config.token.as_deref()) {
315        return error(401, "Unauthorized bearer credential");
316    }
317    for name in [
318        "content-encoding",
319        "content-type",
320        "user-agent",
321        "x-prometheus-remote-write-version",
322    ] {
323        if parts.headers.get_all(name).iter().count() != 1 {
324            return error(400, "One required protocol header expected");
325        }
326    }
327    let text = |name| {
328        parts
329            .headers
330            .get(name)
331            .and_then(|v| v.to_str().ok())
332            .unwrap_or("")
333    };
334    if !text("content-encoding")
335        .trim()
336        .eq_ignore_ascii_case("snappy")
337        || !request_encoding(text("content-type"))
338    {
339        return error(
340            415,
341            "Expected Snappy block and prometheus.WriteRequest protobuf",
342        );
343    }
344    if text("x-prometheus-remote-write-version") != "0.1.0" {
345        return error(
346            415,
347            "Only published remote write 1.0 is supported (wire header0.1.0)",
348        );
349    }
350    if text("user-agent").is_empty() || text("user-agent").len() > 512 {
351        return error(400, "User-Agent required/byte limit");
352    }
353    let decoded = match codec::decode_batch(&bytes) {
354        Ok(v) => v,
355        Err(_) => return error(400, "Malformed, invalid or over-limit remote write samples"),
356    };
357    let event = Event::new(
358        &actions::REMOTE_WRITE_EVENT,
359        json!({
360            "series": decoded.series,
361            "ignored_fields": decoded.ignored_fields,
362            "authenticated": true,
363            "auth_required": config.token.is_some(),
364            "source_addr": peer.to_string(),
365            "version": "1.0",
366            "series_count": decoded.series.len(),
367            "sample_count": decoded.series.iter().map(|s|s.samples.len()).sum::<usize>(),
368            "durable_storage": false,
369        }),
370    );
371    let configured = ctx
372        .state
373        .get_event_handler_config(ctx.server_id)
374        .await
375        .is_some_and(|c| c.find_handler("remote_write_request").is_some());
376    let decision = if configured || config.llm_fallback {
377        match crate::llm::action_helper::call_llm(
378            &ctx.llm_client,
379            &ctx.state,
380            ctx.server_id,
381            Some(id),
382            &event,
383            &actions::PrometheusRemoteWriteProtocol::new(),
384        )
385        .await
386        {
387            Ok(result) if result.failures.is_empty() => {
388                let decisions = result
389                    .protocol_results
390                    .iter()
391                    .filter_map(|r| match r {
392                        crate::llm::ActionResult::Custom { name, data }
393                            if name == "remote_write_decision" =>
394                        {
395                            Some(data.clone())
396                        }
397                        _ => None,
398                    })
399                    .collect::<Vec<_>>();
400                if decisions.len() > 1 {
401                    return failed(ctx, id, event.data, "fail_closed_multiple_decisions").await;
402                }
403                match decisions
404                    .first()
405                    .map(|d| serde_json::from_value::<actions::Decision>(d.clone()))
406                    .transpose()
407                {
408                    Ok(Some(d)) => d,
409                    Ok(None) => actions::Decision::Accept,
410                    Err(_) => {
411                        return failed(ctx, id, event.data, "fail_closed_invalid_decision").await
412                    }
413                }
414            }
415            Ok(_) => return failed(ctx, id, event.data, "fail_closed_handler_action_error").await,
416            Err(_) => return failed(ctx, id, event.data, "fail_closed_handler_error").await,
417        }
418    } else {
419        ctx.state
420            .record_access_log(
421                AccessLogOwner::Server(ctx.server_id.as_u32()),
422                "PrometheusRemoteWrite",
423                Some(id.as_u32()),
424                "remote_write_request",
425                event.data.clone(),
426                vec![json!({"type":"accept_remote_write_samples"})],
427            )
428            .await;
429        actions::Decision::Accept
430    };
431    match decision {
432        actions::Decision::Accept => response(204, "", None),
433        actions::Decision::Reject {
434            status,
435            message,
436            retry_after_seconds,
437        } => response(status, &message, retry_after_seconds),
438    }
439}
440async fn failed(
441    ctx: &SpawnContext,
442    id: ConnectionId,
443    data: serde_json::Value,
444    decision: &str,
445) -> Response<Full<Bytes>> {
446    console_error!(ctx.status_tx, "PrometheusRemoteWrite decision={}", decision);
447    ctx.state
448        .record_access_log(
449            AccessLogOwner::Server(ctx.server_id.as_u32()),
450            "PrometheusRemoteWrite",
451            Some(id.as_u32()),
452            "remote_write_handler_failed",
453            data,
454            vec![json!({"decision":decision})],
455        )
456        .await;
457    error(503, "Write handler failed; no acceptance confirmed")
458}
netget/server/prometheus_remote_write/mod.rs

netget/server/prometheus_remote_write/
mod.rs
