9c340813a8
本期只在新增文件里做，**未改动 Windows agent/ 一个字节**（与 NVENC 线并行，避免争抢）。
- 新增 agent-core：平台无关会话层 —— QUIC/TLS 握手、传输调优（沿用 512KB 发送窗口
结论）、验证码认证、会话/select 循环、latest-wins 单槽、输入分发，以及软件编码
视频管线（openh264）。同步跟进协议新增的 Pong{lost,sent} 下行丢包统计字段
- 新增 linux-agent：X11 GetImage 采集 + 整帧变更检测；XTest 注入 + 扫描码→X11
keycode 映射（主键区 = scancode+8，E0 扩展键必须查表，含单元测试）
- 协议零改动：复用 Set-1 扫描码与虚拟桌面归一化坐标，viewer 无需任何修改
- 新增端到端冒烟测试 tests/smoke.rs 与 scripts/linux-smoke.sh（自建 Xvfb，无需
图形桌面）：采集→编码→QUIC 解帧、认证（含错误码拒绝）、Ping/Pong（含丢包字段）、
静止桌面关键帧恢复、输入注入实证（注入后从另一条 X11 连接读回指针真实坐标）
- 该关键帧用例做过变异测试：把修复退回会失败、恢复后通过，确认断言有效
- .gitattributes：shell 脚本强制 LF（Windows checkout 成 CRLF 会导致容器内无法执行）
- 手册：修正已失效的工具链路径；补充 cargo 镜像；新增第 8 节（Linux）与 6.1 节
（三条待修问题，含两条 Windows 侧仍存在的关键帧缺陷）
验证：
- Linux(容器, USTC 源 + Xvfb): cargo test -p linux-agent 通过，2 passed
- Windows: 本期未改 agent/，default-members 未变，`cargo build --release` 行为不变
408 lines
15 KiB
Rust
408 lines
15 KiB
Rust
//! 平台无关的被控端会话层：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, 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 {
|
||
// 对标 UU：屏幕内容实时档；关键帧周期 300 帧 ≈ 10s
|
||
Self {
|
||
bitrate_bps: 5_000_000,
|
||
max_fps: 30.0,
|
||
intra_period: 300,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 解析 `--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 表示通道已关闭（会话结束）。
|
||
fn encode_send(
|
||
encoder: &mut Encoder,
|
||
raw: &[u8],
|
||
(w, h): (u32, u32),
|
||
ts_ms: u64,
|
||
tx: &watch::Sender<Option<Vec<u8>>>,
|
||
last_frame: &mut Option<Vec<u8>>,
|
||
) -> bool {
|
||
let header = VideoFrameHeader {
|
||
width: w,
|
||
height: h,
|
||
ts_ms,
|
||
};
|
||
let yuv = YUVBuffer::from_bgra8_source(BgraSliceU8::new(raw, (w as usize, h as usize)));
|
||
match encoder.encode(&yuv) {
|
||
Ok(ebs) => {
|
||
let frame = encode_video_frame(&header, &ebs.to_vec());
|
||
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);
|
||
|
||
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, tx, &mut last_frame) {
|
||
return LoopEnd::ChannelClosed; // 会话结束
|
||
}
|
||
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) => {
|
||
let force = keyreq.swap(false, Ordering::Relaxed);
|
||
if force && last_raw.is_some() {
|
||
// 静止桌面下的关键帧恢复：画面没变不代表可以不给关键帧——
|
||
// 解码端一旦出错，若只会收到旧 P 帧就永远恢复不了
|
||
encoder.force_intra_frame();
|
||
let raw = last_raw.as_ref().expect("上面已判断 is_some");
|
||
let ts = start.elapsed().as_millis() as u64;
|
||
if !encode_send(&mut encoder, raw, (w, h), ts, tx, &mut last_frame) {
|
||
return LoopEnd::ChannelClosed;
|
||
}
|
||
last_sent = Instant::now();
|
||
} else if last_sent.elapsed() > Duration::from_millis(500) && last_frame.is_some() {
|
||
// 桌面静止：低频重发最后一帧（重复 P 帧，解码端画面不变）
|
||
let f = last_frame.as_ref().expect("上面已判断 is_some");
|
||
if tx.send(Some(f.clone())).is_err() {
|
||
return LoopEnd::ChannelClosed;
|
||
}
|
||
last_sent = Instant::now();
|
||
} 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 code_ok = matches!(&first, ControlMsg::Auth { code: c } if *c == cfg.code);
|
||
if !code_ok {
|
||
write_control(&mut send, &ControlMsg::AuthFail("bad code".into())).await?;
|
||
send.finish()?;
|
||
// 必须等对端读到 AuthFail 并挂断再退出：直接 return 会让进程带走 QUIC
|
||
// 缓冲区里尚未发出的数据，客户端只会看到 "connection lost" 而不是 AuthFail
|
||
let _ = tokio::time::timeout(Duration::from_secs(5), conn.closed()).await;
|
||
anyhow::bail!("auth failed");
|
||
}
|
||
|
||
// 采集后端先建一次：既拿桌面尺寸回给主控端，也避免白建一个 duplication 只为读尺寸
|
||
let first_cap = capture_factory()?;
|
||
let (w, h) = first_cap.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();
|
||
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(())
|
||
}
|