remote/linux-agent/tests/smoke.rs

//! Linux 被控端端到端冒烟测试。
//!
//! 需要 X server（Xvfb 即可），通过 `DISPLAY` 指定；推荐用 `scripts/linux-smoke.sh` 一键跑。
//!
//! 覆盖链路：X11 采集 → openh264 编码 → QUIC/TLS 下行 → 主控端解帧；
//! 外加验证码认证、Ping/Pong、关键帧请求，以及**输入注入的实证断言**
//! （注入鼠标移动后，从另一条 X11 连接查询指针真实落点，而不是只看进程没崩）。

use anyhow::{bail, Context, Result};
use common::{read_control, read_video_frame, write_control, ControlMsg, InputEvent};
use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};
use rustls::pki_types::{CertificateDer, ServerName, UnixTime};
use rustls::{DigitallySignedStruct, SignatureScheme};
use std::io::{BufRead, BufReader};
use std::net::{TcpListener, ToSocketAddrs};
use std::process::{Child, Command, Stdio};
use std::sync::{mpsc, Arc};
use std::time::{Duration, Instant};
use x11rb::connection::Connection;
use x11rb::protocol::xproto::{ConnectionExt as _, Window};
use x11rb::rust_connection::RustConnection;

/// 与 viewer 一致：MVP 阶段跳过自签名证书校验
#[derive(Debug)]
struct SkipVerify;

