diff --git a/agent-core/src/lib.rs b/agent-core/src/lib.rs index c623a91..d986129 100644 --- a/agent-core/src/lib.rs +++ b/agent-core/src/lib.rs @@ -192,6 +192,13 @@ fn capture_session( let mut seq: u64 = 0; loop { + // 会话结束守卫:发送端已 drop 接收器时立即退出。 + // 与 agent/src/capture.rs 同源——静止桌面分支不发送,只靠 send 失败退出 + // 会让线程在画面不动时永不退出(X11 无 DXGI 式独占,危害较小但同样泄漏线程)。 + if tx.receiver_count() == 0 { + eprintln!("[capture] no receiver, stop"); + return LoopEnd::ChannelClosed; + } let iter_start = Instant::now(); match cap.acquire(80) { Ok(true) => match cap.copy_frame() { diff --git a/agent/src/capture.rs b/agent/src/capture.rs index b760b50..a96d591 100644 --- a/agent/src/capture.rs +++ b/agent/src/capture.rs @@ -387,6 +387,14 @@ pub fn run(tx: Sender>>, keyreq: Arc, congestion: Arc } loop { + // 会话结束守卫:发送端已 drop 接收器时立即退出。 + // 只靠 `tx.send` 失败退出是不够的——桌面完全静止时这里根本不发送, + // 线程会永远持有 duplication,导致下一个会话的 DuplicateOutput 失败 + // (DXGI 每个输出仅允许一个活动实例)。 + if tx.receiver_count() == 0 { + eprintln!("[capture] no receiver, stop"); + return; + } let iter_start = Instant::now(); match dup.acquire(80) { Ok(true) => { diff --git a/agent/src/main.rs b/agent/src/main.rs index 56055a6..837f135 100644 --- a/agent/src/main.rs +++ b/agent/src/main.rs @@ -1,6 +1,8 @@ -//! 被控端 Agent:DXGI 桌面复制采集 → LZ4 压缩 → QUIC 下行视频;上行接收输入事件注入。 +//! 被控端 Agent:DXGI 桌面复制采集 → H.264 编码 → QUIC 下行视频;上行接收输入事件注入。 //! P0 优化:视频帧经 watch 单槽通道(latest-wins,网络慢时只发最新帧,不积压); //! 键鼠/控制走独立 QUIC 双向流,不与视频互相阻塞。 +//! 会话模型:**串行多会话**——accept 循环一次只服务一个主控,会话结束后回到 accept +//! 等下一个连接(进程不退出)。改成并发多会话需要先解决 DXGI 单实例独占,暂不做。 mod capture; mod input; @@ -10,6 +12,7 @@ use common::{read_control, write_control, ControlMsg}; use rustls::pki_types::PrivatePkcs8KeyDer; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; +use std::time::Duration; use tokio::io::AsyncWriteExt; use tokio::sync::{mpsc, watch}; @@ -59,13 +62,31 @@ async fn main() -> Result<()> { )?; println!("[agent] listening 0.0.0.0:{port}, verify code = {code}"); - let conn = endpoint - .accept() - .await - .ok_or_else(|| anyhow::anyhow!("endpoint closed"))? - .await?; - println!("[agent] peer connected: {}", conn.remote_address()); + // 串行多会话:一次只服务一个主控,结束后回到 accept。 + // 不做并发多会话——DXGI Desktop Duplication 每个输出仅允许一个活动实例, + // 第二个 DuplicateOutput 会让第一个收到 ACCESS_LOST。 + // 会话失败(含认证失败)只记日志,不能让进程退出——否则每次连接都要重启被控端。 + loop { + let Some(incoming) = endpoint.accept().await else { + println!("[agent] endpoint closed, exit"); + return Ok(()); + }; + match incoming.await { + Ok(conn) => { + println!("[agent] peer connected: {}", conn.remote_address()); + if let Err(e) = serve_session(conn, &code).await { + println!("[agent] session ended: {e:#}"); + } + println!("[agent] waiting for next connection"); + } + Err(e) => println!("[agent] handshake failed: {e}"), + } + } +} +/// 单次会话:认证 → 采集/发送/输入 → 会话结束。 +/// 返回 Err 表示会话异常结束(含认证失败),调用方继续 accept。 +async fn serve_session(conn: quinn::Connection, code: &str) -> Result<()> { // 流 1:认证 + 视频下行 let (mut send, mut recv) = conn.accept_bi().await?; @@ -93,6 +114,10 @@ async fn main() -> Result<()> { if let Some(msg) = auth_err { write_control(&mut send, &ControlMsg::AuthFail(msg.clone())).await?; send.finish()?; + // 必须等对端读到 AuthFail 再关连接:直接返回会让进程带走 QUIC 缓冲区里 + // 尚未发出的数据,客户端只会看到 "connection lost" 而不是 AuthFail。 + // (agent-core 已有同款修复,此处同源补齐——手册 §6.1 第 1 条) + let _ = tokio::time::timeout(Duration::from_secs(5), conn.closed()).await; anyhow::bail!(msg); } let (w, h) = capture::desktop_size()?; @@ -117,7 +142,14 @@ async fn main() -> Result<()> { // 拥塞反馈标志:发送侧 write 超过单帧预算(20ms) → 置位,采集线程跳帧降注入率 let congestion = Arc::new(AtomicBool::new(false)); let congestion_capture = congestion.clone(); - std::thread::spawn(move || capture::run(tx, keyreq_capture, congestion_capture)); + // 采集线程退出信号:会话结束必须等它真正退出才能 accept 下一个连接—— + // DXGI Desktop Duplication 每个输出只允许一个活动实例,旧线程还持有 duplication + // 时新线程的 DuplicateOutput 会让旧实例收到 ACCESS_LOST(表现为下一个会话黑屏)。 + let (done_tx, done_rx) = std::sync::mpsc::channel::<()>(); + std::thread::spawn(move || { + capture::run(tx, keyreq_capture, congestion_capture); + let _ = done_tx.send(()); + }); // 上行控制消息(Pong)经通道交给主循环写流 2,避免多任务争抢 let (ctl_tx, mut ctl_rx) = mpsc::unbounded_channel::(); @@ -182,6 +214,11 @@ async fn main() -> Result<()> { } } reader.abort(); + // 关掉流并丢弃接收端,让采集线程下一次 send 失败而退出 + let _ = send.finish(); + drop(frames_rx); + // 等采集线程真正退出后再返回,调用方才会 accept 下一个连接(DXGI 独占,见上) + let _ = tokio::task::spawn_blocking(move || done_rx.recv_timeout(Duration::from_secs(2))).await; println!("[agent] session closed"); Ok(()) }