wl-webrtc/src/cap_portal/pipewire_thread.rs

use std::os::fd::{AsRawFd, FromRawFd, OwnedFd};
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::thread::{self, JoinHandle};

use anyhow::Result;
use crossbeam_channel::Sender;

use super::fourcc::spa_to_drm_fourcc;
use super::types::{PortalFormatInfo, PwCtrlEvent, PwDmaBufFrame};

/// PipeWire 捕获线程的上下文数据
///
/// 从主线程传递给 PipeWire 捕获线程的所有必要资源。
/// 该结构体在线程创建时一次性 move 到线程中使用。
struct PwThreadCtx {
    frame_tx: Sender<PwDmaBufFrame>,
    event_tx: Sender<PwCtrlEvent>,
    dropped: Arc<AtomicU64>,
    shutdown_read: OwnedFd,
    pw_fd: OwnedFd,
    node_id: u32,
}

fn pipewire_thread(ctx: PwThreadCtx) {
    use pipewire as pw;
    use pw::properties::properties;
    use pw::spa::param::video::VideoInfoRaw;
    use pw::stream::{StreamBox, StreamFlags};
    use std::cell::Cell;
    use std::rc::Rc;

    // 初始化 PipeWire 进程全局库。
    //
    // pipewire-rs 内部使用 OnceCell 保护 pw::init()，确保只调用一次。
    // pw::deinit() 是 unsafe 且要求"进程生命周期内仅调用一次，且所有
    // PipeWire 使用已停止"。由于 CapPortal 可被多次创建销毁，此函数
    // 不调用 pw::deinit()——进程退出时全局状态由 OS 回收。
    pw::init();

    let PwThreadCtx {
        frame_tx,
        event_tx,
        dropped,
        shutdown_read,
        pw_fd,
        node_id,
    } = ctx;

    let mainloop = match pw::main_loop::MainLoopBox::new(None) {
        Ok(ml) => ml,
        Err(e) => {
            if let Err(e) =
                event_tx.try_send(PwCtrlEvent::Error(format!("MainLoop::new failed: {e}")))
            {
                tracing::error!("MainLoop::new failed and error channel also failed: {e}");
            }
            return;
        }
    };

    let context = match pw::context::ContextBox::new(mainloop.loop_(), None) {
        Ok(c) => c,
        Err(e) => {
            if let Err(e) =
                event_tx.try_send(PwCtrlEvent::Error(format!("Context::new failed: {e}")))
            {
                tracing::error!("Context::new failed and error channel also failed: {e}");
            }
            return;
        }
    };

    let core = match context.connect_fd(pw_fd, None) {
        Ok(c) => c,
        Err(e) => {
            if let Err(e) = event_tx.try_send(PwCtrlEvent::Error(format!("connect_fd failed: {e}")))
            {
                tracing::error!("connect_fd failed and error channel also failed: {e}");
            }
            return;
        }
    };

    // 创建 PipeWire 视频流
    // 属性配置:
    //   - MEDIA_TYPE = "Video": 媒体类型为视频
    //   - MEDIA_CATEGORY = "Capture": 类别为捕获（而非回放）
    //   - MEDIA_ROLE = "Screen": 角色为屏幕（用于策略管理）
    let stream = match StreamBox::new(
        &core,
        "wl-webrtc",
        properties! {
            *pw::keys::MEDIA_TYPE => "Video",
            *pw::keys::MEDIA_CATEGORY => "Capture",
            *pw::keys::MEDIA_ROLE => "Screen",
            *pw::keys::NODE_FORCE_QUANTUM => "512",
        },
    ) {
        Ok(s) => s,
        Err(e) => {
            if let Err(e) =
                event_tx.try_send(PwCtrlEvent::Error(format!("Stream::new failed: {e}")))
            {
                tracing::error!("Stream::new failed and error channel also failed: {e}");
            }
            return;
        }
    };

    let format_info: Rc<Cell<Option<PortalFormatInfo>>> = Rc::new(Cell::new(None));

    let event_tx_state = event_tx.clone();
    let _listener = stream
        .add_local_listener::<()>()
        .state_changed(move |_, _, old, new| {
            tracing::info!("PipeWire stream state: {old:?} -> {new:?}");
            match new {
                pw::stream::StreamState::Error(e) => {
                    tracing::error!("PipeWire stream error: {e}");
                    let _ = event_tx_state.try_send(PwCtrlEvent::StreamEnded);
                }
                pw::stream::StreamState::Unconnected => {
                    let _ = event_tx_state.try_send(PwCtrlEvent::StreamEnded);
                }
                pw::stream::StreamState::Paused => {
                    tracing::warn!("PipeWire stream paused (compositor may be switching content)");
                }
                pw::stream::StreamState::Streaming => {
                    tracing::info!("PipeWire stream (re)started");
                }
                pw::stream::StreamState::Connecting => {}
            }
        })
        // 参数变化回调（格式协商）
        // PipeWire 在流格式协商完成后触发此回调
        // id 为参数类型，param 包含具体的格式参数（分辨率、像素格式等）
        .param_changed({
            let format_info = format_info.clone();
            let event_tx = event_tx.clone();
            move |_, _, id, param| {
                // 仅处理 Format 类型的参数变化
                let Some(param) = param else { return };
                if id != pw::spa::param::ParamType::Format.as_raw() {
                    return;
                }
                // 解析视频格式信息（分辨率、像素格式、修饰符等）
                let mut info = VideoInfoRaw::new();
                if let Err(e) = info.parse(param) {
                    tracing::warn!("Failed to parse video format: {e}");
                    return;
                }
                let width = info.size().width;
                let height = info.size().height;
                // 将 SPA 视频格式转换为 DRM FourCC 格式标识符
                let drm_format = spa_to_drm_fourcc(info.format());
                // 获取 DRM 修饰符，描述 GPU buffer 的内存布局（如 tiling 模式）
                let modifier = info.modifier();
                let framerate = info.framerate();
                let max_framerate = info.max_framerate();
                // 保存协商后的格式信息，供 process 回调读取
                let previous_format = format_info.get();
                format_info.set(Some(PortalFormatInfo {
                    width,
                    height,
                    drm_format,
                    modifier,
                }));
                if let Some(prev) = previous_format {
                    if width != prev.width || height != prev.height {
                        tracing::warn!(
                            "PipeWire dimensions changed: {}x{} (format renegotiation)",
                            width,
                            height
                        );
                        let _ = event_tx.try_send(PwCtrlEvent::FormatChanged { width, height });
                    }
                }
                tracing::info!(
                    "PipeWire format negotiated: {width}x{height}, \
                     drm_format={drm_format:#010x}, modifier={modifier:#x}, \
                     framerate={}/{}, max_framerate={}/{}",
                    framerate.num,
                    framerate.denom,
                    max_framerate.num,
                    max_framerate.denom,
                );
            }
        })
        // 帧处理回调 —— 这是核心的数据路径
        // 每当 PipeWire 有新的帧数据可用时触发
        // 关键操作: 从 buffer 中提取 DMA-BUF fd，dup 后通过 channel 发送给消费者
        .process({
            let format_info = format_info.clone();
            let frame_tx = frame_tx.clone();
            move |stream, _| {
                // SAFETY: raw_buf ownership invariant — PipeWire's process callback
                // contract requires that every buffer acquired via `dequeue_raw_buffer`
                // is returned to the queue EXACTLY ONCE via `queue_raw_buffer` before
                // the callback returns — on every exit path, success or error. Failure
                // to requeue leaks the buffer slot and eventually stalls the stream.
                //
                // Audit map of this closure (verified 2026-06-28):
                //   - null raw_buf  (dequeue returned NULL) → nothing to requeue, return.
                //   - null spa_buf / no data / bad fd / null chunk / no format_info /
                //     invalid dims / dup_fd < 0 → all requeue before early-return.
                //   - success (try_send Ok / Full / Disconnected) → final requeue at end.
                // The fd ownership is independent: dup() creates a fresh fd that lives
                // inside PwDmaBufFrame; on try_send error the frame Drops and closes it.
                let raw_buf = unsafe { stream.dequeue_raw_buffer() };
                if raw_buf.is_null() {
                    tracing::trace!("process: null raw_buf");
                    return;
                }

                // 获取 SPA buffer 结构体，包含数据数组、元数据等
                // SAFETY: raw_buf was checked non-null above. `pw_buffer.buffer` is a
                // valid raw pointer for the lifetime of raw_buf (PipeWire keeps the
                // buffer alive until we queue it back).
                let spa_buf = unsafe { (*raw_buf).buffer };
                if spa_buf.is_null() {
                    tracing::trace!("process: null spa_buf");
                    // SAFETY: raw_buf is the non-null buffer we still own; returning it.
                    unsafe { stream.queue_raw_buffer(raw_buf) };
                    return;
                }

                // 获取 buffer 中的数据项数量和数据指针
                // 对于 DMA-BUF 帧，通常只有 1 个数据项（包含 fd）
                // SAFETY: spa_buf checked non-null above; `n_datas` is a plain u32 field.
                let n_datas = unsafe { (*spa_buf).n_datas };
                // SAFETY: same as above; `datas` is a raw pointer field, may be null.
                let datas_ptr = unsafe { (*spa_buf).datas };
                if n_datas == 0 || datas_ptr.is_null() {
                    tracing::trace!("process: no data (n_datas={n_datas})");
                    // SAFETY: raw_buf still owned, returning it.
                    unsafe { stream.queue_raw_buffer(raw_buf) };
                    return;
                }

                // 从第一个数据项中获取 DMA-BUF 文件描述符
                // 通过 libspa 的 Data 包装类型安全地访问 SPA 数据结构
                // SAFETY: datas_ptr is non-null and n_datas > 0 (checked above). We cast
                // to pw::spa::buffer::Data and take a shared borrow; PipeWire does not
                // mutate the data array during a process cycle, so a shared reference
                // for the duration of this callback is sound.
                let data_ref: &pw::spa::buffer::Data =
                    unsafe { &*(datas_ptr as *const pw::spa::buffer::Data) };
                let fd = data_ref.fd();
                if fd < 0 {
                    tracing::trace!("process: invalid fd={fd}");
                    // SAFETY: raw_buf still owned, returning it.
                    unsafe { stream.queue_raw_buffer(raw_buf) };
                    return;
                }

                if data_ref.as_raw().chunk.is_null() {
                    tracing::trace!("process: null chunk");
                    // SAFETY: raw_buf still owned, returning it.
                    unsafe { stream.queue_raw_buffer(raw_buf) };
                    return;
                }
                let chunk = data_ref.chunk();
                let offset = chunk.offset() as u64;
                let stride = chunk.stride() as u32;

                // 从 SPA_META_Header 元数据中提取 PTS (显示时间戳)
                // 遍历 buffer 的所有元数据项，查找 Header 类型的元数据
                // PTS 可用于音视频同步和帧率控制
                // SAFETY: spa_buf is non-null. `metas` is checked for null before
                // iteration. We iterate `i in 0..n_metas` reading shared POD fields
                // (type_, size, data) — PipeWire keeps the meta array immutable during
                // a process cycle. The size guard (`meta.size >= size_of::<spa_meta_header>()`)
                // and null-data check before reading ensure we never read past the
                // meta's actual extent.
                let pts: i64 = unsafe {
                    let mut pts_val: i64 = 0;
                    let n_metas = (*spa_buf).n_metas;
                    let metas = (*spa_buf).metas;
                    if !metas.is_null() {
                        for i in 0..n_metas {
                            let meta = &*metas.add(i as usize);
                            if meta.type_ == libspa::sys::SPA_META_Header
                                && meta.size as usize
                                    >= std::mem::size_of::<libspa::sys::spa_meta_header>()
                                && !meta.data.is_null()
                            {
                                let header = &*(meta.data as *const libspa::sys::spa_meta_header);
                                pts_val = header.pts;
                                break;
                            }
                        }
                    }
                    pts_val
                };

                // 验证格式信息已协商完成，且分辨率和格式有效
                let Some(fmt) = format_info.get() else {
                    // SAFETY: raw_buf still owned, returning it.
                    unsafe { stream.queue_raw_buffer(raw_buf) };
                    return;
                };
                let PortalFormatInfo {
                    width,
                    height,
                    drm_format: format,
                    modifier,
                } = fmt;
                if width == 0 || height == 0 || format == 0 {
                    tracing::trace!("process: invalid dimensions {width}x{height} format={format}");
                    // SAFETY: raw_buf still owned, returning it.
                    unsafe { stream.queue_raw_buffer(raw_buf) };
                    return;
                }

                // 复制 DMA-BUF 文件描述符
                // 必须 dup，因为原始 fd 由 PipeWire 管理，我们不能持有它
                // dup 后的 fd 由 PwDmaBufFrame 持有，生命周期独立于 PipeWire buffer
                // SAFETY: `fd` is the open DMA-BUF fd reported by PipeWire (>= 0 checked
                // above). libc::dup is the standard POSIX fd duplication call. The
                // original `fd` remains owned by PipeWire (returned with raw_buf later).
                let dup_fd = unsafe { libc::dup(fd) };
                if dup_fd < 0 {
                    // SAFETY: raw_buf still owned, returning it. No fd cleanup needed
                    // because dup() failed and never returned a new fd.
                    unsafe { stream.queue_raw_buffer(raw_buf) };
                    return;
                }

                // 构建帧数据对象，所有必要的帧信息已收集完毕
                // SAFETY: `dup_fd` is a freshly-dup'd open file descriptor (>= 0 checked
                // above) and we are its sole owner. OwnedFd::from_raw_fd takes ownership
                // and will close() it on Drop. The fd's lifecycle is independent of
                // raw_buf: whether try_send succeeds (frame moves into the channel) or
                // fails (Full/Disconnected — the error payload owns the frame and drops
                // it at the end of the match arm), exactly one close() occurs per dup().
                let frame_fd = unsafe { OwnedFd::from_raw_fd(dup_fd) };
                let frame = PwDmaBufFrame {
                    fd: frame_fd,
                    offset,
                    stride,
                    modifier,
                    width,
                    height,
                    format,
                    pts,
                };

                match frame_tx.try_send(frame) {
                    Ok(()) => {}
                    Err(crossbeam_channel::TrySendError::Full(_)) => {
                        dropped.fetch_add(1, Ordering::Relaxed);
                    }
                    Err(crossbeam_channel::TrySendError::Disconnected(_)) => {}
                }
                // SAFETY: final exactly-once requeue of raw_buf. Every path above
                // either returned early with its own requeue, or falls through to here.
                unsafe { stream.queue_raw_buffer(raw_buf) };
            }
        })
        .register();

    let mut params: [&pw::spa::pod::Pod; 0] = [];

    if let Err(e) = stream.connect(
        pw::spa::utils::Direction::Input,
        Some(node_id),
        StreamFlags::AUTOCONNECT | StreamFlags::MAP_BUFFERS,
        &mut params,
    ) {
        if let Err(e) = event_tx.try_send(PwCtrlEvent::Error(format!("stream.connect failed: {e}")))
        {
            tracing::error!("stream.connect failed and error channel also failed: {e}");
        }
        return;
    }

    let loop_ = mainloop.loop_();

    // Register the shutdown eventfd on the PipeWire loop.
    //
    // When CapPortal::drop writes to the eventfd, the loop wakes up and
    // dispatches this callback on the loop thread. Because the callback
    // only fires while mainloop.run() is blocking this thread, mainloop
    // is guaranteed alive — eliminating the UAF that existed with the
    // previous detached helper thread approach.
    // 保存 mainloop 的原始指针，用于在 shutdown 回调中调用 pw_main_loop_quit
    // 这是安全的，因为回调只在 mainloop.run() 阻塞期间执行
    let mainloop_ptr = mainloop.as_raw_ptr();

    let _shutdown_source = loop_.add_io(
        shutdown_read,
        libspa::support::system::IoFlags::IN,
        move |fd| {
            // Drain the eventfd so it doesn't re-trigger
            let mut buf: u64 = 0;
            // SAFETY: `fd` is the registered eventfd owned by the mainloop source; the
            // buffer is a stack u64 of 8 bytes matching the count argument. POSIX
            // read(2) is the standard fd-read syscall; eventfd semantics require the
            // 8-byte buffer.
            let _ = unsafe {
                libc::read(
                    fd.as_raw_fd(),
                    &mut buf as *mut u64 as *mut _,
                    std::mem::size_of::<u64>(),
                )
            };
            // SAFETY: This callback only executes while mainloop.run() is
            // blocking this thread, so mainloop is guaranteed alive.
            unsafe { pipewire::sys::pw_main_loop_quit(mainloop_ptr) };
        },
    );

    // 启动 PipeWire 主事件循环
    // 此调用会阻塞当前线程，直到 mainloop.quit() 被调用
    // quit() 由 shutdown eventfd 的 IO 回调触发
    mainloop.run();

    // run() returned — _shutdown_source drops first (reverse declaration order),
    // which unregisters the callback from the loop. Then mainloop drops.
    // No dangling raw pointers are possible.
    // PipeWire global state is intentionally not deinitialized here — see pw::init() comment above.
}

pub(super) fn spawn_pipewire_thread(
    frame_tx: Sender<PwDmaBufFrame>,
    event_tx: Sender<PwCtrlEvent>,
    dropped: Arc<AtomicU64>,
    shutdown_read: OwnedFd,
    pw_fd: OwnedFd,
    node_id: u32,
) -> Result<JoinHandle<()>> {
    let ctx = PwThreadCtx {
        frame_tx,
        event_tx,
        dropped,
        shutdown_read,
        pw_fd,
        node_id,
    };
    thread::Builder::new()
        .name("pipewire-capture".into())
        .spawn(move || pipewire_thread(ctx))
        .map_err(|e| anyhow::anyhow!("thread spawn failed: {e}"))
}