remote/agent-core/src/lib.rs

//! 平台无关的被控端会话层：QUIC/TLS 握手、验证码认证、H.264 编码下行、输入分发。
//!
//! 从 `agent/`（Windows）抽出来，供 Windows 与 Linux 被控端共用，避免两端各写一份
//! QUIC 握手/认证/丢帧逻辑造成协议漂移。平台差异只通过两个 trait 注入：
//!
//! - [`ScreenCapturer`]：屏幕采集（Windows = DXGI Desktop Duplication，Linux = X11）
//! - [`InputInjector`]：输入注入（Windows = SendInput，Linux = XTest）
//!
//! 线格式完全复用 `common`，**不改协议**：键盘沿用 Set-1 扫描码 + extended 标志，
//! 鼠标沿用虚拟桌面归一化坐标 0..=65535。

use anyhow::Result;
use common::{
    encode_video_frame, h264_contains_idr, read_control, write_control, ControlMsg, InputEvent,
    MouseButton, VideoFrameHeader,
};
use openh264::encoder::{
    BitRate, Encoder, EncoderConfig, FrameRate, IntraFramePeriod, RateControlMode, UsageType,
};
use openh264::formats::{BgraSliceU8, YUVBuffer};
use openh264::OpenH264API;
use rustls::pki_types::PrivatePkcs8KeyDer;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::{mpsc, watch};

/// 屏幕采集后端。
///
/// 语义刻意对齐 DXGI Desktop Duplication，使 Linux 侧只需实现同一套模型：
/// `acquire()` 阻塞至多 `timeout_ms`，返回是否拿到**新**帧；`copy_frame()` 再取出像素。
pub trait ScreenCapturer: Send {
    /// 桌面尺寸（物理像素）
    fn size(&self) -> (u32, u32);

    /// 阻塞至多 `timeout_ms`。
    /// `Ok(true)` = 有新帧（已拷进内部缓冲，可由 `copy_frame` 取出）；
    /// `Ok(false)` = 超时且画面未变（静止桌面）。
    fn acquire(&mut self, timeout_ms: u32) -> Result<bool>;

    /// 取出最近一帧，BGRA 紧凑排列（已去掉 row pitch），供 openh264 直接编码。
    fn copy_frame(&mut self) -> Result<Vec<u8>>;
}

/// 输入注入后端。实现需可在异步运行时中被多任务共享，故要求 `Send + Sync`。
pub trait InputInjector: Send + Sync + 'static {
    fn mouse_move(&self, x: u16, y: u16);
    fn mouse_button(&self, b: MouseButton, down: bool);
    fn wheel(&self, dy: i32);
    fn key(&self, scancode: u16, extended: bool, down: bool);
}

/// 把一条协议输入事件分发到注入后端。
pub fn dispatch_input(inj: &dyn InputInjector, ev: InputEvent) {
    match ev {
        InputEvent::MouseMove { x, y } => inj.mouse_move(x, y),
        InputEvent::MouseDown(b) => inj.mouse_button(b, true),
        InputEvent::MouseUp(b) => inj.mouse_button(b, false),
        InputEvent::Wheel { dy } => inj.wheel(dy),
        InputEvent::KeyDown { scancode, extended } => inj.key(scancode, extended, true),
        InputEvent::KeyUp { scancode, extended } => inj.key(scancode, extended, false),
    }
}

/// 会话参数
#[derive(Debug, Clone)]
pub struct SessionConfig {
    pub port: u16,
    pub code: String,
}

/// 编码档位。默认值即 Windows 现有档位，改动需两端同步评估。
#[derive(Debug, Clone, Copy)]
pub struct EncoderSettings {
    pub bitrate_bps: u32,
    pub max_fps: f32,
    pub intra_period: u32,
}

