pub mod actions;
pub use actions::FingerClientProtocol;
use anyhow::{Context, Result};
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
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::Event;
use crate::state::app_state::AppState;
use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
use crate::state::{AccessLogOwner, ClientId, ClientStatus};
use crate::utils::truncate_for_log;
use actions::{
best_effort_fields, sanitize_response, FingerQuerySpec, DEFAULT_FINGER_PORT,
FINGER_CLIENT_CONNECTED_EVENT, FINGER_CLIENT_RESPONSE_RECEIVED_EVENT, MAX_FOLLOWUP_DEPTH,
MAX_RESPONSE_BYTES,
};
#[derive(Clone, Copy, Debug)]
pub struct FingerClientConfig {
pub default_port: u16,
pub allow_forwarding: bool,
pub response_timeout_secs: u64,
}
impl Default for FingerClientConfig {
fn default() -> Self {
Self {
default_port: DEFAULT_FINGER_PORT,
allow_forwarding: actions::DEFAULT_ALLOW_FORWARDING,
response_timeout_secs: actions::DEFAULT_RESPONSE_TIMEOUT_SECS,
}
}
}
impl FingerClientConfig {
fn response_timeout(&self) -> Option<Duration> {
if self.response_timeout_secs == 0 {
None
} else {
Some(Duration::from_secs(self.response_timeout_secs))
}
}
}
type SentQuery = Arc<std::sync::Mutex<Option<FingerQuerySpec>>>;
enum Applied {
Sent(usize),
Refused(String),
Disconnect,
}
#[derive(Clone, Debug)]
struct Response {
text: String,
bytes: usize,
eof: bool,
truncated: bool,
}
struct SessionCtx {
protocol: Arc<FingerClientProtocol>,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
client_id: ClientId,
target: String,
instruction: String,
config: FingerClientConfig,
}
pub fn resolve_target(remote_addr: &str, default_port: u16) -> String {
let addr = remote_addr.trim();
if let Some(rest) = addr.strip_prefix('[') {
return match rest.find(']') {
Some(close) if addr.len() > close + 2 && addr.as_bytes()[close + 2] == b':' => {
addr.to_string()
}
Some(_) => format!("{addr}:{default_port}"),
None => format!("{addr}:{default_port}"),
};
}
match addr.matches(':').count() {
0 => format!("{addr}:{default_port}"),
1 => addr.to_string(),
_ => format!("[{addr}]:{default_port}"),
}
}
pub struct FingerClient;
impl FingerClient {
pub async fn connect_with_llm_actions(
remote_addr: String,
config: FingerClientConfig,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
client_id: ClientId,
) -> Result<SocketAddr> {
let target = resolve_target(&remote_addr, config.default_port);
let stream = TcpStream::connect(&target)
.await
.with_context(|| format!("Failed to connect to Finger server at {target}"))?;
let local_addr = stream.local_addr()?;
let remote_sock_addr = stream.peer_addr()?;
info!(
"Finger client {} connected to {} (local: {})",
client_id, remote_sock_addr, local_addr
);
if config.allow_forwarding {
warn!(
"Finger client {} allow_forwarding=true: 'user@host' queries will be put on the \
wire, asking {} to relay on our behalf (RFC 1288 3.2.1 calls this a security \
risk)",
client_id, target
);
let _ = status_tx.send(format!(
"[CLIENT] ⚠ Finger client {} will send forwarding queries (allow_forwarding=true)",
client_id
));
}
app_state
.update_client_status(client_id, ClientStatus::Connected)
.await;
let _ = status_tx.send(format!("[CLIENT] Finger client {} connected", client_id));
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(FingerClientProtocol::new());
let sent_query: SentQuery = Arc::new(std::sync::Mutex::new(None));
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(),
write_half.clone(),
sent_query.clone(),
config,
client_id,
app_state.clone(),
status_tx.clone(),
));
app_state.register_client_task(client_id, cmd_task).await;
let task_registrar = app_state.clone();
let session_state = app_state.clone();
let session_status = status_tx.clone();
let task_handle = tokio::spawn(async move {
Self::session(
read_half,
write_half,
protocol,
sent_query,
target,
config,
llm_client,
session_state.clone(),
session_status.clone(),
client_id,
)
.await;
session_state.remove_client_handle(client_id).await;
session_state
.update_client_status(client_id, ClientStatus::Disconnected)
.await;
let _ =
session_status.send(format!("[CLIENT] Finger client {} disconnected", client_id));
let _ = session_status.send("__UPDATE_UI__".to_string());
});
task_registrar
.register_client_task(client_id, task_handle)
.await;
Ok(local_addr)
}
#[allow(clippy::too_many_arguments)]
async fn session<R, W>(
read_half: R,
write_half: Arc<Mutex<W>>,
protocol: Arc<FingerClientProtocol>,
sent_query: SentQuery,
target: String,
config: FingerClientConfig,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
client_id: ClientId,
) where
R: tokio::io::AsyncRead + Unpin,
W: tokio::io::AsyncWrite + Unpin,
{
let Some(instruction) = app_state.get_instruction_for_client(client_id).await else {
warn!(
"Finger client {} has no instruction; nothing will be asked",
client_id
);
return;
};
let ctx = Arc::new(SessionCtx {
protocol: protocol.clone(),
llm_client,
app_state: app_state.clone(),
status_tx: status_tx.clone(),
client_id,
target: target.clone(),
instruction,
config,
});
let event = Event::new(
&FINGER_CLIENT_CONNECTED_EVENT,
serde_json::json!({
"remote_addr": target,
"allow_forwarding": config.allow_forwarding,
}),
);
let memory = app_state
.get_memory_for_client(client_id)
.await
.unwrap_or_default();
let mut hung_up = false;
match call_llm_for_client(
&ctx.llm_client,
&app_state,
client_id.to_string(),
&ctx.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;
}
for action in actions {
let result = match protocol.execute_action(action.clone()) {
Ok(result) => result,
Err(e) => {
error!("Finger client {} rejected action: {}", client_id, e);
let _ = status_tx.send(format!(
"[CLIENT] ✖ Finger client {} rejected an action: {}",
client_id, e
));
continue;
}
};
match Self::apply_action(result, &write_half, &sent_query, config, client_id)
.await
{
Ok(Applied::Sent(_)) => {}
Ok(Applied::Refused(reason)) => {
warn!("Finger client {} {}", client_id, reason);
let _ = status_tx
.send(format!("[CLIENT] ⚠ Finger client {} {}", client_id, reason));
}
Ok(Applied::Disconnect) => {
info!(
"Finger client {} disconnected before sending a query",
client_id
);
hung_up = true;
break;
}
Err(e) => {
error!("Finger client {} failed to send query: {}", client_id, 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) => {
error!("LLM error for Finger client {}: {}", client_id, e);
let _ = status_tx.send(format!(
"[CLIENT] ⚠ Finger client {} LLM error: {} (still connected; a query can be \
injected)",
client_id, e
));
}
}
if hung_up {
let _ = write_half.lock().await.shutdown().await;
return;
}
if sent_query.lock().map(|q| q.is_none()).unwrap_or(true) {
info!(
"Finger client {} has no query on the wire yet; waiting for an injected one or \
for {} to close",
client_id, target
);
}
let response = Self::read_response(read_half, config.response_timeout(), client_id).await;
debug!(
"Finger client {} read {} bytes (eof={}, truncated={})",
client_id, response.bytes, response.eof, response.truncated
);
trace!(
"Finger client {} response: {}",
client_id,
truncate_for_log(&response.text, 4096)
);
app_state.remove_client_handle(client_id).await;
let _ = status_tx.send("__UPDATE_UI__".to_string());
let spec = sent_query.lock().ok().and_then(|q| q.clone());
match spec {
Some(spec) => {
let _ = status_tx.send(format!(
"[CLIENT] Finger client {} received {} bytes for {:?}",
client_id,
response.bytes,
spec.query_text()
));
Self::handle_response(ctx, spec, response, 0).await;
}
None => {
info!(
"Finger client {} closed with no query ever sent ({} bytes read)",
client_id, response.bytes
);
}
}
}
fn handle_response(
ctx: Arc<SessionCtx>,
spec: FingerQuerySpec,
response: Response,
depth: usize,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>> {
Box::pin(async move {
let client_id = ctx.client_id;
let event = Event::new(
&FINGER_CLIENT_RESPONSE_RECEIVED_EVENT,
serde_json::json!({
"response": response.text,
"query": spec.query_text(),
"username": spec.username,
"verbose": spec.verbose,
"forward_host": spec.forward_host,
"bytes": response.bytes,
"eof": response.eof,
"truncated": response.truncated,
"best_effort": best_effort_fields(&response.text),
}),
);
let memory = ctx
.app_state
.get_memory_for_client(client_id)
.await
.unwrap_or_default();
let result = 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 result {
Ok(result) => result,
Err(e) => {
error!(
"LLM error for Finger client {} on response event: {}",
client_id, e
);
return;
}
};
if let Some(mem) = memory_updates {
ctx.app_state.set_memory_for_client(client_id, mem).await;
}
let mut followed = false;
for action in actions {
let result = match ctx.protocol.execute_action(action.clone()) {
Ok(result) => result,
Err(e) => {
error!("Finger client {} rejected action: {}", client_id, e);
continue;
}
};
match result {
ClientActionResult::Custom { name, data } if name == "finger_query" => {
let next = match FingerQuerySpec::from_json(&data) {
Ok(next) => next,
Err(e) => {
error!("Finger client {} built an invalid query: {}", client_id, e);
continue;
}
};
if next.is_forwarding() && !ctx.config.allow_forwarding {
warn!(
"Finger client {} decision=forward_refused query={:?} \
forward_host={:?}: RFC 1288 3.2.1 calls forwarding a security \
risk; start the client with allow_forwarding=true to permit it",
client_id,
next.query_text(),
next.forward_host
);
let _ = ctx.status_tx.send(format!(
"[CLIENT] ⚠ Finger client {} refused a forwarding follow-up",
client_id
));
continue;
}
if followed {
warn!(
"Finger client {} ignoring extra follow-up query {:?}: one query \
per turn, because each needs its own connection",
client_id,
next.query_text()
);
continue;
}
if depth + 1 >= MAX_FOLLOWUP_DEPTH {
warn!(
"Finger client {} reached MAX_FOLLOWUP_DEPTH ({}); not following \
up with {:?}",
client_id,
MAX_FOLLOWUP_DEPTH,
next.query_text()
);
let _ = ctx.status_tx.send(format!(
"[CLIENT] ⚠ Finger client {} hit the follow-up depth limit",
client_id
));
continue;
}
followed = true;
info!(
"Finger client {} following up with {:?} on a new connection to {} \
(depth {})",
client_id,
next.query_text(),
ctx.target,
depth + 1
);
match Self::run_query(&ctx.target, &next, ctx.config.response_timeout())
.await
{
Ok(next_response) => {
let _ = ctx.status_tx.send(format!(
"[CLIENT] Finger client {} received {} bytes for {:?}",
client_id,
next_response.bytes,
next.query_text()
));
Self::handle_response(ctx.clone(), next, next_response, depth + 1)
.await;
}
Err(e) => {
error!(
"Finger client {} follow-up query {:?} failed: {}",
client_id,
next.query_text(),
e
);
let _ = ctx.status_tx.send(format!(
"[CLIENT] ✖ Finger client {} follow-up failed: {}",
client_id, e
));
}
}
}
ClientActionResult::Disconnect => {
debug!(
"Finger client {} asked to disconnect; the server already closed",
client_id
);
}
ClientActionResult::WaitForMore => {
debug!(
"Finger client {} waiting; the server already closed this connection",
client_id
);
}
other => {
trace!(
"Finger client {} produced an inert result: {:?}",
client_id,
other
);
}
}
}
})
}
async fn run_query(
target: &str,
spec: &FingerQuerySpec,
timeout: Option<Duration>,
) -> Result<Response> {
let stream = TcpStream::connect(target)
.await
.with_context(|| format!("Finger follow-up could not reach {target}"))?;
let (read_half, mut write_half) = tokio::io::split(stream);
write_half
.write_all(spec.to_wire_line().as_bytes())
.await
.context("Finger follow-up failed to write the query")?;
write_half
.flush()
.await
.context("Finger follow-up failed to flush the query")?;
Ok(Self::read_response_inner(read_half, timeout).await)
}
async fn read_response<R>(
read_half: R,
timeout: Option<Duration>,
client_id: ClientId,
) -> Response
where
R: tokio::io::AsyncRead + Unpin,
{
let response = Self::read_response_inner(read_half, timeout).await;
if !response.eof && !response.truncated {
warn!(
"Finger client {} stopped reading after {} bytes without seeing EOF; the server \
did not close",
client_id, response.bytes
);
}
if response.truncated {
warn!(
"Finger client {} truncated the response at {} bytes",
client_id, MAX_RESPONSE_BYTES
);
}
response
}
async fn read_response_inner<R>(mut read_half: R, timeout: Option<Duration>) -> Response
where
R: tokio::io::AsyncRead + Unpin,
{
let mut raw: Vec<u8> = Vec::new();
let mut truncated = false;
let read_loop = async {
let mut buf = vec![0u8; 4096];
loop {
match read_half.read(&mut buf).await {
Ok(0) => return true,
Ok(n) => {
let room = MAX_RESPONSE_BYTES.saturating_sub(raw.len());
if n >= room {
raw.extend_from_slice(&buf[..room]);
truncated = true;
return false;
}
raw.extend_from_slice(&buf[..n]);
}
Err(e) => {
debug!("Finger client read error: {}", e);
return false;
}
}
}
};
let eof = match timeout {
Some(duration) => tokio::time::timeout(duration, read_loop)
.await
.unwrap_or(false),
None => read_loop.await,
};
let bytes = raw.len();
Response {
text: sanitize_response(&String::from_utf8_lossy(&raw)),
bytes,
eof,
truncated,
}
}
async fn apply_action<W>(
result: ClientActionResult,
write_half: &Arc<Mutex<W>>,
sent_query: &SentQuery,
config: FingerClientConfig,
client_id: ClientId,
) -> Result<Applied>
where
W: tokio::io::AsyncWrite + Unpin,
{
match result {
ClientActionResult::Custom { name, data } if name == "finger_query" => {
let spec = FingerQuerySpec::from_json(&data)?;
if spec.is_forwarding() && !config.allow_forwarding {
let reason = format!(
"decision=forward_refused query={:?} forward_host={:?}: RFC 1288 3.2.1 \
calls 'user@host' forwarding a security risk, and this client was not \
started with allow_forwarding=true. Nothing was sent.",
spec.query_text(),
spec.forward_host
);
return Ok(Applied::Refused(reason));
}
let line = spec.to_wire_line();
debug!(
"Finger client {} querying: {:?}",
client_id,
spec.query_text()
);
{
let mut writer = write_half.lock().await;
writer.write_all(line.as_bytes()).await?;
writer.flush().await?;
}
if let Ok(mut slot) = sent_query.lock() {
slot.get_or_insert(spec);
}
Ok(Applied::Sent(line.len()))
}
ClientActionResult::Disconnect => {
debug!("Finger client {} disconnecting", client_id);
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<FingerClientProtocol>,
write_half: Arc<Mutex<W>>,
sent_query: SentQuery,
config: FingerClientConfig,
client_id: ClientId,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
) where
W: tokio::io::AsyncWrite + Unpin,
{
while let Some(command) = command_rx.recv().await {
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, &sent_query, config, client_id)
.await
.map(|applied| match applied {
Applied::Disconnect => ClientSendOutcome::Disconnected,
Applied::Refused(reason) => {
ClientSendOutcome::Rejected { error: reason }
}
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));
match &outcome {
Ok(ClientSendOutcome::Rejected { error }) => {
warn!(
"Finger client {} refused an injected action: {}",
client_id, error
);
let _ = status_tx.send(format!(
"[CLIENT] ⚠ Finger client {} refused an injected action: {}",
client_id, error
));
}
Err(e) => {
error!("Finger client {} injected action failed: {}", client_id, e);
let _ = status_tx.send(format!(
"[CLIENT] ✖ Finger client {} injected action failed: {}",
client_id, 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;
}
}
}
}