1pub 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";
23pub 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
149fn 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
184async 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
226async 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 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}