netget/server/loki/
mod.rs

1//! One-request-per-connection HTTP/1.1 push 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_ENCODING, 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    require_tenant: bool,
40}
41pub struct LokiServer;
42impl LokiServer {
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 require_tenant = ctx
61            .startup_params
62            .as_ref()
63            .map(|p| p.get_optional_bool("require_tenant"))
64            .transpose()?
65            .flatten()
66            .unwrap_or(codec::DEFAULT_REQUIRE_TENANT);
67        let config = Arc::new(Config {
68            token,
69            llm_fallback,
70            require_tenant,
71        });
72        let listener = tokio::net::TcpListener::bind(ctx.legacy_listen_addr()).await?;
73        let local = listener.local_addr()?;
74        console_info!(
75            ctx.status_tx,
76            "Loki push listening on {} (auth_required={},llm_fallback={})",
77            local,
78            config.token.is_some(),
79            llm_fallback
80        );
81        let owner = ctx.state.clone();
82        let server_id = ctx.server_id;
83        owner.spawn_server_task(server_id,async move {
84            let limiter=ConnectionLimiter::new(DEFAULT_MAX_CONNECTIONS);
85            loop {
86                let (stream,peer,permit)=match accept_bounded(&listener,&limiter,CAP_REFUSAL,"Loki",Some(&ctx.status_tx)).await {
87                    Ok(value)=>value,
88                    Err(error)=>{console_error!(ctx.status_tx,"Loki accept failed: {}",error);break;}
89                };
90                let id=ConnectionId::new(ctx.state.get_next_unified_id().await);
91                let now=crate::utils::clock::Instant::now();
92                ctx.state.add_connection_to_server(server_id,ConnectionState {id,remote_addr:peer,local_addr:local,bytes_sent:0,bytes_received:0,packets_sent:0,packets_received:0,last_activity:now,status:ConnectionStatus::Active,status_changed_at:now,protocol_info:ProtocolConnectionInfo::empty()}).await;
93                let child_ctx=ctx.clone(); let config=config.clone();
94                ctx.state.spawn_server_task(server_id,async move {
95                    let _permit=permit;
96                    let write_ready=Arc::new(Notify::new());
97                    let service_ready=write_ready.clone();
98                    let request_ctx=child_ctx.clone();
99                    let service=service_fn(move |request| {
100                        let ctx=request_ctx.clone(); let config=config.clone(); let ready=service_ready.clone();
101                        async move {
102                            let response=handle_request(request,peer,id,&ctx,&config).await;
103                            ready.notify_one();
104                            Ok::<_,Infallible>(response)
105                        }
106                    });
107                    let mut builder=http1::Builder::new();
108                    builder.keep_alive(false).timer(TokioTimer::new()).header_read_timeout(HEADER_TIMEOUT).max_headers(MAX_HEADERS).max_buf_size(MAX_HEADER_BYTES);
109                    let connection=builder.serve_connection(TokioIo::new(stream),service);
110                    tokio::pin!(connection);
111                    tokio::select! {
112                        result=&mut connection=>{if let Err(error)=result {console_error!(child_ctx.status_tx,"Loki HTTP framing/header error: {}",error);}}
113                        _=async {write_ready.notified().await;tokio::time::sleep(WRITE_TIMEOUT).await;}=>{console_error!(child_ctx.status_tx,"Loki decision=fail_closed_response_write_timeout");}
114                    }
115                    child_ctx.state.remove_connection_from_server(server_id,id).await;
116                    let _=child_ctx.status_tx.send("__UPDATE_UI__".into());
117                }).await;
118            }
119        }).await;
120        Ok(local)
121    }
122}
123fn response(status: u16, message: &str, retry: Option<u16>) -> Response<Full<Bytes>> {
124    let bytes = if status == 204 {
125        Vec::new()
126    } else {
127        message.as_bytes().to_vec()
128    };
129    let mut r = Response::new(Full::new(Bytes::from(bytes)));
130    *r.status_mut() = StatusCode::from_u16(status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
131    r.headers_mut().insert(
132        CONTENT_TYPE,
133        HeaderValue::from_static("text/plain; charset=utf-8"),
134    );
135    if let Some(n) = retry {
136        if let Ok(v) = HeaderValue::from_str(&n.to_string()) {
137            r.headers_mut().insert(RETRY_AFTER, v);
138        }
139    }
140    r
141}
142fn error(status: u16, message: &str) -> Response<Full<Bytes>> {
143    response(status, message, None)
144}
145fn authorized(headers: &hyper::HeaderMap, expected: Option<&str>) -> bool {
146    let Some(expected) = expected else {
147        return true;
148    };
149    let values = headers.get_all(AUTHORIZATION).iter().collect::<Vec<_>>();
150    if values.len() != 1 {
151        return false;
152    }
153    let Some((scheme, token)) = values[0].to_str().ok().and_then(|v| v.split_once(' ')) else {
154        return false;
155    };
156    let token = token.trim_start_matches(' ');
157    if !scheme.eq_ignore_ascii_case("Bearer") || codec::validate_token(token).is_err() {
158        return false;
159    }
160    let mut difference = expected.len() ^ token.len();
161    for at in 0..expected.len().max(token.len()) {
162        difference |= usize::from(
163            expected.as_bytes().get(at).copied().unwrap_or(0)
164                ^ token.as_bytes().get(at).copied().unwrap_or(0),
165        );
166    }
167    difference == 0
168}
169// RFC9110: media types, parameter names, auth schemes and content codings
170// are case-insensitive tokens. JSON accepts an optional UTF-8 charset.
171fn request_encoding(content_type: &str, content_encoding: &str) -> Option<codec::Encoding> {
172    let mut parts = content_type.split(';');
173    let media = parts.next()?.trim();
174    let json = media.eq_ignore_ascii_case("application/json");
175    if !json && !media.eq_ignore_ascii_case("application/x-protobuf") {
176        return None;
177    }
178    let mut charset_seen = false;
179    for parameter in parts {
180        let parameter = parameter.trim();
181        if parameter.is_empty() {
182            continue;
183        }
184        let (name, value) = parameter.split_once('=')?;
185        if !json || charset_seen || !name.trim().eq_ignore_ascii_case("charset") {
186            return None;
187        }
188        let value = value.trim();
189        let value = if value.starts_with('"') {
190            value.strip_prefix('"')?.strip_suffix('"')?
191        } else {
192            value
193        };
194        if !value.eq_ignore_ascii_case("utf-8") {
195            return None;
196        }
197        charset_seen = true;
198    }
199    let coding = content_encoding.trim();
200    if json && coding.eq_ignore_ascii_case("identity") {
201        Some(codec::Encoding::Json)
202    } else if json && coding.eq_ignore_ascii_case("gzip") {
203        Some(codec::Encoding::GzipJson)
204    } else if !json
205        && (coding.eq_ignore_ascii_case("identity") || coding.eq_ignore_ascii_case("snappy"))
206    {
207        Some(codec::Encoding::SnappyProtobuf)
208    } else {
209        None
210    }
211}
212async fn handle_request(
213    request: Request<Incoming>,
214    peer: SocketAddr,
215    id: ConnectionId,
216    ctx: &SpawnContext,
217    config: &Config,
218) -> Response<Full<Bytes>> {
219    let (parts, body) = request.into_parts();
220    if parts.headers.len() > MAX_HEADERS
221        || parts
222            .headers
223            .iter()
224            .map(|(k, v)| k.as_str().len() + v.as_bytes().len() + 4)
225            .sum::<usize>()
226            > MAX_HEADER_BYTES
227    {
228        return error(413, "Request header byte/count limit");
229    }
230    let bytes = match tokio::time::timeout(
231        BODY_TIMEOUT,
232        Limited::new(body, codec::MAX_BODY_BYTES).collect(),
233    )
234    .await
235    {
236        Err(_) => return error(408, "Request body deadline exceeded"),
237        Ok(Err(error)) => {
238            return if error
239                .downcast_ref::<http_body_util::LengthLimitError>()
240                .is_some()
241            {
242                self::error(413, "Request body limit exceeded")
243            } else {
244                self::error(400, "Incomplete HTTP body")
245            };
246        }
247        Ok(Ok(body)) => body.to_bytes(),
248    };
249    ctx.state
250        .update_connection_stats(
251            ctx.server_id,
252            id,
253            Some(bytes.len() as u64),
254            None,
255            Some(1),
256            None,
257        )
258        .await;
259    if parts.uri.path() != "/loki/api/v1/push" {
260        return error(404, "Only the push endpoint is implemented");
261    }
262    if parts.method != Method::POST {
263        let mut r = error(405, "Only POST is supported");
264        r.headers_mut()
265            .insert(ALLOW, HeaderValue::from_static("POST"));
266        return r;
267    }
268    if parts.uri.query().is_some() {
269        return error(400, "Push query parameters unsupported");
270    }
271    if !authorized(&parts.headers, config.token.as_deref()) {
272        return error(401, "Unauthorized bearer credential");
273    }
274    let tenants = parts
275        .headers
276        .get_all("X-Scope-OrgID")
277        .iter()
278        .collect::<Vec<_>>();
279    if tenants.len() > 1 {
280        return error(400, "Duplicate tenant header");
281    }
282    let tenant = match tenants.first() {
283        Some(v) => match v.to_str() {
284            Ok(t) if codec::validate_tenant(t).is_ok() => Some(t.to_owned()),
285            _ => return error(400, "Invalid single tenant ID"),
286        },
287        None => None,
288    };
289    if config.require_tenant && tenant.is_none() {
290        return error(401, "Tenant ID required");
291    }
292    if parts.headers.get_all(CONTENT_ENCODING).iter().count() > 1
293        || parts.headers.get_all(CONTENT_TYPE).iter().count() > 1
294    {
295        return error(400, "Duplicate encoding/content type headers");
296    }
297    let content_type = parts
298        .headers
299        .get(CONTENT_TYPE)
300        .and_then(|v| v.to_str().ok())
301        .unwrap_or("");
302    let content_encoding = match parts
303        .headers
304        .get(CONTENT_ENCODING)
305        .map(|v| v.to_str())
306        .transpose()
307    {
308        Ok(Some(v)) => v,
309        Ok(None) => "identity",
310        Err(_) => return error(415, "Invalid content encoding"),
311    };
312    let encoding = match request_encoding(content_type, content_encoding) {
313        Some(encoding) => encoding,
314        None => return error(415, "Expected JSON, gzip JSON or Snappy protobuf"),
315    };
316    let streams = match codec::decode_batch(&bytes, encoding) {
317        Ok(v) => v,
318        Err(e) => {
319            return error(
320                if e.to_string().contains("limit") {
321                    413
322                } else {
323                    400
324                },
325                "Malformed/oversized Loki push",
326            )
327        }
328    };
329    let event = Event::new(
330        &actions::LOKI_PUSH_EVENT,
331        json!({"tenant_id":tenant.as_deref().unwrap_or("fake"),"tenant_provided":tenant.is_some(),"encoding":encoding,"streams":streams,"authenticated":true,"auth_required":config.token.is_some(),"source_addr":peer.to_string()}),
332    );
333    let configured = ctx
334        .state
335        .get_event_handler_config(ctx.server_id)
336        .await
337        .is_some_and(|c| c.find_handler("loki_push").is_some());
338    let decision = if configured || config.llm_fallback {
339        match crate::llm::action_helper::call_llm(
340            &ctx.llm_client,
341            &ctx.state,
342            ctx.server_id,
343            Some(id),
344            &event,
345            &actions::LokiProtocol::new(),
346        )
347        .await
348        {
349            Ok(result) if result.failures.is_empty() => {
350                let decisions = result
351                    .protocol_results
352                    .iter()
353                    .filter_map(|r| match r {
354                        crate::llm::ActionResult::Custom { name, data }
355                            if name == "loki_push_decision" =>
356                        {
357                            Some(data.clone())
358                        }
359                        _ => None,
360                    })
361                    .collect::<Vec<_>>();
362                if decisions.len() > 1 {
363                    return failed(ctx, id, event.data, "fail_closed_multiple_decisions").await;
364                }
365                match decisions
366                    .first()
367                    .map(|d| serde_json::from_value::<actions::Decision>(d.clone()))
368                    .transpose()
369                {
370                    Ok(Some(d)) => d,
371                    Ok(None) => actions::Decision::Accept,
372                    Err(_) => {
373                        return failed(ctx, id, event.data, "fail_closed_invalid_decision").await
374                    }
375                }
376            }
377            Ok(_) => return failed(ctx, id, event.data, "fail_closed_handler_action_error").await,
378            Err(_) => return failed(ctx, id, event.data, "fail_closed_handler_error").await,
379        }
380    } else {
381        ctx.state
382            .record_access_log(
383                AccessLogOwner::Server(ctx.server_id.as_u32()),
384                "Loki",
385                Some(id.as_u32()),
386                "loki_push",
387                event.data.clone(),
388                vec![json!({"type":"accept_loki_entries"})],
389            )
390            .await;
391        actions::Decision::Accept
392    };
393    match decision {
394        actions::Decision::Accept => response(204, "", None),
395        actions::Decision::Reject {
396            status,
397            message,
398            retry_after_seconds,
399        } => response(status, &message, retry_after_seconds),
400    }
401}
402async fn failed(
403    ctx: &SpawnContext,
404    id: ConnectionId,
405    data: serde_json::Value,
406    decision: &str,
407) -> Response<Full<Bytes>> {
408    console_error!(ctx.status_tx, "Loki decision={}", decision);
409    ctx.state
410        .record_access_log(
411            AccessLogOwner::Server(ctx.server_id.as_u32()),
412            "Loki",
413            Some(id.as_u32()),
414            "loki_handler_failed",
415            data,
416            vec![json!({"decision":decision})],
417        )
418        .await;
419    error(503, "Push handler failed; no acceptance confirmed")
420}
netget/server/loki/mod.rs

netget/server/loki/
mod.rs
