pub mod actions;
pub mod wire;
pub use actions::NetbiosNsClientProtocol;
use crate::utils::clock::Instant;
use std::collections::VecDeque;
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use std::time::Duration;
use anyhow::{Context, Result};
use tokio::net::UdpSocket;
use tokio::sync::{mpsc, Mutex};
use tracing::{debug, error, info, trace, warn};
use crate::client::llm_budget::call_llm_for_client;
use crate::llm::actions::client_trait::{Client, ClientActionResult};
use crate::llm::actions::protocol_trait::Protocol;
use crate::llm::ollama_client::OllamaClient;
use crate::llm::ClientLlmResult;
use crate::protocol::{ConnectContext, Event, EventType};
use crate::server::netbios_ns::packet as pkt;
use crate::state::app_state::AppState;
use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
use crate::state::{AccessLogOwner, ClientId, ClientStatus};
use self::actions::{
NETBIOS_NAME_RESPONSE_EVENT, NETBIOS_NEGATIVE_RESPONSE_EVENT,
NETBIOS_NODE_STATUS_RESPONSE_EVENT, NETBIOS_NS_CONNECTED_EVENT, NETBIOS_QUERY_TIMEOUT_EVENT,
};
const DEFAULT_NBNS_PORT: u16 = 137;
const DEFAULT_QUERY_TIMEOUT_SECS: u64 = 3;
const MAX_FOLLOWUP_DEPTH: usize = 6;
struct Transport {
socket: Mutex<UdpSocket>,
default_target: SocketAddr,
query_timeout: Duration,
}
enum Applied {
Answered {
event_type: &'static EventType,
event_data: serde_json::Value,
detail: String,
},
Nothing(String),
Disconnect,
}
enum QueryOutcome {
Reply {
response: wire::NbnsResponse,
from: SocketAddr,
},
TimedOut { ignored: usize },
}
pub struct NetbiosNsClient;
impl NetbiosNsClient {
pub async fn connect_with_llm_actions(ctx: ConnectContext) -> Result<SocketAddr> {
let ConnectContext {
remote_addr,
llm_client,
state: app_state,
status_tx,
client_id,
startup_params,
} = ctx;
let port_override = match &startup_params {
Some(params) => params.get_optional_u32("port")?,
None => None,
};
let port_override = match port_override {
Some(p) => Some(u16::try_from(p).map_err(|_| {
anyhow::anyhow!("startup parameter 'port' must be between 0 and 65535, got {p}")
})?),
None => None,
};
let query_timeout_secs = match &startup_params {
Some(params) => params
.get_optional_u64("query_timeout_secs")?
.unwrap_or(DEFAULT_QUERY_TIMEOUT_SECS),
None => DEFAULT_QUERY_TIMEOUT_SECS,
};
if query_timeout_secs == 0 {
anyhow::bail!(
"startup parameter 'query_timeout_secs' must be at least 1: a zero deadline \
discards every reply before it can arrive"
);
}
let default_target = resolve_target(&remote_addr, port_override).await?;
let socket = UdpSocket::bind(("0.0.0.0", 0))
.await
.context("failed to bind the NetBIOS-NS client's UDP socket")?;
let local_addr = socket
.local_addr()
.context("bound NetBIOS-NS socket has no local address")?;
info!(
"NetBIOS-NS client {} querying from {} to {}",
client_id, local_addr, default_target
);
app_state
.update_client_status(client_id, ClientStatus::Connected)
.await;
let _ = status_tx.send(format!(
"[CLIENT] NetBIOS-NS client {} ready ({} -> {})",
client_id, local_addr, default_target
));
let _ = status_tx.send("__UPDATE_UI__".to_string());
let transport = Arc::new(Transport {
socket: Mutex::new(socket),
default_target,
query_timeout: Duration::from_secs(query_timeout_secs),
});
let protocol = Arc::new(NetbiosNsClientProtocol::new());
let command_rx =
crate::client::command_support::register_command_channel(&app_state, client_id).await;
let cmd_task = tokio::spawn(Self::command_loop(
command_rx,
protocol.clone(),
transport.clone(),
client_id,
app_state.clone(),
llm_client.clone(),
status_tx.clone(),
));
app_state.register_client_task(client_id, cmd_task).await;
let conversation = tokio::spawn(Self::conversation(
protocol.clone(),
transport.clone(),
client_id,
app_state.clone(),
llm_client,
status_tx,
default_target,
local_addr,
));
app_state
.register_client_task(client_id, conversation)
.await;
Ok(local_addr)
}
#[allow(clippy::too_many_arguments)]
async fn conversation(
protocol: Arc<NetbiosNsClientProtocol>,
transport: Arc<Transport>,
client_id: ClientId,
app_state: Arc<AppState>,
llm_client: OllamaClient,
status_tx: mpsc::UnboundedSender<String>,
default_target: SocketAddr,
local_addr: SocketAddr,
) {
let Some(instruction) = app_state.get_instruction_for_client(client_id).await else {
debug!(
"NetBIOS-NS client {} has no instruction; nothing to drive",
client_id
);
return;
};
let event = Event::new(
&NETBIOS_NS_CONNECTED_EVENT,
serde_json::json!({
"remote_addr": default_target.to_string(),
"local_addr": local_addr.to_string(),
}),
);
let memory = app_state
.get_memory_for_client(client_id)
.await
.unwrap_or_default();
match call_llm_for_client(
&llm_client,
&app_state,
client_id.to_string(),
&instruction,
&memory,
Some(&event),
protocol.as_ref(),
&status_tx,
)
.await
{
Ok(ClientLlmResult {
actions,
memory_updates,
}) => {
if let Some(mem) = memory_updates {
app_state.set_memory_for_client(client_id, mem).await;
}
Self::run_actions(
&protocol,
&transport,
actions,
0,
client_id,
&app_state,
&llm_client,
&status_tx,
)
.await;
}
Err(e) => {
error!("LLM error for NetBIOS-NS client {}: {}", client_id, e);
let _ = status_tx.send(format!(
"[CLIENT] ✖ NetBIOS-NS client {} could not start: {}",
client_id, e
));
}
}
debug!("NetBIOS-NS client {} conversation task finished", client_id);
}
#[allow(clippy::too_many_arguments)]
async fn run_actions(
protocol: &Arc<NetbiosNsClientProtocol>,
transport: &Arc<Transport>,
initial_actions: Vec<serde_json::Value>,
depth: usize,
client_id: ClientId,
app_state: &Arc<AppState>,
llm_client: &OllamaClient,
status_tx: &mpsc::UnboundedSender<String>,
) {
let mut pending: VecDeque<(serde_json::Value, usize)> =
initial_actions.into_iter().map(|a| (a, depth)).collect();
while let Some((action, depth)) = pending.pop_front() {
let applied =
match Self::apply_action(protocol, transport, action, client_id, status_tx).await {
Ok(applied) => applied,
Err(e) => {
error!("NetBIOS-NS client {} action error: {}", client_id, e);
let _ = status_tx.send(format!(
"[WARN] NetBIOS-NS client {} action failed: {}",
client_id, e
));
continue;
}
};
match applied {
Applied::Disconnect => {
Self::finish(client_id, app_state, status_tx).await;
return;
}
Applied::Nothing(detail) => {
debug!("NetBIOS-NS client {}: {}", client_id, detail);
}
Applied::Answered {
event_type,
event_data,
..
} => {
if depth >= MAX_FOLLOWUP_DEPTH {
warn!(
"NetBIOS-NS client {} reached the follow-up depth cap ({}); the \
'{}' answer is not being handed back to the model",
client_id, MAX_FOLLOWUP_DEPTH, event_type.id
);
let _ = status_tx.send(format!(
"[CLIENT] ⚠ NetBIOS-NS client {} stopped following up at depth {}",
client_id, MAX_FOLLOWUP_DEPTH
));
continue;
}
let follow_ups = Self::report(
event_type, event_data, protocol, client_id, app_state, llm_client,
status_tx,
)
.await;
pending.extend(follow_ups.into_iter().map(|a| (a, depth + 1)));
}
}
}
}
async fn report(
event_type: &'static EventType,
event_data: serde_json::Value,
protocol: &Arc<NetbiosNsClientProtocol>,
client_id: ClientId,
app_state: &Arc<AppState>,
llm_client: &OllamaClient,
status_tx: &mpsc::UnboundedSender<String>,
) -> Vec<serde_json::Value> {
let Some(instruction) = app_state.get_instruction_for_client(client_id).await else {
return Vec::new();
};
let event = Event::new(event_type, event_data);
let memory = app_state
.get_memory_for_client(client_id)
.await
.unwrap_or_default();
match call_llm_for_client(
llm_client,
app_state,
client_id.to_string(),
&instruction,
&memory,
Some(&event),
protocol.as_ref(),
status_tx,
)
.await
{
Ok(ClientLlmResult {
actions,
memory_updates,
}) => {
if let Some(mem) = memory_updates {
app_state.set_memory_for_client(client_id, mem).await;
}
actions
}
Err(e) => {
error!("LLM error for NetBIOS-NS client {}: {}", client_id, e);
Vec::new()
}
}
}
async fn apply_action(
protocol: &Arc<NetbiosNsClientProtocol>,
transport: &Arc<Transport>,
action: serde_json::Value,
client_id: ClientId,
status_tx: &mpsc::UnboundedSender<String>,
) -> Result<Applied> {
match protocol.execute_action(action)? {
ClientActionResult::Custom { name, data } if name == "netbios_name_query" => {
Self::query(transport, &data, pkt::QTYPE_NB, client_id, status_tx).await
}
ClientActionResult::Custom { name, data } if name == "netbios_node_status_query" => {
Self::query(transport, &data, pkt::QTYPE_NBSTAT, client_id, status_tx).await
}
ClientActionResult::Disconnect => Ok(Applied::Disconnect),
ClientActionResult::WaitForMore => {
Ok(Applied::Nothing("wait_for_more: no query sent".to_string()))
}
ClientActionResult::Custom { name, .. } => Ok(Applied::Nothing(format!(
"custom result '{name}' is not a NetBIOS-NS query; nothing sent"
))),
other => Ok(Applied::Nothing(format!(
"action result {other:?} produced no datagram"
))),
}
}
async fn query(
transport: &Arc<Transport>,
data: &serde_json::Value,
qtype: u16,
client_id: ClientId,
status_tx: &mpsc::UnboundedSender<String>,
) -> Result<Applied> {
let name = data
.get("name")
.and_then(|v| v.as_str())
.context("query is missing 'name'")?
.to_string();
let suffix = u8::try_from(data.get("suffix").and_then(|v| v.as_u64()).unwrap_or(0))
.context("'suffix' must be between 0 and 255")?;
let broadcast = data
.get("broadcast")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let target = target_for(transport.default_target, data)?;
let trn_id: u16 = rand::random();
let datagram = if qtype == pkt::QTYPE_NBSTAT {
wire::encode_node_status_query(trn_id, &name, suffix)?
} else {
wire::encode_name_query(trn_id, &name, suffix, broadcast)?
};
let question_type = pkt::qtype_name(qtype);
debug!(
"NetBIOS-NS client {} sending {} query for {}<0x{:02x}> to {} (trn_id 0x{:04x}, \
{} octets)",
client_id,
question_type,
name,
suffix,
target,
trn_id,
datagram.len()
);
let outcome = Self::run_query(transport, &datagram, target, trn_id, client_id).await?;
match outcome {
QueryOutcome::TimedOut { ignored } => {
info!(
"NetBIOS-NS client {} got no matching answer for {}<0x{:02x}> from {} in {}s \
({} datagram(s) discarded)",
client_id,
name,
suffix,
target,
transport.query_timeout.as_secs(),
ignored
);
let _ = status_tx.send(format!(
"[CLIENT] NetBIOS-NS {} query for {} timed out (normal when the node does \
not hold the name)",
question_type, name
));
Ok(Applied::Answered {
event_type: &NETBIOS_QUERY_TIMEOUT_EVENT,
event_data: serde_json::json!({
"name": name,
"suffix": suffix,
"suffix_label": wire::suffix_label(suffix),
"question_type": question_type,
"target": target.to_string(),
"timeout_secs": transport.query_timeout.as_secs(),
"ignored_datagrams": ignored,
}),
detail: format!(
"{question_type} query for {name} timed out after {}s ({ignored} \
datagram(s) discarded)",
transport.query_timeout.as_secs()
),
})
}
QueryOutcome::Reply { response, from } => {
Ok(Self::describe_reply(response, from, client_id, status_tx))
}
}
}
fn describe_reply(
response: wire::NbnsResponse,
from: SocketAddr,
client_id: ClientId,
status_tx: &mpsc::UnboundedSender<String>,
) -> Applied {
let responder = from.to_string();
match response.answer {
wire::NbnsAnswer::Name(answer) => {
let addresses: Vec<String> =
answer.addresses.iter().map(|a| a.to_string()).collect();
info!(
"NetBIOS-NS client {} resolved {}<0x{:02x}> to {} (ttl {}, group {})",
client_id,
answer.name,
answer.suffix,
addresses.join(", "),
answer.ttl,
answer.group
);
let _ = status_tx.send(format!(
"[CLIENT] NetBIOS-NS {}<0x{:02x}> = {}",
answer.name,
answer.suffix,
addresses.join(", ")
));
let detail = format!(
"{}<0x{:02x}> resolved to {}",
answer.name,
answer.suffix,
addresses.join(", ")
);
Applied::Answered {
event_type: &NETBIOS_NAME_RESPONSE_EVENT,
event_data: serde_json::json!({
"name": answer.name,
"suffix": answer.suffix,
"suffix_label": wire::suffix_label(answer.suffix),
"addresses": addresses,
"ttl": answer.ttl,
"group": answer.group,
"node_type": answer.node_type,
"responder": responder,
}),
detail,
}
}
wire::NbnsAnswer::NodeStatus(answer) => {
let names: Vec<serde_json::Value> = answer
.names
.iter()
.map(|entry| {
serde_json::json!({
"name": entry.name,
"suffix": entry.suffix,
"suffix_label": wire::suffix_label(entry.suffix),
"group": entry.group,
"active": entry.active,
})
})
.collect();
info!(
"NetBIOS-NS client {} node status from {}: {} name(s), adapter {}",
client_id,
responder,
names.len(),
answer.mac_address
);
let _ = status_tx.send(format!(
"[CLIENT] NetBIOS-NS node status: {} name(s), adapter {}",
names.len(),
answer.mac_address
));
let detail = format!(
"node status listed {} name(s), adapter {}",
names.len(),
answer.mac_address
);
Applied::Answered {
event_type: &NETBIOS_NODE_STATUS_RESPONSE_EVENT,
event_data: serde_json::json!({
"name": answer.name,
"suffix": answer.suffix,
"suffix_label": wire::suffix_label(answer.suffix),
"names": names,
"mac_address": answer.mac_address,
"responder": responder,
}),
detail,
}
}
wire::NbnsAnswer::Negative(answer) => {
info!(
"NetBIOS-NS client {} refused for {}<0x{:02x}>: rcode {} ({})",
client_id, answer.name, answer.suffix, answer.rcode, answer.rcode_name
);
let _ = status_tx.send(format!(
"[CLIENT] NetBIOS-NS {} -> {}",
answer.name, answer.rcode_name
));
let detail = format!(
"{}<0x{:02x}> refused: {}",
answer.name, answer.suffix, answer.rcode_name
);
Applied::Answered {
event_type: &NETBIOS_NEGATIVE_RESPONSE_EVENT,
event_data: serde_json::json!({
"name": answer.name,
"suffix": answer.suffix,
"suffix_label": wire::suffix_label(answer.suffix),
"rcode": answer.rcode,
"rcode_name": answer.rcode_name,
"responder": responder,
}),
detail,
}
}
}
}
async fn run_query(
transport: &Arc<Transport>,
datagram: &[u8],
target: SocketAddr,
trn_id: u16,
client_id: ClientId,
) -> Result<QueryOutcome> {
let deadline = Instant::now() + transport.query_timeout;
let mut ignored = 0usize;
let mut buffer = vec![0u8; wire::RECV_BUFFER];
let socket = transport.socket.lock().await;
socket
.send_to(datagram, target)
.await
.with_context(|| format!("failed to send a NetBIOS-NS query to {target}"))?;
loop {
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return Ok(QueryOutcome::TimedOut { ignored });
}
let received =
match tokio::time::timeout(remaining, socket.recv_from(&mut buffer)).await {
Err(_) => return Ok(QueryOutcome::TimedOut { ignored }),
Ok(Ok(received)) => received,
Ok(Err(e)) => {
return Err(anyhow::anyhow!(
"NetBIOS-NS client {client_id} socket read failed: {e}"
))
}
};
let (n, from) = received;
trace!(
"NetBIOS-NS client {} read {} octets from {}",
client_id,
n,
from
);
match wire::parse_response(&buffer[..n]) {
Ok(response) if response.trn_id == trn_id => {
return Ok(QueryOutcome::Reply { response, from })
}
Ok(response) => {
ignored += 1;
warn!(
"NetBIOS-NS client {} discarded a datagram from {}: transaction id \
0x{:04x} does not match the outstanding 0x{:04x}",
client_id, from, response.trn_id, trn_id
);
}
Err(e) => {
ignored += 1;
debug!(
"NetBIOS-NS client {} discarded a datagram from {}: {}",
client_id, from, e
);
}
}
}
}
#[allow(clippy::too_many_arguments)]
async fn command_loop(
mut command_rx: mpsc::Receiver<ClientCommand>,
protocol: Arc<NetbiosNsClientProtocol>,
transport: Arc<Transport>,
client_id: ClientId,
app_state: Arc<AppState>,
llm_client: OllamaClient,
status_tx: mpsc::UnboundedSender<String>,
) {
while let Some(command) = command_rx.recv().await {
let action = command.action.clone();
let mut answered: Option<(&'static EventType, serde_json::Value)> = None;
let outcome = match Self::apply_action(
&protocol,
&transport,
action.clone(),
client_id,
&status_tx,
)
.await
{
Ok(Applied::Answered {
event_type,
event_data,
detail,
}) => {
answered = Some((event_type, event_data));
Ok(ClientSendOutcome::Executed { detail })
}
Ok(Applied::Nothing(detail)) => Ok(ClientSendOutcome::Executed { detail }),
Ok(Applied::Disconnect) => Ok(ClientSendOutcome::Disconnected),
Err(e)
if e.to_string()
.starts_with("Unknown NetBIOS-NS client action") =>
{
Ok(ClientSendOutcome::Rejected {
error: e.to_string(),
})
}
Err(e) => Err(e),
};
let outcome_json = match &outcome {
Ok(outcome) => serde_json::to_value(outcome).unwrap_or(serde_json::Value::Null),
Err(e) => serde_json::json!({"error": e.to_string()}),
};
app_state
.record_access_log(
AccessLogOwner::Client(client_id.as_u32()),
protocol.protocol_name(),
None,
"injected_action",
action,
vec![outcome_json],
)
.await;
let disconnect = matches!(outcome, Ok(ClientSendOutcome::Disconnected));
if let Err(e) = &outcome {
error!(
"NetBIOS-NS client {} injected action failed: {}",
client_id, e
);
let _ = status_tx.send(format!(
"[WARN] Client {} injected action failed: {}",
client_id, e
));
}
let _ = status_tx.send("__UPDATE_UI__".to_string());
crate::client::command_support::reply(command, outcome);
if disconnect {
Self::finish(client_id, &app_state, &status_tx).await;
break;
}
if let Some((event_type, event_data)) = answered {
let follow_ups = Self::report(
event_type,
event_data,
&protocol,
client_id,
&app_state,
&llm_client,
&status_tx,
)
.await;
Self::run_actions(
&protocol,
&transport,
follow_ups,
1,
client_id,
&app_state,
&llm_client,
&status_tx,
)
.await;
}
}
}
async fn finish(
client_id: ClientId,
app_state: &Arc<AppState>,
status_tx: &mpsc::UnboundedSender<String>,
) {
info!("NetBIOS-NS client {} disconnecting", client_id);
app_state.remove_client_handle(client_id).await;
app_state
.update_client_status(client_id, ClientStatus::Disconnected)
.await;
let _ = status_tx.send("__UPDATE_UI__".to_string());
}
}
async fn resolve_target(remote_addr: &str, port_override: Option<u16>) -> Result<SocketAddr> {
let remote_addr = remote_addr.trim();
if let Ok(addr) = remote_addr.parse::<SocketAddr>() {
return Ok(match port_override {
Some(port) => SocketAddr::new(addr.ip(), port),
None => addr,
});
}
if let Ok(ip) = remote_addr.parse::<IpAddr>() {
return Ok(SocketAddr::new(
ip,
port_override.unwrap_or(DEFAULT_NBNS_PORT),
));
}
let (host, port) = match remote_addr.rsplit_once(':') {
Some((host, port)) if port.chars().all(|c| c.is_ascii_digit()) && !port.is_empty() => (
host,
port.parse::<u16>()
.with_context(|| format!("'{port}' is not a port number"))?,
),
_ => (remote_addr, DEFAULT_NBNS_PORT),
};
let port = port_override.unwrap_or(port);
tokio::net::lookup_host((host, port))
.await
.with_context(|| format!("could not resolve NetBIOS-NS target '{remote_addr}'"))?
.find(|addr| addr.is_ipv4())
.with_context(|| format!("'{remote_addr}' has no IPv4 address; NetBIOS-NS is IPv4 only"))
}
fn target_for(default_target: SocketAddr, data: &serde_json::Value) -> Result<SocketAddr> {
let ip = match data.get("target_address").and_then(|v| v.as_str()) {
Some(text) => text.trim().parse::<IpAddr>().with_context(|| {
format!("'target_address' must be a dotted-quad IPv4 address, got '{text}'")
})?,
None => default_target.ip(),
};
let port = match data.get("target_port").and_then(|v| v.as_u64()) {
Some(port) => u16::try_from(port)
.with_context(|| format!("'target_port' must be between 0 and 65535, got {port}"))?,
None => default_target.port(),
};
Ok(SocketAddr::new(ip, port))
}