pub mod actions;
use crate::llm::action_helper::call_llm;
use crate::llm::ollama_client::OllamaClient;
use crate::logging::emit::Log;
use crate::protocol::{Event, SpawnContext};
use crate::server::connection::ConnectionId;
use crate::state::app_state::AppState;
use crate::state::ServerId;
use actions::{LlmnrProtocol, LLMNR_IPV4_GROUP, LLMNR_IPV6_GROUP, LLMNR_QUERY_EVENT};
use anyhow::{Context, Result};
use hickory_proto::op::{Message as DnsMessage, MessageType, OpCode};
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream, UdpSocket};
use tokio::sync::mpsc;
use tracing::{debug, error, info, trace, warn};
pub const MAX_MESSAGE_LEN: usize = 4096;
const TCP_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
const MAX_CONNECTIONS: usize = crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS;
const CONNECTION_CAP_REFUSAL: &[u8] = b"";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Decision {
ModelResponse,
ModelSilent,
ModelReject,
ModelRejectSuppressed,
FailClosedNoAction,
FailClosedActionError,
FailClosedLlmError,
}
impl Decision {
pub fn as_str(&self) -> &'static str {
match self {
Decision::ModelResponse => "model_response",
Decision::ModelSilent => "model_silent",
Decision::ModelReject => "model_reject",
Decision::ModelRejectSuppressed => "model_reject_suppressed_udp",
Decision::FailClosedNoAction => "fail_closed_no_action",
Decision::FailClosedActionError => "fail_closed_action_error",
Decision::FailClosedLlmError => "fail_closed_llm_error",
}
}
pub fn is_fail_closed(&self) -> bool {
matches!(
self,
Decision::FailClosedNoAction
| Decision::FailClosedActionError
| Decision::FailClosedLlmError
)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Transport {
Udp,
Tcp,
}
impl Transport {
fn as_str(self) -> &'static str {
match self {
Transport::Udp => "udp",
Transport::Tcp => "tcp",
}
}
}
#[derive(Clone)]
struct Responder {
llm_client: OllamaClient,
state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
server_id: ServerId,
local_addr: SocketAddr,
}
pub struct LlmnrServer;
impl LlmnrServer {
pub async fn spawn_with_llm_actions(ctx: SpawnContext) -> Result<SocketAddr> {
let listen_addr = ctx.legacy_listen_addr();
let SpawnContext {
llm_client,
state,
status_tx,
server_id,
startup_params,
..
} = ctx;
let (join_multicast, multicast_interface, enable_tcp) = match startup_params {
Some(params) => (
params.get_optional_bool("join_multicast")?.unwrap_or(true),
params.get_optional_string("multicast_interface")?,
params.get_optional_bool("enable_tcp")?.unwrap_or(true),
),
None => (true, None, true),
};
let socket = Arc::new(
UdpSocket::bind(listen_addr)
.await
.with_context(|| format!("LLMNR failed to bind UDP {}", listen_addr))?,
);
let local_addr = socket
.local_addr()
.context("LLMNR could not read its own UDP address")?;
let log = Log::new(Some(&status_tx));
log.info(format!("LLMNR responder listening on {}", local_addr));
info!("LLMNR responder listening on UDP {}", local_addr);
if join_multicast {
join_llmnr_group(
&socket,
local_addr,
multicast_interface.as_deref(),
&status_tx,
);
} else {
debug!("LLMNR multicast join skipped (join_multicast=false)");
}
let responder = Responder {
llm_client,
state: state.clone(),
status_tx: status_tx.clone(),
server_id,
local_addr,
};
let udp_responder = responder.clone();
let udp_socket = socket.clone();
let udp_handle = tokio::spawn(async move {
let mut buffer = vec![0u8; MAX_MESSAGE_LEN];
loop {
let (n, peer_addr) = match udp_socket.recv_from(&mut buffer).await {
Ok(pair) => pair,
Err(e) => {
Log::new(Some(&udp_responder.status_tx))
.error(format!("LLMNR UDP receive error: {}", e));
break;
}
};
let data = buffer[..n].to_vec();
let sink = UdpSink {
socket: udp_socket.clone(),
peer: peer_addr,
};
let r = udp_responder.clone();
let task = tokio::spawn(async move {
r.handle_query(data, peer_addr, Transport::Udp, Some(sink))
.await;
});
udp_responder
.state
.register_server_task(server_id, task)
.await;
}
});
state.register_server_task(server_id, udp_handle).await;
if enable_tcp {
match TcpListener::bind(SocketAddr::new(local_addr.ip(), local_addr.port())).await {
Ok(listener) => {
info!("LLMNR also listening on TCP {}", local_addr);
let tcp_responder = responder.clone();
let limiter =
crate::server::accept_bounded::ConnectionLimiter::new(MAX_CONNECTIONS);
let tcp_handle = tokio::spawn(async move {
loop {
let (stream, peer_addr, permit) =
match crate::server::accept_bounded::accept_bounded(
&listener,
&limiter,
CONNECTION_CAP_REFUSAL,
"LLMNR",
Some(&tcp_responder.status_tx),
)
.await
{
Ok(triple) => triple,
Err(e) => {
Log::new(Some(&tcp_responder.status_tx))
.error(format!("LLMNR TCP accept error: {}", e));
break;
}
};
let r = tcp_responder.clone();
let task = tokio::spawn(async move {
let _permit = permit;
r.serve_tcp_connection(stream, peer_addr).await;
});
tcp_responder
.state
.register_server_task(server_id, task)
.await;
}
});
state.register_server_task(server_id, tcp_handle).await;
}
Err(e) => {
warn!(
"LLMNR could not bind TCP {}: {}; UDP queries are still answered, \
TCP queries are not",
local_addr, e
);
log.warn(format!(
"LLMNR TCP listener unavailable on {} ({}); UDP only",
local_addr, e
));
}
}
} else {
debug!("LLMNR TCP listener disabled (enable_tcp=false)");
}
Ok(local_addr)
}
}
struct UdpSink {
socket: Arc<UdpSocket>,
peer: SocketAddr,
}
impl Responder {
async fn serve_tcp_connection(&self, mut stream: TcpStream, peer_addr: SocketAddr) {
loop {
let mut len_buf = [0u8; 2];
match tokio::time::timeout(TCP_IDLE_TIMEOUT, stream.read_exact(&mut len_buf)).await {
Ok(Ok(_)) => {}
Ok(Err(_)) => break,
Err(_) => {
debug!(
"LLMNR closing an idle TCP connection from {} after {:?}",
peer_addr, TCP_IDLE_TIMEOUT
);
break;
}
}
let len = u16::from_be_bytes(len_buf) as usize;
if len == 0 || len > MAX_MESSAGE_LEN {
warn!(
"LLMNR dropping TCP query from {}: length prefix {} is out of range",
peer_addr, len
);
break;
}
let mut data = vec![0u8; len];
match tokio::time::timeout(TCP_IDLE_TIMEOUT, stream.read_exact(&mut data)).await {
Ok(Ok(_)) => {}
Ok(Err(_)) => break,
Err(_) => {
warn!(
"LLMNR closing a TCP connection from {}: {} octets were announced but \
did not arrive within {:?}",
peer_addr, len, TCP_IDLE_TIMEOUT
);
break;
}
}
let outcome = self
.handle_query(data, peer_addr, Transport::Tcp, None)
.await;
let Some((reply, connection_id)) = outcome else {
continue;
};
let mut framed = Vec::with_capacity(reply.len() + 2);
let Ok(prefix) = u16::try_from(reply.len()) else {
error!(
"LLMNR TCP response to {} is {} bytes, too long to frame",
peer_addr,
reply.len()
);
break;
};
framed.extend_from_slice(&prefix.to_be_bytes());
framed.extend_from_slice(&reply);
if let Err(e) = stream.write_all(&framed).await {
Log::new(Some(&self.status_tx)).error(format!(
"LLMNR failed to reply over TCP to {}: {}",
peer_addr, e
));
break;
}
trace!(
"LLMNR sent {} framed bytes to {} over TCP",
framed.len(),
peer_addr
);
self.count_sent_for(connection_id, framed.len()).await;
}
}
async fn handle_query(
&self,
data: Vec<u8>,
peer_addr: SocketAddr,
transport: Transport,
sink: Option<UdpSink>,
) -> Option<(Vec<u8>, ConnectionId)> {
let log = Log::new(Some(&self.status_tx));
trace!(
"LLMNR {} bytes from {} over {}: {}",
data.len(),
peer_addr,
transport.as_str(),
hex::encode(&data)
);
let query = match DnsMessage::from_vec(&data) {
Ok(m) => m,
Err(e) => {
debug!(
"LLMNR dropped unparseable message from {}: {}",
peer_addr, e
);
return None;
}
};
if query.message_type() != MessageType::Query {
debug!("LLMNR dropped a response (not a query) from {}", peer_addr);
return None;
}
if query.op_code() != OpCode::Query {
debug!(
"LLMNR dropped query from {}: unsupported OPCODE {:?}",
peer_addr,
query.op_code()
);
return None;
}
if query.queries().len() != 1 {
debug!(
"LLMNR dropped query from {}: QDCOUNT is {}, must be 1",
peer_addr,
query.queries().len()
);
return None;
}
let question = &query.queries()[0];
let name = question.name().to_string();
let query_type = question.query_type().to_string();
let query_class = question.query_class().to_string();
let transaction_id = query.id();
let connection_id = self.record_connection(peer_addr, data.len()).await;
let _ = self.status_tx.send("__UPDATE_UI__".to_string());
debug!(
"LLMNR query id={} {} {} {} from {} over {}",
transaction_id,
name,
query_class,
query_type,
peer_addr,
transport.as_str()
);
let event = Event::new(
&LLMNR_QUERY_EVENT,
serde_json::json!({
"transaction_id": transaction_id,
"name": name,
"query_type": query_type,
"query_class": query_class,
"source_address": peer_addr.to_string(),
"transport": transport.as_str(),
"conflict": actions::conflict_bit(&query),
"tentative": actions::tentative_bit(&query),
}),
);
let protocol = LlmnrProtocol::new();
let llm_outcome = call_llm(
&self.llm_client,
&self.state,
self.server_id,
None,
&event,
&protocol,
)
.await;
let (decision, reply) = self.decide(llm_outcome, transport, peer_addr);
let summary = format!(
"LLMNR {} {} from {} decision={}",
query_type,
name,
peer_addr,
decision.as_str()
);
if decision.is_fail_closed() {
error!(
"{} (nothing sent: no usable answer was produced; a guessed LLMNR record \
would poison the querier's name cache)",
summary
);
log.error(format!(
"{} - staying silent rather than guessing a name binding",
summary
));
} else {
log.info(&summary);
}
let reply = reply?;
let Some(udp) = sink else {
return Some((reply, connection_id));
};
match udp.socket.send_to(&reply, udp.peer).await {
Ok(sent) => {
trace!(
"LLMNR sent {} bytes to {}: {}",
sent,
udp.peer,
hex::encode(&reply)
);
self.count_sent_for(connection_id, sent).await;
log.debug(format!("LLMNR sent {} bytes to {}", sent, udp.peer));
}
Err(e) => {
Log::new(Some(&self.status_tx))
.error(format!("LLMNR failed to reply to {}: {}", udp.peer, e));
}
}
None
}
fn decide(
&self,
llm_outcome: Result<crate::llm::ExecutionResult>,
transport: Transport,
peer_addr: SocketAddr,
) -> (Decision, Option<Vec<u8>>) {
let log = Log::new(Some(&self.status_tx));
let execution = match llm_outcome {
Ok(result) => {
for message in &result.messages {
log.info(message);
}
result
}
Err(e) => {
let category = match crate::utils::wire_failure::WireFailure::classify(&e) {
crate::utils::wire_failure::WireFailure::Overloaded => "overloaded",
crate::utils::wire_failure::WireFailure::Unavailable => "unavailable",
};
error!(
"LLMNR LLM call failed for {} (category={}): {}",
peer_addr, category, e
);
return (Decision::FailClosedLlmError, None);
}
};
let mut chosen: Option<Vec<u8>> = None;
let mut extra = 0usize;
for protocol_result in &execution.protocol_results {
for output in protocol_result.get_all_output() {
if chosen.is_none() {
chosen = Some(output);
} else {
extra += 1;
}
}
}
if extra > 0 {
warn!(
"LLMNR ignored {} extra output(s) for the query from {}; one query gets one \
response",
extra, peer_addr
);
}
if let Some(bytes) = chosen {
let rcode = bytes.get(3).map(|b| b & 0x0F).unwrap_or(0);
if rcode == 0 {
return (Decision::ModelResponse, Some(bytes));
}
if transport == Transport::Udp {
warn!(
"LLMNR suppressed an RCODE {} response to {}: a response to a UDP query \
must carry RCODE 0 (RFC 4795 2.1.1); use no_response instead",
rcode, peer_addr
);
return (Decision::ModelRejectSuppressed, None);
}
return (Decision::ModelReject, Some(bytes));
}
if execution.raw_actions.is_empty() {
return (Decision::FailClosedNoAction, None);
}
if execution
.raw_actions
.iter()
.any(|a| a.get("type").and_then(|t| t.as_str()) == Some("no_response"))
{
return (Decision::ModelSilent, None);
}
(Decision::FailClosedActionError, None)
}
async fn record_connection(&self, peer_addr: SocketAddr, bytes: usize) -> ConnectionId {
use crate::state::server::{
ConnectionState as ServerConnectionState, ConnectionStatus, ProtocolConnectionInfo,
};
let connection_id = ConnectionId::new(self.state.get_next_unified_id().await);
let now = crate::utils::clock::Instant::now();
let conn_state = ServerConnectionState {
id: connection_id,
remote_addr: peer_addr,
local_addr: self.local_addr,
bytes_sent: 0,
bytes_received: bytes as u64,
packets_sent: 0,
packets_received: 1,
last_activity: now,
status: ConnectionStatus::Active,
status_changed_at: now,
protocol_info: ProtocolConnectionInfo::empty(),
};
self.state
.add_connection_to_server(self.server_id, conn_state)
.await;
connection_id
}
async fn count_sent_for(&self, connection_id: ConnectionId, bytes: usize) {
self.state
.update_connection_stats(
self.server_id,
connection_id,
None,
Some(bytes as u64),
None,
Some(1),
)
.await;
}
}
fn join_llmnr_group(
socket: &UdpSocket,
local_addr: SocketAddr,
multicast_interface: Option<&str>,
status_tx: &mpsc::UnboundedSender<String>,
) {
let log = Log::new(Some(status_tx));
match local_addr.ip() {
IpAddr::V4(_) => {
let interface = match multicast_interface {
Some(value) => match value.parse::<Ipv4Addr>() {
Ok(ip) => ip,
Err(e) => {
warn!(
"LLMNR multicast_interface '{}' is not an IPv4 address ({}); \
letting the host choose",
value, e
);
Ipv4Addr::UNSPECIFIED
}
},
None => Ipv4Addr::UNSPECIFIED,
};
match socket.join_multicast_v4(LLMNR_IPV4_GROUP, interface) {
Ok(()) => {
info!(
"LLMNR joined {} on interface {}",
LLMNR_IPV4_GROUP, interface
);
log.info(format!("LLMNR joined multicast group {}", LLMNR_IPV4_GROUP));
}
Err(e) => {
warn!(
"LLMNR could not join {} on interface {}: {}; queries sent directly to \
this port are still answered",
LLMNR_IPV4_GROUP, interface, e
);
log.warn(format!(
"LLMNR multicast join failed ({}); answering direct queries only",
e
));
}
}
}
IpAddr::V6(_) => {
match socket.join_multicast_v6(&LLMNR_IPV6_GROUP, 0) {
Ok(()) => {
info!("LLMNR joined {}", LLMNR_IPV6_GROUP);
log.info(format!("LLMNR joined multicast group {}", LLMNR_IPV6_GROUP));
}
Err(e) => {
warn!(
"LLMNR could not join {}: {}; queries sent directly to this port are \
still answered",
LLMNR_IPV6_GROUP, e
);
log.warn(format!(
"LLMNR multicast join failed ({}); answering direct queries only",
e
));
}
}
}
}
}