1pub 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 write(ctx, id, &mut write_half, "ERR DATA-STALE\n").await?;
327 return Ok(());
328 }
329 }
330 }
331}