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