impl ServerCertVerifier for SkipVerify {
    fn verify_server_cert(
        &self,
        _end_entity: &CertificateDer<'_>,
        _intermediates: &[CertificateDer<'_>],
        _server_name: &ServerName<'_>,
        _ocsp_response: &[u8],
        _now: UnixTime,
    ) -> Result<ServerCertVerified, rustls::Error> {
        Ok(ServerCertVerified::assertion())
    }
    fn verify_tls12_signature(
        &self,
        _message: &[u8],
        _certificate: &CertificateDer<'_>,
        _dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, rustls::Error> {
        Ok(HandshakeSignatureValid::assertion())
    }
    fn verify_tls13_signature(
        &self,
        _message: &[u8],
        _certificate: &CertificateDer<'_>,
        _dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, rustls::Error> {
        Ok(HandshakeSignatureValid::assertion())
    }
    fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
        vec![
            SignatureScheme::RSA_PKCS1_SHA256,
            SignatureScheme::ECDSA_NISTP256_SHA256,
            SignatureScheme::ED25519,
            SignatureScheme::RSA_PSS_SHA256,
        ]
    }
}

/// 被测进程的 RAII 句柄：测试无论成功失败都会收尸
struct AgentProc {
    child: Child,
}

impl Drop for AgentProc {
    fn drop(&mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}

fn expected_screen() -> (u32, u32) {
    let s = std::env::var("SMOKE_SCREEN").unwrap_or_else(|_| "1280x720".to_string());
    let (w, h) = s
        .split_once('x')
        .unwrap_or_else(|| panic!("SMOKE_SCREEN 需形如 1280x720，实际 {s:?}"));
    (w.parse().expect("宽"), h.parse().expect("高"))
}

fn free_port() -> Result<u16> {
    let l = TcpListener::bind("127.0.0.1:0")?;
    Ok(l.local_addr()?.port())
}

fn spawn_agent(port: u16, code: &str) -> Result<AgentProc> {
    let mut child = Command::new(env!("CARGO_BIN_EXE_linux-agent"))
        .args(["--port", &port.to_string(), "--code", code])
        .stdout(Stdio::piped())
        .stderr(Stdio::inherit())
        .spawn()
        .context("启动 linux-agent 失败")?;

    let stdout = child.stdout.take().expect("stdout 已 piped");
    let (tx, rx) = mpsc::channel();
    std::thread::spawn(move || {
        for line in BufReader::new(stdout).lines().map_while(|l| l.ok()) {
            println!("[agent] {line}");
            let _ = tx.send(line);
        }
    });

    // 等 QUIC 端口就绪
    let deadline = Instant::now() + Duration::from_secs(30);
    loop {
        let left = deadline.saturating_duration_since(Instant::now());
        if left.is_zero() {
            bail!("等待 agent 监听超时（DISPLAY 是否已设置？）");
        }
        match rx.recv_timeout(left) {
            Ok(line) if line.contains("listening") => return Ok(AgentProc { child }),
            Ok(_) => continue,
            Err(_) => bail!("等待 agent 监听超时，进程可能已退出"),
        }
    }
}

async fn connect(port: u16) -> Result<(quinn::Endpoint, quinn::Connection)> {
    let tls = rustls::ClientConfig::builder()
        .dangerous()
        .with_custom_certificate_verifier(Arc::new(SkipVerify))
        .with_no_client_auth();
    let qc: quinn::crypto::rustls::QuicClientConfig = Arc::new(tls).try_into()?;

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

    let mut client_config = quinn::ClientConfig::new(Arc::new(qc));
    client_config.transport_config(Arc::new(transport));

    let mut endpoint = quinn::Endpoint::client("0.0.0.0:0".parse()?)?;
    endpoint.set_default_client_config(client_config);

    let addr = ("127.0.0.1", port)
        .to_socket_addrs()?
        .next()
        .ok_or_else(|| anyhow::anyhow!("地址解析失败"))?;
    let conn = endpoint.connect(addr, "localhost")?.await?;
    // endpoint 必须与 conn 同生命周期，故一并返回
    Ok((endpoint, conn))
}

fn has_start_code(payload: &[u8]) -> bool {
    let head = &payload[..payload.len().min(64)];
    head.windows(4).any(|w| w == [0, 0, 0, 1]) || head.windows(3).any(|w| w == [0, 0, 1])
}

/// H.264 NAL 类型：5 = IDR 片（关键帧）
const NAL_IDR: u8 = 5;

/// 按 Annex-B 起始码扫描，判断载荷里是否含指定类型的 NAL
fn contains_nal(payload: &[u8], nal_type: u8) -> bool {
    let mut i = 0usize;
    while i + 3 <= payload.len() {
        let sc_len = if payload[i..].starts_with(&[0, 0, 0, 1]) {
            4
        } else if payload[i..].starts_with(&[0, 0, 1]) {
            3
        } else {
            i += 1;
            continue;
        };
        if i + sc_len >= payload.len() {
            return false;
        }
        if payload[i + sc_len] & 0x1F == nal_type {
            return true;
        }
        i += sc_len + 1;
    }
    false
}

/// 在根窗口左上角画一个**小**矩形，制造一次细微画面变化。
///
/// 刻意不用「整屏刷白」那类大变化：大变化会触发编码器的场景切换自动 IDR，
/// 于是「重发旧帧」和「真重编关键帧」拿到同样的字节，用例就失去区分度了。
fn perturb_screen(conn: &RustConnection, root: Window) -> Result<()> {
    use x11rb::protocol::xproto::{CreateGCAux, Rectangle};
    let gc = conn.generate_id()?;
    conn.create_gc(
        gc,
        root,
        &CreateGCAux::new().foreground(0x00FF_FFFF).graphics_exposures(0),
    )?;
    conn.poly_fill_rectangle(
        root,
        gc,
        &[Rectangle {
            x: 32,
            y: 32,
            width: 64,
            height: 64,
        }],
    )?;
    conn.free_gc(gc)?;
    conn.flush()?;
    println!("[smoke] 已扰动屏幕（根窗口画 64x64 白块）");
    Ok(())
}

/// 读帧直到拿到一个非关键帧（跳过静止桌面的低频重发），返回该帧载荷
async fn read_until_p_frame(recv: &mut quinn::RecvStream, timeout: Duration) -> Result<Vec<u8>> {
    let deadline = Instant::now() + timeout;
    while Instant::now() < deadline {
        let (_h, p) = tokio::time::timeout(Duration::from_secs(5), read_video_frame(recv))
            .await
            .context("等待扰动帧超时")??;
        if !contains_nal(&p, NAL_IDR) {
            return Ok(p);
        }
    }
    bail!("扰动后未拿到非关键帧，用例前置条件不成立（扰动量是否太大？）")
}

/// 请求关键帧后，容错地等待一个 IDR：期间允许出现静止桌面的 P 帧重发。
/// 返回是否在超时内拿到 IDR。
async fn wait_for_idr(recv: &mut quinn::RecvStream, timeout: Duration) -> Result<bool> {
    let deadline = Instant::now() + timeout;
    while Instant::now() < deadline {
        match tokio::time::timeout(Duration::from_secs(2), read_video_frame(recv)).await {
            Ok(Ok((_h, p))) => {
                if contains_nal(&p, NAL_IDR) {
                    return Ok(true);
                }
            }
            _ => break,
        }
    }
    Ok(false)
}

/// 归一化坐标 → 根窗口像素（与 linux-agent/src/input.rs 的 scale 完全一致）
fn scale(v: u16, extent: u32) -> i16 {
    if extent <= 1 {
        return 0;
    }
    ((v as i64 * (extent as i64 - 1)) / 65535) as i16
}

fn open_display() -> Result<(RustConnection, Window)> {
    let (conn, screen_num) = x11rb::connect(None)
        .map_err(|e| anyhow::anyhow!("连接 X11 失败（DISPLAY 是否已设置？）: {e}"))?;
    let root = conn.setup().roots[screen_num].root;
    Ok((conn, root))
}

fn pointer_position(conn: &RustConnection, root: Window) -> Result<(i16, i16)> {
    let r = conn.query_pointer(root)?.reply()?;
    Ok((r.root_x, r.root_y))
}

fn wait_pointer(
    conn: &RustConnection,
    root: Window,
    want: (i16, i16),
    timeout: Duration,
) -> Result<(i16, i16)> {
    let deadline = Instant::now() + timeout;
    let mut last = pointer_position(conn, root)?;
    while Instant::now() < deadline {
        last = pointer_position(conn, root)?;
        if (last.0 - want.0).abs() <= 2 && (last.1 - want.1).abs() <= 2 {
            return Ok(last);
        }
        std::thread::sleep(Duration::from_millis(50));
    }
    Ok(last)
}

#[tokio::test]
async fn end_to_end_capture_encode_transport_input() -> Result<()> {
    rustls::crypto::ring::default_provider()
        .install_default()
        .ok();

    let (sw, sh) = expected_screen();
    let port = free_port()?;
    let _agent = spawn_agent(port, "smoke-code")?;
    let (_endpoint, conn) = connect(port).await?;

    // 流 1：认证，并核对被控端上报的桌面尺寸就是 Xvfb 的尺寸
    let (mut send, mut recv) = conn.open_bi().await?;
    write_control(
        &mut send,
        &ControlMsg::Auth {
            code: "smoke-code".into(),
            proto_version: common::PROTO_VERSION,
        },
    )
    .await?;
    match read_control(&mut recv).await? {
        ControlMsg::AuthOk {
            width,
            height,
            proto_version,
        } => {
            assert_eq!(proto_version, common::PROTO_VERSION, "被控端协议版本不符");
            assert_eq!(
                (width, height),
                (sw, sh),
                "agent 上报桌面尺寸与实际屏幕不符"
            );
        }
        other => bail!("认证应成功，却收到 {other:?}"),
    }

    // 流 2：输入 + 控制
    let (mut isend, mut irecv) = conn.open_bi().await?;

    // Ping/Pong（延迟测量通道 + 下行丢包统计）
    write_control(&mut isend, &ControlMsg::Ping(4242)).await?;
    match read_control(&mut irecv).await? {
        ControlMsg::Pong {
            t,
            agent_ts,
            lost,
            sent,
        } => {
            assert_eq!(t, 4242, "Pong 回显的时间戳不一致");
            assert!(sent > 0, "下行发包计数应大于 0（视频流已在发）");
            println!("[smoke] Pong 收到，agent_ts={agent_ts}，下行 sent={sent} lost={lost}");
        }
        other => bail!("期望 Pong，收到 {other:?}"),
    }

    // 首帧：证明采集 + 编码 + 线格式整条链路通了
    let (h1, p1) = tokio::time::timeout(Duration::from_secs(15), read_video_frame(&mut recv))
        .await
        .context("等待首个视频帧超时")??;
    assert_eq!((h1.width, h1.height), (sw, sh), "视频帧头尺寸不符");
    assert!(!p1.is_empty(), "视频载荷为空");
    println!(
        "[smoke] 首帧 {}x{} ts={} 编码后 {} 字节，头部 {:02X?}",
        h1.width,
        h1.height,
        h1.ts_ms,
        p1.len(),
        &p1[..p1.len().min(12)]
    );
    assert!(
        has_start_code(&p1),
        "载荷不像 H.264（前 16 字节 {:02X?}）",
        &p1[..p1.len().min(16)]
    );

    // 关键帧恢复：先在**静止**桌面上请求关键帧。
    // 为了有区分度，先扰动一下屏幕，让缓存里落入一个非关键帧（P 帧），否则
    // 「重发旧帧」和「真重编关键帧」两种行为拿到的字节是一样的，测不出来。
    let (probe, root) = open_display()?;
    perturb_screen(&probe, root)?;
    let pp = read_until_p_frame(&mut recv, Duration::from_secs(10)).await?;
    println!(
        "[smoke] 扰动帧 {} 字节，头部 {:02X?}（应为 P 帧，无 IDR）",
        pp.len(),
        &pp[..pp.len().min(12)]
    );
    // 前置条件自检：没有这一步，用例可能因为「扰动本身产生了 IDR」而虚假通过
    assert!(
        !contains_nal(&pp, NAL_IDR),
        "扰动帧仍是 IDR，本用例无法区分重发旧帧与真重编关键帧"
    );
    tokio::time::sleep(Duration::from_millis(800)).await; // 等画面重新静止

    write_control(&mut isend, &ControlMsg::KeyFrameReq).await?;
    assert!(
        wait_for_idr(&mut recv, Duration::from_secs(3)).await?,
        "静止桌面请求关键帧后 3s 内未收到 IDR：解码端一旦出错将永远无法恢复"
    );
    println!("[smoke] 静止桌面请求关键帧 → 收到 IDR");

    // 输入注入实证：注入到屏幕中心，再从独立 X11 连接查询指针真实落点
    let before = pointer_position(&probe, root)?;
    write_control(
        &mut isend,
        &ControlMsg::Input(InputEvent::MouseMove {
            x: 32768,
            y: 32768,
        }),
    )
    .await?;
    let want = (scale(32768, sw), scale(32768, sh));
    let got = wait_pointer(&probe, root, want, Duration::from_secs(5))?;
    println!("[smoke] 指针 {before:?} → {got:?}（期望 ≈{want:?}）");
    assert!(
        (got.0 - want.0).abs() <= 2 && (got.1 - want.1).abs() <= 2,
        "XTest 注入未生效：期望指针 ≈{want:?}，实际 {got:?}"
    );

    // 键盘：链路不得因此崩溃（无窗口可断言落点）
    write_control(
        &mut isend,
        &ControlMsg::Input(InputEvent::KeyDown {
            scancode: 0x1E,
            extended: false,
        }),
    )
    .await?;
    write_control(
        &mut isend,
        &ControlMsg::Input(InputEvent::KeyUp {
            scancode: 0x1E,
            extended: false,
        }),
    )
    .await?;

    tokio::time::sleep(Duration::from_millis(300)).await;
    assert!(
        conn.close_reason().is_none(),
        "会话被意外关闭：{:?}",
        conn.close_reason()
    );

    println!("[smoke] 全部断言通过");
    Ok(())
}

#[tokio::test]
async fn rejects_wrong_verification_code() -> Result<()> {
    rustls::crypto::ring::default_provider()
        .install_default()
        .ok();

    let port = free_port()?;
    let _agent = spawn_agent(port, "right-code")?;
    let (_endpoint, conn) = connect(port).await?;

    let (mut send, mut recv) = conn.open_bi().await?;
    write_control(
        &mut send,
        &ControlMsg::Auth {
            code: "wrong-code".into(),
            proto_version: common::PROTO_VERSION,
        },
    )
    .await?;

    match read_control(&mut recv).await? {
        ControlMsg::AuthFail(_) => Ok(()),
        other => bail!("错误验证码竟被接受：{other:?}"),
    }
}