impl Default for EncoderSettings {
    fn default() -> Self {
        // 关键帧周期：以**时间**为准而非固定帧数，避免帧率变化后恢复窗口漂移。
        // 1 秒 @30fps = 30 帧。主恢复路径是 KeyFrameReq（协议 v2 后已确认生效），
        // 自动 IDR 只是兜底；此前 300 帧 @30fps 意味着兜底窗口长达 10 秒。
        let max_fps = 30.0f32;
        Self {
            bitrate_bps: 5_000_000,
            max_fps,
            intra_period: (max_fps as u32).max(1),
        }
    }
}

/// 解析 `--port` / `--code`（两端共用，保持命令行一致）
pub 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 + 1 < a.len() {
        match a[i].as_str() {
            "--port" => port = a[i + 1].parse()?,
            "--code" => code = a[i + 1].clone(),
            _ => {}
        }
        i += 2;
    }
    Ok((port, code))
}

fn make_encoder(enc: &EncoderSettings) -> Result<Encoder> {
    Ok(Encoder::with_api_config(
        OpenH264API::from_source(),
        EncoderConfig::new()
            .bitrate(BitRate::from_bps(enc.bitrate_bps))
            .max_frame_rate(FrameRate::from_hz(enc.max_fps))
            .rate_control_mode(RateControlMode::Bitrate)
            .usage_type(UsageType::ScreenContentRealTime)
            .intra_frame_period(IntraFramePeriod::from_num_frames(enc.intra_period)),
    )?)
}

enum LoopEnd {
    /// 下行通道关闭 —— 会话结束，采集线程退出
    ChannelClosed,
    /// 采集后端失效（Windows 的 ACCESS_LOST / X11 连接断开）—— 重建后继续
    Rebuild,
}

/// 编码一帧并推入下行通道。返回 false 表示通道已关闭（会话结束）。
///
/// `seq` 由调用方持有并自增：帧序号必须覆盖**每一次实际发送**（含静止桌面重发），
/// 不能由编码器计数（编码器只知道编码了多少帧，不知道发了多少帧）。
fn encode_send(
    encoder: &mut Encoder,
    raw: &[u8],
    (w, h): (u32, u32),
    ts_ms: u64,
    seq: u64,
    tx: &watch::Sender<Option<Vec<u8>>>,
    last_frame: &mut Option<Vec<u8>>,
) -> bool {
    let yuv = YUVBuffer::from_bgra8_source(BgraSliceU8::new(raw, (w as usize, h as usize)));
    match encoder.encode(&yuv) {
        Ok(ebs) => {
            let payload = ebs.to_vec();
            let header = VideoFrameHeader {
                width: w,
                height: h,
                ts_ms,
                seq,
                keyframe: h264_contains_idr(&payload),
            };
            let frame = encode_video_frame(&header, &payload);
            if tx.send(Some(frame.clone())).is_err() {
                return false;
            }
            *last_frame = Some(frame);
            true
        }
        Err(e) => {
            eprintln!("[capture] encode failed: {e}");
            true
        }
    }
}

