pub mod actions;
pub use actions::IdentClientProtocol;
use std::future::Future;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use anyhow::{Context, Result};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::sync::{mpsc, Mutex};
use crate::client::llm_budget::call_llm_for_client;
use crate::llm::actions::client_trait::{Client, ClientActionResult};
use crate::llm::ollama_client::OllamaClient;
use crate::llm::ClientLlmResult;
use crate::logging::emit::Log;
use crate::protocol::{Event, StartupParams};
use crate::state::app_state::AppState;
use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
use crate::state::{AccessLogOwner, ClientId, ClientStatus};
use actions::{
IDENT_CLIENT_CONNECTED_EVENT, IDENT_CLIENT_ERROR_RECEIVED_EVENT,
IDENT_CLIENT_REPLY_MISMATCH_EVENT, IDENT_CLIENT_RESPONSE_RECEIVED_EVENT,
};
pub const DEFAULT_IDENT_PORT: u16 = 113;
const DEFAULT_RESPONSE_TIMEOUT_SECS: u64 = 30;
pub const MAX_REPLY_BYTES: usize = 1024;
pub const MAX_FOLLOWUP_DEPTH: usize = 4;
const IDENT_ERROR_TOKENS: [&str; 4] = ["INVALID-PORT", "NO-USER", "HIDDEN-USER", "UNKNOWN-ERROR"];
type PendingQuery = Arc<std::sync::Mutex<Option<(u16, u16)>>>;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum IdentReply {
Userid {
server_port: u16,
client_port: u16,
opsys: String,
charset: Option<String>,
userid: String,
},
Error {
server_port: u16,
client_port: u16,
token: String,
},
Malformed { reason: &'static str },
}
impl IdentReply {
pub fn port_pair(&self) -> Option<(u16, u16)> {
match self {
IdentReply::Userid {
server_port,
client_port,
..
}
| IdentReply::Error {
server_port,
client_port,
..
} => Some((*server_port, *client_port)),
IdentReply::Malformed { .. } => None,
}
}
}
fn parse_port_pair(text: &str) -> Option<(u16, u16)> {
let (left, right) = text.split_once(',')?;
let server_port: u32 = left.trim().parse().ok()?;
let client_port: u32 = right.trim().parse().ok()?;
if !(1..=65535).contains(&server_port) || !(1..=65535).contains(&client_port) {
return None;
}
Some((server_port as u16, client_port as u16))
}
fn is_error_token(token: &str) -> bool {
if IDENT_ERROR_TOKENS.contains(&token) {
return true;
}
token.starts_with('X')
&& token.len() > 1
&& token
.chars()
.all(|c| c.is_ascii_uppercase() || c.is_ascii_digit() || c == '-')
}
pub fn parse_ident_reply(line: &str) -> IdentReply {
if line.len() > MAX_REPLY_BYTES {
return IdentReply::Malformed {
reason: "oversized",
};
}
let line = line.trim();
let Some((pair_text, rest)) = line.split_once(':') else {
return IdentReply::Malformed {
reason: "no_response_type",
};
};
let Some((server_port, client_port)) = parse_port_pair(pair_text) else {
return IdentReply::Malformed {
reason: "port_pair",
};
};
let Some((kind, addl)) = rest.split_once(':') else {
return IdentReply::Malformed {
reason: "no_addl_info",
};
};
match kind.trim().to_ascii_uppercase().as_str() {
"USERID" => {
let Some((opsys_field, userid)) = addl.split_once(':') else {
return IdentReply::Malformed {
reason: "userid_fields",
};
};
let userid = userid.trim();
if userid.is_empty() {
return IdentReply::Malformed {
reason: "empty_userid",
};
}
let (opsys, charset) = match opsys_field.split_once(',') {
Some((opsys, charset)) => {
let charset = charset.trim();
(
opsys.trim().to_string(),
(!charset.is_empty()).then(|| charset.to_string()),
)
}
None => (opsys_field.trim().to_string(), None),
};
if opsys.is_empty() {
return IdentReply::Malformed {
reason: "empty_opsys",
};
}
IdentReply::Userid {
server_port,
client_port,
opsys,
charset,
userid: userid.to_string(),
}
}
"ERROR" => {
let token = addl.trim().to_ascii_uppercase();
if !is_error_token(&token) {
return IdentReply::Malformed {
reason: "error_token",
};
}
IdentReply::Error {
server_port,
client_port,
token,
}
}
_ => IdentReply::Malformed {
reason: "response_type",
},
}
}
fn sanitize_line(line: &str) -> String {
let cleaned = crate::utils::sanitize::strip_controls(line);
crate::utils::truncate_for_log(cleaned.trim(), 256)
}
fn split_host_port(addr: &str) -> (String, Option<u16>) {
if let Some(rest) = addr.strip_prefix('[') {
if let Some((host, tail)) = rest.split_once(']') {
let port = tail.strip_prefix(':').and_then(|p| p.trim().parse().ok());
return (host.to_string(), port);
}
}
match addr.rsplit_once(':') {
Some((head, tail)) if !head.contains(':') => match tail.trim().parse::<u16>() {
Ok(port) => (head.to_string(), Some(port)),
Err(_) => (addr.to_string(), None),
},
_ => (addr.to_string(), None),
}
}
pub fn resolve_target(remote_addr: &str, ident_port: Option<u16>) -> Result<String> {
let trimmed = remote_addr.trim();
let trimmed = trimmed
.strip_prefix("ident://")
.unwrap_or(trimmed)
.trim_end_matches('/');
let (host, addr_port) = split_host_port(trimmed);
if host.is_empty() {
anyhow::bail!("remote_addr {remote_addr:?} names no host to query");
}
let port = ident_port.or(addr_port).unwrap_or(DEFAULT_IDENT_PORT);
Ok(if host.contains(':') {
format!("[{host}]:{port}")
} else {
format!("{host}:{port}")
})
}
enum Applied {
Sent(usize),
Disconnect,
}
enum ReplyRead {
Line(String),
Oversize(String),
Closed,
TimedOut,
Failed(String),
}
struct SessionCtx {
protocol: Arc<IdentClientProtocol>,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
client_id: ClientId,
instruction: String,
target: String,
response_timeout: Duration,
}
pub struct IdentClient;
impl IdentClient {
pub async fn connect_with_llm_actions(
remote_addr: String,
startup_params: Option<StartupParams>,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
client_id: ClientId,
) -> Result<SocketAddr> {
let ident_port = match &startup_params {
Some(params) => params.get_optional_u64("ident_port")?,
None => None,
};
let ident_port = match ident_port {
Some(port) if (1..=65535).contains(&port) => Some(port as u16),
Some(port) => anyhow::bail!("ident_port must be a TCP port in 1-65535, got {port}"),
None => None,
};
let response_timeout = match &startup_params {
Some(params) => params.get_optional_u64("response_timeout_secs")?,
None => None,
}
.filter(|secs| *secs > 0)
.unwrap_or(DEFAULT_RESPONSE_TIMEOUT_SECS);
let response_timeout = Duration::from_secs(response_timeout);
let target = resolve_target(&remote_addr, ident_port)?;
let stream = TcpStream::connect(&target)
.await
.with_context(|| format!("Failed to connect to ident server at {target}"))?;
let local_addr = stream.local_addr()?;
let peer_addr = stream.peer_addr()?;
Log::new(Some(&status_tx)).info(format!(
"Ident client {client_id} connected to {peer_addr} (local: {local_addr})"
));
app_state
.update_client_status(client_id, ClientStatus::Connected)
.await;
let _ = status_tx.send("__UPDATE_UI__".to_string());
let (read_half, write_half) = tokio::io::split(stream);
let write_half = Arc::new(Mutex::new(write_half));
let protocol = Arc::new(IdentClientProtocol::new());
let pending: PendingQuery = Arc::new(std::sync::Mutex::new(None));
let command_rx =
crate::client::command_support::register_command_channel(&app_state, client_id).await;
let command_task = tokio::spawn(Self::command_loop(
command_rx,
protocol.clone(),
write_half.clone(),
pending.clone(),
client_id,
app_state.clone(),
status_tx.clone(),
));
app_state
.register_client_task(client_id, command_task)
.await;
let session_registrar = app_state.clone();
let session_task = tokio::spawn(async move {
Self::session(
read_half,
write_half,
protocol,
pending,
target,
response_timeout,
llm_client,
app_state.clone(),
status_tx.clone(),
client_id,
)
.await;
app_state.remove_client_handle(client_id).await;
let _ = status_tx.send("__UPDATE_UI__".to_string());
});
session_registrar
.register_client_task(client_id, session_task)
.await;
Ok(local_addr)
}
#[allow(clippy::too_many_arguments)]
async fn session(
mut read_half: tokio::io::ReadHalf<TcpStream>,
write_half: Arc<Mutex<tokio::io::WriteHalf<TcpStream>>>,
protocol: Arc<IdentClientProtocol>,
pending: PendingQuery,
target: String,
response_timeout: Duration,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
client_id: ClientId,
) {
let log = Log::new(Some(&status_tx));
let Some(instruction) = app_state.get_instruction_for_client(client_id).await else {
return;
};
let ctx = Arc::new(SessionCtx {
protocol: protocol.clone(),
llm_client: llm_client.clone(),
app_state: app_state.clone(),
status_tx: status_tx.clone(),
client_id,
instruction: instruction.clone(),
target: target.clone(),
response_timeout,
});
let event = Event::new(
&IDENT_CLIENT_CONNECTED_EVENT,
serde_json::json!({ "remote_addr": target }),
);
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(memory) = memory_updates {
app_state.set_memory_for_client(client_id, memory).await;
}
for action in actions {
let result = match protocol.execute_action(action) {
Ok(result) => result,
Err(e) => {
log.error(format!("Ident client {client_id} rejected action: {e}"));
continue;
}
};
match Self::apply_action(result, &write_half, &pending, client_id, &log).await {
Ok(Applied::Sent(_)) => {}
Ok(Applied::Disconnect) => {
log.info(format!(
"Ident client {client_id} disconnected before asking anything"
));
app_state
.update_client_status(client_id, ClientStatus::Disconnected)
.await;
let _ = status_tx.send("__UPDATE_UI__".to_string());
return;
}
Err(e) => {
log.error(format!(
"Ident client {client_id} failed to send query: {e}"
));
app_state
.update_client_status(client_id, ClientStatus::Error(e.to_string()))
.await;
let _ = status_tx.send("__UPDATE_UI__".to_string());
return;
}
}
}
}
Err(e) => {
log.error(format!(
"Ident client {client_id} LLM error on connect: {e} (still connected; a \
query can be injected)"
));
}
}
match Self::read_reply(&mut read_half, response_timeout).await {
ReplyRead::Line(line) => {
log.trace(format!("Ident client {client_id} reply line: {line:?}"));
let asked = pending.lock().ok().and_then(|p| *p);
let reply = parse_ident_reply(&line);
Self::handle_reply(ctx, 0, reply, asked, line).await;
}
ReplyRead::Oversize(line) => {
let asked = pending.lock().ok().and_then(|p| *p);
log.warn(format!(
"Ident client {client_id} reply exceeded {MAX_REPLY_BYTES} bytes with no \
newline; rejected without parsing"
));
Self::handle_reply(
ctx,
0,
IdentReply::Malformed {
reason: "oversized",
},
asked,
line,
)
.await;
}
ReplyRead::Closed => {
let asked = pending.lock().ok().and_then(|p| *p);
match asked {
Some((server_port, client_port)) => log.info(format!(
"Ident client {client_id} got no answer for {server_port},{client_port}: \
the server closed first"
)),
None => log.info(format!(
"Ident client {client_id} was closed by the server before any query went \
out"
)),
}
}
ReplyRead::TimedOut => log.warn(format!(
"Ident client {client_id} gave up after {}s with no reply",
response_timeout.as_secs()
)),
ReplyRead::Failed(e) => {
log.error(format!("Ident client {client_id} read error: {e}"));
app_state
.update_client_status(client_id, ClientStatus::Error(e))
.await;
let _ = status_tx.send("__UPDATE_UI__".to_string());
return;
}
}
let _ = write_half.lock().await.shutdown().await;
app_state
.update_client_status(client_id, ClientStatus::Disconnected)
.await;
log.info(format!("Ident client {client_id} disconnected"));
let _ = status_tx.send("__UPDATE_UI__".to_string());
}
fn handle_reply(
ctx: Arc<SessionCtx>,
depth: usize,
reply: IdentReply,
asked: Option<(u16, u16)>,
line: String,
) -> Pin<Box<dyn Future<Output = ()> + Send>> {
Box::pin(async move {
let log = Log::new(Some(&ctx.status_tx));
let client_id = ctx.client_id;
let event = match (&reply, asked, reply.port_pair()) {
(IdentReply::Malformed { reason }, _, _) => {
log.warn(format!(
"Ident client {client_id} rejected a reply ({reason}): {}",
sanitize_line(&line)
));
Self::mismatch_event(reason, asked, None, &line)
}
(_, None, got) => {
log.warn(format!(
"Ident client {client_id} received a reply with no query outstanding; \
rejected"
));
Self::mismatch_event("unsolicited", None, got, &line)
}
(_, Some(query), Some(got)) if query != got => {
log.warn(format!(
"Ident client {client_id} asked about {},{} and was answered about \
{},{}; rejected (RFC 1413 matches a reply by its port pair)",
query.0, query.1, got.0, got.1
));
Self::mismatch_event("port_pair_mismatch", asked, Some(got), &line)
}
(
IdentReply::Userid {
server_port,
client_port,
opsys,
charset,
userid,
},
_,
_,
) => {
log.info(format!(
"Ident client {client_id} {server_port},{client_port} USERID on {opsys}"
));
let mut data = serde_json::json!({
"server_port": server_port,
"client_port": client_port,
"userid": userid,
"opsys": opsys,
});
if let Some(charset) = charset {
data["charset"] = serde_json::Value::String(charset.clone());
}
Event::new(&IDENT_CLIENT_RESPONSE_RECEIVED_EVENT, data)
}
(
IdentReply::Error {
server_port,
client_port,
token,
},
_,
_,
) => {
log.info(format!(
"Ident client {client_id} {server_port},{client_port} ERROR {token}"
));
Event::new(
&IDENT_CLIENT_ERROR_RECEIVED_EVENT,
serde_json::json!({
"server_port": server_port,
"client_port": client_port,
"error_token": token,
}),
)
}
};
let memory = ctx
.app_state
.get_memory_for_client(client_id)
.await
.unwrap_or_default();
let outcome = call_llm_for_client(
&ctx.llm_client,
&ctx.app_state,
client_id.to_string(),
&ctx.instruction,
&memory,
Some(&event),
ctx.protocol.as_ref(),
&ctx.status_tx,
)
.await;
let ClientLlmResult {
actions,
memory_updates,
} = match outcome {
Ok(result) => result,
Err(e) => {
log.error(format!(
"Ident client {client_id} LLM error on {}: {e}",
event.id()
));
return;
}
};
if let Some(memory) = memory_updates {
ctx.app_state.set_memory_for_client(client_id, memory).await;
}
for action in actions {
match ctx.protocol.execute_action(action) {
Ok(ClientActionResult::Custom { name, data }) if name == "ident_query" => {
let (Some(server_port), Some(client_port)) = (
data.get("server_port").and_then(|v| v.as_u64()),
data.get("client_port").and_then(|v| v.as_u64()),
) else {
continue;
};
if depth + 1 > MAX_FOLLOWUP_DEPTH {
log.warn(format!(
"Ident client {client_id} stopped at follow-up depth \
{MAX_FOLLOWUP_DEPTH}; the request for {server_port},\
{client_port} was not sent"
));
continue;
}
Self::followup_query(
ctx.clone(),
depth + 1,
server_port as u16,
client_port as u16,
)
.await;
}
Ok(ClientActionResult::Disconnect) => {
log.debug(format!(
"Ident client {client_id} was asked to disconnect; the connection is \
already finished"
));
break;
}
Ok(ClientActionResult::WaitForMore) => log.debug(format!(
"Ident client {client_id} asked to wait, but RFC 1413 sends one line and \
closes; nothing more will arrive"
)),
Ok(_) => {}
Err(e) => log.error(format!("Ident client {client_id} rejected action: {e}")),
}
}
})
}
async fn followup_query(
ctx: Arc<SessionCtx>,
depth: usize,
server_port: u16,
client_port: u16,
) {
let log = Log::new(Some(&ctx.status_tx));
let client_id = ctx.client_id;
let mut stream = match TcpStream::connect(&ctx.target).await {
Ok(stream) => stream,
Err(e) => {
log.error(format!(
"Ident client {client_id} follow-up could not reach {}: {e}",
ctx.target
));
return;
}
};
let query = format!("{server_port} , {client_port}\r\n");
if let Err(e) = stream.write_all(query.as_bytes()).await {
log.error(format!(
"Ident client {client_id} follow-up write failed: {e}"
));
return;
}
if let Err(e) = stream.flush().await {
log.error(format!(
"Ident client {client_id} follow-up flush failed: {e}"
));
return;
}
log.debug(format!(
"Ident client {client_id} follow-up #{depth} asked about {server_port},{client_port}"
));
let (mut read_half, write_half) = tokio::io::split(stream);
let outcome = Self::read_reply(&mut read_half, ctx.response_timeout).await;
let mut write_half = write_half;
let _ = write_half.shutdown().await;
match outcome {
ReplyRead::Line(line) => {
let reply = parse_ident_reply(&line);
Self::handle_reply(ctx, depth, reply, Some((server_port, client_port)), line).await;
}
ReplyRead::Oversize(line) => {
Self::handle_reply(
ctx,
depth,
IdentReply::Malformed {
reason: "oversized",
},
Some((server_port, client_port)),
line,
)
.await;
}
ReplyRead::Closed => log.info(format!(
"Ident client {client_id} follow-up for {server_port},{client_port} was closed \
unanswered"
)),
ReplyRead::TimedOut => log.warn(format!(
"Ident client {client_id} follow-up for {server_port},{client_port} timed out"
)),
ReplyRead::Failed(e) => log.error(format!(
"Ident client {client_id} follow-up read error: {e}"
)),
}
}
fn mismatch_event(
reason: &str,
asked: Option<(u16, u16)>,
got: Option<(u16, u16)>,
line: &str,
) -> Event {
let mut data = serde_json::json!({
"reason": reason,
"reply_line": sanitize_line(line),
});
if let Some((server_port, client_port)) = asked {
data["queried_server_port"] = serde_json::json!(server_port);
data["queried_client_port"] = serde_json::json!(client_port);
}
if let Some((server_port, client_port)) = got {
data["reply_server_port"] = serde_json::json!(server_port);
data["reply_client_port"] = serde_json::json!(client_port);
}
Event::new(&IDENT_CLIENT_REPLY_MISMATCH_EVENT, data)
}
async fn read_reply<R>(reader: &mut R, timeout: Duration) -> ReplyRead
where
R: tokio::io::AsyncRead + Unpin,
{
let mut accumulated: Vec<u8> = Vec::new();
let mut buffer = vec![0u8; 512];
loop {
let read = match tokio::time::timeout(timeout, reader.read(&mut buffer)).await {
Err(_) => return ReplyRead::TimedOut,
Ok(read) => read,
};
match read {
Ok(0) => {
if accumulated.is_empty() {
return ReplyRead::Closed;
}
break;
}
Ok(n) => {
accumulated.extend_from_slice(&buffer[..n]);
if accumulated.contains(&b'\n') {
break;
}
if accumulated.len() > MAX_REPLY_BYTES {
let text = String::from_utf8_lossy(&accumulated).into_owned();
return ReplyRead::Oversize(text);
}
}
Err(e) => return ReplyRead::Failed(e.to_string()),
}
}
let text = String::from_utf8_lossy(&accumulated).into_owned();
ReplyRead::Line(text.lines().next().unwrap_or("").to_string())
}
async fn apply_action<W>(
result: ClientActionResult,
write_half: &Arc<Mutex<W>>,
pending: &PendingQuery,
client_id: ClientId,
log: &Log<'_>,
) -> Result<Applied>
where
W: tokio::io::AsyncWrite + Unpin,
{
match result {
ClientActionResult::Custom { name, data } if name == "ident_query" => {
let server_port = data
.get("server_port")
.and_then(|v| v.as_u64())
.context("ident_query without a server_port")?;
let server_port =
u16::try_from(server_port).context("ident server_port exceeds u16")?;
let client_port = data
.get("client_port")
.and_then(|v| v.as_u64())
.context("ident_query without a client_port")?;
let client_port =
u16::try_from(client_port).context("ident client_port exceeds u16")?;
let query = format!("{server_port} , {client_port}\r\n");
{
let mut writer = write_half.lock().await;
writer.write_all(query.as_bytes()).await?;
writer.flush().await?;
}
log.debug(format!(
"Ident client {client_id} asked about {server_port},{client_port}"
));
if let Ok(mut slot) = pending.lock() {
slot.get_or_insert((server_port, client_port));
}
Ok(Applied::Sent(query.len()))
}
ClientActionResult::Disconnect => {
let _ = write_half.lock().await.shutdown().await;
Ok(Applied::Disconnect)
}
_ => Ok(Applied::Sent(0)),
}
}
#[allow(clippy::too_many_arguments)]
async fn command_loop<W>(
mut command_rx: mpsc::Receiver<ClientCommand>,
protocol: Arc<IdentClientProtocol>,
write_half: Arc<Mutex<W>>,
pending: PendingQuery,
client_id: ClientId,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
) where
W: tokio::io::AsyncWrite + Unpin,
{
use crate::llm::actions::protocol_trait::Protocol;
while let Some(command) = command_rx.recv().await {
let log = Log::new(Some(&status_tx));
let action = command.action.clone();
let outcome = match protocol.execute_action(action.clone()) {
Err(e) => Ok(ClientSendOutcome::Rejected {
error: e.to_string(),
}),
Ok(result) => Self::apply_action(result, &write_half, &pending, client_id, &log)
.await
.map(|applied| match applied {
Applied::Disconnect => ClientSendOutcome::Disconnected,
Applied::Sent(0) => ClientSendOutcome::Executed {
detail: "executed (nothing to write)".to_string(),
},
Applied::Sent(bytes_sent) => ClientSendOutcome::Sent { bytes_sent },
}),
};
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 {
log.error(format!(
"Ident client {client_id} injected action failed: {e}"
));
}
let _ = status_tx.send("__UPDATE_UI__".to_string());
crate::client::command_support::reply(command, outcome);
if disconnect {
app_state.remove_client_handle(client_id).await;
let _ = status_tx.send("__UPDATE_UI__".to_string());
break;
}
}
}
}