netget-tool 0.3.4

Network protocol stubbing toolkit: ~280 protocol servers and clients whose responses come from an LLM, a script, a static rule or a person.
Documentation
//! Receive/parse ownership stays independent of parked common handlers.
pub mod actions;
pub mod codec;
use crate::protocol::{Event, SpawnContext};
use crate::server::connection::ConnectionId;
use crate::state::{
    server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo},
    AccessLogOwner,
};
use crate::{console_error, console_info};
use anyhow::{ensure, Result};
use serde_json::json;
use std::{net::SocketAddr, time::Duration};
pub const MAX_QUEUED_EVENTS: usize = 32;
pub struct SflowServer;
pub fn duration(value: u64) -> Result<Duration> {
    ensure!(
        (1..=86400).contains(&value),
        "sFlow duration must be1..86400seconds"
    );
    Ok(Duration::from_secs(value))
}
impl SflowServer {
    pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
        let params = ctx.startup_params.as_ref();
        let idle = duration(
            params
                .map(|p| p.get_optional_u64("session_idle_seconds"))
                .transpose()?
                .flatten()
                .unwrap_or(codec::DEFAULT_SESSION_IDLE),
        )?;
        let llm = params
            .map(|p| p.get_optional_bool("llm_fallback"))
            .transpose()?
            .flatten()
            .unwrap_or(codec::DEFAULT_LLM_FALLBACK);
        let socket = tokio::net::UdpSocket::bind(ctx.legacy_listen_addr()).await?;
        let local = socket.local_addr()?;
        console_info!(
            ctx.status_tx,
            "sFlow UDP listening on {} (llm_fallback={})",
            local,
            llm
        );
        let (tx, mut rx) =
            tokio::sync::mpsc::channel::<(SocketAddr, usize, codec::Message)>(MAX_QUEUED_EVENTS);
        let dispatcher = ctx.clone();
        let owner = ctx.state.clone();
        let id = ctx.server_id;
        owner
            .spawn_server_task(id, async move {
                while let Some((peer, n, message)) = rx.recv().await {
                    let cid = ConnectionId::new(dispatcher.state.get_next_unified_id().await);
                    let now = crate::utils::clock::Instant::now();
                    dispatcher
                        .state
                        .add_connection_to_server(
                            id,
                            ConnectionState {
                                id: cid,
                                remote_addr: peer,
                                local_addr: local,
                                bytes_sent: 0,
                                bytes_received: n as u64,
                                packets_sent: 0,
                                packets_received: 1,
                                last_activity: now,
                                status: ConnectionStatus::Active,
                                status_changed_at: now,
                                protocol_info: ProtocolConnectionInfo::empty(),
                            },
                        )
                        .await;
                    let event = Event::new(
                        &actions::SFLOW_MESSAGE_EVENT,
                        json!({"message":message, "source_addr":peer.to_string()}),
                    );
                    let configured = dispatcher
                        .state
                        .get_event_handler_config(id)
                        .await
                        .is_some_and(|c| c.find_handler("sflow_message").is_some());
                    if llm || configured {
                        match crate::llm::action_helper::call_llm(
                            &dispatcher.llm_client,
                            &dispatcher.state,
                            id,
                            Some(cid),
                            &event,
                            &actions::SflowProtocol::new(),
                        )
                        .await
                        {
                            Ok(result) if result.failures.is_empty() => {
                                let decision = if result
                                    .protocol_result_actions
                                    .iter()
                                    .any(|action| action == "collect_sflow_samples")
                                {
                                    "handler_collect"
                                } else {
                                    "handler_silent"
                                };
                                for m in result.messages {
                                    console_info!(dispatcher.status_tx, "{}", m);
                                }
                                console_info!(
                                    dispatcher.status_tx,
                                    "sFlow source={} decision={} udp_silent=true",
                                    peer,
                                    decision
                                );
                                dispatcher
                                    .state
                                    .record_access_log(
                                        AccessLogOwner::Server(id.as_u32()),
                                        "sFlow",
                                        Some(cid.as_u32()),
                                        "sflow_handler_decision",
                                        event.data,
                                        vec![json!({"decision":decision,"udp_silent":true})],
                                    )
                                    .await;
                            }
                            result => {
                                let (decision, reason) = match result {
                                    Ok(r) => (
                                        "fail_closed_action_error",
                                        format!("{} failed actions", r.failures.len()),
                                    ),
                                    Err(e) => ("fail_closed_dispatch_error", e.to_string()),
                                };
                                console_error!(
                                    dispatcher.status_tx,
                                    "sFlow source={} decision={} udp_silent=true error={}",
                                    peer,
                                    decision,
                                    reason
                                );
                                dispatcher
                                    .state
                                    .record_access_log(
                                        AccessLogOwner::Server(id.as_u32()),
                                        "sFlow",
                                        Some(cid.as_u32()),
                                        "sflow_handler_failed",
                                        event.data,
                                        vec![json!({"decision":decision, "error":reason,"terminal_decision":decision,"udp_silent":true})],
                                    )
                                    .await;
                            }
                        }
                    } else {
                        console_info!(
                            dispatcher.status_tx,
                            "sFlow source={} decision=default_collect udp_silent=true",
                            peer
                        );
                        dispatcher
                            .state
                            .record_access_log(
                                AccessLogOwner::Server(id.as_u32()),
                                "sFlow",
                                Some(cid.as_u32()),
                                "sflow_message",
                                event.data,
                                vec![json!({"type":"collect_sflow_samples","decision":"default_collect","udp_silent":true})],
                            )
                            .await;
                    }
                    dispatcher
                        .state
                        .remove_connection_from_server(id, cid)
                        .await;
                    let _ = dispatcher.status_tx.send("__UPDATE_UI__".into());
                }
            })
            .await;
        owner
            .spawn_server_task(id, async move {
                let mut cache = codec::SequenceCache::new(idle);
                let mut buffer = vec![0; codec::MAX_MESSAGE_BYTES + 1];
                let mut expiry = tokio::time::interval(Duration::from_secs(1));
                expiry.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
                loop {
                    let (n, peer) = tokio::select! {
                        _ = expiry.tick() => {
                            cache.expire(tokio::time::Instant::now());
                            continue;
                        },
                        r = socket.recv_from(&mut buffer) => match r {
                            Ok(v) => v,
                            Err(e) => {
                                console_error!(ctx.status_tx, "sFlow decision=fail_closed_receive_error udp_silent=true error={}", e);
                                break;
                            }
                        }
                    };
                    match cache.ingest(peer, &buffer[..n], tokio::time::Instant::now()) {
                        Ok(message) => {
                            if let Err(error) = tx.try_send((peer, n, message)) {
                                let closed = matches!(
                                    error,
                                    tokio::sync::mpsc::error::TrySendError::Closed(_)
                                );
                                let decision = if closed {
                                    "fail_closed_dispatch_closed"
                                } else {
                                    "fail_closed_event_capacity"
                                };
                                console_error!(
                                    ctx.status_tx,
                                    "sFlow source={} decision={} udp_silent=true event_cap={}",
                                    peer,
                                    decision,
                                    MAX_QUEUED_EVENTS
                                );
                                ctx.state
                                    .record_access_log(
                                        AccessLogOwner::Server(id.as_u32()),
                                        "sFlow",
                                        None,
                                        "sflow_event_capacity",
                                        json!({"source_addr":peer.to_string()}),
                                        vec![json!({"decision":"fail_closed_event_capacity", "event_cap":MAX_QUEUED_EVENTS,"terminal_decision":decision,"udp_silent":true})],
                                    )
                                    .await;
                                if closed {
                                    break;
                                }
                            }
                        }
                        Err(error) => {
                            console_error!(
                                ctx.status_tx,
                                "sFlow source={} decision=fail_closed_invalid_datagram udp_silent=true received_bytes={} error={}",
                                peer,
                                n,
                                error
                            );
                            ctx.state
                                .record_access_log(
                                    AccessLogOwner::Server(id.as_u32()),
                                    "sFlow",
                                    None,
                                    "sflow_invalid_datagram",
                                    json!({"source_addr":peer.to_string(), "received_bytes":n}),
                                    vec![json!({"decision":"fail_closed_invalid_datagram", "error":error.to_string()})],
                                )
                                .await;
                        }
                    }
                }
            })
            .await;
        Ok(local)
    }
}