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 流窗口）
    // - 发送缓冲上限 2MB：QUIC 内部排队有界（约 1 帧多），配合上游 latest-wins 丢帧
    //   把端到端排队延迟压在 ~0.4s 内
    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(2 * 1024 * 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 code_ok = matches!(&first, ControlMsg::Auth { code: c } if *c == code);
    if !code_ok {
        write_control(&mut send, &ControlMsg::AuthFail("bad code".into())).await?;
        send.finish()?;
        anyhow::bail!("auth failed");
    }
    let (w, h) = capture::desktop_size()?;
    write_control(&mut send, &ControlMsg::AuthOk { width: w, height: h }).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();
    std::thread::spawn(move || capture::run(tx, keyreq_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() {
                    if send.write_all(&data).await.is_err() {
                        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(())
}