remote/agent/src/main.rs

//! 被控端 Agent：DXGI 桌面复制采集 → H.264 编码 → QUIC 下行视频；上行接收输入事件注入。
//! P0 优化：视频帧经 watch 单槽通道（latest-wins，网络慢时只发最新帧，不积压）；
//! 键鼠/控制走独立 QUIC 双向流，不与视频互相阻塞。
//! 会话模型：**串行多会话**——accept 循环一次只服务一个主控，会话结束后回到 accept
//! 等下一个连接（进程不退出）。改成并发多会话需要先解决 DXGI 单实例独占，暂不做。

mod capture;
mod input;

use anyhow::Result;
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};

fn parse_args() -> Result<(u16, String)> {
    let mut port = 47800u16;
    let mut code = "123456".to_string();
    let a: Vec<String> = std::env::args().collect();
    let mut i = 1;
    while i < a.len() {
        match a[i].as_str() {
            "--port" => port = a[i + 1].parse()?,
            "--code" => code = a[i + 1].clone(),
            _ => {}
        }
        i += 2;
    }
    Ok((port, code))
}

/// 显式声明进程 DPI 感知（**必须在任何采集/注入调用之前**，越早越好）。
///
/// 为什么不能省（3b 开工实测，见个人层 `docs/output/2026-10-31-3b-开工前实测.md` §4）：
/// 相对鼠标走「`GetCursorPos` + 增量」，而 `GetCursorPos` 的坐标空间取决于进程的
/// DPI 感知模式——UNAWARE 进程拿到的是**逻辑坐标**（本机 125% 缩放下是 1536x864），
/// 而 DXGI 采集与 `AuthOk` 下发的是**物理像素**（1920x1080）。两者混用会让
/// clamp 上界偏小、增量尺度偏小：指针到不了屏幕右/下边缘，且移动距离只有 80%。
///
/// 为什么不能依赖现状：本机 `agent.exe` 实测确实是 `PER_MONITOR_AWARE(2)`，但那是
/// `HKCU\...\AppCompatFlags\Layers` 里一条**路径绑定**的 `HIGHDPIAWARE` 覆盖造成的
/// ——同一个 exe 复制到 `%TEMP%` 启动就变成 `UNAWARE(0)`。换机器、换构建路径、
/// CI 产物、用户机器上都不成立，所以必须由代码显式设置。
///
/// 失败时只告警不退出：DPI 感知不可用是老系统（< Win10 1703）的情况，
/// 此时相对鼠标会退化（精度下降），但绝对注入与其余功能不受影响——
/// 让整个被控端起不来是更差的选择。
fn init_dpi_awareness() {
    use windows::Win32::UI::HiDpi::{
        GetThreadDpiAwarenessContext, SetProcessDpiAwarenessContext, GetAwarenessFromDpiAwarenessContext,
        DPI_AWARENESS_CONTEXT_PER_MONITOR_AWARE_V2,
    };
    match unsafe { SetProcessDpiAwarenessContext(DPI_AWARENESS_CONTEXT_PER_MONITOR_AWARE_V2) } {
        Ok(()) => println!("[agent] DPI awareness: PER_MONITOR_AWARE_V2 (set by this process)"),
        Err(e) => {
            // 失败有两种完全不同的原因，**必须分开报告**，否则日志会误导排查方向：
            //  a) 进程的感知模式已被别处设定（注册表 AppCompat 覆盖、清单、父进程）——
            //     ERROR_ACCESS_DENIED(0x80070005)，此时若已经是 per-monitor 就完全没问题；
            //  b) 系统过旧不支持该 API——这才是真正的降级。
            // 本机实测是 (a)：HKCU 的 AppCompatFlags\Layers 里有一条路径绑定的 HIGHDPIAWARE。
            let actual = unsafe {
                GetAwarenessFromDpiAwarenessContext(GetThreadDpiAwarenessContext())
            };
            // DPI_AWARENESS（WinUser.h）：INVALID=-1、UNAWARE=0、SYSTEM_AWARE=1、
            // PER_MONITOR_AWARE=2。注意 PER_MONITOR 是 **2** 而不是 3。
            let is_per_monitor = actual.0 == 2;
            if is_per_monitor {
                println!(
                    "[agent] DPI awareness: 已由外部设定为 PER_MONITOR_AWARE（无需本进程设置；{e}）"
                );
            } else {
                eprintln!(
                    "[agent] 警告：DPI 感知为 {}（非 per-monitor），且本进程设置失败（{e}）。\
                     相对鼠标的位移比例会与被采集桌面不一致；绝对触摸不受影响。",
                    actual.0
                );
            }
        }
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    // 第一件事：定下 DPI 感知，后续所有坐标运算都依赖它
    init_dpi_awareness();
    let (port, code) = parse_args()?;
    rustls::crypto::ring::default_provider()
        .install_default()
        .ok();

    // 自签名证书（MVP：客户端跳过校验；正式版改为固定公钥/账号体系）
    let ck = rcgen::generate_simple_self_signed(vec!["localhost".into()])?;
    let cert = ck.cert.der().clone();
    let key = PrivatePkcs8KeyDer::from(ck.key_pair.serialize_der());

    // 传输参数调优：
    // - 收发窗口放大，避免高 RTT 下吞吐被 窗口/RTT 死锁（默认 512KB 流窗口）
    // - 发送缓冲上限 64KB：QUIC 内部排队硬上限（实测 WiFi ~3Mbps 吞吐 → E2E ≤ ~170ms；
    //   突发吞吐 = 64KB/7ms ≈ 72Mbps，不限制带宽）。注意此窗口为连接级，流 2 输入/Ping
    //   数据量极小不受影响。配合下方拥塞反馈跳帧，把排队压在数帧以内
    let mut transport = quinn::TransportConfig::default();
    transport.receive_window(quinn::VarInt::from_u32(16 * 1024 * 1024));
    transport.stream_receive_window(quinn::VarInt::from_u32(8 * 1024 * 1024));
    transport.send_window(64 * 1024);
    // 空闲超时：主控端被强杀（非正常 close）时收不到它的 Ping，连接会一直占住
    // 串行 accept，下一个连接要等几十秒才能建立。10s = 20 个 Ping 周期，
    // 正常会话（Ping 每 500ms）永不触发，只用来回收死连接。
    transport.max_idle_timeout(Some(
        Duration::from_secs(10)
            .try_into()
            .expect("10s fits in VarInt"),
    ));
    let mut server_config = quinn::ServerConfig::with_single_cert(vec![cert], key.into())?;
    server_config.transport_config(Arc::new(transport));

    let endpoint = quinn::Endpoint::server(
        server_config,
        format!("0.0.0.0:{port}").parse()?,
    )?;
    println!("[agent] listening 0.0.0.0:{port}, verify code = {code}");

    // 串行多会话：一次只服务一个主控，结束后回到 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?;

    // 认证：先校验版本，再校验验证码。
    // 顺序刻意如此——版本不一致时字段布局可能不同，验证码的解析结果不可信。
    let first = read_control(&mut recv).await?;
    let auth_err: Option<String> = match &first {
        ControlMsg::Auth {
            code: c,
            proto_version,
        } => {
            if *proto_version != common::PROTO_VERSION {
                Some(format!(
                    "protocol version mismatch: client={proto_version}, server={}",
                    common::PROTO_VERSION
                ))
            } else if *c != code {
                Some("bad code".into())
            } else {
                None
            }
        }
        other => Some(format!("expected Auth, got {other:?}")),
    };
    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()?;
    write_control(
        &mut send,
        &ControlMsg::AuthOk {
            width: w,
            height: h,
            proto_version: common::PROTO_VERSION,
        },
    )
    .await?;
    println!("[agent] auth ok, desktop {w}x{h}");

    // 流 2：键鼠输入 + 控制上行（与视频分离，避免被视频积压阻塞）
    let (mut isend, mut irecv) = conn.accept_bi().await?;

    // 采集线程 → 视频帧（watch 单槽：新帧覆盖旧帧，latest-wins）
    let (tx, mut frames_rx) = watch::channel::<Option<Vec<u8>>>(None);
    let keyreq = Arc::new(AtomicBool::new(false));
    let keyreq_capture = keyreq.clone();
    // 拥塞反馈标志：发送侧 write 超过单帧预算(20ms) → 置位，采集线程跳帧降注入率
    let congestion = Arc::new(AtomicBool::new(false));
    let congestion_capture = congestion.clone();
    // 采集线程退出信号：会话结束必须等它真正退出才能 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::<ControlMsg>();
    let t0 = std::time::Instant::now();

    // 读循环（流 2）：输入事件 + Ping；Pong 附带视频下行方向的 QUIC 丢包统计
    // （conn.stats().path 统计的是本端发送方向，agent 的发送方向即视频下行）
    let keyreq_reader = keyreq.clone();
    let conn_stats = conn.clone();
    let reader = tokio::spawn(async move {
        loop {
            match read_control(&mut irecv).await {
                Ok(ControlMsg::Input(ev)) => input::handle(ev),
                Ok(ControlMsg::KeyFrameReq) => {
                    keyreq_reader.store(true, Ordering::Relaxed);
                }
                Ok(ControlMsg::Ping(t)) => {
                    let agent_ts = t0.elapsed().as_millis() as u64;
                    let p = conn_stats.stats().path;
                    let _ = ctl_tx.send(ControlMsg::Pong {
                        t,
                        agent_ts,
                        lost: p.lost_packets,
                        sent: p.sent_packets,
                    });
                }
                Ok(_) => {}
                Err(_) => break,
            }
        }
    });

    loop {
        tokio::select! {
            changed = frames_rx.changed() => {
                if changed.is_err() {
                    break; // 采集线程退出
                }
                if let Some(data) = frames_rx.borrow().clone() {
                    let wt = std::time::Instant::now();
                    let ok = send.write_all(&data).await.is_ok();
                    // write 耗时 > 单帧预算(16.7ms 的 1.2 倍) → QUIC 窗口排队深，
                    // 通知采集线程跳帧，把注入率降到网络实际吞吐水平
                    if wt.elapsed() > std::time::Duration::from_millis(20) {
                        congestion.store(true, Ordering::Relaxed);
                    }
                    if !ok {
                        break;
                    }
                }
            }
            ctl = ctl_rx.recv() => {
                match ctl {
                    Some(m) => {
                        if write_control(&mut isend, &m).await.is_err() {
                            break;
                        }
                    }
                    None => break,
                }
            }
        }
    }
    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(())
}