netget/server/icap/
mod.rs

1//! ICAP (RFC 3507) server. Rust owns framing, OPTIONS, preview continuation and response
2//! assembly; the handler decides each REQMOD/RESPMOD.
3pub mod actions;
4pub mod wire;
5
6use crate::llm::action_helper::call_llm;
7use crate::llm::actions::protocol_trait::ActionResult;
8use crate::logging::emit::Log;
9use crate::protocol::{Event, SpawnContext};
10use crate::server::accept_bounded::{accept_bounded, ConnectionLimiter, DEFAULT_MAX_CONNECTIONS};
11use crate::server::connection::ConnectionId;
12use crate::state::server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo};
13use anyhow::{bail, ensure, Context, Result};
14use serde_json::{json, Value};
15use std::net::SocketAddr;
16use std::sync::Arc;
17use std::time::Duration;
18use tokio::io::BufReader;
19use wire::HttpHead;
20
21pub const IDLE_TIMEOUT: Duration = Duration::from_secs(300);
22pub const DEFAULT_SERVICE: &str = "netget";
23pub const DEFAULT_PREVIEW: u64 = 1024;
24const ISTAG: &str = "\"NetGet-1\"";
25
26struct Service {
27    name: String,
28    methods: Vec<String>,
29}
30
31struct Shared {
32    ctx: SpawnContext,
33    services: Vec<Service>,
34    preview: u64,
35    idle: Duration,
36}
37
38pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
39    let p = ctx.startup_params.as_ref();
40    let services = match p
41        .map(|p| p.get_optional_array("services"))
42        .transpose()?
43        .flatten()
44    {
45        None => vec![Service {
46            name: DEFAULT_SERVICE.into(),
47            methods: vec!["REQMOD".into(), "RESPMOD".into()],
48        }],
49        Some(list) => {
50            ensure!(
51                !list.is_empty() && list.len() <= 32,
52                "services must list 1..=32 services"
53            );
54            let mut out = Vec::new();
55            for s in list {
56                let name = s["name"].as_str().context("each service needs a name")?;
57                ensure!(
58                    !name.is_empty()
59                        && name.len() <= 64
60                        && name
61                            .bytes()
62                            .all(|b| b.is_ascii_alphanumeric() || b"-_.".contains(&b)),
63                    "service name '{name}' must be letters, digits, '-', '_' or '.'"
64                );
65                let methods: Vec<String> = s["methods"]
66                    .as_array()
67                    .context("each service needs methods")?
68                    .iter()
69                    .map(|m| m.as_str().unwrap_or("").to_owned())
70                    .collect();
71                ensure!(
72                    !methods.is_empty() && methods.iter().all(|m| m == "REQMOD" || m == "RESPMOD"),
73                    "methods must be REQMOD and/or RESPMOD"
74                );
75                out.push(Service {
76                    name: name.into(),
77                    methods,
78                });
79            }
80            out
81        }
82    };
83    let preview = p
84        .map(|p| p.get_optional_u64("preview_bytes"))
85        .transpose()?
86        .flatten()
87        .unwrap_or(DEFAULT_PREVIEW);
88    ensure!(preview <= 65536, "preview_bytes must be 0..=65536");
89    let idle = p
90        .map(|p| p.get_optional_u64("idle_timeout_secs"))
91        .transpose()?
92        .flatten()
93        .unwrap_or(IDLE_TIMEOUT.as_secs());
94    ensure!(
95        (1..=86400).contains(&idle),
96        "idle_timeout_secs must be 1..=86400"
97    );
98    let listener = tokio::net::TcpListener::bind(ctx.legacy_listen_addr()).await?;
99    let local = listener.local_addr()?;
100    Log::new(Some(&ctx.status_tx)).info(format!(
101        "ICAP listening on {local}; services: {}",
102        services
103            .iter()
104            .map(|s| s.name.as_str())
105            .collect::<Vec<_>>()
106            .join(", ")
107    ));
108    let shared = Arc::new(Shared {
109        ctx: ctx.clone(),
110        services,
111        preview,
112        idle: Duration::from_secs(idle),
113    });
114    let server_id = ctx.server_id;
115    let accept = tokio::spawn(async move {
116        let limiter = ConnectionLimiter::new(DEFAULT_MAX_CONNECTIONS);
117        let refusal = b"ICAP/1.0 503 Service Unavailable\r\nISTag: \"NetGet-1\"\r\nEncapsulated: null-body=0\r\n\r\n";
118        loop {
119            let (socket, peer, permit) = match accept_bounded(
120                &listener,
121                &limiter,
122                refusal,
123                "ICAP",
124                Some(&shared.ctx.status_tx),
125            )
126            .await
127            {
128                Ok(v) => v,
129                Err(_) => break,
130            };
131            let id = ConnectionId::new(shared.ctx.state.get_next_unified_id().await);
132            let now = crate::utils::clock::Instant::now();
133            shared
134                .ctx
135                .state
136                .add_connection_to_server(
137                    server_id,
138                    ConnectionState {
139                        id,
140                        remote_addr: peer,
141                        local_addr: local,
142                        bytes_sent: 0,
143                        bytes_received: 0,
144                        packets_sent: 0,
145                        packets_received: 0,
146                        last_activity: now,
147                        status: ConnectionStatus::Active,
148                        status_changed_at: now,
149                        protocol_info: ProtocolConnectionInfo::empty(),
150                    },
151                )
152                .await;
153            let child = shared.clone();
154            shared
155                .ctx
156                .state
157                .spawn_server_task(server_id, async move {
158                    let _permit = permit;
159                    if let Err(e) = connection(&child, id, socket).await {
160                        Log::new(Some(&child.ctx.status_tx))
161                            .debug(format!("ICAP connection {id} ended: {e}"));
162                    }
163                    child
164                        .ctx
165                        .state
166                        .update_connection_status(server_id, id, ConnectionStatus::Closed)
167                        .await;
168                    let _ = child.ctx.status_tx.send("__UPDATE_UI__".into());
169                })
170                .await;
171        }
172    });
173    ctx.state.register_server_task(server_id, accept).await;
174    Ok(local)
175}
176
177fn outcome(ctx: &SpawnContext, id: ConnectionId, operation: &str, decision: &str) {
178    let summary = format!("ICAP connection {id} request={operation} decision={decision}");
179    let log = Log::new(Some(&ctx.status_tx));
180    if decision.starts_with("fail_closed_") || decision == "model_silent" {
181        log.error(summary);
182    } else {
183        log.info(summary);
184    }
185}
186
187fn status_only(code: u16, reason: &str) -> Result<Vec<u8>> {
188    wire::message(
189        &format!("ICAP/1.0 {code} {reason}"),
190        vec![("ISTag".into(), ISTAG.into())],
191        None,
192        None,
193        None,
194        "",
195    )
196}
197
198async fn decide(shared: &Shared, id: ConnectionId, event: Event) -> Result<Value> {
199    let ctx = &shared.ctx;
200    let result = match call_llm(
201        &ctx.llm_client,
202        &ctx.state,
203        ctx.server_id,
204        Some(id),
205        &event,
206        &actions::IcapProtocol,
207    )
208    .await
209    {
210        Ok(r) => r,
211        Err(e) => {
212            outcome(ctx, id, event.id(), "fail_closed_llm_error");
213            return Err(e);
214        }
215    };
216    if !result.failures.is_empty() {
217        outcome(ctx, id, event.id(), "fail_closed_invalid_reply");
218        bail!("ICAP handler supplied an invalid action");
219    }
220    let mut found = None;
221    let mut pending = result.protocol_results;
222    while let Some(r) = pending.pop() {
223        match r {
224            ActionResult::Custom { name, data } if name == "icap_response" => {
225                if found.replace(data).is_some() {
226                    outcome(ctx, id, event.id(), "fail_closed_invalid_reply");
227                    bail!("ICAP handler supplied more than one response");
228                }
229            }
230            ActionResult::Multiple(items) => pending.extend(items),
231            _ => {}
232        }
233    }
234    found
235        .context("ICAP handler did not answer")
236        .inspect_err(|_| outcome(ctx, id, event.id(), "model_silent"))
237}
238
239struct Request {
240    method: String,
241    service: String,
242    headers: Vec<(String, String)>,
243    req: Option<HttpHead>,
244    res: Option<HttpHead>,
245    body: Option<Vec<u8>>,
246}
247
248async fn connection(
249    shared: &Shared,
250    id: ConnectionId,
251    socket: tokio::net::TcpStream,
252) -> Result<()> {
253    let ctx = &shared.ctx;
254    let (read, mut writer) = tokio::io::split(socket);
255    let mut reader = BufReader::new(read);
256    loop {
257        let head = match tokio::time::timeout(
258            shared.idle,
259            wire::read_head(&mut reader, wire::MAX_HEAD_BYTES),
260        )
261        .await
262        {
263            Err(_) => return Ok(()),
264            Ok(r) => r,
265        };
266        let Some(head) = (match head {
267            Ok(h) => h,
268            Err(e) => {
269                let _ = wire::write_all(&mut writer, &status_only(400, "Bad Request")?).await;
270                return Err(e);
271            }
272        }) else {
273            return Ok(());
274        };
275        ctx.state
276            .update_connection_stats(
277                ctx.server_id,
278                id,
279                Some(head.len() as u64),
280                None,
281                Some(1),
282                None,
283            )
284            .await;
285        let parsed = match wire::parse_head(&head) {
286            Ok(p) => p,
287            Err(e) => {
288                let _ = wire::write_all(&mut writer, &status_only(400, "Bad Request")?).await;
289                return Err(e);
290            }
291        };
292        let close = wire::header(&parsed.headers, "Connection")
293            .is_some_and(|v| v.eq_ignore_ascii_case("close"));
294        let reply = tokio::time::timeout(
295            wire::IO_TIMEOUT * 4,
296            handle(shared, id, parsed, &mut reader, &mut writer),
297        )
298        .await
299        .context("ICAP request deadline")??;
300        wire::write_all(&mut writer, &reply).await?;
301        ctx.state
302            .update_connection_stats(
303                ctx.server_id,
304                id,
305                None,
306                Some(reply.len() as u64),
307                None,
308                Some(1),
309            )
310            .await;
311        // After an error the request's body may still be unread; close rather than parse it
312        // as the next request.
313        let failed = reply.starts_with(b"ICAP/1.0 4") || reply.starts_with(b"ICAP/1.0 5");
314        if close || failed {
315            return Ok(());
316        }
317    }
318}
319
320async fn handle<R, W>(
321    shared: &Shared,
322    id: ConnectionId,
323    head: HttpHead,
324    reader: &mut R,
325    writer: &mut W,
326) -> Result<Vec<u8>>
327where
328    R: tokio::io::AsyncBufRead + Unpin,
329    W: tokio::io::AsyncWrite + Unpin,
330{
331    let ctx = &shared.ctx;
332    let [method, uri, version] = &head.start;
333    if version != "ICAP/1.0" {
334        return status_only(505, "ICAP Version Not Supported");
335    }
336    let Some(rest) = uri.strip_prefix("icap://") else {
337        return status_only(400, "Bad Request");
338    };
339    let service = rest
340        .split_once('/')
341        .map(|(_, s)| s.split('?').next().unwrap_or(""))
342        .unwrap_or("")
343        .to_owned();
344    let Some(svc) = shared.services.iter().find(|s| s.name == service) else {
345        return status_only(404, "ICAP Service Not Found");
346    };
347    if method == "OPTIONS" {
348        let headers = vec![
349            ("Methods".into(), svc.methods.join(", ")),
350            ("Service".into(), "NetGet ICAP".into()),
351            ("ISTag".into(), ISTAG.into()),
352            (
353                "Max-Connections".into(),
354                crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS.to_string(),
355            ),
356            ("Options-TTL".into(), "3600".into()),
357            ("Allow".into(), "204".into()),
358            ("Preview".into(), shared.preview.to_string()),
359            ("Transfer-Preview".into(), "*".into()),
360        ];
361        outcome(ctx, id, "OPTIONS", "protocol_options");
362        return wire::message("ICAP/1.0 200 OK", headers, None, None, None, "");
363    }
364    if method != "REQMOD" && method != "RESPMOD" {
365        return status_only(501, "Method Not Implemented");
366    }
367    if !svc.methods.iter().any(|m| m == method) {
368        return status_only(405, "Method Not Allowed For Service");
369    }
370    let encapsulated = wire::parse_encapsulated(
371        wire::header(&head.headers, "Encapsulated").context("Encapsulated header is required")?,
372    )?;
373    let mut req = None;
374    let mut res = None;
375    let mut body = None;
376    for (i, (entity, offset)) in encapsulated.iter().enumerate() {
377        let next = encapsulated.get(i + 1).map(|(_, o)| *o);
378        match entity.as_str() {
379            "req-hdr" => {
380                req = Some(
381                    wire::read_exact_head(
382                        reader,
383                        next.context("req-hdr must be followed by another entity")? - offset,
384                    )
385                    .await?,
386                )
387            }
388            "res-hdr" => {
389                res = Some(
390                    wire::read_exact_head(
391                        reader,
392                        next.context("res-hdr must be followed by another entity")? - offset,
393                    )
394                    .await?,
395                )
396            }
397            "req-body" | "res-body" => {
398                let first = wire::read_chunks(reader, 0).await?;
399                let mut data = first.data;
400                if wire::header(&head.headers, "Preview").is_some() && !first.ieof {
401                    wire::write_all(writer, b"ICAP/1.0 100 Continue\r\n\r\n").await?;
402                    let rest = wire::read_chunks(reader, data.len()).await?;
403                    data.extend(rest.data);
404                }
405                body = Some(data);
406            }
407            _ => {}
408        }
409    }
410    ensure!(
411        method == "RESPMOD" || res.is_none(),
412        "REQMOD carries no response head"
413    );
414    let request = Request {
415        method: method.clone(),
416        service,
417        headers: head.headers.clone(),
418        req,
419        res,
420        body,
421    };
422    let allow_204 = wire::header(&request.headers, "Allow")
423        .is_some_and(|v| v.split(',').any(|t| t.trim() == "204"));
424    let mut data = json!({
425        "method": request.method,
426        "service": request.service,
427        "allow_204": allow_204,
428        "icap_headers": request.headers.iter().map(|(k, v)| json!([k, v])).collect::<Vec<_>>(),
429    });
430    if let Some(h) = &request.req {
431        data["http_request"] = wire::head_json(h, true);
432    }
433    if let Some(h) = &request.res {
434        data["http_response"] = wire::head_json(h, false);
435    }
436    if let Some(b) = &request.body {
437        for (k, v) in wire::body_json(b).as_object().expect("object") {
438            data[k] = v.clone();
439        }
440    }
441    let label = request.method.clone();
442    let decision = decide(shared, id, Event::new(&actions::REQUEST_EVENT, data)).await;
443    let decision = match decision {
444        Err(_) => return status_only(500, "Server Error"),
445        Ok(v) => match actions::validate(&v) {
446            Ok(()) => v,
447            Err(_) => {
448                outcome(ctx, id, &label, "fail_closed_invalid_reply");
449                return status_only(500, "Server Error");
450            }
451        },
452    };
453    let verdict = decision["verdict"].as_str().unwrap_or("error");
454    let istag = vec![("ISTag".to_owned(), ISTAG.to_owned())];
455    let text_body = decision["body_text"]
456        .as_str()
457        .map(|s| s.as_bytes().to_vec());
458    let reply = match verdict {
459        "no_modification" => {
460            outcome(ctx, id, &label, "model_pass");
461            if allow_204 {
462                status_only(204, "No Content")
463            } else if request.method == "REQMOD" {
464                wire::message(
465                    "ICAP/1.0 200 OK",
466                    istag,
467                    request.req.as_ref(),
468                    None,
469                    request.body.as_deref(),
470                    "req-body",
471                )
472            } else {
473                wire::message(
474                    "ICAP/1.0 200 OK",
475                    istag,
476                    None,
477                    request.res.as_ref(),
478                    request.body.as_deref(),
479                    "res-body",
480                )
481            }
482        }
483        "block" => {
484            outcome(ctx, id, &label, "model_block");
485            let page = text_body.unwrap_or_else(|| b"Blocked by policy".to_vec());
486            let mut head = HttpHead {
487                start: ["HTTP/1.1".into(), "403".into(), "Forbidden".into()],
488                headers: vec![("Content-Type".into(), "text/plain; charset=utf-8".into())],
489            };
490            wire::set_length(&mut head, page.len());
491            wire::message(
492                "ICAP/1.0 200 OK",
493                istag,
494                None,
495                Some(&head),
496                Some(&page),
497                "res-body",
498            )
499        }
500        // RFC 3507 §4.4.1: a RESPMOD response carries only [res-hdr] res-body.
501        "modify" if request.method == "RESPMOD" && decision["http_response"].is_null() => {
502            outcome(ctx, id, &label, "fail_closed_invalid_reply");
503            status_only(500, "Server Error")
504        }
505        "modify" => {
506            outcome(ctx, id, &label, "model_modify");
507            if !decision["http_response"].is_null() {
508                let mut head = wire::response_head(&decision["http_response"])?;
509                let body = text_body
510                    .or_else(|| request.body.clone())
511                    .unwrap_or_default();
512                wire::set_length(&mut head, body.len());
513                wire::message(
514                    "ICAP/1.0 200 OK",
515                    istag,
516                    None,
517                    Some(&head),
518                    Some(&body),
519                    "res-body",
520                )
521            } else {
522                let mut head = wire::request_head(&decision["http_request"])?;
523                let body = text_body.or_else(|| request.body.clone());
524                if let Some(b) = &body {
525                    wire::set_length(&mut head, b.len());
526                }
527                wire::message(
528                    "ICAP/1.0 200 OK",
529                    istag,
530                    Some(&head),
531                    None,
532                    body.as_deref(),
533                    "req-body",
534                )
535            }
536        }
537        _ => {
538            outcome(ctx, id, &label, "model_error");
539            let code = decision["status"].as_u64().unwrap_or(500) as u16;
540            let reason = match code {
541                400 => "Bad Request",
542                403 => "Forbidden",
543                404 => "ICAP Service Not Found",
544                503 => "Service Unavailable",
545                _ => "Server Error",
546            };
547            status_only(code, reason)
548        }
549    };
550    reply
551}
netget/server/icap/mod.rs

netget/server/icap/
mod.rs
