Skip to main content

netget/server/lpd/
mod.rs

1//! LPD (RFC 1179) print server. Rust owns the one-command-per-connection framing, the
2//! receive-job subcommands and their acknowledgements, and every bound; handlers decide
3//! whether each complete job is queued, what lpq is told and what lprm removed. No spool:
4//! nothing received is kept after the handler has seen it.
5pub mod actions;
6pub mod wire;
7
8use crate::llm::action_helper::call_llm;
9use crate::llm::actions::protocol_trait::ActionResult;
10use crate::logging::emit::Log;
11use crate::protocol::{Event, SpawnContext};
12use crate::server::connection::ConnectionId;
13use crate::state::server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo};
14use anyhow::{Context, Result};
15use serde_json::{json, Value};
16use std::{collections::HashMap, net::SocketAddr, time::Duration};
17use tokio::{
18    io::{AsyncWriteExt, BufReader},
19    net::{TcpListener, TcpStream},
20};
21
22/// Deadline for each command line and each file transfer.
23pub const IDLE_TIMEOUT: Duration = Duration::from_secs(300);
24pub const MAX_IDLE_TIMEOUT_SECS: u64 = 3600;
25/// Job-list operands accepted on a queue or removal command.
26pub const MAX_LIST_OPERANDS: usize = 64;
27
28#[derive(Clone)]
29struct Config {
30    queues: Option<Vec<String>>,
31    max_job: u64,
32    idle: Duration,
33}
34
35fn config(ctx: &SpawnContext) -> Result<Config> {
36    let params = ctx.startup_params.as_ref();
37    let queues = match params
38        .map(|p| p.get_optional_array("queues"))
39        .transpose()?
40        .flatten()
41    {
42        None => None,
43        Some(list) => {
44            let mut queues = Vec::new();
45            for queue in list {
46                let queue = queue
47                    .as_str()
48                    .filter(|q| wire::valid_token(q))
49                    .ok_or_else(|| {
50                        anyhow::anyhow!("queues must be printable names without spaces")
51                    })?;
52                queues.push(queue.to_string());
53            }
54            anyhow::ensure!(
55                !queues.is_empty() && queues.len() <= 256,
56                "queues must name 1 to 256 queues"
57            );
58            Some(queues)
59        }
60    };
61    let max_job = params
62        .map(|p| p.get_optional_u64("max_job_bytes"))
63        .transpose()?
64        .flatten()
65        .unwrap_or(wire::DEFAULT_MAX_JOB_BYTES);
66    anyhow::ensure!(
67        (1..=wire::MAX_JOB_BYTES_LIMIT).contains(&max_job),
68        "max_job_bytes must be between 1 and {}",
69        wire::MAX_JOB_BYTES_LIMIT
70    );
71    let idle = params
72        .map(|p| p.get_optional_u64("idle_timeout_secs"))
73        .transpose()?
74        .flatten()
75        .unwrap_or(IDLE_TIMEOUT.as_secs());
76    anyhow::ensure!(
77        (1..=MAX_IDLE_TIMEOUT_SECS).contains(&idle),
78        "idle_timeout_secs must be between 1 and {MAX_IDLE_TIMEOUT_SECS}"
79    );
80    Ok(Config {
81        queues,
82        max_job,
83        idle: Duration::from_secs(idle),
84    })
85}
86
87pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
88    let cfg = config(&ctx)?;
89    let listener = TcpListener::bind(ctx.legacy_listen_addr()).await?;
90    let addr = listener.local_addr()?;
91    Log::new(Some(&ctx.status_tx)).info(format!("LPD listening on {addr}"));
92    let state = ctx.state.clone();
93    let server_id = ctx.server_id;
94    let accept = tokio::spawn(async move {
95        let limiter = crate::server::accept_bounded::ConnectionLimiter::new(
96            crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS,
97        );
98        loop {
99            // LPD has no greeting or error line; a refused connection is simply closed.
100            let (socket, peer, permit) = match crate::server::accept_bounded::accept_bounded(
101                &listener,
102                &limiter,
103                b"",
104                "LPD",
105                Some(&ctx.status_tx),
106            )
107            .await
108            {
109                Ok(v) => v,
110                Err(_) => break,
111            };
112            let id = ConnectionId::new(ctx.state.get_next_unified_id().await);
113            let now = crate::utils::clock::Instant::now();
114            ctx.state
115                .add_connection_to_server(
116                    server_id,
117                    ConnectionState {
118                        id,
119                        remote_addr: peer,
120                        local_addr: addr,
121                        bytes_sent: 0,
122                        bytes_received: 0,
123                        packets_sent: 0,
124                        packets_received: 0,
125                        last_activity: now,
126                        status: ConnectionStatus::Active,
127                        status_changed_at: now,
128                        protocol_info: ProtocolConnectionInfo::empty(),
129                    },
130                )
131                .await;
132            let child = ctx.clone();
133            let cfg = cfg.clone();
134            ctx.state
135                .spawn_server_task(server_id, async move {
136                    let _permit = permit;
137                    if let Err(e) = session(&child, id, socket, &cfg).await {
138                        Log::new(Some(&child.status_tx))
139                            .warn(format!("LPD connection {id} ended: {e}"));
140                    }
141                    child
142                        .state
143                        .update_connection_status(server_id, id, ConnectionStatus::Closed)
144                        .await;
145                    let _ = child.status_tx.send("__UPDATE_UI__".into());
146                })
147                .await;
148        }
149    });
150    state.register_server_task(server_id, accept).await;
151    Ok(addr)
152}
153
154async fn write<W: tokio::io::AsyncWrite + Unpin>(
155    ctx: &SpawnContext,
156    id: ConnectionId,
157    w: &mut W,
158    bytes: &[u8],
159) -> Result<()> {
160    tokio::time::timeout(wire::IO_TIMEOUT, w.write_all(bytes))
161        .await
162        .context("LPD write deadline")??;
163    ctx.state
164        .update_connection_stats(
165            ctx.server_id,
166            id,
167            None,
168            Some(bytes.len() as u64),
169            None,
170            Some(1),
171        )
172        .await;
173    Ok(())
174}
175
176fn outcome(ctx: &SpawnContext, id: ConnectionId, operation: &str, decision: &str) {
177    let summary = format!("LPD connection {id} operation={operation} decision={decision}");
178    let log = Log::new(Some(&ctx.status_tx));
179    if decision.starts_with("fail_closed_") || decision == "model_silent" {
180        log.error(summary);
181    } else {
182        log.info(summary);
183    }
184}
185
186async fn decision(
187    ctx: &SpawnContext,
188    id: ConnectionId,
189    event: Event,
190    expected: &str,
191) -> Result<Value> {
192    let result = match call_llm(
193        &ctx.llm_client,
194        &ctx.state,
195        ctx.server_id,
196        Some(id),
197        &event,
198        &actions::LpdProtocol,
199    )
200    .await
201    {
202        Ok(result) => result,
203        Err(error) => {
204            outcome(ctx, id, event.id(), "fail_closed_llm_error");
205            return Err(error);
206        }
207    };
208    if !result.failures.is_empty() {
209        outcome(ctx, id, event.id(), "fail_closed_invalid_reply");
210        anyhow::bail!("LPD handler supplied an invalid action");
211    }
212    let mut found = None;
213    let mut pending = result.protocol_results;
214    while let Some(result) = pending.pop() {
215        match result {
216            ActionResult::Custom { name, data } if name == expected => {
217                if found.is_some() {
218                    outcome(ctx, id, event.id(), "fail_closed_invalid_reply");
219                    anyhow::bail!("Multiple LPD answers to one event");
220                }
221                found = Some(data);
222            }
223            ActionResult::Multiple(items) => pending.extend(items),
224            _ => {}
225        }
226    }
227    if found.is_none() {
228        outcome(ctx, id, event.id(), "model_silent");
229    }
230    found.context("Handler did not answer the LPD event")
231}
232
233async fn session(
234    ctx: &SpawnContext,
235    id: ConnectionId,
236    socket: TcpStream,
237    cfg: &Config,
238) -> Result<()> {
239    let (read, mut out) = tokio::io::split(socket);
240    let mut reader = BufReader::new(read);
241    let Some(line) = wire::read_line(&mut reader, cfg.idle).await? else {
242        return Ok(());
243    };
244    ctx.state
245        .update_connection_stats(
246            ctx.server_id,
247            id,
248            Some(line.len() as u64 + 1),
249            None,
250            Some(1),
251            None,
252        )
253        .await;
254    let Some((&code, rest)) = line.split_first() else {
255        anyhow::bail!("Empty LPD command");
256    };
257    let operands = wire::operands(rest);
258    let Some(queue) = operands.first().filter(|q| wire::valid_token(q)).cloned() else {
259        anyhow::bail!("LPD command without a valid queue name");
260    };
261    let list: Vec<String> = operands.iter().skip(1).cloned().collect();
262    anyhow::ensure!(
263        list.len() <= MAX_LIST_OPERANDS + 1 && list.iter().all(|t| wire::valid_token(t)),
264        "LPD command operands out of bounds"
265    );
266    match code {
267        wire::CMD_PRINT_WAITING => {
268            // RFC 1179 5.1 defines no reply; the request is logged and the connection closed.
269            Log::new(Some(&ctx.status_tx))
270                .info(format!("LPD print-waiting-jobs for queue {queue}"));
271            Ok(())
272        }
273        wire::CMD_RECEIVE_JOB => receive_job(ctx, id, &mut reader, &mut out, cfg, &queue).await,
274        wire::CMD_QUEUE_SHORT | wire::CMD_QUEUE_LONG => {
275            let long = code == wire::CMD_QUEUE_LONG;
276            let event = Event::new(
277                &actions::QUEUE_EVENT,
278                json!({"queue": queue, "long": long, "list": list}),
279            );
280            let text = match decision(ctx, id, event, "lpd_queue_status").await {
281                Ok(answer) => {
282                    outcome(ctx, id, "lpd_queue_query", "model_answer");
283                    let jobs = answer["jobs"].as_array().cloned().unwrap_or_default();
284                    wire::render_queue(&queue, long, answer["status"].as_str(), &jobs)
285                }
286                Err(_) => format!("{queue}: queue status unavailable\n"),
287            };
288            write(ctx, id, &mut out, text.as_bytes()).await?;
289            let _ = out.shutdown().await;
290            Ok(())
291        }
292        wire::CMD_REMOVE_JOBS => {
293            let Some((agent, jobs)) = list.split_first() else {
294                anyhow::bail!("LPD remove command without an agent");
295            };
296            let event = Event::new(
297                &actions::REMOVE_EVENT,
298                json!({"queue": queue, "agent": agent, "list": jobs}),
299            );
300            // On failure nothing was removed, and nothing is reported as removed.
301            let text = match decision(ctx, id, event, "lpd_remove_result").await {
302                Ok(answer) => {
303                    let removed = answer["removed"].as_array().cloned().unwrap_or_default();
304                    outcome(
305                        ctx,
306                        id,
307                        "lpd_remove_request",
308                        if removed.is_empty() {
309                            "model_reject"
310                        } else {
311                            "model_answer"
312                        },
313                    );
314                    removed
315                        .iter()
316                        .map(|job| match job {
317                            Value::Number(n) => match n.as_u64() {
318                                Some(job) => format!("job {job:03} dequeued\n"),
319                                None => format!("job {n} dequeued\n"),
320                            },
321                            other => format!(
322                                "job {} dequeued\n",
323                                wire::clean(other.as_str().unwrap_or(""), 64)
324                            ),
325                        })
326                        .collect::<String>()
327                }
328                Err(_) => String::new(),
329            };
330            write(ctx, id, &mut out, text.as_bytes()).await?;
331            let _ = out.shutdown().await;
332            Ok(())
333        }
334        other => anyhow::bail!("Unknown LPD command code {other:#04x}"),
335    }
336}
337
338async fn receive_job<R, W>(
339    ctx: &SpawnContext,
340    id: ConnectionId,
341    reader: &mut R,
342    out: &mut W,
343    cfg: &Config,
344    queue: &str,
345) -> Result<()>
346where
347    R: tokio::io::AsyncBufRead + Unpin,
348    W: tokio::io::AsyncWrite + Unpin,
349{
350    if cfg
351        .queues
352        .as_ref()
353        .is_some_and(|q| !q.iter().any(|name| name == queue))
354    {
355        Log::new(Some(&ctx.status_tx)).info(format!("LPD refused job for unknown queue {queue}"));
356        write(ctx, id, out, &[1]).await?;
357        return Ok(());
358    }
359    write(ctx, id, out, &[0]).await?;
360    let mut data: HashMap<String, Vec<u8>> = HashMap::new();
361    let mut control: Option<(String, wire::ControlFile)> = None;
362    let mut job_bytes: u64 = 0;
363    loop {
364        let Some(line) = wire::read_line(reader, cfg.idle).await? else {
365            if control.is_some() || !data.is_empty() {
366                Log::new(Some(&ctx.status_tx)).warn(format!(
367                    "LPD connection {id} closed with an incomplete job; discarded"
368                ));
369            }
370            return Ok(());
371        };
372        let Some((&sub, rest)) = line.split_first() else {
373            anyhow::bail!("Empty LPD subcommand");
374        };
375        if sub == wire::SUB_ABORT {
376            data.clear();
377            control = None;
378            job_bytes = 0;
379            continue;
380        }
381        anyhow::ensure!(
382            sub == wire::SUB_CONTROL_FILE || sub == wire::SUB_DATA_FILE,
383            "Unknown LPD receive-job subcommand {sub:#04x}"
384        );
385        let operands = wire::operands(rest);
386        let count = operands.first().and_then(|c| c.parse::<u64>().ok());
387        let name = operands.get(1).filter(|n| wire::valid_token(n)).cloned();
388        let is_control = sub == wire::SUB_CONTROL_FILE;
389        let admissible = match (count, &name) {
390            (Some(count), Some(name)) if is_control => {
391                name.starts_with("cf") && count <= wire::MAX_CONTROL_FILE && control.is_none()
392            }
393            (Some(count), Some(name)) => {
394                name.starts_with("df")
395                    && job_bytes.saturating_add(count) <= cfg.max_job
396                    && data.len() < wire::MAX_FILES_PER_JOB
397                    && !data.contains_key(name)
398            }
399            _ => false,
400        };
401        let (Some(count), Some(name), true) = (count, name, admissible) else {
402            // Refused before a byte of the file is read; the client will not send it.
403            write(ctx, id, out, &[1]).await?;
404            return Ok(());
405        };
406        write(ctx, id, out, &[0]).await?;
407        let body = wire::read_file(reader, count, cfg.idle).await?;
408        ctx.state
409            .update_connection_stats(ctx.server_id, id, Some(count + 1), None, Some(1), None)
410            .await;
411        if is_control {
412            match wire::ControlFile::parse(&body) {
413                Ok(parsed) => control = Some((name, parsed)),
414                Err(e) => {
415                    Log::new(Some(&ctx.status_tx))
416                        .warn(format!("LPD connection {id} refused a control file: {e}"));
417                    write(ctx, id, out, &[1]).await?;
418                    return Ok(());
419                }
420            }
421        } else {
422            job_bytes += count;
423            data.insert(name, body);
424        }
425        let complete = control
426            .as_ref()
427            .is_some_and(|(_, c)| c.data_files().iter().all(|n| data.contains_key(n)));
428        if !complete {
429            write(ctx, id, out, &[0]).await?;
430            continue;
431        }
432        let (control_name, parsed) = control.take().expect("complete implies a control file");
433        let accepted = decide_job(ctx, id, queue, &control_name, &parsed, &data).await;
434        write(ctx, id, out, &[if accepted { 0 } else { 1 }]).await?;
435        data.clear();
436        job_bytes = 0;
437    }
438}
439
440async fn decide_job(
441    ctx: &SpawnContext,
442    id: ConnectionId,
443    queue: &str,
444    control_name: &str,
445    control: &wire::ControlFile,
446    data: &HashMap<String, Vec<u8>>,
447) -> bool {
448    let (job_id, origin) = wire::job_number(control_name).unwrap_or_default();
449    let files: Vec<Value> = control
450        .prints
451        .iter()
452        .map(|(format, name, source)| {
453            let bytes = data.get(name).map(Vec::as_slice).unwrap_or_default();
454            let (text, _binary) = wire::preview(bytes);
455            json!({
456                "name": name,
457                "source_name": source,
458                "format": format.to_string(),
459                "size": bytes.len(),
460                "text": text,
461            })
462        })
463        .collect();
464    let event = Event::new(
465        &actions::JOB_EVENT,
466        json!({
467            "queue": queue,
468            "job_id": job_id,
469            "host": control.host.clone().or(Some(origin).filter(|o| !o.is_empty())),
470            "user": control.user,
471            "job_name": control.job_name,
472            "title": control.title,
473            "class": control.class,
474            "mail": control.mail,
475            "files": files,
476        }),
477    );
478    match decision(ctx, id, event, "lpd_job_reply").await {
479        Ok(answer) if answer["accept"] == true => {
480            outcome(ctx, id, "lpd_print_job", "model_answer");
481            true
482        }
483        Ok(_) => {
484            outcome(ctx, id, "lpd_print_job", "model_reject");
485            false
486        }
487        Err(_) => false,
488    }
489}