1pub mod actions;
2pub mod transport;
3use crate::llm::actions::client_trait::{Client, ClientActionResult};
4use crate::protocol::{ConnectContext, Event};
5use crate::server::ipfix::codec::Batch;
6use crate::state::{
7 client_handles::{ClientCommand, ClientSendOutcome},
8 AccessLogOwner, ClientStatus,
9};
10use crate::{console_error, console_info};
11pub use actions::IpfixClientProtocol;
12use anyhow::{Context, Result};
13use serde_json::{json, Value};
14use std::{collections::VecDeque, future::Future, pin::Pin, sync::Arc, time::Duration};
15pub const MAX_QUEUED_EVENTS: usize = 32;
16pub const MAX_HANDLER_ACTIONS: usize = 32;
17pub const MAX_FOLLOWUP_DEPTH: usize = 8;
18type Handler = Pin<Box<dyn Future<Output = Result<crate::llm::ClientLlmResult>> + Send>>;
19type SendFuture = Pin<Box<dyn Future<Output = Result<usize>> + Send>>;
20struct Inflight {
21 send: SendFuture,
22 prepared: Option<transport::Prepared>,
23 command: Option<ClientCommand>,
24 depth: usize,
25}
26fn handler(ctx: ConnectContext, event: Event) -> Handler {
27 Box::pin(async move {
28 let instruction = ctx
29 .state
30 .get_instruction_for_client(ctx.client_id)
31 .await
32 .unwrap_or_default();
33 let configured = ctx
34 .state
35 .get_client_event_handler_config(ctx.client_id)
36 .await
37 .is_some_and(|c| c.find_handler(event.id()).is_some());
38 if instruction.is_empty() && !configured {
39 ctx.state
40 .record_access_log(
41 AccessLogOwner::Client(ctx.client_id.as_u32()),
42 "IPFIX",
43 None,
44 event.id(),
45 event.data.clone(),
46 vec![],
47 )
48 .await;
49 return Ok(crate::llm::ClientLlmResult {
50 actions: vec![],
51 memory_updates: None,
52 });
53 }
54 let memory = ctx
55 .state
56 .get_memory_for_client(ctx.client_id)
57 .await
58 .unwrap_or_default();
59 crate::client::llm_budget::call_llm_for_client(
60 &ctx.llm_client,
61 &ctx.state,
62 ctx.client_id.to_string(),
63 &instruction,
64 &memory,
65 Some(&event),
66 &IpfixClientProtocol::new(),
67 &ctx.status_tx,
68 )
69 .await
70 })
71}
72pub struct IpfixClient;
73impl IpfixClient {
74 pub async fn connect(ctx: ConnectContext) -> Result<std::net::SocketAddr> {
75 anyhow::ensure!(
76 ctx.remote_addr.len() <= 1024,
77 "IPFIX destination length bound"
78 );
79 let refresh = ctx
80 .startup_params
81 .as_ref()
82 .map(|p| p.get_optional_u64("template_refresh_seconds"))
83 .transpose()?
84 .flatten()
85 .unwrap_or(transport::DEFAULT_REFRESH_SECONDS);
86 anyhow::ensure!(
87 (1..=3600).contains(&refresh),
88 "template refresh must be1..3600seconds"
89 );
90 let peer = tokio::time::timeout(
91 transport::IO_TIMEOUT,
92 tokio::net::lookup_host(&ctx.remote_addr),
93 )
94 .await
95 .context("IPFIX resolve deadline")??
96 .next()
97 .context("collector resolved to no address")?;
98 let socket = Arc::new(
99 tokio::net::UdpSocket::bind(if peer.is_ipv4() {
100 "0.0.0.0:0"
101 } else {
102 "[::]:0"
103 })
104 .await?,
105 );
106 socket.connect(peer).await?;
107 let local = socket.local_addr()?;
108 let mut commands =
109 crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id)
110 .await;
111 ctx.state
112 .update_client_status(ctx.client_id, ClientStatus::Connected)
113 .await;
114 let owner = ctx.state.clone();
115 let id = ctx.client_id;
116 console_info!(
117 ctx.status_tx,
118 "IPFIX UDP exporter {} ready for {}",
119 id,
120 peer
121 );
122 owner
123 .spawn_client_task(id, async move {
124 let protocol = IpfixClientProtocol::new();
125 let mut catalog = transport::Catalog::default();
126 let mut current = Some((
127 handler(ctx.clone(), Event::new(
128 &actions::IPFIX_CONNECTED_EVENT,
129 json!({"remote_addr": peer.to_string(), "local_addr": local.to_string()})
130 )),
131 0
132 ));
133 let mut events = VecDeque::new();
134 let mut actions = VecDeque::new();
135 let mut pending: Option<Inflight> = None;
136 let period = Duration::from_secs(refresh);
137 let mut refresh_tick = tokio::time::interval_at(
138 tokio::time::Instant::now() + period, period
139 );
140 refresh_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
141 let mut incoming = [0u8; 1];
142 'session: loop {
143 if current.is_none() {
144 if let Some((event, depth)) = events.pop_front() {
145 current = Some((handler(ctx.clone(), event), depth));
146 }
147 }
148 if pending.is_none() {
149 if let Some((action, depth)) = actions.pop_front() {
150 match prepare(&protocol, action, depth, &catalog) {
151 Ok(Prepared::Write(p)) => {
152 pending = Some(inflight(socket.clone(), p, None, depth));
153 },
154 Ok(Prepared::Disconnect) => break,
155 Err(e) => console_error!(
156 ctx.status_tx, "IPFIX handler action rejected: {}", e
157 )
158 }
159 continue;
160 }
161 }
162 tokio::select! {
163 result = socket.recv(&mut incoming) => {
164 match result {
165 Ok(_) => console_error!(
166 ctx.status_tx,
167 "IPFIX collector sent unexpected UDP reply; closing"
168 ),
169 Err(e) => console_error!(
170 ctx.status_tx, "IPFIX UDP transport failed: {}", e
171 )
172 }
173 break;
174 },
175 result = async {
176 pending.as_mut().unwrap().send.as_mut().await
177 }, if pending.is_some() => {
178 let p = pending.take().unwrap();
179 match result {
180 Ok(count) => if let Some(prepared) = p.prepared {
181 catalog.domains.insert(
182 prepared.info.observation_domain_id, prepared.domain
183 );
184 if let Some(command) = p.command {
185 finish(&ctx, command, Ok(ClientSendOutcome::Executed {
186 detail: format!(
187 "IPFIX local UDP send: {} records",
188 prepared.info.record_count
189 )
190 })).await;
191 }
192 if events.len() == MAX_QUEUED_EVENTS {
193 console_error!(ctx.status_tx, "IPFIX export-event queue limit");
194 break;
195 }
196 events.push_back((Event::new(
197 &actions::IPFIX_EXPORTED_EVENT,
198 serde_json::to_value(prepared.info).unwrap_or(Value::Null)
199 ), p.depth));
200 } else {
201 ctx.state.record_access_log(
202 AccessLogOwner::Client(id.as_u32()),
203 "IPFIX",
204 None,
205 "ipfix_template_refresh",
206 json!({"message_count": count, "local_transport_only": true}),
207 vec![]
208 ).await;
209 },
210 Err(e) => {
211 if let Some(command) = p.command {
212 finish(&ctx, command, Err(anyhow::anyhow!(
213 "IPFIX UDP send failed"
214 ))).await;
215 }
216 console_error!(ctx.status_tx, "IPFIX UDP send failed: {}", e);
217 break;
218 }
219 }
220 },
221 result = async {
222 current.as_mut().unwrap().0.as_mut().await
223 }, if current.is_some() => {
224 let (_, depth) = current.take().unwrap();
225 match result {
226 Ok(result) => {
227 if let Some(memory) = result.memory_updates {
228 ctx.state.set_memory_for_client(id, memory).await;
229 }
230 if result.actions.len() > MAX_HANDLER_ACTIONS
231 || actions.len() + result.actions.len() > MAX_HANDLER_ACTIONS
232 {
233 console_error!(ctx.status_tx, "IPFIX handler action count limit");
234 break;
235 }
236 actions.extend(result.actions.into_iter().map(|a| (a, depth + 1)));
237 },
238 Err(e) => console_error!(ctx.status_tx, "IPFIX handler failed: {}", e)
239 }
240 },
241 command = commands.recv() => {
242 let Some(command) = command else { break; };
243 match prepare(&protocol, command.action.clone(), 0, &catalog) {
244 Ok(Prepared::Disconnect) => {
245 finish(&ctx, command, Ok(ClientSendOutcome::Disconnected)).await;
246 break 'session;
247 },
248 Ok(Prepared::Write(p)) if pending.is_none() =>
249 pending = Some(inflight(socket.clone(), p, Some(command), 0)),
250 Ok(Prepared::Write(_)) => finish(
251 &ctx, command, Ok(ClientSendOutcome::Rejected {
252 error: "One UDP send is already in flight".into()
253 })
254 ).await,
255 Err(e) => finish(
256 &ctx, command, Ok(ClientSendOutcome::Rejected { error: e.to_string() })
257 ).await
258 }
259 },
260 _ = refresh_tick.tick(), if pending.is_none() && !catalog.domains.is_empty() => {
261 match catalog.refresh() {
262 Ok(messages) => pending = Some(Inflight {
263 send: Box::pin(transport::send(socket.clone(), messages)),
264 prepared: None,
265 command: None,
266 depth: 0
267 }),
268 Err(e) => {
269 console_error!(ctx.status_tx, "IPFIX template refresh failed: {}", e);
270 break;
271 }
272 }
273 },
274 }
275 let _ = ctx.status_tx.send("__UPDATE_UI__".into());
276 }
277 if let Some(p) = pending.take() {
278 if let Some(command) = p.command {
279 finish(&ctx, command, Err(anyhow::anyhow!("IPFIX UDP send cancelled"))).await;
280 }
281 }
282 current.take();
283 events.clear();
284 actions.clear();
285 ctx.state.remove_client_handle(id).await;
286 ctx.state.update_client_status(id, ClientStatus::Disconnected).await;
287 let _ = ctx.status_tx.send("__UPDATE_UI__".into());
288 }).await;
289 Ok(local)
290 }
291}
292enum Prepared {
293 Disconnect,
294 Write(transport::Prepared),
295}
296fn prepare(
297 protocol: &IpfixClientProtocol,
298 action: Value,
299 depth: usize,
300 catalog: &transport::Catalog,
301) -> Result<Prepared> {
302 match protocol.execute_action(action)? {
303 ClientActionResult::Disconnect => Ok(Prepared::Disconnect),
304 ClientActionResult::Custom { name, data } if name == "export_ipfix_records" => {
305 anyhow::ensure!(depth <= MAX_FOLLOWUP_DEPTH, "IPFIX followup depth bound8");
306 let batch: Batch = serde_json::from_value(data)?;
307 Ok(Prepared::Write(catalog.prepare(&batch)?))
308 }
309 _ => anyhow::bail!("unsupported IPFIX action result"),
310 }
311}
312fn inflight(
313 socket: Arc<tokio::net::UdpSocket>,
314 p: transport::Prepared,
315 command: Option<ClientCommand>,
316 depth: usize,
317) -> Inflight {
318 let bytes = p.bytes.clone();
319 Inflight {
320 send: Box::pin(transport::send(socket, vec![bytes])),
321 prepared: Some(p),
322 command,
323 depth,
324 }
325}
326async fn finish(ctx: &ConnectContext, command: ClientCommand, result: Result<ClientSendOutcome>) {
327 let outcome = match &result {
328 Ok(outcome) => serde_json::to_value(outcome).unwrap_or(Value::Null),
329 Err(e) => json!({"error":e.to_string()}),
330 };
331 ctx.state
332 .record_access_log(
333 AccessLogOwner::Client(ctx.client_id.as_u32()),
334 "IPFIX",
335 None,
336 "injected_action",
337 command.action.clone(),
338 vec![outcome],
339 )
340 .await;
341 crate::client::command_support::reply(command, result);
342}