From 9e50bc62377c3eab4c69ca1d85d5125025e03e38 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=83=91=E5=8D=9A=E6=96=87?= <1323515022@qq.com> Date: Wed, 28 Oct 2026 16:01:09 +0800 Subject: [PATCH] =?UTF-8?q?feat(common):=20=E5=8D=8F=E8=AE=AE=20v2?= =?UTF-8?q?=E2=80=94=E2=80=94=E7=89=88=E6=9C=AC=E5=8F=B7=20+=20=E5=B8=A7?= =?UTF-8?q?=E5=A4=B4=20seq/keyframe=20+=20=E9=95=BF=E5=BA=A6=E4=B8=8B?= =?UTF-8?q?=E7=95=8C=E6=A0=A1=E9=AA=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 为接入安卓主控做准备，先把三端共享契约改到能承载多客户端。改动集中在 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 新增） --- agent-core/src/lib.rs | 112 ++++++++++--- agent/src/capture.rs | 113 +++++++++++-- agent/src/main.rs | 38 ++++- common/src/lib.rs | 321 ++++++++++++++++++++++++++++++++++++- linux-agent/tests/smoke.rs | 9 +- viewer/src/main.rs | 22 ++- 6 files changed, 560 insertions(+), 55 deletions(-) diff --git a/agent-core/src/lib.rs b/agent-core/src/lib.rs index 713a3fc..c6c6f67 100644 --- a/agent-core/src/lib.rs +++ b/agent-core/src/lib.rs @@ -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>>, last_frame: &mut Option>, ) -> 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> = 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 = 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：键鼠输入 + 控制上行（与视频分离，避免被视频积压阻塞） diff --git a/agent/src/capture.rs b/agent/src/capture.rs index 9489495..a01c6f7 100644 --- a/agent/src/capture.rs +++ b/agent/src/capture.rs @@ -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>>, keyreq: Arc, congestion: Arc } }; - let mut last_frame: Option> = None; // 软编静止重发 + // 最近一帧的**裸载荷**与关键帧标志：静止重发时要用它们重建帧头， + // 不能直接重发整帧（那样 seq/ts_ms 会停在上次发送时刻） + let mut last_payload: Option> = 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 = 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 = $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>>, keyreq: Arc, 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>>, keyreq: Arc, 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>>, keyreq: Arc, 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>>, keyreq: Arc, 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) } } } diff --git a/agent/src/main.rs b/agent/src/main.rs index c96df0e..56055a6 100644 --- a/agent/src/main.rs +++ b/agent/src/main.rs @@ -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 = 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：键鼠输入 + 控制上行（与视频分离，避免被视频积压阻塞） diff --git a/common/src/lib.rs b/common/src/lib.rs index 8a32db2..f3d8891 100644 --- a/common/src/lib.rs +++ b/common/src/lib.rs @@ -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: &mut R) -> Result { @@ -81,7 +108,9 @@ pub async fn read_control(r: &mut R) -> Result Vec { 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 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: &mut R, ) -> Result<(VideoFrameHeader, Vec)> { 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 字节 + } +} diff --git a/linux-agent/tests/smoke.rs b/linux-agent/tests/smoke.rs index f40d539..acfffb0 100644 --- a/linux-agent/tests/smoke.rs +++ b/linux-agent/tests/smoke.rs @@ -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?; diff --git a/viewer/src/main.rs b/viewer/src/main.rs index 8521538..d81a04e 100644 --- a/viewer/src/main.rs +++ b/viewer/src/main.rs @@ -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?;