Skip to main content

netget/client/inetd/
mod.rs

1//! Clients for the classic inetd services (RFC 862-868). Each query is one exchange: a fresh
2//! TCP connection or one UDP datagram. `connect` opens nothing, because connecting is itself
3//! a request to Daytime, QOTD, Time and Chargen.
4pub mod actions;
5use crate::client::llm_budget::call_llm_for_client;
6use crate::llm::actions::client_trait::ClientActionResult;
7use crate::logging::emit::Log;
8use crate::protocol::{ConnectContext, Event};
9use crate::server::inetd::{actions::Service, wire};
10use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
11use crate::state::{AccessLogOwner, ClientStatus};
12use actions::Query;
13use anyhow::{Context, Result};
14use serde_json::{json, Value};
15use std::{net::SocketAddr, time::Duration};
16use tokio::{
17    io::{AsyncReadExt, AsyncWriteExt},
18    net::{TcpStream, UdpSocket},
19    sync::mpsc,
20};
21
22pub const DEFAULT_TRANSPORT: &str = "tcp";
23/// How long a UDP query waits for its reply.
24pub const UDP_REPLY_TIMEOUT: Duration = Duration::from_secs(5);
25
26pub async fn connect(ctx: ConnectContext, service: Service) -> Result<SocketAddr> {
27    let udp = match ctx
28        .startup_params
29        .as_ref()
30        .map(|p| p.get_optional_string("transport"))
31        .transpose()?
32        .flatten()
33        .as_deref()
34        .unwrap_or(DEFAULT_TRANSPORT)
35    {
36        "tcp" => false,
37        "udp" => true,
38        other => anyhow::bail!("transport must be tcp or udp, not {other}"),
39    };
40    let remote: SocketAddr = tokio::net::lookup_host(&ctx.remote_addr)
41        .await
42        .context("resolve the service address")?
43        .next()
44        .context("the service address resolved to nothing")?;
45    ctx.state
46        .update_client_status(ctx.client_id, ClientStatus::Connected)
47        .await;
48    let external =
49        crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id).await;
50    let (internal_tx, internal_rx) = mpsc::channel::<Value>(16);
51    let (event_tx, mut event_rx) = mpsc::channel::<Event>(16);
52    event_tx.try_send(Event::new(
53        actions::ready_event(service),
54        json!({"remote_addr": ctx.remote_addr, "transport": if udp { "udp" } else { "tcp" }}),
55    ))?;
56    let events_ctx = ctx.clone();
57    let dispatcher = tokio::spawn(async move {
58        while let Some(event) = event_rx.recv().await {
59            let instruction = events_ctx
60                .state
61                .get_instruction_for_client(events_ctx.client_id)
62                .await
63                .unwrap_or_default();
64            let memory = events_ctx
65                .state
66                .get_memory_for_client(events_ctx.client_id)
67                .await
68                .unwrap_or_default();
69            match call_llm_for_client(
70                &events_ctx.llm_client,
71                &events_ctx.state,
72                events_ctx.client_id.to_string(),
73                &instruction,
74                &memory,
75                Some(&event),
76                actions::protocol(service),
77                &events_ctx.status_tx,
78            )
79            .await
80            {
81                Ok(result) => {
82                    if let Some(memory) = result.memory_updates {
83                        events_ctx
84                            .state
85                            .set_memory_for_client(events_ctx.client_id, memory)
86                            .await;
87                    }
88                    for action in result.actions {
89                        if internal_tx.send(action).await.is_err() {
90                            return;
91                        }
92                    }
93                }
94                Err(e) => Log::new(Some(&events_ctx.status_tx))
95                    .warn(format!("{} client handler: {e}", service.name())),
96            }
97        }
98    });
99    let dispatcher_abort = dispatcher.abort_handle();
100    ctx.state
101        .register_client_task(ctx.client_id, dispatcher)
102        .await;
103    let session_ctx = ctx.clone();
104    let task = tokio::spawn(async move {
105        let result = session(
106            &session_ctx,
107            service,
108            remote,
109            udp,
110            external,
111            internal_rx,
112            event_tx,
113        )
114        .await;
115        if result.is_err() {
116            dispatcher_abort.abort();
117        }
118        let status = match result {
119            Ok(()) => ClientStatus::Disconnected,
120            Err(e) => {
121                Log::new(Some(&session_ctx.status_tx))
122                    .warn(format!("{} client ended: {e}", service.name()));
123                ClientStatus::Error(e.to_string())
124            }
125        };
126        session_ctx
127            .state
128            .update_client_status(session_ctx.client_id, status)
129            .await;
130        session_ctx
131            .state
132            .remove_client_handle(session_ctx.client_id)
133            .await;
134        let _ = session_ctx.status_tx.send("__UPDATE_UI__".into());
135    });
136    ctx.state.register_client_task(ctx.client_id, task).await;
137    Ok(SocketAddr::new(
138        std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED),
139        0,
140    ))
141}
142
143fn answer(injected: &mut Option<ClientCommand>, outcome: Result<ClientSendOutcome>) {
144    if let Some(command) = injected.take() {
145        crate::client::command_support::reply(command, outcome);
146    }
147}
148
149/// Turn received bytes into the response event's data.
150fn response(service: Service, transport: &str, sent: &[u8], received: &[u8]) -> Result<Value> {
151    Ok(match service {
152        Service::Echo => {
153            let (data, encoding) = wire::encode(received);
154            json!({"transport": transport, "data": data, "encoding": encoding, "bytes": received.len(), "matches": received == sent})
155        }
156        Service::Discard => json!({"transport": transport, "bytes": sent.len()}),
157        Service::Daytime => {
158            let text = String::from_utf8_lossy(received);
159            json!({"transport": transport, "text": text.trim_end_matches(['\r', '\n'])})
160        }
161        Service::Qotd => {
162            let quote = String::from_utf8_lossy(received).replace("\r\n", "\n");
163            json!({"transport": transport, "quote": quote.trim_end_matches('\n')})
164        }
165        Service::Chargen => {
166            let text = String::from_utf8_lossy(received).to_string();
167            let conforms = wire::conforms_to_chargen(&text, &wire::default_charset());
168            json!({"transport": transport, "text": text, "bytes": received.len(), "conforms": conforms})
169        }
170        Service::Time => {
171            anyhow::ensure!(
172                received.len() == 4,
173                "a Time reply is 4 bytes, got {}",
174                received.len()
175            );
176            let mut value =
177                wire::describe_time([received[0], received[1], received[2], received[3]]);
178            value["transport"] = json!(transport);
179            value
180        }
181    })
182}
183
184/// One TCP exchange. `injected` is answered once the request is on the wire.
185async fn over_tcp(
186    service: Service,
187    remote: SocketAddr,
188    query: &Query,
189    injected: &mut Option<ClientCommand>,
190) -> Result<Option<Value>> {
191    let mut stream = tokio::time::timeout(wire::IO_TIMEOUT, TcpStream::connect(remote))
192        .await
193        .context("connect deadline")??;
194    if !query.payload.is_empty() {
195        tokio::time::timeout(wire::IO_TIMEOUT, stream.write_all(&query.payload))
196            .await
197            .context("write deadline")??;
198    }
199    answer(
200        injected,
201        Ok(ClientSendOutcome::Sent {
202            bytes_sent: query.payload.len(),
203        }),
204    );
205    let limit = match service {
206        Service::Discard => {
207            let _ = stream.shutdown().await;
208            return Ok(Some(response(service, "tcp", &query.payload, &[])?));
209        }
210        Service::Echo => {
211            stream.shutdown().await?;
212            query.payload.len() as u64
213        }
214        Service::Chargen => query.read_bytes,
215        Service::Time => 4,
216        Service::Daytime | Service::Qotd => wire::MAX_STREAM_REPLY as u64,
217    };
218    let mut received = Vec::new();
219    let mut bounded = (&mut stream).take(limit);
220    tokio::time::timeout(wire::IO_TIMEOUT, bounded.read_to_end(&mut received))
221        .await
222        .context("reply deadline")??;
223    Ok(Some(response(service, "tcp", &query.payload, &received)?))
224}
225
226/// One UDP exchange; `None` when the reply never came.
227async fn over_udp(
228    service: Service,
229    remote: SocketAddr,
230    query: &Query,
231    injected: &mut Option<ClientCommand>,
232) -> Result<Option<Value>> {
233    let local: SocketAddr = if remote.is_ipv4() {
234        "0.0.0.0:0"
235    } else {
236        "[::]:0"
237    }
238    .parse()?;
239    let socket = UdpSocket::bind(local).await?;
240    socket.connect(remote).await?;
241    // RFC 864-868 ignore a request datagram's contents; one newline keeps it non-empty.
242    let request: &[u8] = if query.payload.is_empty() {
243        b"\n"
244    } else {
245        &query.payload
246    };
247    socket.send(request).await?;
248    answer(
249        injected,
250        Ok(ClientSendOutcome::Sent {
251            bytes_sent: request.len(),
252        }),
253    );
254    if service == Service::Discard {
255        return Ok(Some(response(service, "udp", &query.payload, &[])?));
256    }
257    let mut buf = vec![0u8; wire::READ_CHUNK];
258    match tokio::time::timeout(UDP_REPLY_TIMEOUT, socket.recv(&mut buf)).await {
259        Ok(Ok(n)) => Ok(Some(response(service, "udp", &query.payload, &buf[..n])?)),
260        Ok(Err(e)) => Err(e.into()),
261        Err(_) => Ok(None),
262    }
263}
264
265async fn session(
266    ctx: &ConnectContext,
267    service: Service,
268    remote: SocketAddr,
269    udp: bool,
270    mut external: mpsc::Receiver<ClientCommand>,
271    mut internal: mpsc::Receiver<Value>,
272    events: mpsc::Sender<Event>,
273) -> Result<()> {
274    loop {
275        let (action, mut injected) = tokio::select! {
276            command = external.recv() => match command {
277                Some(c) => (c.action.clone(), Some(c)),
278                None => return Ok(()),
279            },
280            action = internal.recv() => match action {
281                Some(a) => (a, None),
282                None => return Ok(()),
283            },
284        };
285        let query = match actions::execute(service, action.clone()) {
286            Ok(ClientActionResult::Disconnect) => {
287                answer(&mut injected, Ok(ClientSendOutcome::Disconnected));
288                return Ok(());
289            }
290            Ok(_) => actions::parse_query(service, &action),
291            Err(e) => Err(e),
292        };
293        let query = match query {
294            Ok(query) => query,
295            Err(e) => {
296                answer(
297                    &mut injected,
298                    Ok(ClientSendOutcome::Rejected {
299                        error: e.to_string(),
300                    }),
301                );
302                continue;
303            }
304        };
305        if injected.is_some() {
306            ctx.state
307                .record_access_log(
308                    AccessLogOwner::Client(ctx.client_id.as_u32()),
309                    service.name(),
310                    None,
311                    "injected_action",
312                    json!({"type": action["type"]}),
313                    vec![],
314                )
315                .await;
316        }
317        let result = if udp {
318            over_udp(service, remote, &query, &mut injected).await
319        } else {
320            over_tcp(service, remote, &query, &mut injected).await
321        };
322        match result {
323            Ok(Some(data)) => events
324                .try_send(Event::new(actions::response_event(service), data))
325                .context("event queue full; consumer stalled")?,
326            Ok(None) => Log::new(Some(&ctx.status_tx)).warn(format!(
327                "{} query over udp got no reply within {:?}",
328                service.name(),
329                UDP_REPLY_TIMEOUT
330            )),
331            Err(e) => {
332                answer(&mut injected, Err(anyhow::anyhow!(e.to_string())));
333                Log::new(Some(&ctx.status_tx))
334                    .warn(format!("{} query failed: {e:#}", service.name()));
335            }
336        }
337    }
338}