1pub mod actions;
2
3use crate::client::llm_budget::call_llm_for_client;
4use crate::client::pop3::actions::{
5 POP3_CLIENT_CONNECTED_EVENT, POP3_CLIENT_RESPONSE_RECEIVED_EVENT,
6};
7use crate::llm::actions::client_trait::Client;
8use crate::llm::ollama_client::OllamaClient;
9use crate::protocol::Event;
10use crate::state::app_state::AppState;
11use crate::state::client::{ClientId, ClientStatus};
12use anyhow::Result;
13use serde_json::json;
14use std::net::SocketAddr;
15use std::sync::Arc;
16use tokio::io::{AsyncWriteExt, BufReader};
17use tokio::net::TcpStream;
18use tokio::sync::mpsc;
19use tracing::{debug, error, info};
20
21pub use actions::Pop3ClientProtocol;
22
23pub struct CommandWriter<W> {
26 writer: W,
27 replies: std::collections::VecDeque<bool>,
28 failed: bool,
29}
30
31impl<W: tokio::io::AsyncWrite + Unpin> CommandWriter<W> {
32 pub fn new(writer: W) -> Self {
33 Self {
34 writer,
35 replies: Default::default(),
36 failed: false,
37 }
38 }
39
40 pub async fn send(&mut self, command: &str) -> Result<usize> {
41 if self.failed {
42 anyhow::bail!("POP3 writer is unusable after an interrupted write");
43 }
44 if command.is_empty() || command.contains(['\r', '\n']) {
45 anyhow::bail!("POP3 command must be one nonempty line");
46 }
47 if self.replies.len() >= 1024 {
48 anyhow::bail!("too many outstanding POP3 commands");
49 }
50 let mut words = command.split_ascii_whitespace();
51 let verb = words
52 .next()
53 .ok_or_else(|| anyhow::anyhow!("empty POP3 command"))?;
54 let multiline = match verb.to_ascii_uppercase().as_str() {
55 "RETR" | "TOP" | "CAPA" => true,
56 "LIST" | "UIDL" => words.next().is_none(),
57 _ => false,
58 };
59 self.failed = true;
60 self.replies.push_back(multiline);
61 tokio::time::timeout(std::time::Duration::from_secs(30), async {
62 self.writer.write_all(command.as_bytes()).await?;
63 self.writer.write_all(b"\r\n").await?;
64 self.writer.flush().await
65 })
66 .await
67 .map_err(|_| anyhow::anyhow!("POP3 command write deadline exceeded"))??;
68 self.failed = false;
69 Ok(command.len() + 2)
70 }
71
72 pub fn next_response_is_multiline(&mut self, response: &str) -> Result<bool> {
73 if self.failed {
74 anyhow::bail!("POP3 command write failed before its reply");
75 }
76 let multiline = self
77 .replies
78 .pop_front()
79 .ok_or_else(|| anyhow::anyhow!("unsolicited POP3 response"))?;
80 if !(response == "+OK"
81 || response.starts_with("+OK ")
82 || response == "-ERR"
83 || response.starts_with("-ERR "))
84 {
85 anyhow::bail!("invalid POP3 response status");
86 }
87 Ok(multiline && response.starts_with("+OK"))
88 }
89
90 async fn shutdown(&mut self) -> std::io::Result<()> {
91 self.writer.shutdown().await
92 }
93}
94
95pub struct Pop3Client;
96
97impl Pop3Client {
98 pub async fn connect_with_llm_actions(
110 remote_addr: String,
111 llm_client: OllamaClient,
112 app_state: Arc<AppState>,
113 status_tx: mpsc::UnboundedSender<String>,
114 client_id: ClientId,
115 startup_params: Option<crate::protocol::StartupParams>,
116 ) -> Result<SocketAddr> {
117 if let Some(params) = &startup_params {
118 if params.get_optional_bool("use_tls")?.unwrap_or(false) {
119 return Err(anyhow::anyhow!(
120 "POP3 client: `use_tls: true` was requested, but this client does not \
121 implement TLS - it speaks POP3 over a plain TCP socket. Connecting \
122 anyway would put the USER/PASS exchange for {remote_addr} on the wire \
123 in cleartext while reporting an encrypted session. Pass \
124 `use_tls: false` to accept a plaintext connection deliberately, or \
125 terminate TLS in front of the server."
126 ));
127 }
128 }
129 Self::connect_plain(remote_addr, llm_client, app_state, status_tx, client_id).await
130 }
131
132 async fn connect_plain(
133 remote_addr: String,
134 llm_client: OllamaClient,
135 app_state: Arc<AppState>,
136 status_tx: mpsc::UnboundedSender<String>,
137 client_id: ClientId,
138 ) -> Result<SocketAddr> {
139 let stream = TcpStream::connect(&remote_addr).await?;
140 let local_addr = stream.local_addr()?;
141
142 info!("POP3 client {} connected to {}", client_id, remote_addr);
143
144 app_state
146 .update_client_status(client_id, ClientStatus::Connected)
147 .await;
148
149 let (read_half, write_half) = tokio::io::split(stream);
150 let reader = BufReader::new(read_half);
151 let write_half = Arc::new(tokio::sync::Mutex::new(CommandWriter::new(write_half)));
152
153 let protocol = Arc::new(Pop3ClientProtocol);
154
155 let command_rx =
159 crate::client::command_support::register_command_channel(&app_state, client_id).await;
160
161 let cmd_state = app_state.clone();
164 let cmd_tx = status_tx.clone();
165 let cmd_write = write_half.clone();
166 let cmd_protocol = protocol.clone();
167 let cmd_task = tokio::spawn(async move {
168 Self::command_loop(
169 command_rx,
170 cmd_protocol,
171 cmd_write,
172 client_id,
173 cmd_state,
174 cmd_tx,
175 )
176 .await;
177 });
178 app_state.register_client_task(client_id, cmd_task).await;
179
180 let task_registrar = app_state.clone();
184 let task_handle = tokio::spawn(async move {
185 let cleanup_state = app_state.clone();
186 let cleanup_status = status_tx.clone();
187 let result = Self::read_loop(
188 reader,
189 write_half,
190 llm_client,
191 app_state,
192 status_tx,
193 client_id,
194 protocol,
195 remote_addr,
196 )
197 .await;
198 let status = match result {
199 Ok(()) => ClientStatus::Disconnected,
200 Err(e) => {
201 error!("POP3 client {} read loop error: {}", client_id, e);
202 ClientStatus::Error(e.to_string())
203 }
204 };
205 cleanup_state.update_client_status(client_id, status).await;
206 cleanup_state.remove_client_handle(client_id).await;
207 let _ = cleanup_status.send("__UPDATE_UI__".to_string());
208 });
209 task_registrar
210 .register_client_task(client_id, task_handle)
211 .await;
212
213 Ok(local_addr)
214 }
215
216 async fn command_loop<W>(
224 mut command_rx: tokio::sync::mpsc::Receiver<crate::state::client_handles::ClientCommand>,
225 protocol: Arc<Pop3ClientProtocol>,
226 write_half: Arc<tokio::sync::Mutex<CommandWriter<W>>>,
227 client_id: ClientId,
228 app_state: Arc<AppState>,
229 status_tx: mpsc::UnboundedSender<String>,
230 ) where
231 W: tokio::io::AsyncWrite + Unpin + Send + 'static,
232 {
233 use crate::llm::actions::protocol_trait::Protocol;
234 use crate::state::client_handles::ClientSendOutcome;
235 use crate::state::AccessLogOwner;
236
237 while let Some(command) = command_rx.recv().await {
238 let action = command.action.clone();
239 let outcome = match protocol.as_ref().execute_action(action.clone()) {
240 Err(e) => Ok(ClientSendOutcome::Rejected {
241 error: e.to_string(),
242 }),
243 Ok(result) => Self::apply_action(result, &write_half, client_id)
244 .await
245 .map(|applied| match applied {
246 Applied::Disconnect => ClientSendOutcome::Disconnected,
247 Applied::Sent(0) => ClientSendOutcome::Executed {
248 detail: "executed (nothing to write)".to_string(),
249 },
250 Applied::Sent(bytes_sent) => ClientSendOutcome::Sent { bytes_sent },
251 }),
252 };
253
254 let outcome_json = match &outcome {
255 Ok(outcome) => serde_json::to_value(outcome).unwrap_or(serde_json::Value::Null),
256 Err(e) => json!({"error": e.to_string()}),
257 };
258 app_state
259 .record_access_log(
260 AccessLogOwner::Client(client_id.as_u32()),
261 protocol.protocol_name(),
262 None,
263 "injected_action",
264 action,
265 vec![outcome_json],
266 )
267 .await;
268
269 let disconnect = matches!(outcome, Ok(ClientSendOutcome::Disconnected));
270 if let Err(e) = &outcome {
271 error!("POP3 client {} injected action failed: {}", client_id, e);
272 let _ = status_tx.send(format!(
273 "[WARN] Client {} injected action failed: {}",
274 client_id, e
275 ));
276 }
277 let _ = status_tx.send("__UPDATE_UI__".to_string());
278 crate::client::command_support::reply(command, outcome);
279
280 if disconnect {
281 let _ = write_half.lock().await.shutdown().await;
284 break;
285 }
286 }
287 }
288
289 async fn read_loop<R, W>(
290 mut reader: BufReader<R>,
291 write_half: Arc<tokio::sync::Mutex<CommandWriter<W>>>,
292 llm_client: OllamaClient,
293 app_state: Arc<AppState>,
294 status_tx: mpsc::UnboundedSender<String>,
295 client_id: ClientId,
296 protocol: Arc<Pop3ClientProtocol>,
297 remote_addr: String,
298 ) -> Result<()>
299 where
300 R: tokio::io::AsyncRead + Unpin,
301 W: tokio::io::AsyncWrite + Unpin + Send + 'static,
302 {
303 let mut line = String::new();
305 if crate::client::response_reader::read_expected_response_line(&mut reader, &mut line)
306 .await?
307 == 0
308 {
309 anyhow::bail!("POP3 server closed before its greeting");
310 }
311 let greeting = line.trim().to_string();
312
313 debug!("POP3 client {} received greeting: {}", client_id, greeting);
314
315 let is_ok = greeting.starts_with("+OK");
316
317 let (instruction, memory) = app_state
319 .with_client_mut(client_id, |client| {
320 (client.instruction.to_string(), client.memory.clone())
321 })
322 .await
323 .unwrap_or_default();
324
325 let event = Event::new(
327 &POP3_CLIENT_CONNECTED_EVENT,
328 json!({
329 "pop3_server": remote_addr,
330 "greeting": greeting,
331 "is_ok": is_ok,
332 }),
333 );
334
335 if let Err(e) = Self::handle_llm_response(
337 &event,
338 &llm_client,
339 &app_state,
340 &status_tx,
341 client_id,
342 &protocol,
343 &write_half,
344 &instruction,
345 &memory,
346 )
347 .await
348 {
349 error!(
350 "POP3 client {} failed to process greeting: {}",
351 client_id, e
352 );
353 return Err(e);
354 }
355
356 loop {
358 line.clear();
359 match crate::client::response_reader::read_response_line(&mut reader, &mut line).await {
360 Ok(0) => {
361 debug!("POP3 client {} connection closed by server", client_id);
362 break;
363 }
364 Ok(_) => {
365 let response = line.trim().to_string();
366 if response.is_empty() {
367 continue;
368 }
369
370 debug!("POP3 client {} received response: {}", client_id, response);
371
372 let is_multiline = write_half
374 .lock()
375 .await
376 .next_response_is_multiline(&response)?;
377
378 let full_response = if is_multiline {
379 crate::client::response_reader::read_dot_response(&mut reader, response)
380 .await?
381 } else {
382 response
383 };
384
385 let is_ok = full_response.starts_with("+OK");
386
387 let (instruction, memory) = app_state
389 .with_client_mut(client_id, |client| {
390 (client.instruction.to_string(), client.memory.clone())
391 })
392 .await
393 .unwrap_or_default();
394
395 let event = Event::new(
396 &POP3_CLIENT_RESPONSE_RECEIVED_EVENT,
397 json!({
398 "response": full_response,
399 "is_ok": is_ok,
400 }),
401 );
402
403 if let Err(e) = Self::handle_llm_response(
404 &event,
405 &llm_client,
406 &app_state,
407 &status_tx,
408 client_id,
409 &protocol,
410 &write_half,
411 &instruction,
412 &memory,
413 )
414 .await
415 {
416 error!(
417 "POP3 client {} failed to process response: {}",
418 client_id, e
419 );
420 break;
421 }
422 }
423 Err(e) => {
424 return Err(e.into());
425 }
426 }
427 }
428
429 info!("POP3 client {} disconnected", client_id);
430 Ok(())
431 }
432
433 async fn handle_llm_response<W>(
434 event: &Event,
435 llm_client: &OllamaClient,
436 app_state: &Arc<AppState>,
437 status_tx: &mpsc::UnboundedSender<String>,
438 client_id: ClientId,
439 protocol: &Arc<Pop3ClientProtocol>,
440 write_half: &Arc<tokio::sync::Mutex<CommandWriter<W>>>,
441 instruction: &str,
442 memory: &str,
443 ) -> Result<()>
444 where
445 W: tokio::io::AsyncWrite + Unpin,
446 {
447 let llm_result = call_llm_for_client(
449 llm_client,
450 app_state,
451 client_id.to_string(),
452 instruction,
453 memory,
454 Some(event),
455 protocol.as_ref(),
456 status_tx,
457 )
458 .await?;
459
460 if let Some(new_memory) = llm_result.memory_updates {
462 app_state
463 .with_client_mut(client_id, |client| {
464 client.memory = new_memory.clone();
465 })
466 .await;
467 }
468
469 for action in llm_result.actions {
471 let action_result = protocol.as_ref().execute_action(action)?;
472 if let Applied::Disconnect =
473 Self::apply_action(action_result, write_half, client_id).await?
474 {
475 return Ok(());
476 }
477 }
478
479 Ok(())
480 }
481
482 async fn apply_action<W>(
485 action_result: crate::llm::actions::client_trait::ClientActionResult,
486 write_half: &Arc<tokio::sync::Mutex<CommandWriter<W>>>,
487 client_id: ClientId,
488 ) -> Result<Applied>
489 where
490 W: tokio::io::AsyncWrite + Unpin,
491 {
492 use crate::llm::actions::client_trait::ClientActionResult;
493
494 match action_result {
495 ClientActionResult::Custom { name, data } => {
496 if name != "pop3_command" {
497 return Ok(Applied::Sent(0));
498 }
499 let command = data["command"]
500 .as_str()
501 .ok_or_else(|| anyhow::anyhow!("Missing command in action data"))?;
502
503 debug!("POP3 client {} sending command: {}", client_id, command);
504
505 let mut writer = write_half.lock().await;
506 Ok(Applied::Sent(writer.send(command).await?))
507 }
508 ClientActionResult::Disconnect => {
509 debug!("POP3 client {} disconnecting", client_id);
510 let mut writer = write_half.lock().await;
512 writer.send("QUIT").await?;
513 Ok(Applied::Disconnect)
514 }
515 _ => Ok(Applied::Sent(0)),
517 }
518 }
519}
520
521enum Applied {
523 Sent(usize),
525 Disconnect,
527}