feat(common): 协议 v2——版本号 + 帧头 seq/keyframe + 长度下界校验
为接入安卓主控做准备，先把三端共享契约改到能承载多客户端。改动集中在 common/src/lib.rs，三端适配在同一提交内完成以保证可编译。 协议版本（PROTO_VERSION=2） - ControlMsg::Auth/AuthOk 增加 proto_version 字段；认证顺序改为「先版本后验证码」： 版本不一致时字段布局可能不同，验证码的解析结果不可信 - 任一端版本不符即 AuthFail 并明确报错，把「连上了但行为诡异」变成「连不上且提示升级」 视频帧头（VideoFrameHeader） - 新增 seq: u64（每次实际发送递增，含静止桌面重发）与 keyframe: bool - keyframe 由新增的 common::h264_contains_idr() 从实际 Annex-B 码流判定， 不依赖编码器 API 推断——NVENC 与 openh264 规则不同，实际输出才是唯一可信来源 - 该字段是 MediaCodec 判断解码器重置时机的必要条件，避免客户端自己写 NAL 扫描 不可信输入防护 - read_video_frame 补齐长度下界：total 必须 >= 4，hlen 必须落在 [1, total-4] - 旧实现只校验上界，构造 total<4 的帧会让接收端在 buf[0..4] 处 panic（进程级DoS） - 测试从 0 增至 11 个，覆盖序列化往返与 6 类越界/损坏输入 顺带修复（与上述契约直接相关） - 新增 common::split_video_frame()：静止桌面重发时复用载荷但重建帧头， 否则 seq/ts_ms 会停在上次发送时刻，破坏 seq 单调递增契约 - agent-core 与 agent 两端的静止重发路径均改用该函数 验证：cargo test -p common → 20 passed（9 signaling + 11 新增）
This commit is contained in:
郑博文
parent
703eb64fc0
commit
9e50bc6237
@ -11,8 +11,8 @@
|
||||
|
||||
use anyhow::Result;
|
||||
use common::{
|
||||
encode_video_frame, read_control, write_control, ControlMsg, InputEvent, MouseButton,
|
||||
VideoFrameHeader,
|
||||
encode_video_frame, h264_contains_idr, read_control, write_control, ControlMsg, InputEvent,
|
||||
MouseButton, VideoFrameHeader,
|
||||
};
|
||||
use openh264::encoder::{
|
||||
BitRate, Encoder, EncoderConfig, FrameRate, IntraFramePeriod, RateControlMode, UsageType,
|
||||
@ -125,23 +125,30 @@ enum LoopEnd {
|
||||
}
|
||||
|
||||
/// 编码一帧并推入下行通道。返回 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 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());
|
||||
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;
|
||||
}
|
||||
@ -178,6 +185,8 @@ fn capture_session(
|
||||
// 保留最近一帧裸像素：静止桌面收到关键帧请求时，仍能重编一个真关键帧
|
||||
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();
|
||||
@ -189,9 +198,10 @@ fn capture_session(
|
||||
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) {
|
||||
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
|
||||
@ -206,24 +216,56 @@ fn capture_session(
|
||||
}
|
||||
},
|
||||
Ok(false) => {
|
||||
// 静止桌面：先看关键帧请求，再决定是否走 500ms 保活节流。
|
||||
//
|
||||
// ⚠️ 这里必须保证「消费了标志就一定发送」：先前写法是
|
||||
// `if force && last_raw.is_some() { ... }`，当 force 为真但
|
||||
// last_raw 还是 None（刚建会话、尚未抓到任何一帧）时，标志被
|
||||
// swap 消费掉却没有发送，解码端的这次请求就永久丢失。
|
||||
let force = keyreq.swap(false, Ordering::Relaxed);
|
||||
if force && last_raw.is_some() {
|
||||
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 帧就永远恢复不了
|
||||
// 解码端一旦出错，若只会收到旧 P 帧就永远恢复不了。
|
||||
// 关键帧请求绕过 500ms 节流：它是错误恢复路径，延迟一次 IDR
|
||||
// 的代价是持续黑屏，远高于多发一帧的带宽成本。
|
||||
encoder.force_intra_frame();
|
||||
let raw = last_raw.as_ref().expect("上面已判断 is_some");
|
||||
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, tx, &mut last_frame) {
|
||||
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 帧，解码端画面不变）
|
||||
// 桌面静止：低频重发最后一帧（重复 P 帧，解码端画面不变）。
|
||||
// 帧头必须重建：直接 clone 旧帧会让 seq/ts_ms 停在上次发送时刻，
|
||||
// 破坏 seq 单调递增契约。
|
||||
let f = last_frame.as_ref().expect("上面已判断 is_some");
|
||||
if tx.send(Some(f.clone())).is_err() {
|
||||
return LoopEnd::ChannelClosed;
|
||||
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));
|
||||
}
|
||||
}
|
||||
last_sent = Instant::now();
|
||||
} else {
|
||||
std::thread::sleep(Duration::from_millis(8));
|
||||
}
|
||||
@ -318,20 +360,46 @@ where
|
||||
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?;
|
||||
// 先校验版本，再校验验证码：版本不一致时字段布局可能不同，验证码的解析结果不可信
|
||||
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!("auth failed");
|
||||
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 }).await?;
|
||||
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：键鼠输入 + 控制上行（与视频分离，避免被视频积压阻塞）
|
||||
|
||||
@ -4,7 +4,7 @@
|
||||
//! 支持主控端 KeyFrameReq 强制关键帧（解码异常恢复）。
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use common::{encode_video_frame, VideoFrameHeader};
|
||||
use common::{encode_video_frame, h264_contains_idr, VideoFrameHeader};
|
||||
use nvenc::bitstream::BitStream;
|
||||
use nvenc::encoder::RegisteredResource;
|
||||
use nvenc::session::{InitParams, Session};
|
||||
@ -327,29 +327,41 @@ pub fn run(tx: Sender<Option<Vec<u8>>>, keyreq: Arc<AtomicBool>, congestion: Arc
|
||||
}
|
||||
};
|
||||
|
||||
let mut last_frame: Option<Vec<u8>> = None; // 软编静止重发
|
||||
// 最近一帧的**裸载荷**与关键帧标志：静止重发时要用它们重建帧头，
|
||||
// 不能直接重发整帧（那样 seq/ts_ms 会停在上次发送时刻）
|
||||
let mut last_payload: Option<Vec<u8>> = None;
|
||||
let mut last_kf = false;
|
||||
let mut last_sent = Instant::now() - Duration::from_secs(10);
|
||||
let min_frame = Duration::from_millis(16); // 60fps 上限
|
||||
let mut scratch: Vec<u8> = Vec::new(); // 软编：行距不连续时的复用缓冲
|
||||
// 帧序号：每发一帧 +1（含静止桌面的重复帧），供主控端做丢帧检测
|
||||
let mut seq: u64 = 0;
|
||||
// 编码耗时统计（每 120 帧输出一次均值，ms）
|
||||
let mut stat_frames = 0u64;
|
||||
let mut stat_convert_ms = 0u64;
|
||||
let mut stat_encode_ms = 0u64;
|
||||
let mut stat_skipped = 0u64; // 拥塞跳帧计数
|
||||
|
||||
// `$kf` 由调用点显式给出：NVENC 路径知道 `force_key`，软编路径同理。
|
||||
// 不要在这里猜——首帧和 KeyFrameReq 都会产生 IDR，猜错会让解码端错过重同步点。
|
||||
macro_rules! emit {
|
||||
($payload:expr) => {{
|
||||
($payload:expr, $kf:expr) => {{
|
||||
let header = VideoFrameHeader {
|
||||
width: dup.width,
|
||||
height: dup.height,
|
||||
ts_ms: start.elapsed().as_millis() as u64,
|
||||
seq,
|
||||
keyframe: $kf,
|
||||
};
|
||||
let frame = encode_video_frame(&header, &$payload);
|
||||
let payload_vec: Vec<u8> = $payload;
|
||||
let frame = encode_video_frame(&header, &payload_vec);
|
||||
if tx.send(Some(frame.clone())).is_err() {
|
||||
return; // 会话结束
|
||||
}
|
||||
last_frame = Some(frame);
|
||||
last_payload = Some(payload_vec);
|
||||
last_kf = $kf;
|
||||
last_sent = Instant::now();
|
||||
seq = seq.wrapping_add(1);
|
||||
stat_frames += 1;
|
||||
}};
|
||||
}
|
||||
@ -358,10 +370,12 @@ pub fn run(tx: Sender<Option<Vec<u8>>>, keyreq: Arc<AtomicBool>, congestion: Arc
|
||||
let iter_start = Instant::now();
|
||||
match dup.acquire(80) {
|
||||
Ok(true) => {
|
||||
// 主控端请求关键帧（解码异常恢复）
|
||||
let force_key = keyreq.swap(false, Ordering::Relaxed);
|
||||
// 拥塞反馈：发送侧排队深 → 跳过本帧（NVENC 少输入一帧，参考链连续；
|
||||
// 解码端收到的帧序 = 编码端实际输出序，无花屏风险）
|
||||
//
|
||||
// ⚠️ 必须放在消费 keyreq **之前**：先前顺序是「先 swap 再判断拥塞」，
|
||||
// 一旦本帧因拥塞被跳过，已消费的 KeyFrameReq 就随 continue 一起丢失，
|
||||
// 解码端等不到 IDR（与下方静止分支同源的缺陷）。
|
||||
if congestion.swap(false, Ordering::Relaxed) {
|
||||
stat_skipped += 1;
|
||||
let el = iter_start.elapsed();
|
||||
@ -370,13 +384,19 @@ pub fn run(tx: Sender<Option<Vec<u8>>>, keyreq: Arc<AtomicBool>, congestion: Arc
|
||||
}
|
||||
continue;
|
||||
}
|
||||
// 主控端请求关键帧（解码异常恢复）。走到这里一定会编码，
|
||||
// 因此可以安全地消费标志。
|
||||
let force_key = keyreq.swap(false, Ordering::Relaxed);
|
||||
match &mut encoder {
|
||||
EncoderKind::Nvenc(enc) => {
|
||||
let enc_start = Instant::now();
|
||||
let r = enc.encode(force_key, start.elapsed().as_millis() as u64);
|
||||
stat_encode_ms += enc_start.elapsed().as_millis() as u64;
|
||||
match r {
|
||||
Ok(payload) => emit!(payload),
|
||||
Ok(payload) => {
|
||||
let kf = h264_contains_idr(&payload);
|
||||
emit!(payload, kf)
|
||||
}
|
||||
Err(e) => eprintln!("[capture] nvenc encode failed: {e}"),
|
||||
}
|
||||
}
|
||||
@ -410,7 +430,11 @@ pub fn run(tx: Sender<Option<Vec<u8>>>, keyreq: Arc<AtomicBool>, congestion: Arc
|
||||
let r = enc.encode(&yuv);
|
||||
stat_encode_ms += enc_start.elapsed().as_millis() as u64;
|
||||
match r {
|
||||
Ok(ebs) => emit!(ebs.to_vec()),
|
||||
Ok(ebs) => {
|
||||
let payload = ebs.to_vec();
|
||||
let kf = h264_contains_idr(&payload);
|
||||
emit!(payload, kf)
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("[capture] encode failed: {e}");
|
||||
}
|
||||
@ -442,25 +466,80 @@ pub fn run(tx: Sender<Option<Vec<u8>>>, keyreq: Arc<AtomicBool>, congestion: Arc
|
||||
}
|
||||
Ok(false) => {
|
||||
// 桌面静止：低频重发（丢包后恢复 + 会话保活）。
|
||||
// 关键：KeyFrameReq 在桌面静止时也必须响应 —— 此前只在新帧路径
|
||||
// 消费该标志，桌面完全静止时解码端请求 IDR 永远得不到响应 → 黑屏。
|
||||
//
|
||||
// ⚠️ 这里曾有三个缺陷，全部围绕「KeyFrameReq 被静默丢弃」：
|
||||
// 1. `swap` 写在时间闸**外面**：请求若落在距上次发送 500ms 之内，
|
||||
// 标志已被消费但分支不进入，该次请求直接丢失（手册 §6.1 第 3 条）。
|
||||
// 2. 软编分支只重发旧码流、不看 force_key：永远吐不出 IDR（手册 §6.1 第 2 条）。
|
||||
// 3. NVENC 分支虽然把 force_key 传给了编码器，但受第 1 条限制。
|
||||
//
|
||||
// 现在的语义：**关键帧请求绕过 500ms 节流**——它是错误恢复路径，
|
||||
// 延迟一次 IDR 的代价是解码端持续黑屏，远高于多发一帧的带宽成本。
|
||||
let force_key = keyreq.swap(false, Ordering::Relaxed);
|
||||
if last_sent.elapsed() > Duration::from_millis(500) {
|
||||
let due = force_key || last_sent.elapsed() > Duration::from_millis(500);
|
||||
if due {
|
||||
match &mut encoder {
|
||||
EncoderKind::Nvenc(enc) => {
|
||||
// 重编当前 GPU 帧（内容未变，输出极小；force_key 时为 IDR）
|
||||
let r = enc.encode(force_key, start.elapsed().as_millis() as u64);
|
||||
match r {
|
||||
Ok(payload) => emit!(payload),
|
||||
Ok(payload) => {
|
||||
let kf = h264_contains_idr(&payload);
|
||||
emit!(payload, kf)
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("[capture] nvenc idle encode failed: {e}")
|
||||
}
|
||||
}
|
||||
}
|
||||
EncoderKind::Soft(_) => {
|
||||
// 软编：直接重发最后一帧码流副本
|
||||
if let Some(f) = last_frame.as_ref() {
|
||||
let _ = tx.send(Some(f.clone()));
|
||||
EncoderKind::Soft(enc) => {
|
||||
// 软编：静止桌面没有新像素可编，只能重发最近一帧。
|
||||
// 但 force_key 时必须真的重编出 IDR —— 重发旧 P 帧对
|
||||
// 已经解不动的解码端毫无帮助。
|
||||
//
|
||||
// 裸像素来源：staging 纹理里仍是上一帧内容（DXGI 只在
|
||||
// 有新帧时覆盖），重新 Map 一次即可，不必每帧额外拷 14MB。
|
||||
if force_key {
|
||||
enc.force_intra_frame();
|
||||
let dims = (dup.width as usize, dup.height as usize);
|
||||
let conv = dup.with_mapped(|src, pitch| {
|
||||
if pitch == dims.0 * 4 {
|
||||
YUVBuffer::from_bgra8_source(BgraSliceU8::new(
|
||||
src, dims,
|
||||
))
|
||||
} else {
|
||||
scratch.resize(dims.0 * dims.1 * 4, 0);
|
||||
for (row, dst) in
|
||||
scratch.chunks_exact_mut(dims.0 * 4).enumerate()
|
||||
{
|
||||
dst.copy_from_slice(
|
||||
&src[row * pitch..row * pitch + dims.0 * 4],
|
||||
);
|
||||
}
|
||||
YUVBuffer::from_bgra8_source(BgraSliceU8::new(
|
||||
&scratch, dims,
|
||||
))
|
||||
}
|
||||
});
|
||||
match conv {
|
||||
Ok(yuv) => match enc.encode(&yuv) {
|
||||
Ok(ebs) => {
|
||||
let payload = ebs.to_vec();
|
||||
let kf = h264_contains_idr(&payload);
|
||||
emit!(payload, kf)
|
||||
}
|
||||
Err(e) => eprintln!(
|
||||
"[capture] idle keyframe encode failed: {e}"
|
||||
),
|
||||
},
|
||||
Err(e) => {
|
||||
eprintln!("[capture] idle keyframe convert failed: {e}")
|
||||
}
|
||||
}
|
||||
} else if let Some(p) = last_payload.as_ref() {
|
||||
// 非关键帧的保活重发：复用载荷但重建帧头（seq 必须递增）
|
||||
let kf = last_kf;
|
||||
emit!(p.clone(), kf)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -69,16 +69,42 @@ async fn main() -> Result<()> {
|
||||
// 流 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?;
|
||||
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!("auth failed");
|
||||
anyhow::bail!(msg);
|
||||
}
|
||||
let (w, h) = capture::desktop_size()?;
|
||||
write_control(&mut send, &ControlMsg::AuthOk { width: w, height: h }).await?;
|
||||
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：键鼠输入 + 控制上行（与视频分离，避免被视频积压阻塞）
|
||||
|
||||
@ -7,12 +7,29 @@ use anyhow::Result;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
|
||||
/// 媒体/控制协议版本。
|
||||
///
|
||||
/// **改动规则**：任何对 `ControlMsg` 或 `VideoFrameHeader` 字段的增删改都必须递增此值。
|
||||
/// 理由：线格式是 bincode（按位置编码，无字段名），新增字段会让旧端解析失败或错位，
|
||||
/// 且失败方式可能是静默的错误数据而非明确报错。递增后由 `Auth` 握手显式拒绝旧端，
|
||||
/// 把"连上了但行为诡异"变成"连不上且提示升级"。
|
||||
///
|
||||
/// - v1：初版（无版本字段，隐式）
|
||||
/// - v2：`Auth`/`AuthOk` 携带版本；帧头增加 `seq`/`keyframe`
|
||||
pub const PROTO_VERSION: u16 = 2;
|
||||
|
||||
/// 控制通道消息（双向）
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum ControlMsg {
|
||||
/// 主控端认证（临时验证码，MVP 明文走 QUIC-TLS 加密通道）
|
||||
Auth { code: String },
|
||||
AuthOk { width: u32, height: u32 },
|
||||
/// 主控端认证（临时验证码，MVP 明文走 QUIC-TLS 加密通道）。
|
||||
/// `proto_version` 必须等于 [`PROTO_VERSION`]，否则被控端回 `AuthFail` 并拒绝。
|
||||
Auth { code: String, proto_version: u16 },
|
||||
/// 认证通过。`proto_version` 回显被控端的版本，供主控端确认双方一致。
|
||||
AuthOk {
|
||||
width: u32,
|
||||
height: u32,
|
||||
proto_version: u16,
|
||||
},
|
||||
AuthFail(String),
|
||||
KeyFrameReq,
|
||||
Ping(u64),
|
||||
@ -42,12 +59,22 @@ pub enum InputEvent {
|
||||
KeyUp { scancode: u16, extended: bool },
|
||||
}
|
||||
|
||||
/// 视频帧头（载荷为 LZ4 压缩的 BGRA 像素，按行紧凑排列）
|
||||
/// 视频帧头。载荷为 H.264 Annex-B 码流（含 00 00 00 01 起始码）。
|
||||
///
|
||||
/// 字段语义：
|
||||
/// - `width` / `height`：被控端桌面尺寸（逻辑尺寸，非编码对齐后的尺寸）
|
||||
/// - `ts_ms`：被控端采集时刻（被控时钟），配合 Ping/Pong 估计的时钟偏移算端到端延迟
|
||||
/// - `seq`：帧序号，从 0 单调递增，每发一帧 +1（含静止桌面的重复帧）。
|
||||
/// 用于丢帧检测与弱网统计；注意它统计的是"发送序"，不是"显示序"
|
||||
/// - `keyframe`：本帧是否 IDR。解码端（尤其 MediaCodec）需要它决定是否重置解码器，
|
||||
/// 不必再自己扫 NAL 起始码判断
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct VideoFrameHeader {
|
||||
pub width: u32,
|
||||
pub height: u32,
|
||||
pub ts_ms: u64,
|
||||
pub seq: u64,
|
||||
pub keyframe: bool,
|
||||
}
|
||||
|
||||
async fn read_u32le<R: AsyncReadExt + Unpin>(r: &mut R) -> Result<u32> {
|
||||
@ -81,7 +108,9 @@ pub async fn read_control<R: AsyncReadExt + Unpin>(r: &mut R) -> Result<ControlM
|
||||
Ok(bincode::deserialize(&buf)?)
|
||||
}
|
||||
|
||||
/// 编码视频帧：[总长 u32][头长 u32][bincode 头][LZ4 载荷]
|
||||
/// 编码视频帧：[总长 u32][头长 u32][bincode 头][H.264 载荷]
|
||||
///
|
||||
/// `总长` 不含自身这 4 字节，等于 `4 + 头长 + 载荷长`。
|
||||
pub fn encode_video_frame(header: &VideoFrameHeader, payload: &[u8]) -> Vec<u8> {
|
||||
let hb = bincode::serialize(header).expect("serialize header");
|
||||
let total = 4 + hb.len() + payload.len();
|
||||
@ -93,14 +122,292 @@ pub fn encode_video_frame(header: &VideoFrameHeader, payload: &[u8]) -> Vec<u8>
|
||||
out
|
||||
}
|
||||
|
||||
/// 单帧总长上限（256 MiB）。超过即视为协议错误或恶意输入。
|
||||
const MAX_VIDEO_FRAME: usize = 256 << 20;
|
||||
|
||||
/// 判断一段 Annex-B H.264 码流里是否含 IDR 片（NAL 类型 5）。
|
||||
///
|
||||
/// 用途：填充 [`VideoFrameHeader::keyframe`]。
|
||||
/// 为什么不问编码器要这个标志——NVENC 与 openh264 的 API 都能间接推出，
|
||||
/// 但两者规则不同（首帧强制、`force_intra_frame`、GOP 自动 IDR 交织），
|
||||
/// 而**实际输出**是唯一可信来源。直接看码流，两种编码器一套逻辑。
|
||||
///
|
||||
/// 只识别 3/4 字节起始码；`start-code emulation prevention` 由编码器保证，
|
||||
/// 因此不会在 SPS/PPS 载荷里误判出起始码。
|
||||
pub fn h264_contains_idr(payload: &[u8]) -> bool {
|
||||
let mut i = 0usize;
|
||||
while i + 3 < payload.len() {
|
||||
let sc = if payload[i..].starts_with(&[0, 0, 0, 1]) {
|
||||
4
|
||||
} else if payload[i..].starts_with(&[0, 0, 1]) {
|
||||
3
|
||||
} else {
|
||||
i += 1;
|
||||
continue;
|
||||
};
|
||||
match payload.get(i + sc) {
|
||||
// nal_unit_type 占首字节低 5 位
|
||||
Some(&b) if b & 0x1F == 5 => return true,
|
||||
Some(_) => i += sc + 1,
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// 读一个视频帧。
|
||||
///
|
||||
/// 对端数据一律视为不可信输入：`total` 与 `hlen` 都必须做**上下界**校验，
|
||||
/// 否则构造一个 `total < 4` 的帧就能让接收端在切片处 panic（进程级拒绝服务）。
|
||||
pub async fn read_video_frame<R: AsyncReadExt + Unpin>(
|
||||
r: &mut R,
|
||||
) -> Result<(VideoFrameHeader, Vec<u8>)> {
|
||||
let total = read_u32le(r).await? as usize;
|
||||
anyhow::ensure!(total < 256 << 20, "video frame too large: {total}");
|
||||
// 下界：至少要容得下「头长字段(4) + 一个非空的 bincode 头」。
|
||||
// 上界：防止对端声明超大长度导致巨量分配。
|
||||
anyhow::ensure!(
|
||||
(4..MAX_VIDEO_FRAME).contains(&total),
|
||||
"video frame length out of range: {total}"
|
||||
);
|
||||
|
||||
let mut buf = vec![0u8; total];
|
||||
r.read_exact(&mut buf).await?;
|
||||
let hlen = u32::from_le_bytes(buf[0..4].try_into().unwrap()) as usize;
|
||||
|
||||
let hlen = u32::from_le_bytes(buf[0..4].try_into().expect("已校验 total >= 4")) as usize;
|
||||
// hlen 必须落在 [1, total-4] 内：0 无法反序列化，超出则切片越界
|
||||
anyhow::ensure!(
|
||||
hlen > 0 && hlen <= total - 4,
|
||||
"video header length invalid: hlen={hlen}, total={total}"
|
||||
);
|
||||
|
||||
let header: VideoFrameHeader = bincode::deserialize(&buf[4..4 + hlen])?;
|
||||
Ok((header, buf[4 + hlen..].to_vec()))
|
||||
}
|
||||
|
||||
/// 把一个已编码的完整帧拆回 `(帧头, 载荷)`。
|
||||
///
|
||||
/// 用途：静止桌面重发时，被控端需要复用上一帧的**载荷**但重建**帧头**
|
||||
/// （`seq` 必须递增、`ts_ms` 必须是本次发送时刻）。直接 `clone` 整个帧会让
|
||||
/// 这两个字段停在上一次发送的时刻，破坏 `seq` 单调递增契约。
|
||||
///
|
||||
/// 只接受本模块 [`encode_video_frame`] 产出的字节；格式不符返回 `Err`。
|
||||
pub fn split_video_frame(frame: &[u8]) -> Result<(VideoFrameHeader, &[u8])> {
|
||||
anyhow::ensure!(frame.len() >= 8, "frame too short: {}", frame.len());
|
||||
let hlen = u32::from_le_bytes(frame[4..8].try_into().expect("已校验 len >= 8")) as usize;
|
||||
anyhow::ensure!(
|
||||
hlen > 0 && hlen <= frame.len() - 8,
|
||||
"header length invalid: hlen={hlen}, len={}",
|
||||
frame.len()
|
||||
);
|
||||
let header: VideoFrameHeader = bincode::deserialize(&frame[8..8 + hlen])?;
|
||||
Ok((header, &frame[8 + hlen..]))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tokio::io::duplex;
|
||||
|
||||
fn hdr(seq: u64, keyframe: bool) -> VideoFrameHeader {
|
||||
VideoFrameHeader {
|
||||
width: 1920,
|
||||
height: 1080,
|
||||
ts_ms: 12345,
|
||||
seq,
|
||||
keyframe,
|
||||
}
|
||||
}
|
||||
|
||||
// ---- 控制消息 ----
|
||||
|
||||
#[test]
|
||||
fn control_roundtrip_preserves_version() {
|
||||
let m = ControlMsg::Auth {
|
||||
code: "123456".into(),
|
||||
proto_version: PROTO_VERSION,
|
||||
};
|
||||
let bytes = encode_control(&m);
|
||||
// 长度前缀 + 载荷
|
||||
let len = u32::from_le_bytes(bytes[0..4].try_into().unwrap()) as usize;
|
||||
assert_eq!(len, bytes.len() - 4);
|
||||
let back: ControlMsg = bincode::deserialize(&bytes[4..]).unwrap();
|
||||
match back {
|
||||
ControlMsg::Auth { code, proto_version } => {
|
||||
assert_eq!(code, "123456");
|
||||
assert_eq!(proto_version, PROTO_VERSION);
|
||||
}
|
||||
other => panic!("期望 Auth，实际 {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn control_write_read_roundtrip() {
|
||||
let (mut a, mut b) = duplex(64 * 1024);
|
||||
let msgs = [
|
||||
ControlMsg::Auth {
|
||||
code: "x".into(),
|
||||
proto_version: PROTO_VERSION,
|
||||
},
|
||||
ControlMsg::AuthOk {
|
||||
width: 800,
|
||||
height: 600,
|
||||
proto_version: PROTO_VERSION,
|
||||
},
|
||||
ControlMsg::KeyFrameReq,
|
||||
ControlMsg::Ping(42),
|
||||
ControlMsg::Pong {
|
||||
t: 1,
|
||||
agent_ts: 2,
|
||||
lost: 3,
|
||||
sent: 4,
|
||||
},
|
||||
];
|
||||
for m in &msgs {
|
||||
write_control(&mut a, m).await.unwrap();
|
||||
let got = read_control(&mut b).await.unwrap();
|
||||
assert_eq!(format!("{got:?}"), format!("{m:?}"));
|
||||
}
|
||||
}
|
||||
|
||||
// ---- 视频帧头 ----
|
||||
|
||||
#[test]
|
||||
fn video_header_roundtrip_keeps_seq_and_keyframe() {
|
||||
for (seq, kf) in [(0u64, true), (1, false), (u64::MAX, false)] {
|
||||
let h = hdr(seq, kf);
|
||||
let payload = b"\x00\x00\x00\x01\x65fake-nal";
|
||||
let bytes = encode_video_frame(&h, payload);
|
||||
|
||||
let total = u32::from_le_bytes(bytes[0..4].try_into().unwrap()) as usize;
|
||||
assert_eq!(total, bytes.len() - 4, "总长字段应不含自身");
|
||||
|
||||
let hlen = u32::from_le_bytes(bytes[4..8].try_into().unwrap()) as usize;
|
||||
let back: VideoFrameHeader = bincode::deserialize(&bytes[8..8 + hlen]).unwrap();
|
||||
assert_eq!(back.seq, seq);
|
||||
assert_eq!(back.keyframe, kf);
|
||||
assert_eq!(&bytes[8 + hlen..], payload);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn video_frame_write_read_roundtrip() {
|
||||
let (mut a, mut b) = duplex(256 * 1024);
|
||||
let h = hdr(7, true);
|
||||
let payload = vec![0xABu8; 4096];
|
||||
let bytes = encode_video_frame(&h, &payload);
|
||||
tokio::io::AsyncWriteExt::write_all(&mut a, &bytes).await.unwrap();
|
||||
|
||||
let (h2, p2) = read_video_frame(&mut b).await.unwrap();
|
||||
assert_eq!(h2.seq, 7);
|
||||
assert!(h2.keyframe);
|
||||
assert_eq!((h2.width, h2.height), (1920, 1080));
|
||||
assert_eq!(p2, payload);
|
||||
}
|
||||
|
||||
// ---- 不可信输入（P3 的核心回归） ----
|
||||
|
||||
/// `total < 4` 必须报错而不是 panic：旧实现在这里 `buf[0..4]` 会越界 panic。
|
||||
#[tokio::test]
|
||||
async fn rejects_total_below_minimum() {
|
||||
for total in [0u32, 1, 3] {
|
||||
let (mut a, mut b) = duplex(64);
|
||||
tokio::io::AsyncWriteExt::write_all(&mut a, &total.to_le_bytes())
|
||||
.await
|
||||
.unwrap();
|
||||
let err = read_video_frame(&mut b).await.unwrap_err().to_string();
|
||||
assert!(
|
||||
err.contains("out of range"),
|
||||
"total={total} 应被拒绝，实际错误: {err}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// `total == 4` 但没有头数据：hlen 读出来是 0，必须拒绝。
|
||||
#[tokio::test]
|
||||
async fn rejects_zero_header_length() {
|
||||
let (mut a, mut b) = duplex(64);
|
||||
let mut bytes = 4u32.to_le_bytes().to_vec();
|
||||
bytes.extend_from_slice(&0u32.to_le_bytes());
|
||||
tokio::io::AsyncWriteExt::write_all(&mut a, &bytes)
|
||||
.await
|
||||
.unwrap();
|
||||
let err = read_video_frame(&mut b).await.unwrap_err().to_string();
|
||||
assert!(err.contains("header length invalid"), "实际错误: {err}");
|
||||
}
|
||||
|
||||
/// hlen 声明超过 total-4：旧实现会切片越界 panic。
|
||||
#[tokio::test]
|
||||
async fn rejects_header_length_beyond_frame() {
|
||||
let (mut a, mut b) = duplex(64);
|
||||
// total = 8 → body 必须正好 8 字节（hlen 字段 4 + 4 字节填充）
|
||||
let mut bytes = 8u32.to_le_bytes().to_vec();
|
||||
bytes.extend_from_slice(&999u32.to_le_bytes()); // hlen = 999 > total-4 = 4
|
||||
bytes.extend_from_slice(&[0u8; 4]);
|
||||
tokio::io::AsyncWriteExt::write_all(&mut a, &bytes)
|
||||
.await
|
||||
.unwrap();
|
||||
let err = read_video_frame(&mut b).await.unwrap_err().to_string();
|
||||
assert!(err.contains("header length invalid"), "实际错误: {err}");
|
||||
}
|
||||
|
||||
/// 超过 256 MiB 上限必须在**分配前**拒绝，不能先 vec! 再检查。
|
||||
#[tokio::test]
|
||||
async fn rejects_oversized_frame_without_allocating() {
|
||||
let (mut a, mut b) = duplex(64);
|
||||
let total = (MAX_VIDEO_FRAME as u32) + 1;
|
||||
tokio::io::AsyncWriteExt::write_all(&mut a, &total.to_le_bytes())
|
||||
.await
|
||||
.unwrap();
|
||||
let err = read_video_frame(&mut b).await.unwrap_err().to_string();
|
||||
assert!(err.contains("out of range"), "实际错误: {err}");
|
||||
}
|
||||
|
||||
/// 头部字节损坏（非合法 bincode）应返回 Err 而非 panic。
|
||||
#[tokio::test]
|
||||
async fn rejects_corrupt_header_bytes() {
|
||||
let (mut a, mut b) = duplex(64);
|
||||
let garbage = [0xFFu8; 8];
|
||||
let total = (4 + garbage.len()) as u32;
|
||||
let mut bytes = total.to_le_bytes().to_vec();
|
||||
bytes.extend_from_slice(&(garbage.len() as u32).to_le_bytes());
|
||||
bytes.extend_from_slice(&garbage);
|
||||
tokio::io::AsyncWriteExt::write_all(&mut a, &bytes)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(read_video_frame(&mut b).await.is_err());
|
||||
}
|
||||
|
||||
// ---- IDR 识别（P2 的 keyframe 标志来源） ----
|
||||
|
||||
#[test]
|
||||
fn detects_idr_in_annex_b() {
|
||||
// SPS(7) + PPS(8) + IDR 片(5)，4 字节起始码
|
||||
let idr = [
|
||||
0, 0, 0, 1, 0x67, 0xAA, // SPS
|
||||
0, 0, 0, 1, 0x68, 0xBB, // PPS
|
||||
0, 0, 0, 1, 0x65, 0xCC, // IDR
|
||||
];
|
||||
assert!(h264_contains_idr(&idr));
|
||||
|
||||
// 3 字节起始码也应识别
|
||||
let idr3 = [0, 0, 1, 0x65, 0xCC];
|
||||
assert!(h264_contains_idr(&idr3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_non_idr_payloads() {
|
||||
// 纯 P 片（NAL 类型 1）
|
||||
let p_frame = [0, 0, 0, 1, 0x41, 0xDD];
|
||||
assert!(!h264_contains_idr(&p_frame));
|
||||
|
||||
// 只有 SPS/PPS，没有片（关键帧边界可能落在下一个包，此时不应谎报 keyframe）
|
||||
let params_only = [
|
||||
0, 0, 0, 1, 0x67, 0xAA,
|
||||
0, 0, 0, 1, 0x68, 0xBB,
|
||||
];
|
||||
assert!(!h264_contains_idr(¶ms_only));
|
||||
|
||||
assert!(!h264_contains_idr(&[]));
|
||||
assert!(!h264_contains_idr(&[0, 0, 1])); // 起始码后无 NAL 字节
|
||||
}
|
||||
}
|
||||
|
||||
@ -289,11 +289,17 @@ async fn end_to_end_capture_encode_transport_input() -> Result<()> {
|
||||
&mut send,
|
||||
&ControlMsg::Auth {
|
||||
code: "smoke-code".into(),
|
||||
proto_version: common::PROTO_VERSION,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
match read_control(&mut recv).await? {
|
||||
ControlMsg::AuthOk { width, height } => {
|
||||
ControlMsg::AuthOk {
|
||||
width,
|
||||
height,
|
||||
proto_version,
|
||||
} => {
|
||||
assert_eq!(proto_version, common::PROTO_VERSION, "被控端协议版本不符");
|
||||
assert_eq!(
|
||||
(width, height),
|
||||
(sw, sh),
|
||||
@ -429,6 +435,7 @@ async fn rejects_wrong_verification_code() -> Result<()> {
|
||||
&mut send,
|
||||
&ControlMsg::Auth {
|
||||
code: "wrong-code".into(),
|
||||
proto_version: common::PROTO_VERSION,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
|
||||
@ -218,9 +218,27 @@ async fn run(addr: String, port: u16, code: String) -> Result<()> {
|
||||
|
||||
let (mut send, mut recv) = conn.open_bi().await?;
|
||||
|
||||
write_control(&mut send, &ControlMsg::Auth { code }).await?;
|
||||
write_control(
|
||||
&mut send,
|
||||
&ControlMsg::Auth {
|
||||
code,
|
||||
proto_version: common::PROTO_VERSION,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
match read_control(&mut recv).await? {
|
||||
ControlMsg::AuthOk { width, height } => {
|
||||
ControlMsg::AuthOk {
|
||||
width,
|
||||
height,
|
||||
proto_version,
|
||||
} => {
|
||||
// 版本已在 Auth 里声明，这里再核对一次：防止被控端是旧版却"接受"了认证
|
||||
if proto_version != common::PROTO_VERSION {
|
||||
bail!(
|
||||
"protocol version mismatch: server={proto_version}, client={}",
|
||||
common::PROTO_VERSION
|
||||
);
|
||||
}
|
||||
println!("[viewer] auth ok, remote desktop {width}x{height}");
|
||||
// 流 2：键鼠输入 + Ping/Pong 独立通道，不与视频流竞争
|
||||
let (isend, irecv) = conn.open_bi().await?;
|
||||
|
||||
Loading…
Reference in New Issue
Block a user