/// 单次采集会话：抓帧 → 编码 → 推入 watch 单槽。
fn capture_session(
    cap: &mut dyn ScreenCapturer,
    tx: &watch::Sender<Option<Vec<u8>>>,
    keyreq: &AtomicBool,
    enc: &EncoderSettings,
) -> LoopEnd {
    let mut encoder = match make_encoder(enc) {
        Ok(e) => e,
        Err(e) => {
            eprintln!("[capture] encoder init failed: {e}");
            return LoopEnd::ChannelClosed;
        }
    };

    let (w, h) = cap.size();
    eprintln!("[capture] started {w}x{h}");
    let start = Instant::now();
    let min_frame = Duration::from_secs_f32(1.0 / enc.max_fps);
    let mut last_frame: Option<Vec<u8>> = None;
    // 保留最近一帧裸像素：静止桌面收到关键帧请求时，仍能重编一个真关键帧
    let mut last_raw: Option<Vec<u8>> = None;
    let mut last_sent = Instant::now() - Duration::from_secs(10);
    // 帧序号：覆盖每一次实际发送（含静止桌面重发），供主控端做丢帧检测
    let mut seq: u64 = 0;

    loop {
        let iter_start = Instant::now();
        match cap.acquire(80) {
            Ok(true) => match cap.copy_frame() {
                Ok(raw) => {
                    // 主控端请求关键帧（解码异常恢复）
                    if keyreq.swap(false, Ordering::Relaxed) {
                        encoder.force_intra_frame();
                    }
                    let ts = start.elapsed().as_millis() as u64;
                    if !encode_send(&mut encoder, &raw, (w, h), ts, seq, tx, &mut last_frame) {
                        return LoopEnd::ChannelClosed; // 会话结束
                    }
                    seq = seq.wrapping_add(1);
                    last_raw = Some(raw);
                    last_sent = Instant::now();
                    // 帧间隔节流到 max_fps
                    let el = iter_start.elapsed();
                    if el < min_frame {
                        std::thread::sleep(min_frame - el);
                    }
                }
                Err(e) => {
                    eprintln!("[capture] frame copy failed: {e}");
                    return LoopEnd::Rebuild;
                }
            },
            Ok(false) => {
                // 静止桌面：先看关键帧请求，再决定是否走 500ms 保活节流。
                //
                // ⚠️ 这里必须保证「消费了标志就一定发送」：先前写法是
                // `if force && last_raw.is_some() { ... }`，当 force 为真但
                // last_raw 还是 None（刚建会话、尚未抓到任何一帧）时，标志被
                // swap 消费掉却没有发送，解码端的这次请求就永久丢失。
                let force = keyreq.swap(false, Ordering::Relaxed);
                let can_key = last_raw.is_some();
                if force && !can_key {
                    // 还没有裸帧可重编：把请求还回去，让下一轮（拿到首帧后）继续生效
                    keyreq.store(true, Ordering::Relaxed);
                    eprintln!("[capture] KeyFrameReq 暂存：尚无裸帧可重编");
                    std::thread::sleep(Duration::from_millis(8));
                } else if force {
                    // 静止桌面下的关键帧恢复：画面没变不代表可以不给关键帧——
                    // 解码端一旦出错，若只会收到旧 P 帧就永远恢复不了。
                    // 关键帧请求绕过 500ms 节流：它是错误恢复路径，延迟一次 IDR
                    // 的代价是持续黑屏，远高于多发一帧的带宽成本。
                    encoder.force_intra_frame();
                    let raw = last_raw.as_ref().expect("can_key 已保证");
                    let ts = start.elapsed().as_millis() as u64;
                    if !encode_send(&mut encoder, raw, (w, h), ts, seq, tx, &mut last_frame) {
                        return LoopEnd::ChannelClosed;
                    }
                    seq = seq.wrapping_add(1);
                    last_sent = Instant::now();
                } else if last_sent.elapsed() > Duration::from_millis(500) && last_frame.is_some() {
                    // 桌面静止：低频重发最后一帧（重复 P 帧，解码端画面不变）。
                    // 帧头必须重建：直接 clone 旧帧会让 seq/ts_ms 停在上次发送时刻，
                    // 破坏 seq 单调递增契约。
                    let f = last_frame.as_ref().expect("上面已判断 is_some");
                    match common::split_video_frame(f) {
                        Ok((mut hdr, payload)) => {
                            hdr.seq = seq;
                            hdr.ts_ms = start.elapsed().as_millis() as u64;
                            let frame = encode_video_frame(&hdr, payload);
                            if tx.send(Some(frame)).is_err() {
                                return LoopEnd::ChannelClosed;
                            }
                            seq = seq.wrapping_add(1);
                            last_sent = Instant::now();
                        }
                        Err(e) => {
                            // 自己产出的帧解析失败 = 编码/封装路径有 bug。
                            // 丢弃缓存而不是每轮重复报错：静止桌面下这里会以 ~125Hz 触发。
                            eprintln!("[capture] 静止重发时帧格式异常，丢弃缓存: {e}");
                            last_frame = None;
                            std::thread::sleep(Duration::from_millis(8));
                        }
                    }
                } else {
                    std::thread::sleep(Duration::from_millis(8));
                }
            }
            Err(e) => {
                eprintln!("[capture] acquire failed: {e}");
                return LoopEnd::Rebuild;
            }
        }
    }
}

