1pub mod actions;
2pub mod transport;
3use crate::llm::actions::client_trait::{Client, ClientActionResult};
4use crate::protocol::{ConnectContext, Event};
5use crate::server::sflow::codec::Batch;
6use crate::state::{
7 client_handles::{ClientCommand, ClientSendOutcome},
8 AccessLogOwner, ClientStatus,
9};
10use crate::{console_error, console_info};
11pub use actions::SflowClientProtocol;
12use anyhow::{Context, Result};
13use serde_json::{json, Value};
14use std::{collections::VecDeque, future::Future, pin::Pin, sync::Arc};
15
16pub const MAX_QUEUED_EVENTS: usize = 32;
17pub const MAX_HANDLER_ACTIONS: usize = 32;
18pub const MAX_FOLLOWUP_DEPTH: usize = 8;
19type Handler = Pin<Box<dyn Future<Output = Result<crate::llm::ClientLlmResult>> + Send>>;
20type SendFuture = Pin<Box<dyn Future<Output = Result<usize>> + Send>>;
21struct Inflight {
22 send: SendFuture,
23 prepared: transport::Prepared,
24 command: Option<ClientCommand>,
25 depth: usize,
26}
27fn handler(ctx: ConnectContext, event: Event) -> Handler {
28 Box::pin(async move {
29 let instruction = ctx
30 .state
31 .get_instruction_for_client(ctx.client_id)
32 .await
33 .unwrap_or_default();
34 let configured = ctx
35 .state
36 .get_client_event_handler_config(ctx.client_id)
37 .await
38 .is_some_and(|c| c.find_handler(event.id()).is_some());
39 if instruction.is_empty() && !configured {
40 ctx.state
41 .record_access_log(
42 AccessLogOwner::Client(ctx.client_id.as_u32()),
43 "sFlow",
44 None,
45 event.id(),
46 event.data.clone(),
47 vec![],
48 )
49 .await;
50 return Ok(crate::llm::ClientLlmResult {
51 actions: vec![],
52 memory_updates: None,
53 });
54 }
55 let memory = ctx
56 .state
57 .get_memory_for_client(ctx.client_id)
58 .await
59 .unwrap_or_default();
60 crate::client::llm_budget::call_llm_for_client(
61 &ctx.llm_client,
62 &ctx.state,
63 ctx.client_id.to_string(),
64 &instruction,
65 &memory,
66 Some(&event),
67 &SflowClientProtocol::new(),
68 &ctx.status_tx,
69 )
70 .await
71 })
72}
73pub struct SflowClient;
74impl SflowClient {
75 pub async fn connect(ctx: ConnectContext) -> Result<std::net::SocketAddr> {
76 anyhow::ensure!(
77 ctx.remote_addr.len() <= 1024,
78 "sFlow destination length bound"
79 );
80 let peer = tokio::time::timeout(
81 transport::IO_TIMEOUT,
82 tokio::net::lookup_host(&ctx.remote_addr),
83 )
84 .await
85 .context("sFlow resolve deadline")??
86 .next()
87 .context("collector resolved to no address")?;
88 let socket = Arc::new(
89 tokio::net::UdpSocket::bind(if peer.is_ipv4() {
90 "0.0.0.0:0"
91 } else {
92 "[::]:0"
93 })
94 .await?,
95 );
96 socket.connect(peer).await?;
97 let local = socket.local_addr()?;
98 let mut commands =
99 crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id)
100 .await;
101 ctx.state
102 .update_client_status(ctx.client_id, ClientStatus::Connected)
103 .await;
104 let owner = ctx.state.clone();
105 let id = ctx.client_id;
106 console_info!(
107 ctx.status_tx,
108 "sFlow UDP exporter {} ready for {}",
109 id,
110 peer
111 );
112 owner.spawn_client_task(id, async move {
113 let protocol = SflowClientProtocol::new();
114 let started = tokio::time::Instant::now();
115 let mut sequences = transport::Sequences::default();
116 let mut current = Some((
117 handler(ctx.clone(), Event::new(
118 &actions::SFLOW_CONNECTED_EVENT,
119 json!({"remote_addr":peer.to_string(),"local_addr":local.to_string()})
120 )), 0
121 ));
122 let mut events = VecDeque::new();
123 let mut actions = VecDeque::new();
124 let mut pending: Option<Inflight> = None;
125 let mut incoming = [0u8; 1];
126 'session: loop {
127 if current.is_none() {
128 if let Some((event, depth)) = events.pop_front() {
129 current = Some((handler(ctx.clone(), event), depth));
130 }
131 }
132 if pending.is_none() {
133 if let Some((action, depth)) = actions.pop_front() {
134 match prepare(
135 &protocol, action, depth, &sequences, started.elapsed().as_millis() as u32
136 ) {
137 Ok(Prepared::Write(p)) => pending = Some(inflight(socket.clone(), p, None, depth)),
138 Ok(Prepared::Disconnect) => break,
139 Err(e) => console_error!(ctx.status_tx, "sFlow decision=fail_closed_action_error error={}", e),
140 }
141 continue;
142 }
143 }
144 tokio::select! {
145 result = socket.recv(&mut incoming) => {
146 match result {
147 Ok(_) => console_error!(ctx.status_tx, "sFlow decision=fail_closed_unexpected_reply closing UDP exporter"),
148 Err(e) => console_error!(ctx.status_tx, "sFlow decision=fail_closed_receive_error error={}", e),
149 }
150 break;
151 },
152 result = async {
153 pending.as_mut().unwrap().send.as_mut().await
154 }, if pending.is_some() => {
155 let p = pending.take().unwrap();
156 match result {
157 Ok(byte_count) => {
158 let info = json!({"agent_address":p.prepared.agent,"sub_agent_id":p.prepared.sub_agent,
159 "sequence_number":p.prepared.sequence,"uptime_ms":p.prepared.uptime,"sample_count":p.prepared.sample_count,
160 "record_count":p.prepared.records,"byte_count":byte_count,"local_transport_only":true});
161 sequences.commit(p.prepared);
162 if let Some(command) = p.command {
163 finish(&ctx, command, Ok(ClientSendOutcome::Executed {
164 detail: "sFlow datagram accepted by local UDP transport".into()
165 })).await;
166 }
167 console_info!(ctx.status_tx, "sFlow decision=local_export local_transport_only=true byte_count={}", byte_count);
168 if events.len() == MAX_QUEUED_EVENTS {
169 console_error!(ctx.status_tx, "sFlow decision=fail_closed_event_capacity event_cap={}", MAX_QUEUED_EVENTS);
170 break;
171 }
172 events.push_back((Event::new(&actions::SFLOW_EXPORTED_EVENT, info), p.depth));
173 },
174 Err(e) => {
175 if let Some(command) = p.command {
176 finish(&ctx, command, Err(anyhow::anyhow!("sFlow UDP send failed"))).await;
177 }
178 console_error!(ctx.status_tx, "sFlow decision=fail_closed_send_error error={}", e);
179 break;
180 },
181 }
182 },
183 result = async {
184 current.as_mut().unwrap().0.as_mut().await
185 }, if current.is_some() => {
186 let (_, depth) = current.take().unwrap();
187 match result {
188 Ok(result) => {
189 if let Some(memory) = result.memory_updates {
190 ctx.state.set_memory_for_client(id, memory).await;
191 }
192 if result.actions.len() > MAX_HANDLER_ACTIONS
193 || actions.len() + result.actions.len() > MAX_HANDLER_ACTIONS
194 {
195 console_error!(ctx.status_tx, "sFlow decision=fail_closed_action_capacity action_cap={}", MAX_HANDLER_ACTIONS);
196 break;
197 }
198 actions.extend(result.actions.into_iter().map(|a| (a, depth + 1)));
199 },
200 Err(e) => console_error!(ctx.status_tx, "sFlow decision=fail_closed_dispatch_error error={}", e),
201 }
202 },
203 command = commands.recv() => {
204 let Some(command) = command else { break; };
205 match prepare(
206 &protocol, command.action.clone(), 0, &sequences, started.elapsed().as_millis() as u32
207 ) {
208 Ok(Prepared::Disconnect) => {
209 finish(&ctx, command, Ok(ClientSendOutcome::Disconnected)).await;
210 break 'session;
211 },
212 Ok(Prepared::Write(p)) if pending.is_none() => pending = Some(inflight(socket.clone(), p, Some(command), 0)),
213 Ok(Prepared::Write(_)) => finish(&ctx, command, Ok(ClientSendOutcome::Rejected {
214 error: "One UDP send is already in flight".into()
215 })).await,
216 Err(e) => finish(&ctx, command, Ok(ClientSendOutcome::Rejected { error: e.to_string() })).await,
217 }
218 },
219 }
220 let _ = ctx.status_tx.send("__UPDATE_UI__".into());
221 }
222 if let Some(p) = pending.take() {
223 if let Some(command) = p.command {
224 finish(&ctx, command, Err(anyhow::anyhow!("sFlow UDP send cancelled"))).await;
225 }
226 }
227 current.take();
228 events.clear();
229 actions.clear();
230 ctx.state.remove_client_handle(id).await;
231 ctx.state.update_client_status(id, ClientStatus::Disconnected).await;
232 let _ = ctx.status_tx.send("__UPDATE_UI__".into());
233 }).await;
234 Ok(local)
235 }
236}
237enum Prepared {
238 Disconnect,
239 Write(transport::Prepared),
240}
241fn prepare(
242 protocol: &SflowClientProtocol,
243 action: Value,
244 depth: usize,
245 sequences: &transport::Sequences,
246 uptime: u32,
247) -> Result<Prepared> {
248 match protocol.execute_action(action)? {
249 ClientActionResult::Disconnect => Ok(Prepared::Disconnect),
250 ClientActionResult::Custom { name, data } if name == "export_sflow_samples" => {
251 anyhow::ensure!(depth <= MAX_FOLLOWUP_DEPTH, "sFlow followup depth bound8");
252 let batch: Batch = serde_json::from_value(data)?;
253 Ok(Prepared::Write(sequences.prepare(&batch, uptime)?))
254 }
255 _ => anyhow::bail!("unsupported sFlow action result"),
256 }
257}
258fn inflight(
259 socket: Arc<tokio::net::UdpSocket>,
260 prepared: transport::Prepared,
261 command: Option<ClientCommand>,
262 depth: usize,
263) -> Inflight {
264 let bytes = prepared.bytes.clone();
265 Inflight {
266 send: Box::pin(async move {
267 let count = tokio::time::timeout(transport::IO_TIMEOUT, socket.send(&bytes))
268 .await
269 .context("sFlow UDP write deadline")??;
270 anyhow::ensure!(count == bytes.len(), "incomplete UDP datagram send");
271 Ok(count)
272 }),
273 prepared,
274 command,
275 depth,
276 }
277}
278async fn finish(ctx: &ConnectContext, command: ClientCommand, result: Result<ClientSendOutcome>) {
279 let outcome = match &result {
280 Ok(outcome) => serde_json::to_value(outcome).unwrap_or(Value::Null),
281 Err(e) => json!({"error":e.to_string()}),
282 };
283 ctx.state
284 .record_access_log(
285 AccessLogOwner::Client(ctx.client_id.as_u32()),
286 "sFlow",
287 None,
288 "injected_action",
289 command.action.clone(),
290 vec![outcome],
291 )
292 .await;
293 crate::client::command_support::reply(command, result);
294}