init: MVP remote desktop (agent/viewer/common) with H.264 pipeline

agent/src/capture.rs

@@ -0,0 +1,256 @@
+//! DXGI Desktop Duplication 采集（Windows 8+）。
+//! 帧处理：BGRA → openh264 H.264 编码（屏幕内容实时档，5Mbps）→ 线格式。
+//! 支持主控端 KeyFrameReq 强制关键帧（解码异常恢复）。
+
+use anyhow::{anyhow, Result};
+use common::{encode_video_frame, VideoFrameHeader};
+use openh264::encoder::{
+    BitRate, Encoder, EncoderConfig, FrameRate, IntraFramePeriod, RateControlMode, UsageType,
+};
+use openh264::formats::{BgraSliceU8, YUVBuffer};
+use openh264::OpenH264API;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use tokio::sync::watch::Sender;
+use windows::core::Interface;
+use windows::Win32::Foundation::HMODULE;
+use windows::Win32::Graphics::Direct3D::{
+    D3D_DRIVER_TYPE_HARDWARE, D3D_FEATURE_LEVEL_11_0,
+};
+use windows::Win32::Graphics::Direct3D11::{
+    D3D11CreateDevice, ID3D11Device, ID3D11DeviceContext, ID3D11Resource, ID3D11Texture2D,
+    D3D11_CREATE_DEVICE_FLAG, D3D11_CPU_ACCESS_READ, D3D11_MAP, D3D11_MAP_READ,
+    D3D11_MAPPED_SUBRESOURCE, D3D11_SDK_VERSION, D3D11_TEXTURE2D_DESC, D3D11_USAGE_STAGING,
+};
+use windows::Win32::Graphics::Dxgi::Common::{DXGI_FORMAT_B8G8R8A8_UNORM, DXGI_SAMPLE_DESC};
+use windows::Win32::Graphics::Dxgi::{
+    IDXGIDevice, IDXGIOutput1, IDXGIOutputDuplication, IDXGIResource, DXGI_ERROR_WAIT_TIMEOUT,
+    DXGI_OUTDUPL_DESC, DXGI_OUTDUPL_FRAME_INFO,
+};
+
+use std::time::{Duration, Instant};
+
+pub struct Duplication {
+    _device: ID3D11Device,
+    context: ID3D11DeviceContext,
+    staging_res: ID3D11Resource,
+    dup: IDXGIOutputDuplication,
+    pub width: u32,
+    pub height: u32,
+}
+
+impl Duplication {
+    pub fn new() -> Result<Self> {
+        unsafe {
+            let mut device: Option<ID3D11Device> = None;
+            let mut context: Option<ID3D11DeviceContext> = None;
+            D3D11CreateDevice(
+                None,
+                D3D_DRIVER_TYPE_HARDWARE,
+                HMODULE::default(),
+                D3D11_CREATE_DEVICE_FLAG(0),
+                Some(&[D3D_FEATURE_LEVEL_11_0]),
+                D3D11_SDK_VERSION,
+                Some(&mut device),
+                None,
+                Some(&mut context),
+            )
+            .map_err(|e| anyhow!("D3D11CreateDevice failed: {e}"))?;
+            let device = device.ok_or_else(|| anyhow!("no d3d device"))?;
+            let context = context.ok_or_else(|| anyhow!("no d3d context"))?;
+
+            let dxgi: IDXGIDevice = device.cast()?;
+            let adapter = dxgi.GetAdapter()?;
+            let output = adapter
+                .EnumOutputs(0)
+                .map_err(|e| anyhow!("no display output: {e}"))?;
+            let output1: IDXGIOutput1 = output.cast()?;
+            let dup: IDXGIOutputDuplication = output1
+                .DuplicateOutput(&device)
+                .map_err(|e| anyhow!("DuplicateOutput failed: {e}"))?;
+
+            let desc: DXGI_OUTDUPL_DESC = dup.GetDesc();
+            let (width, height) = (desc.ModeDesc.Width, desc.ModeDesc.Height);
+            anyhow::ensure!(width > 0 && height > 0, "bad desktop size");
+
+            let sdesc = D3D11_TEXTURE2D_DESC {
+                Width: width,
+                Height: height,
+                MipLevels: 1,
+                ArraySize: 1,
+                Format: DXGI_FORMAT_B8G8R8A8_UNORM,
+                SampleDesc: DXGI_SAMPLE_DESC {
+                    Count: 1,
+                    Quality: 0,
+                },
+                Usage: D3D11_USAGE_STAGING,
+                BindFlags: 0,
+                CPUAccessFlags: D3D11_CPU_ACCESS_READ.0 as u32,
+                MiscFlags: 0,
+            };
+            let mut staging_opt: Option<ID3D11Texture2D> = None;
+            device
+                .CreateTexture2D(&sdesc, None, Some(&mut staging_opt))
+                .map_err(|e| anyhow!("CreateTexture2D failed: {e}"))?;
+            let staging = staging_opt.ok_or_else(|| anyhow!("no staging texture"))?;
+            let staging_res: ID3D11Resource = staging.cast()?;
+
+            Ok(Self {
+                _device: device,
+                context,
+                staging_res,
+                dup,
+                width,
+                height,
+            })
+        }
+    }
+
+    /// true=拿到新帧并已复制到 staging；false=超时无更新
+    pub fn acquire(&mut self, timeout_ms: u32) -> Result<bool> {
+        unsafe {
+            let mut info = DXGI_OUTDUPL_FRAME_INFO::default();
+            let mut resource: Option<IDXGIResource> = None;
+            match self.dup.AcquireNextFrame(timeout_ms, &mut info, &mut resource) {
+                Ok(()) => {
+                    let res = resource.ok_or_else(|| anyhow!("no desktop resource"))?;
+                    let src: ID3D11Resource = res.cast()?;
+                    self.context.CopyResource(&self.staging_res, &src);
+                    let _ = self.dup.ReleaseFrame();
+                    Ok(true)
+                }
+                Err(e) if e.code() == DXGI_ERROR_WAIT_TIMEOUT => Ok(false),
+                Err(e) => Err(e.into()),
+            }
+        }
+    }
+
+    /// 把 staging 纹理映射到内存并去掉行间距
+    pub fn copy_frame(&mut self) -> Result<Vec<u8>> {
+        unsafe {
+            let mut mapped = D3D11_MAPPED_SUBRESOURCE::default();
+            self.context
+                .Map(&self.staging_res, 0, D3D11_MAP_READ, 0, Some(&mut mapped))
+                .map_err(|e| anyhow!("Map staging failed: {e}"))?;
+            let (w, h) = (self.width as usize, self.height as usize);
+            let pitch = mapped.RowPitch as usize;
+            let src = std::slice::from_raw_parts(mapped.pData as *const u8, pitch * h);
+            let mut out = Vec::with_capacity(w * h * 4);
+            for row in 0..h {
+                out.extend_from_slice(&src[row * pitch..row * pitch + w * 4]);
+            }
+            self.context.Unmap(&self.staging_res, 0);
+            Ok(out)
+        }
+    }
+}
+
+#[allow(dead_code)]
+pub fn desktop_size() -> Result<(u32, u32)> {
+    let d = Duplication::new()?;
+    Ok((d.width, d.height))
+}
+
+/// 采集主循环：独立线程运行；连接断开（通道关闭）时退出。
+/// 通道为 watch 单槽（latest-wins）：发送即覆盖，网络慢时不产生积压。
+pub fn run(tx: Sender<Option<Vec<u8>>>, keyreq: Arc<AtomicBool>) {
+    let start = Instant::now();
+    loop {
+        let mut dup = match Duplication::new() {
+            Ok(d) => d,
+            Err(e) => {
+                eprintln!("[capture] init failed: {e}, retry in 2s");
+                std::thread::sleep(Duration::from_secs(2));
+                continue;
+            }
+        };
+        eprintln!(
+            "[capture] duplication started {}x{}",
+            dup.width, dup.height
+        );
+
+        // H.264 编码器：屏幕内容实时档 + 码率控制（对标 UU：3.2Mbps@60fps）
+        let mut encoder = match Encoder::with_api_config(
+            OpenH264API::from_source(),
+            EncoderConfig::new()
+                .bitrate(BitRate::from_bps(5_000_000))
+                .max_frame_rate(FrameRate::from_hz(30.0))
+                .rate_control_mode(RateControlMode::Bitrate)
+                .usage_type(UsageType::ScreenContentRealTime)
+                .intra_frame_period(IntraFramePeriod::from_num_frames(300)),
+        ) {
+            Ok(e) => e,
+            Err(e) => {
+                eprintln!("[capture] encoder init failed: {e}");
+                return;
+            }
+        };
+
+        let mut last_frame: Option<Vec<u8>> = None;
+        let mut last_sent = Instant::now() - Duration::from_secs(10);
+        let min_frame = Duration::from_millis(33); // 30fps 上限
+        loop {
+            let iter_start = Instant::now();
+            match dup.acquire(80) {
+                Ok(true) => match dup.copy_frame() {
+                    Ok(raw) => {
+                        // 主控端请求关键帧（解码异常恢复）
+                        if keyreq.swap(false, Ordering::Relaxed) {
+                            encoder.force_intra_frame();
+                        }
+                        let header = VideoFrameHeader {
+                            width: dup.width,
+                            height: dup.height,
+                            ts_ms: start.elapsed().as_millis() as u64,
+                        };
+                        let yuv = YUVBuffer::from_bgra8_source(BgraSliceU8::new(
+                            &raw,
+                            (dup.width as usize, dup.height as usize),
+                        ));
+                        match encoder.encode(&yuv) {
+                            Ok(ebs) => {
+                                let payload = ebs.to_vec();
+                                let frame = encode_video_frame(&header, &payload);
+                                if tx.send(Some(frame.clone())).is_err() {
+                                    return; // 会话结束
+                                }
+                                last_frame = Some(frame);
+                                last_sent = Instant::now();
+                            }
+                            Err(e) => {
+                                eprintln!("[capture] encode failed: {e}");
+                            }
+                        }
+                        // 帧间隔节流到 ~30fps
+                        let el = iter_start.elapsed();
+                        if el < min_frame {
+                            std::thread::sleep(min_frame - el);
+                        }
+                    }
+                    Err(e) => {
+                        eprintln!("[capture] frame copy failed: {e}");
+                        break; // 重建 duplication
+                    }
+                },
+                Ok(false) => {
+                    // 桌面静止：低频重发最后一帧（重复 P 帧，解码端画面不变）
+                    if last_sent.elapsed() > Duration::from_millis(500) {
+                        if let Some(f) = last_frame.as_ref() {
+                            if tx.send(Some(f.clone())).is_err() {
+                                return;
+                            }
+                            last_sent = Instant::now();
+                        }
+                    } else {
+                        std::thread::sleep(Duration::from_millis(8));
+                    }
+                }
+                Err(e) => {
+                    eprintln!("[capture] acquire failed: {e}");
+                    break; // ACCESS_LOST 等，重建
+                }
+            }
+        }
+        std::thread::sleep(Duration::from_millis(300));
+    }
+}