1use crate::{
4 llm::actions::{client_trait::Client, protocol_trait::Server},
5 protocol::{ConnectContext, Event, SpawnContext},
6 server::connection::ConnectionId,
7 state::{AccessLogOwner, ClientStatus},
8};
9use anyhow::{ensure, Result};
10use async_trait::async_trait;
11use serde_json::Value;
12use std::{net::SocketAddr, sync::Arc, time::Duration};
13use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
14use tokio::net::TcpStream;
15pub trait AsyncIo: AsyncRead + AsyncWrite + Unpin + Send {}
16impl<T: AsyncRead + AsyncWrite + Unpin + Send> AsyncIo for T {}
17pub type Stream = Box<dyn AsyncIo>;
18pub type ReadStream = tokio::io::ReadHalf<Stream>;
19
20pub const DEFAULT_USE_TLS: bool = false;
21
22pub const FIRST_BYTE_TIMEOUT: Duration = Duration::from_secs(30);
23pub const IDLE_TIMEOUT: Duration = Duration::from_secs(600);
24pub const RESPONSE_TIMEOUT: Duration = Duration::from_secs(10);
25
26pub fn parameter(
27 name: &str,
28 kind: &str,
29 description: &str,
30 required: bool,
31) -> crate::llm::actions::Parameter {
32 crate::llm::actions::Parameter {
33 name: name.into(),
34 type_hint: kind.into(),
35 description: description.into(),
36 required,
37 }
38}
39pub fn action(
40 name: &str,
41 description: &str,
42 parameters: Vec<crate::llm::actions::Parameter>,
43 example: Value,
44) -> crate::llm::actions::ActionDefinition {
45 crate::llm::actions::ActionDefinition {
46 name: name.into(),
47 description: description.into(),
48 parameters,
49 example,
50 log_template: Some(crate::protocol::LogTemplate::new().with_info(name)),
51 }
52}
53pub fn number(v: &Value, key: &str, max: u64) -> Result<u64> {
54 let n = v[key]
55 .as_u64()
56 .ok_or_else(|| anyhow::anyhow!("{key} must be an unsigned integer"))?;
57 ensure!(n <= max, "{key} exceeds {max}");
58 Ok(n)
59}
60pub async fn frame(
62 stream: &mut Stream,
63 header_len: usize,
64 max: usize,
65 length: fn(&[u8]) -> Result<usize>,
66) -> Result<Vec<u8>> {
67 let mut header = vec![0; header_len];
68 stream.read_exact(&mut header).await?;
69 let len = length(&header)?;
70 ensure!(len >= header_len && len <= max, "invalid frame length");
71 header.resize(len, 0);
72 stream.read_exact(&mut header[header_len..]).await?;
73 Ok(header)
74}
75#[async_trait]
76pub trait DeviceSession: Send + 'static {
77 fn greeting(&mut self) -> Result<Vec<u8>> {
78 Ok(vec![])
79 }
80 fn outgoing(&mut self) -> Option<tokio::sync::mpsc::Receiver<Vec<u8>>> {
81 None
82 }
83 async fn read(&mut self, stream: &mut ReadStream) -> Result<Vec<u8>>;
84 fn receive(&mut self, frame: &[u8]) -> Result<(Vec<u8>, Option<Value>)>;
86 fn answer(&mut self, request: &Value, action: Option<&Value>) -> Result<Vec<u8>>;
87}
88#[async_trait]
89pub trait ScannerSession: Send + 'static {
90 fn connected(&self) -> Value {
91 serde_json::json!({})
92 }
93 async fn open(&mut self, stream: &mut Stream) -> Result<()>;
94 async fn exchange(&mut self, stream: &mut Stream, action: &Value) -> Result<Value>;
95 fn outcome(
96 &self,
97 _action: &Value,
98 response: &Value,
99 ) -> crate::state::client_handles::ClientSendOutcome {
100 crate::state::client_handles::ClientSendOutcome::Executed {
101 detail: response.to_string(),
102 }
103 }
104 async fn close(&mut self, _stream: &mut Stream) -> Result<()> {
105 Ok(())
106 }
107 async fn idle(&mut self, stream: &mut Stream) -> Result<Option<Value>> {
108 let mut b = [0];
109 stream.read_exact(&mut b).await?;
110 anyhow::bail!("unexpected data or peer closed")
111 }
112}
113pub async fn spawn<S: DeviceSession, F: Fn() -> S + Send + Sync + Clone + 'static>(
114 ctx: SpawnContext,
115 protocol: Arc<dyn Server>,
116 factory: F,
117 kinds: &'static [crate::protocol::EventType],
118) -> Result<SocketAddr> {
119 let tls = server_tls(ctx.startup_params.as_ref())?;
120 let listener =
121 crate::server::socket_helpers::create_reusable_tcp_listener(ctx.legacy_listen_addr())
122 .await?;
123 spawn_accept_bounded(ctx, protocol, factory, kinds, listener, tls).await
124}
125pub async fn spawn_accept_bounded<
126 S: DeviceSession,
127 F: Fn() -> S + Send + Sync + Clone + 'static,
128>(
129 ctx: SpawnContext,
130 protocol: Arc<dyn Server>,
131 factory: F,
132 kinds: &'static [crate::protocol::EventType],
133 listener: tokio::net::TcpListener,
134 tls: Option<tokio_rustls::TlsAcceptor>,
135) -> Result<SocketAddr> {
136 let local = listener.local_addr()?;
137 let state = ctx.state.clone();
138 let sid = ctx.server_id;
139 let task = tokio::spawn(async move {
140 let cap = crate::server::accept_bounded::ConnectionLimiter::new(
141 crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS,
142 );
143 loop {
144 let Ok((stream, peer, permit)) = crate::server::accept_bounded::accept_bounded(
145 &listener,
146 &cap,
147 b"",
148 protocol.protocol_name(),
149 Some(&ctx.status_tx),
150 )
151 .await
152 else {
153 break;
154 };
155 let id = ConnectionId::new(ctx.state.get_next_unified_id().await);
156 let now = crate::utils::clock::Instant::now();
157 ctx.state
158 .add_connection_to_server(
159 sid,
160 crate::state::server::ConnectionState {
161 id,
162 remote_addr: peer,
163 local_addr: local,
164 bytes_sent: 0,
165 bytes_received: 0,
166 packets_sent: 0,
167 packets_received: 0,
168 last_activity: now,
169 status: crate::state::server::ConnectionStatus::Active,
170 status_changed_at: now,
171 protocol_info: crate::state::server::ProtocolConnectionInfo::empty(),
172 },
173 )
174 .await;
175 let mut commands =
176 crate::server::peer_support::register_peer_channel(&ctx.state, sid, id.as_u32())
177 .await;
178 let factory = factory.clone();
179 let tls = tls.clone();
180 let child = ctx.clone();
181 let p = protocol.clone();
182 ctx.state.spawn_server_task(sid, async move {
183 let _permit = permit;
184 let stream: Stream = if let Some(tls) = tls {
185 let handshake = tokio::time::timeout(RESPONSE_TIMEOUT,tls.accept(stream));
186 tokio::pin!(handshake);
187 let accepted = loop { tokio::select! {
188 result=&mut handshake => break result.ok().and_then(Result::ok),
189 Some(c)=commands.recv() => {
190 let close=c.action["type"]=="disconnect";
191 crate::client::command_support::reply(c, if close {Ok(crate::state::client_handles::ClientSendOutcome::Disconnected)} else {Ok(crate::state::client_handles::ClientSendOutcome::Rejected{error:"TLS handshake pending".into()})});
192 if close {break None;}
193 }
194 }};
195 let Some(stream)=accepted else {
196 child.state.remove_peer_handle(sid,id.as_u32()).await;
197 child.state.update_connection_status(sid,id,crate::state::server::ConnectionStatus::Closed).await;
198 return;
199 }; Box::new(stream)
200 } else {Box::new(stream)};
201 let mut session = factory(); let mut first = true;
202 let mut outgoing=session.outgoing();
203 let (mut read_stream,mut write_stream)=tokio::io::split(stream);
204 let result: Result<()> = async {
205 tokio::time::timeout(RESPONSE_TIMEOUT,write_stream.write_all(&session.greeting()?)).await??;
206 loop {
207 let input = {
208 let read=tokio::time::timeout(if first {FIRST_BYTE_TIMEOUT} else {IDLE_TIMEOUT},session.read(&mut read_stream));
209 tokio::pin!(read);
210 loop {
211 tokio::select! {
212 v=&mut read=>break Some(v??),
213 outbound=async {match outgoing.as_mut(){Some(rx)=>rx.recv().await,None=>std::future::pending().await}}=>{
214 if let Some(bytes)=outbound {tokio::time::timeout(RESPONSE_TIMEOUT,write_stream.write_all(&bytes)).await??;} else {outgoing=None;}
215 }
216 Some(command)=commands.recv()=>{
217 if command.action["type"]=="disconnect" {crate::client::command_support::reply(command,Ok(crate::state::client_handles::ClientSendOutcome::Disconnected));break None;}
218 crate::client::command_support::reply(command,Ok(crate::state::client_handles::ClientSendOutcome::Rejected{error:"Only disconnect is valid outside a pending request".into()}));
219 }
220 }
221 }
222 };
223 let Some(input)=input else {break;};
224 first = false;
225 child.state.update_connection_stats(sid, id, Some(input.len() as u64), None, Some(1), None).await;
226 let (automatic, request) = session.receive(&input)?;
227 let mut reply = automatic;
228 if let Some(request) = request {
229 let event = Event::new(&kinds[0], request.clone());
230 let turn=crate::llm::action_helper::call_llm(&child.llm_client,&child.state,sid,Some(id),&event,p.as_ref());tokio::pin!(turn);
231 let result=loop {tokio::select! {result=&mut turn=>break result,Some(command)=commands.recv()=>{if command.action["type"]=="disconnect" {crate::client::command_support::reply(command,Ok(crate::state::client_handles::ClientSendOutcome::Disconnected));anyhow::bail!("operator disconnected peer");}crate::client::command_support::reply(command,Ok(crate::state::client_handles::ClientSendOutcome::Rejected{error:"Reply to the pending request through its handler; only disconnect is an asynchronous action".into()}));}}};
232 let answer = result.as_ref().ok().and_then(|r| r.protocol_results.iter().find_map(|a| if let crate::llm::ActionResult::Custom {data,..}=a {Some(data)} else {None}));
233 if answer.is_none() { crate::logging::emit::Log::new(Some(&child.status_tx)).warn(format!("{} decision=fail_closed_no_action", p.protocol_name())); }
234 reply.extend(session.answer(&request, answer)?);
235 }
236 if !reply.is_empty() { tokio::time::timeout(RESPONSE_TIMEOUT,write_stream.write_all(&reply)).await??; child.state.update_connection_stats(sid, id, None, Some(reply.len() as u64), None, Some(1)).await; }
237 } Ok(())
238 }.await;
239 if let Err(e) = result { crate::logging::emit::Log::new(Some(&child.status_tx)).debug(format!("{} peer ended: {e}", p.protocol_name())); }
240 let _ = write_stream.shutdown().await;
241 child.state.remove_peer_handle(sid, id.as_u32()).await;
242 child.state.update_connection_status(sid, id, crate::state::server::ConnectionStatus::Closed).await;
243 }).await;
244 }
245 });
246 state.register_server_task(sid, task).await;
247 Ok(local)
248}
249pub async fn connect<S: ScannerSession>(
250 ctx: ConnectContext,
251 protocol: Arc<dyn Client>,
252 mut session: S,
253 kinds: [&'static crate::protocol::EventType; 2],
254) -> Result<SocketAddr> {
255 ensure!(!ctx.remote_addr.is_empty(), "remote_addr is required");
256 let tcp =
257 tokio::time::timeout(RESPONSE_TIMEOUT, TcpStream::connect(&ctx.remote_addr)).await??;
258 let local = tcp.local_addr()?;
259 let mut stream = client_transport(tcp, &ctx.remote_addr, ctx.startup_params.as_ref()).await?;
260 tokio::time::timeout(RESPONSE_TIMEOUT, session.open(&mut stream)).await??;
261 let mut commands =
262 crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id).await;
263 ctx.state
264 .update_client_status(ctx.client_id, ClientStatus::Connected)
265 .await;
266 let state = ctx.state.clone();
267 let cid = ctx.client_id;
268 state
269 .spawn_client_task(cid, async move {
270 let (events_tx, mut events_rx) = tokio::sync::mpsc::channel::<Event>(32);
272 let (actions_tx, mut actions_rx) = tokio::sync::mpsc::channel::<Value>(32);
273 let event_ctx = ctx.clone();
274 let p = protocol.clone();
275 let turns = ctx
276 .state
277 .spawn_client_task(cid, async move {
278 while let Some(event) = events_rx.recv().await {
279 let instruction = event_ctx
280 .state
281 .get_instruction_for_client(cid)
282 .await
283 .unwrap_or_default();
284 let memory = event_ctx
285 .state
286 .get_memory_for_client(cid)
287 .await
288 .unwrap_or_default();
289 match crate::client::llm_budget::call_llm_for_client(
290 &event_ctx.llm_client,
291 &event_ctx.state,
292 cid.to_string(),
293 &instruction,
294 &memory,
295 Some(&event),
296 p.as_ref(),
297 &event_ctx.status_tx,
298 )
299 .await
300 {
301 Ok(result) => {
302 if let Some(memory) = result.memory_updates {
303 event_ctx.state.set_memory_for_client(cid, memory).await;
304 }
305 for a in result.actions {
306 if actions_tx.send(a).await.is_err() {
307 return;
308 }
309 }
310 }
311 Err(e) => crate::logging::emit::Log::new(Some(&event_ctx.status_tx))
312 .warn(format!("{} client handler: {e}", p.protocol_name())),
313 }
314 }
315 })
316 .await;
317 let _ = events_tx.try_send(Event::new(
318 kinds[0],
319 {let mut connected=session.connected();if let Some(fields)=connected.as_object_mut(){fields.insert("remote_addr".into(),serde_json::json!(ctx.remote_addr));}connected},
320 ));
321 loop {
322 let (action, command) = tokio::select! {
323 Some(c) = commands.recv() => (c.action.clone(),Some(c)),
324 Some(a) = actions_rx.recv() => (a,None),
325 incoming = session.idle(&mut stream) => {match incoming {Ok(Some(response))=>{ctx.state.record_access_log(AccessLogOwner::Client(cid.as_u32()),protocol.protocol_name(),None,&kinds[1].id,response.clone(),vec![]).await;if events_tx.try_send(Event::new(kinds[1],response)).is_err(){break;}},Ok(None)=>{},Err(_)=>break};continue;},
326 else => break,
327 };
328 if action["type"] == "disconnect" {
329 if let Some(c) = command {
330 crate::client::command_support::reply(
331 c,
332 Ok(crate::state::client_handles::ClientSendOutcome::Disconnected),
333 );
334 }
335 break;
336 }
337 match protocol.execute_action(action.clone()) {
338 Err(e)=>{if let Some(c)=command{crate::client::command_support::reply(c,Ok(crate::state::client_handles::ClientSendOutcome::Rejected{error:e.to_string()}));}continue;}
339 Ok(crate::llm::actions::client_trait::ClientActionResult::Disconnect)=>{let _=tokio::time::timeout(RESPONSE_TIMEOUT,session.close(&mut stream)).await;if let Some(c)=command{crate::client::command_support::reply(c,Ok(crate::state::client_handles::ClientSendOutcome::Disconnected));}break;}
340 Ok(crate::llm::actions::client_trait::ClientActionResult::WaitForMore)=>{if let Some(c)=command{crate::client::command_support::reply(c,Ok(crate::state::client_handles::ClientSendOutcome::Executed{detail:"Waiting for peer".into()}));}continue;}
341 Ok(_)=>{}
342 }
343 if command.is_some(){ctx.state.record_access_log(AccessLogOwner::Client(cid.as_u32()),protocol.protocol_name(),None,"injected_action",action.clone(),vec![]).await;}
344 let result = match tokio::time::timeout(
345 RESPONSE_TIMEOUT,
346 session.exchange(&mut stream, &action),
347 )
348 .await
349 {
350 Ok(r) => r,
351 Err(_) => Err(anyhow::anyhow!("device response deadline exceeded")),
352 };
353 let terminal = result.is_err();
354 match result {
355 Ok(response) => {
356 ctx.state
357 .record_access_log(
358 AccessLogOwner::Client(cid.as_u32()),
359 protocol.protocol_name(),
360 None,
361 &kinds[1].id,
362 response.clone(),
363 vec![],
364 )
365 .await;
366 let _ = events_tx.try_send(Event::new(kinds[1], response.clone()));
367 if let Some(c) = command {
368 crate::client::command_support::reply(
369 c,
370 Ok(session.outcome(&action,&response)),
371 );
372 }
373 }
374 Err(e) => {
375 if let Some(c) = command {
376 crate::client::command_support::reply(c, Err(e));
377 }
378 }
379 }
380 if terminal {
381 break;
382 }
383 }
384 turns.abort();
385 let _ = stream.shutdown().await;
386 ctx.state.remove_client_handle(cid).await;
387 ctx.state
388 .update_client_status(cid, ClientStatus::Disconnected)
389 .await;
390 })
391 .await;
392 Ok(local)
393}
394
395pub fn tls_parameters(server: bool) -> Vec<crate::llm::actions::ParameterDefinition> {
396 let mut fields = vec![crate::llm::actions::ParameterDefinition {
397 name: "use_tls".into(),
398 type_hint: "boolean".into(),
399 description: "Use implicit TLS (ADCS/NNTPS); TLS 1.2+ with certificate validation".into(),
400 required: false,
401 example: serde_json::json!(true),
402 default: Some(serde_json::json!(DEFAULT_USE_TLS)),
403 }];
404 for (name, description) in if server {
405 vec![
406 (
407 "cert_path",
408 "PEM server certificate chain; required with use_tls",
409 ),
410 ("key_path", "PEM private key; required with use_tls"),
411 ]
412 } else {
413 vec![
414 (
415 "ca_path",
416 "Optional PEM trust anchors, added to public roots",
417 ),
418 (
419 "server_name",
420 "Optional expected TLS DNS name; defaults to remote hostname",
421 ),
422 ]
423 } {
424 fields.push(crate::llm::actions::ParameterDefinition {
425 name: name.into(),
426 type_hint: "string".into(),
427 description: description.into(),
428 required: false,
429 example: serde_json::json!(""),
430 default: None,
431 });
432 }
433 fields
434}
435fn optional(params: Option<&crate::protocol::StartupParams>, key: &str) -> Result<Option<String>> {
436 Ok(params
437 .map(|p| p.get_optional_string(key))
438 .transpose()?
439 .flatten())
440}
441pub fn use_tls(params: Option<&crate::protocol::StartupParams>) -> Result<bool> {
442 Ok(params
443 .map(|p| p.get_optional_bool("use_tls"))
444 .transpose()?
445 .flatten()
446 .unwrap_or(DEFAULT_USE_TLS))
447}
448pub fn server_tls(
449 params: Option<&crate::protocol::StartupParams>,
450) -> Result<Option<tokio_rustls::TlsAcceptor>> {
451 use std::{fs::File, io::BufReader};
452 if !use_tls(params)? {
453 return Ok(None);
454 }
455 let cert =
456 optional(params, "cert_path")?.ok_or_else(|| anyhow::anyhow!("TLS requires cert_path"))?;
457 let key =
458 optional(params, "key_path")?.ok_or_else(|| anyhow::anyhow!("TLS requires key_path"))?;
459 let certs = rustls_pemfile::certs(&mut BufReader::new(File::open(cert)?))
460 .collect::<std::io::Result<Vec<_>>>()?;
461 let key = rustls_pemfile::private_key(&mut BufReader::new(File::open(key)?))?
462 .ok_or_else(|| anyhow::anyhow!("PEM key missing"))?;
463 let _ = rustls::crypto::ring::default_provider().install_default();
464 let config = rustls::ServerConfig::builder()
465 .with_no_client_auth()
466 .with_single_cert(certs, key)?;
467 Ok(Some(tokio_rustls::TlsAcceptor::from(Arc::new(config))))
468}
469pub async fn client_transport(
470 tcp: TcpStream,
471 remote: &str,
472 params: Option<&crate::protocol::StartupParams>,
473) -> Result<Stream> {
474 use std::{fs::File, io::BufReader};
475 if !use_tls(params)? {
476 return Ok(Box::new(tcp));
477 }
478 let _ = rustls::crypto::ring::default_provider().install_default();
479 let mut roots = rustls::RootCertStore {
480 roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(),
481 };
482 if let Some(path) = optional(params, "ca_path")? {
483 for cert in rustls_pemfile::certs(&mut BufReader::new(File::open(path)?)) {
484 roots.add(cert?)?;
485 }
486 }
487 let name = optional(params, "server_name")?.unwrap_or_else(|| {
488 remote
489 .rsplit_once(':')
490 .map_or(remote, |p| p.0)
491 .trim_matches(['[', ']'])
492 .to_string()
493 });
494 let name = rustls::pki_types::ServerName::try_from(name)?;
495 let config = rustls::ClientConfig::builder()
496 .with_root_certificates(roots)
497 .with_no_client_auth();
498 let stream = tokio::time::timeout(
499 RESPONSE_TIMEOUT,
500 tokio_rustls::TlsConnector::from(Arc::new(config)).connect(name, tcp),
501 )
502 .await??;
503 Ok(Box::new(stream))
504}
505
506#[derive(Default)]
508pub struct Framer {
509 pending: Vec<u8>,
510 deadline: Option<tokio::time::Instant>,
511}
512impl Framer {
513 async fn byte(&mut self, stream: &mut (impl AsyncRead + Unpin)) -> Result<u8> {
514 let deadline = self
515 .deadline
516 .unwrap_or_else(|| tokio::time::Instant::now() + IDLE_TIMEOUT);
517 let byte = tokio::time::timeout_at(deadline, stream.read_u8()).await??;
518 self.deadline
519 .get_or_insert_with(|| tokio::time::Instant::now() + RESPONSE_TIMEOUT);
520 Ok(byte)
521 }
522 pub async fn delimited(
523 &mut self,
524 stream: &mut (impl AsyncRead + Unpin),
525 delimiter: u8,
526 max: usize,
527 ) -> Result<Vec<u8>> {
528 loop {
529 ensure!(self.pending.len() < max, "frame exceeds cap");
530 let byte = self.byte(stream).await?;
531 self.pending.push(byte);
532 if byte == delimiter {
533 self.deadline = None;
534 return Ok(std::mem::take(&mut self.pending));
535 }
536 }
537 }
538 pub async fn sized(
539 &mut self,
540 stream: &mut (impl AsyncRead + Unpin),
541 header: usize,
542 max: usize,
543 length: fn(&[u8]) -> Result<usize>,
544 ) -> Result<Vec<u8>> {
545 loop {
546 let wanted = if self.pending.len() < header {
547 header
548 } else {
549 length(&self.pending[..header])?
550 };
551 ensure!((header..=max).contains(&wanted), "invalid frame size");
552 if self.pending.len() == wanted {
553 self.deadline = None;
554 return Ok(std::mem::take(&mut self.pending));
555 }
556 let byte = self.byte(stream).await?;
557 self.pending.push(byte);
558 }
559 }
560}
561
562pub fn base32(bytes: &[u8]) -> String {
563 const ALPHABET: &[u8; 32] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";
564 let mut value = 0u32;
565 let mut bits = 0usize;
566 let mut text = String::new();
567 for &byte in bytes {
568 value = (value << 8) | u32::from(byte);
569 bits += 8;
570 while bits >= 5 {
571 bits -= 5;
572 text.push(ALPHABET[((value >> bits) & 31) as usize] as char);
573 }
574 }
575 if bits > 0 {
576 text.push(ALPHABET[((value << (5 - bits)) & 31) as usize] as char);
577 }
578 text
579}
580
581#[cfg(not(target_arch = "wasm32"))]
582pub mod files;