/// 采集主循环：独立线程运行；后端失效时经 `factory` 重建，通道关闭时退出。
/// 通道为 watch 单槽（latest-wins）：发送即覆盖，网络慢时不产生积压。
pub fn run_capture_loop<F>(
    tx: watch::Sender<Option<Vec<u8>>>,
    keyreq: Arc<AtomicBool>,
    mut cap: Box<dyn ScreenCapturer>,
    mut factory: F,
    enc: EncoderSettings,
) where
    F: FnMut() -> Result<Box<dyn ScreenCapturer>>,
{
    loop {
        match capture_session(cap.as_mut(), &tx, keyreq.as_ref(), &enc) {
            LoopEnd::ChannelClosed => return,
            LoopEnd::Rebuild => {}
        }
        std::thread::sleep(Duration::from_millis(300));
        loop {
            match factory() {
                Ok(c) => {
                    cap = c;
                    break;
                }
                Err(e) => {
                    eprintln!("[capture] init failed: {e}, retry in 2s");
                    std::thread::sleep(Duration::from_secs(2));
                }
            }
        }
    }
}

/// 跑一次完整的被控端会话：监听 → 认证 → 视频下行 + 输入上行。
///
/// `capture_factory` 可被多次调用（采集后端失效时重建），故需为 `FnMut`。
pub async fn serve<F>(
    cfg: SessionConfig,
    enc: EncoderSettings,
    mut capture_factory: F,
    injector: Arc<dyn InputInjector>,
) -> Result<()>
where
    F: FnMut() -> Result<Box<dyn ScreenCapturer>> + Send + 'static,
{
    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 流窗口）
    // - 发送缓冲上限 512KB：QUIC 内部排队 ≈ 3 帧 @5Mbps，E2E 延迟天花板 ~150ms；
    //   配合上游 latest-wins 丢帧，超量部分在进 QUIC 前就被丢弃
    //   （与 Windows 端 agent/src/main.rs 保持一致，沿用同一次调优结论）
    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(512 * 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:{}", cfg.port).parse()?)?;
    println!(
        "[agent] listening 0.0.0.0:{}, verify code = {}",
        cfg.port, cfg.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 != cfg.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 并挂断再退出：直接 return 会让进程带走 QUIC
        // 缓冲区里尚未发出的数据，客户端只会看到 "connection lost" 而不是 AuthFail
        let _ = tokio::time::timeout(Duration::from_secs(5), conn.closed()).await;
        anyhow::bail!(msg);
    }

    // 采集后端先建一次：既拿桌面尺寸回给主控端，也避免白建一个 duplication 只为读尺寸
    let first_cap = capture_factory()?;
    let (w, h) = first_cap.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();
    let enc_settings = enc;
    std::thread::spawn(move || {
        run_capture_loop(tx, keyreq_capture, first_cap, capture_factory, enc_settings)
    });

    // 上行控制消息（Pong）经通道交给主循环写流 2，避免多任务争抢
    let (ctl_tx, mut ctl_rx) = mpsc::unbounded_channel::<ControlMsg>();
    let t0 = Instant::now();

    // 读循环（流 2）：输入事件 + Ping；Pong 附带**视频下行方向**的 QUIC 丢包统计
    // （conn.stats().path 统计的是本端发送方向，被控端的发送方向正是视频下行）
    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)) => dispatch_input(injector.as_ref(), 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(())
}