wl-webrtc/src/avhw/dmabuf.rs

use std::mem;
// AsRawFd is required by `frame.fd.as_raw_fd()` below but rustc emits a false
// "unused_imports" warning because OwnedFd also has an inherent `as_raw_fd`.
// E0599 if removed -> must stay; warning is a known rustc quirk.
use std::os::fd::AsRawFd;
use std::os::raw::c_void;
use std::path::Path;
use std::ptr;

use anyhow::{bail, Result};
use ffmpeg_next as ff;
use ffmpeg_next::ffi;

use crate::cap_portal::PwDmaBufFrame;

use super::{ff_err, AvHwDevCtx, AvHwFrameCtx};

/// Test whether `drm_device` can import the PipeWire DMA-BUF frame via VAAPI.
pub fn test_dma_buf_import(drm_device: &Path, frame: &PwDmaBufFrame) -> Result<()> {
    let hw_dev = AvHwDevCtx::new_vaapi(drm_device)?;
    let frames =
        AvHwFrameCtx::for_capture(&hw_dev, frame.width, frame.height, ff::format::Pixel::BGRA)?;

    // SAFETY: frames is a live VAAPI frames context; frame carries valid DMA-BUF metadata.
    unsafe { import_dma_buf_to_vaapi(frames.as_ptr(), frame) }?;

    Ok(())
}

/// Import a DMA-BUF into a VAAPI hardware frame via zero-copy `av_hwframe_map`.
///
/// # Safety
/// Imports a DMA-BUF frame into a VAAPI hardware frame pool for GPU-side processing.
///
/// Takes the negotiated format/geometry from `frame` (a `PwDmaBufFrame` from
/// PipeWire capture) plus the target `frames_ctx` (VAAPI frame pool from
/// `AvHwFrameCtx`) and returns an `ff::frame::Video` whose data[3] points to
/// the hardware frame.
///
/// # Safety
///
/// - `frames_ctx` must point to an initialized AVHWCramesContext for VAAPI
/// - `frame.fd` must be a valid DMA-BUF file descriptor
pub unsafe fn import_dma_buf_to_vaapi(
    frames_ctx: *mut ffi::AVBufferRef,
    frame: &PwDmaBufFrame,
) -> Result<ff::frame::Video> {
    let duped_fd = libc::dup(frame.fd.as_raw_fd());
    if duped_fd < 0 {
        bail!("dup(fd) failed: {}", std::io::Error::last_os_error());
    }

    let mut desc: ffi::AVDRMFrameDescriptor = mem::zeroed();
    desc.nb_objects = 1;
    desc.objects[0].fd = duped_fd;
    desc.objects[0].size = (frame.height as usize) * (frame.stride as usize);
    desc.objects[0].format_modifier = frame.modifier;
    desc.nb_layers = 1;
    desc.layers[0].format = frame.format;
    desc.layers[0].nb_planes = 1;
    desc.layers[0].planes[0].object_index = 0;
    desc.layers[0].planes[0].offset = frame.offset as isize;
    desc.layers[0].planes[0].pitch = frame.stride as isize;

    let desc_box = Box::new(desc);
    let desc_ptr = Box::into_raw(desc_box);

    let buf_ref = ffi::av_buffer_create(
        desc_ptr as *mut u8,
        std::mem::size_of::<ffi::AVDRMFrameDescriptor>(),
        Some(cleanup_drm_descriptor),
        ptr::null_mut(),
        0,
    );
    if buf_ref.is_null() {
        let desc_box = Box::from_raw(desc_ptr);
        libc::close(desc_box.objects[0].fd);
        bail!("av_buffer_create returned null for DRM descriptor");
    }

    let mut src = ff::frame::Video::empty();
    {
        let sp = src.as_mut_ptr();
        (*sp).format = ffi::AVPixelFormat::AV_PIX_FMT_DRM_PRIME as i32;
        (*sp).width = frame.width as i32;
        (*sp).height = frame.height as i32;
        (*sp).data[0] = (*buf_ref).data;
        (*sp).buf[0] = buf_ref;
    }

    let mut dst = ff::frame::Video::empty();
    // SAFETY: frames_ctx is guaranteed by this unsafe function's contract to be a
    // valid initialized VAAPI frames context; we set format/hw_frames_ctx on a
    // freshly allocated dst frame.
    unsafe {
        let dp = dst.as_mut_ptr();
        (*dp).format = ffi::AVPixelFormat::AV_PIX_FMT_VAAPI as i32;
        (*dp).hw_frames_ctx = ffi::av_buffer_ref(frames_ctx);
        if (*dp).hw_frames_ctx.is_null() {
            bail!("av_buffer_ref(frames_ctx) returned null");
        }
    }
    // SAFETY: src and dst are initialized AVFrames; dst has a valid hw_frames_ctx
    // ref and av_hwframe_map fills dst from src.
    let ret = unsafe {
        ffi::av_hwframe_map(
            dst.as_mut_ptr(),
            src.as_ptr(),
            ffi::AV_HWFRAME_MAP_READ as i32,
        )
    };
    if ret < 0 {
        bail!("av_hwframe_map failed: {}", ff_err(ret));
    }

    Ok(dst)
}

unsafe extern "C" fn cleanup_drm_descriptor(_opaque: *mut c_void, data: *mut u8) {
    let desc = data as *mut ffi::AVDRMFrameDescriptor;
    if !desc.is_null() && (*desc).nb_objects > 0 && (*desc).objects[0].fd >= 0 {
        libc::close((*desc).objects[0].fd);
    }
    let _ = Box::from_raw(data as *mut ffi::AVDRMFrameDescriptor);
}