1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
//! SOCKS5 client implementation
pub mod actions;
pub use actions::Socks5ClientProtocol;
use anyhow::{Context, Result};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::{mpsc, Mutex};
use tokio_socks::tcp::Socks5Stream;
use tracing::{error, info, trace, warn};
use crate::client::llm_budget::call_llm_for_client;
use crate::client::socks5::actions::{
SOCKS5_CLIENT_CONNECTED_EVENT, SOCKS5_CLIENT_DATA_RECEIVED_EVENT,
};
use crate::llm::ollama_client::OllamaClient;
use crate::llm::ClientLlmResult;
use crate::protocol::Event;
use crate::state::app_state::AppState;
use crate::state::{ClientId, ClientStatus};
/// Connection state for LLM processing
#[derive(Debug, Clone, PartialEq)]
enum ConnectionState {
Idle,
Processing,
Accumulating,
}
/// Per-client data for LLM handling
struct ClientData {
state: ConnectionState,
queued_data: Vec<u8>,
memory: String,
}
/// SOCKS5 client that connects to a target server through a SOCKS5 proxy
pub struct Socks5Client;
impl Socks5Client {
/// Connect to a target server through a SOCKS5 proxy with integrated LLM actions
pub async fn connect_with_llm_actions(
proxy_addr: String,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
client_id: ClientId,
startup_params: Option<crate::protocol::StartupParams>,
) -> Result<SocketAddr> {
// Extract target address from startup params
let target_addr = startup_params
.as_ref()
.map(|p| p.get_string("target_addr"))
.transpose()?
.context("Missing required startup parameter 'target_addr'")?;
// Extract optional authentication
let auth_username = startup_params
.as_ref()
.map(|p| p.get_optional_string("auth_username"))
.transpose()?
.flatten();
let auth_password = startup_params
.as_ref()
.map(|p| p.get_optional_string("auth_password"))
.transpose()?
.flatten();
info!(
"SOCKS5 client {} connecting to {} through proxy {}",
client_id, target_addr, proxy_addr
);
// Connect through SOCKS5 proxy
let stream = if let (Some(username), Some(password)) = (auth_username, auth_password) {
// Connect with authentication
Socks5Stream::connect_with_password(
proxy_addr.as_str(),
target_addr.as_str(),
username.as_str(),
password.as_str(),
)
.await
.context(format!(
"Failed to connect to {} through SOCKS5 proxy {} with auth",
target_addr, proxy_addr
))?
} else {
// Connect without authentication
Socks5Stream::connect(proxy_addr.as_str(), target_addr.as_str())
.await
.context(format!(
"Failed to connect to {} through SOCKS5 proxy {}",
target_addr, proxy_addr
))?
};
// Get addresses before consuming the stream
let tcp_stream = stream.into_inner();
let local_addr = tcp_stream.local_addr()?;
let proxy_sock_addr = tcp_stream.peer_addr()?;
info!(
"SOCKS5 client {} connected through proxy {} to target {} (local: {})",
client_id, proxy_sock_addr, target_addr, local_addr
);
// Update client state
app_state
.update_client_status(client_id, ClientStatus::Connected)
.await;
let _ = status_tx.send(format!(
"[CLIENT] SOCKS5 client {} connected to {} through proxy",
client_id, target_addr
));
let _ = status_tx.send("__UPDATE_UI__".to_string());
// Split stream (tcp_stream was already extracted above)
let (mut read_half, write_half) = tokio::io::split(tcp_stream);
let write_half_arc = Arc::new(Mutex::new(write_half));
// Initialize client data
let client_data = Arc::new(Mutex::new(ClientData {
state: ConnectionState::Idle,
queued_data: Vec::new(),
memory: String::new(),
}));
// Call LLM with connected event
let protocol = Arc::new(crate::client::socks5::actions::Socks5ClientProtocol::new());
let connected_event = Event::new(
&SOCKS5_CLIENT_CONNECTED_EVENT,
serde_json::json!({
"proxy_addr": proxy_addr,
"target_addr": target_addr,
}),
);
if let Some(instruction) = app_state.get_instruction_for_client(client_id).await {
// Copy the memory out under its own guard. Passing
// `&client_data.lock().await.memory` straight into the call would keep the guard
// alive for the whole `match` -- temporaries in a match scrutinee live until the
// match ends -- and the `Ok` arm below locks the same mutex again to store the
// model's memory update. `tokio::sync::Mutex` is not reentrant, so that is a
// permanent deadlock, reached only when the model happens to return memory.
let memory = { client_data.lock().await.memory.clone() };
match call_llm_for_client(
&llm_client,
&app_state,
client_id.to_string(),
&instruction,
&memory,
Some(&connected_event),
protocol.as_ref(),
&status_tx,
)
.await
{
Ok(ClientLlmResult {
actions,
memory_updates,
}) => {
// Update memory
if let Some(mem) = memory_updates {
client_data.lock().await.memory = mem;
}
// Execute initial actions
for action in actions {
use crate::llm::actions::client_trait::Client;
match protocol.as_ref().execute_action(action) {
Ok(
crate::llm::actions::client_trait::ClientActionResult::SendData(
bytes,
),
) => {
// The `Err` half used to be discarded outright, so a failed
// write to the tunnel produced no log line at any level, no
// status message and no state change -- the client reported
// success having put nothing on the wire. The protocol's own
// CLAUDE.md claimed these were "logged but not fatal"; only
// the second half was true.
match write_half_arc.lock().await.write_all(&bytes).await {
Ok(()) => trace!(
"SOCKS5 client {} sent {} bytes through tunnel",
client_id,
bytes.len()
),
Err(e) => error!(
"SOCKS5 client {} failed to write {} bytes to the tunnel: {}",
client_id,
bytes.len(),
e
),
}
}
Ok(
crate::llm::actions::client_trait::ClientActionResult::Disconnect,
) => {
warn!("SOCKS5 client {} requested disconnect immediately after connect", client_id);
}
_ => {}
}
}
}
Err(e) => {
error!(
"LLM error for SOCKS5 client {} on connect: {}",
client_id, e
);
}
}
}
// Spawn read loop
// Registered with AppState so stop_client can abort this task —
// dropping a JoinHandle only detaches it in Tokio.
let task_registrar = app_state.clone();
// Command channel: lets the dashboard inject actions into this loop
// via AppState::send_to_client (see client/command_support.rs).
let mut command_rx =
crate::client::command_support::register_command_channel(&app_state, client_id).await;
let task_handle = tokio::spawn(async move {
let mut buffer = vec![0u8; 8192];
loop {
let read_result = tokio::select! {
read = read_half.read(&mut buffer) => read,
Some(cmd) = command_rx.recv() => {
let disconnect = crate::client::command_support::handle_stream_client_command(
&crate::client::socks5::actions::Socks5ClientProtocol,
&write_half_arc,
cmd,
client_id,
&app_state,
&status_tx,
)
.await;
if disconnect {
app_state
.update_client_status(client_id, ClientStatus::Disconnected)
.await;
let _ = status_tx.send(format!(
"[CLIENT] SOCKS5 client {} disconnected (injected action)",
client_id
));
let _ = status_tx.send("__UPDATE_UI__".to_string());
break;
}
continue;
}
};
match read_result {
Ok(0) => {
info!("SOCKS5 client {} disconnected", client_id);
app_state
.update_client_status(client_id, ClientStatus::Disconnected)
.await;
let _ = status_tx
.send(format!("[CLIENT] SOCKS5 client {} disconnected", client_id));
let _ = status_tx.send("__UPDATE_UI__".to_string());
break;
}
Ok(n) => {
let data = buffer[..n].to_vec();
trace!(
"SOCKS5 client {} received {} bytes from target",
client_id,
n
);
// Handle data with LLM
let mut client_data_lock = client_data.lock().await;
match client_data_lock.state {
ConnectionState::Idle => {
// Process immediately
client_data_lock.state = ConnectionState::Processing;
drop(client_data_lock);
// Call LLM
if let Some(instruction) =
app_state.get_instruction_for_client(client_id).await
{
let protocol = Arc::new(
crate::client::socks5::actions::Socks5ClientProtocol::new(),
);
let event = Event::new(
&SOCKS5_CLIENT_DATA_RECEIVED_EVENT,
crate::client::socks5::actions::inbound_event_fields(&data),
);
// Same reasoning as the connected-event call above: the
// guard must not survive into the match arms.
let memory = { client_data.lock().await.memory.clone() };
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,
}) => {
// Update memory
if let Some(mem) = memory_updates {
client_data.lock().await.memory = mem;
}
// Execute actions
for action in actions {
use crate::llm::actions::client_trait::Client;
match protocol.as_ref().execute_action(action) {
Ok(crate::llm::actions::client_trait::ClientActionResult::SendData(bytes)) => {
match write_half_arc.lock().await.write_all(&bytes).await {
Ok(()) => trace!("SOCKS5 client {} sent {} bytes", client_id, bytes.len()),
Err(e) => error!("SOCKS5 client {} failed to write {} bytes to the tunnel: {}", client_id, bytes.len(), e),
}
}
Ok(crate::llm::actions::client_trait::ClientActionResult::Disconnect) => {
info!("SOCKS5 client {} disconnecting", client_id);
break;
}
_ => {}
}
}
}
Err(e) => {
error!(
"LLM error for SOCKS5 client {}: {}",
client_id, e
);
}
}
}
// Process queued data if any
let mut client_data_lock = client_data.lock().await;
if !client_data_lock.queued_data.is_empty() {
client_data_lock.queued_data.clear();
}
client_data_lock.state = ConnectionState::Idle;
}
ConnectionState::Processing => {
// Queue data
client_data_lock.queued_data.extend_from_slice(&data);
client_data_lock.state = ConnectionState::Accumulating;
}
ConnectionState::Accumulating => {
// Continue queuing
client_data_lock.queued_data.extend_from_slice(&data);
}
}
}
Err(e) => {
error!("SOCKS5 client {} read error: {}", client_id, e);
app_state
.update_client_status(client_id, ClientStatus::Error(e.to_string()))
.await;
let _ = status_tx.send("__UPDATE_UI__".to_string());
break;
}
}
}
// The loop owns the only receiver; dropping the registered handle
// makes later send_to_client calls fail fast instead of timing out.
app_state.remove_client_handle(client_id).await;
});
task_registrar
.register_client_task(client_id, task_handle)
.await;
Ok(local_addr)
}
}