Skip to main content

netget/server/nut/
mod.rs

1//! Programmable NUT attachment daemon; application data belongs to event handlers.
2pub mod actions;
3pub mod wire;
4use crate::llm::action_helper::call_llm;
5use crate::llm::actions::protocol_trait::ActionResult;
6use crate::logging::emit::Log;
7use crate::protocol::{Event, SpawnContext};
8use crate::server::connection::ConnectionId;
9use crate::state::server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo};
10use anyhow::{Context, Result};
11use serde_json::{json, Value};
12use std::{net::SocketAddr, time::Duration};
13use tokio::{
14    io::{AsyncWriteExt, BufReader},
15    net::{TcpListener, TcpStream},
16};
17pub const IDLE_TIMEOUT: Duration = Duration::from_secs(300);
18
19pub const MAX_IDLE_TIMEOUT_SECS: u64 = 86400;
20
21pub fn idle_timeout(seconds: Option<u64>) -> Result<Duration> {
22    let seconds = seconds.unwrap_or(IDLE_TIMEOUT.as_secs());
23    anyhow::ensure!(
24        (1..=MAX_IDLE_TIMEOUT_SECS).contains(&seconds),
25        "idle_timeout_secs must be between 1 and 86400"
26    );
27    Ok(Duration::from_secs(seconds))
28}
29
30pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
31    let idle = idle_timeout(
32        ctx.startup_params
33            .as_ref()
34            .map(|p| p.get_optional_u64("idle_timeout_secs"))
35            .transpose()?
36            .flatten(),
37    )?;
38    let listener = TcpListener::bind(ctx.legacy_listen_addr()).await?;
39    let addr = listener.local_addr()?;
40    Log::new(Some(&ctx.status_tx)).info(format!("NUT listening on {addr}"));
41    let state = ctx.state.clone();
42    let server_id = ctx.server_id;
43    let accept = tokio::spawn(async move {
44        let limiter = crate::server::accept_bounded::ConnectionLimiter::new(
45            crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS,
46        );
47        loop {
48            let (socket, peer, permit) = match crate::server::accept_bounded::accept_bounded(
49                &listener,
50                &limiter,
51                b"ERR ACCESS-DENIED\n",
52                "NUT",
53                Some(&ctx.status_tx),
54            )
55            .await
56            {
57                Ok(v) => v,
58                Err(_) => break,
59            };
60            let id = ConnectionId::new(ctx.state.get_next_unified_id().await);
61            let now = crate::utils::clock::Instant::now();
62            ctx.state
63                .add_connection_to_server(
64                    server_id,
65                    ConnectionState {
66                        id,
67                        remote_addr: peer,
68                        local_addr: addr,
69                        bytes_sent: 0,
70                        bytes_received: 0,
71                        packets_sent: 0,
72                        packets_received: 0,
73                        last_activity: now,
74                        status: ConnectionStatus::Active,
75                        status_changed_at: now,
76                        protocol_info: ProtocolConnectionInfo::empty(),
77                    },
78                )
79                .await;
80            let child = ctx.clone();
81            ctx.state
82                .spawn_server_task(server_id, async move {
83                    let _permit = permit;
84                    if let Err(e) = session(&child, id, socket, idle).await {
85                        Log::new(Some(&child.status_tx))
86                            .warn(format!("NUT connection {id} ended: {e}"));
87                    }
88                    child
89                        .state
90                        .update_connection_status(server_id, id, ConnectionStatus::Closed)
91                        .await;
92                    let _ = child.status_tx.send("__UPDATE_UI__".into());
93                })
94                .await;
95        }
96    });
97    state.register_server_task(server_id, accept).await;
98    Ok(addr)
99}
100async fn write<W: tokio::io::AsyncWrite + Unpin>(
101    ctx: &SpawnContext,
102    id: ConnectionId,
103    w: &mut W,
104    reply: &str,
105) -> Result<()> {
106    if let Err(error) = tokio::time::timeout(wire::IO_TIMEOUT, w.write_all(reply.as_bytes()))
107        .await
108        .context("NUT write deadline")
109        .and_then(|result| result.map_err(Into::into))
110    {
111        outcome(ctx, id, "wire_reply", "fail_closed_send_error");
112        return Err(error);
113    }
114    ctx.state
115        .update_connection_stats(
116            ctx.server_id,
117            id,
118            None,
119            Some(reply.len() as u64),
120            None,
121            Some(1),
122        )
123        .await;
124    Ok(())
125}
126async fn decision(
127    ctx: &SpawnContext,
128    id: ConnectionId,
129    event: Event,
130    expected: &str,
131) -> Result<Value> {
132    let result = match call_llm(
133        &ctx.llm_client,
134        &ctx.state,
135        ctx.server_id,
136        Some(id),
137        &event,
138        &actions::NutProtocol,
139    )
140    .await
141    {
142        Ok(result) => result,
143        Err(error) => {
144            outcome(ctx, id, event.id(), "fail_closed_llm_error");
145            return Err(error);
146        }
147    };
148    if !result.failures.is_empty() {
149        outcome(ctx, id, event.id(), "fail_closed_invalid_reply");
150        anyhow::bail!("NUT handler supplied an invalid action");
151    }
152    let mut found = None;
153    let mut pending = result.protocol_results;
154    while let Some(result) = pending.pop() {
155        match result {
156            ActionResult::Custom { name, data } if name == expected => {
157                if found.is_some() {
158                    outcome(ctx, id, event.id(), "fail_closed_invalid_reply");
159                    anyhow::bail!("Multiple NUT replies");
160                }
161                found = Some(data);
162            }
163            ActionResult::Multiple(items) => pending.extend(items),
164            _ => {}
165        }
166    }
167    if found.is_none() {
168        outcome(ctx, id, event.id(), "model_silent");
169    }
170    found.context("Handler did not answer NUT request")
171}
172fn outcome(ctx: &SpawnContext, id: ConnectionId, operation: &str, decision: &str) {
173    let summary = format!("NUT connection {id} operation={operation} decision={decision}");
174    let log = Log::new(Some(&ctx.status_tx));
175    if decision.starts_with("fail_closed_") || decision == "model_silent" {
176        log.error(summary);
177    } else {
178        log.info(summary);
179    }
180}
181
182async fn session(
183    ctx: &SpawnContext,
184    id: ConnectionId,
185    socket: TcpStream,
186    idle: Duration,
187) -> Result<()> {
188    let (read, mut write_half) = tokio::io::split(socket);
189    let mut reader = BufReader::new(read);
190    let mut username: Option<String> = None;
191    let mut password_set = false;
192    let mut authenticated = false;
193    loop {
194        let line = match wire::read_line(&mut reader, idle).await {
195            Ok(Some(line)) => line,
196            Ok(None) => return Ok(()),
197            Err(e) => {
198                let _ = write(ctx, id, &mut write_half, "ERR INVALID-ARGUMENT\n").await;
199                return Err(e);
200            }
201        };
202        ctx.state
203            .update_connection_stats(
204                ctx.server_id,
205                id,
206                Some((line.len() + 1) as u64),
207                None,
208                Some(1),
209                None,
210            )
211            .await;
212        let fixed = match line.as_str() {
213            "STARTTLS" => Some("ERR FEATURE-NOT-SUPPORTED\n"),
214            "VER" => Some("Network UPS Tools netget\n"),
215            "HELP" => Some("Commands: HELP VER GET LIST SET INSTCMD USERNAME PASSWORD LOGOUT\n"),
216            _ => None,
217        };
218        if let Some(reply) = fixed {
219            write(ctx, id, &mut write_half, reply).await?;
220            continue;
221        }
222        let request = match wire::Request::parse(&line) {
223            Ok(r) => r,
224            Err(_) => {
225                write(ctx, id, &mut write_half, "ERR UNKNOWN-COMMAND\n").await?;
226                continue;
227            }
228        };
229        if request.operation == "logout" {
230            write(ctx, id, &mut write_half, "OK Goodbye\n").await?;
231            write_half.shutdown().await?;
232            return Ok(());
233        }
234        if request.operation == "username" {
235            if username.is_some() {
236                write(ctx, id, &mut write_half, "ERR ALREADY-SET\n").await?;
237            } else {
238                username = request.value;
239                write(ctx, id, &mut write_half, "OK\n").await?;
240            }
241            continue;
242        }
243        if request.operation == "password" {
244            let reply = if password_set {
245                "ERR ALREADY-SET\n"
246            } else if username.is_none() {
247                "ERR USERNAME-REQUIRED\n"
248            } else {
249                password_set = true;
250                let result = decision(
251                    ctx,
252                    id,
253                    Event::new(
254                        &actions::AUTH_EVENT,
255                        json!({"username":username,"password":request.value}),
256                    ),
257                    "nut_auth_decision",
258                )
259                .await;
260                authenticated = match result {
261                    Ok(value) => {
262                        let allowed = value["allowed"] == true;
263                        outcome(
264                            ctx,
265                            id,
266                            "nut_auth",
267                            if allowed {
268                                "model_answer"
269                            } else {
270                                "model_reject"
271                            },
272                        );
273                        allowed
274                    }
275                    Err(_) => false,
276                };
277                if authenticated {
278                    "OK\n"
279                } else {
280                    "ERR ACCESS-DENIED\n"
281                }
282            };
283            write(ctx, id, &mut write_half, reply).await?;
284            continue;
285        }
286        if request.is_write() && !authenticated {
287            write(ctx, id, &mut write_half, "ERR ACCESS-DENIED\n").await?;
288            continue;
289        }
290        let mut data = json!(request);
291        data["username"] = json!(if authenticated {
292            username.clone()
293        } else {
294            None
295        });
296        let result = decision(
297            ctx,
298            id,
299            Event::new(&actions::REQUEST_EVENT, data),
300            "nut_reply",
301        )
302        .await
303        .and_then(|value| match wire::render(&request, &value) {
304            Ok(reply) => {
305                outcome(
306                    ctx,
307                    id,
308                    &request.operation,
309                    if value.get("error").is_some() {
310                        "model_reject"
311                    } else {
312                        "model_answer"
313                    },
314                );
315                Ok(reply)
316            }
317            Err(error) => {
318                outcome(ctx, id, &request.operation, "fail_closed_invalid_reply");
319                Err(error)
320            }
321        });
322        match result {
323            Ok(reply) => write(ctx, id, &mut write_half, &reply).await?,
324            Err(_) => {
325                // The terminal decision was logged above without credential or backend details.
326                write(ctx, id, &mut write_half, "ERR DATA-STALE\n").await?;
327                return Ok(());
328            }
329        }
330    }
331}