remote/agent/src/main.rs

//! 被控端 Agent：DXGI 桌面复制采集 → LZ4 压缩 → QUIC 下行视频；上行接收输入事件注入。
//! P0 优化：视频帧经 watch 单槽通道（latest-wins，网络慢时只发最新帧，不积压）；
//! 键鼠/控制走独立 QUIC 双向流，不与视频互相阻塞。

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 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))
}

#[tokio::main]
async fn main() -> Result<()> {
    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);
    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}");

    let conn = endpoint
        .accept()
        .await
        .ok_or_else(|| anyhow::anyhow!("endpoint closed"))?
        .await?;
    println!("[agent] peer connected: {}", conn.remote_address());

    // 流 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()?;
        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();
    std::thread::spawn(move || capture::run(tx, keyreq_capture, congestion_capture));

    // 上行控制消息（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();
    println!("[agent] session closed");
    Ok(())
}