wl-webrtc/src/avhw/software.rs

use std::ffi::CString;
use std::path::Path;
use std::ptr;

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

use super::ff_err;

pub(super) fn create_nv12_to_yuv420p_sws(width: u32, height: u32) -> Result<*mut ffi::SwsContext> {
    // SAFETY: sws_getContext creates an owned scaler context for same-size NV12 -> YUV420P.
    let ctx = unsafe {
        ffi::sws_getContext(
            width as i32,
            height as i32,
            ffi::AVPixelFormat::AV_PIX_FMT_NV12,
            width as i32,
            height as i32,
            ffi::AVPixelFormat::AV_PIX_FMT_YUV420P,
            2,
            ptr::null_mut(),
            ptr::null_mut(),
            ptr::null_mut(),
        )
    };
    if ctx.is_null() {
        bail!("Failed to create NV12 -> YUV420P sws_scale context");
    }
    Ok(ctx)
}

pub(super) fn alloc_yuv420p_frame(width: u32, height: u32) -> Result<*mut ffi::AVFrame> {
    // SAFETY: Allocate an AVFrame, configure format/dimensions, then allocate writable buffers.
    unsafe {
        let mut frame = ffi::av_frame_alloc();
        if frame.is_null() {
            bail!("av_frame_alloc failed");
        }
        (*frame).width = width as i32;
        (*frame).height = height as i32;
        (*frame).format = ffi::AVPixelFormat::AV_PIX_FMT_YUV420P as i32;
        let ret = ffi::av_frame_get_buffer(frame, 0);
        if ret < 0 {
            ffi::av_frame_free(&mut frame);
            bail!("av_frame_get_buffer failed: {}", ff_err(ret));
        }
        Ok(frame)
    }
}

pub(super) fn create_software_h264_muxer(
    output_path: &Path,
    width: u32,
    height: u32,
    fps: u32,
    bitrate: u64,
    gop_size: u32,
) -> Result<(
    ff::codec::encoder::video::Video,
    ff::format::context::Output,
)> {
    let output_cstr = CString::new(output_path.to_str().unwrap())?;
    let codec = ff::encoder::find_by_name("libx264")
        .or_else(|| ff::encoder::find_by_name("libopenh264"))
        .ok_or_else(|| {
            anyhow::anyhow!("No H.264 software encoder found (tried libx264, libopenh264)")
        })?;
    let codec_name = codec.name().to_string();

    let mut enc = {
        let ctx = ff::codec::Context::new_with_codec(codec);
        ctx.encoder().video()?
    };
    enc.set_width(width);
    enc.set_height(height);
    enc.set_format(ff::format::Pixel::YUV420P);
    enc.set_bit_rate(bitrate as usize);
    enc.set_gop(gop_size);
    enc.set_time_base(ff::Rational::new(1, fps as i32));
    enc.set_max_b_frames(3);

    // SAFETY: global headers are needed by MP4 and harmless for other common muxers.
    unsafe {
        (*enc.as_mut_ptr()).flags |= ffi::AV_CODEC_FLAG_GLOBAL_HEADER as i32;
    }

    if codec_name == "libx264" {
        // SAFETY: priv_data and codec context belong to the unopened encoder;
        // strings live for each av_opt_set call.
        unsafe {
            let key = CString::new("preset").unwrap();
            let val = CString::new("fast").unwrap();
            ffi::av_opt_set((*enc.as_mut_ptr()).priv_data, key.as_ptr(), val.as_ptr(), 0);
            let key = CString::new("threads").unwrap();
            let val = CString::new("6").unwrap();
            ffi::av_opt_set((*enc.as_mut_ptr()).priv_data, key.as_ptr(), val.as_ptr(), 0);
            (*enc.as_mut_ptr()).profile = ffi::AV_PROFILE_H264_HIGH;
            // SAFETY: enc is a valid, initialized AVCodecContext from
            // avcodec_alloc_context3. Setting level is a simple i32 field
            // assignment on a properly aligned struct.
            (*enc.as_mut_ptr()).level = 40; // H.264 Level 4.0 (up to 1080p@30)
        }
    }

    let opened = enc
        .open()
        .map_err(|e| anyhow::anyhow!("Failed to open {codec_name} encoder: {e}"))?;
    let enc_video = opened.0;

    let use_null = output_path
        .to_str()
        .map(|s| s.contains("null"))
        .unwrap_or(false);
    let fmt_name = if use_null {
        CString::new("null").unwrap()
    } else {
        CString::new("").unwrap()
    };
    let fmt_name_ptr = if use_null {
        fmt_name.as_ptr()
    } else {
        ptr::null()
    };

    let mut fmt_ctx_ptr: *mut ffi::AVFormatContext = ptr::null_mut();
    // SAFETY: fmt_ctx_ptr is initialized by FFmpeg; C strings live across the call.
    let ret = unsafe {
        ffi::avformat_alloc_output_context2(
            &mut fmt_ctx_ptr,
            ptr::null_mut(),
            fmt_name_ptr,
            output_cstr.as_ptr(),
        )
    };
    if ret < 0 || fmt_ctx_ptr.is_null() {
        bail!("Failed to allocate output format context: {}", ff_err(ret));
    }

    // SAFETY: fmt_ctx_ptr is valid; stream and codec parameters are owned by the format context.
    let stream_ptr = unsafe { ffi::avformat_new_stream(fmt_ctx_ptr, ptr::null()) };
    if stream_ptr.is_null() {
        bail!("Failed to create output stream");
    }

    // SAFETY: stream_ptr and encoder context are valid; parameters are copied into stream.
    let ret =
        unsafe { ffi::avcodec_parameters_from_context((*stream_ptr).codecpar, enc_video.as_ptr()) };
    if ret < 0 {
        bail!("Failed to copy codec parameters to stream: {}", ff_err(ret));
    }
    // SAFETY: stream_ptr is valid and writable during muxer setup.
    unsafe {
        (*stream_ptr).time_base = (*enc_video.as_ptr()).time_base;
    }

    // SAFETY: open an AVIO only for muxers that require files; null muxer advertises NOFILE.
    unsafe {
        if (*(*fmt_ctx_ptr).oformat).flags & ffi::AVFMT_NOFILE == 0 {
            let ret = ffi::avio_open(
                &mut (*fmt_ctx_ptr).pb,
                output_cstr.as_ptr(),
                ffi::AVIO_FLAG_WRITE,
            );
            if ret < 0 {
                bail!(
                    "Failed to open output file '{}': {}",
                    output_path.display(),
                    ff_err(ret)
                );
            }
        }
    }

    // SAFETY: fmt_ctx_ptr is fully configured.
    let ret = unsafe { ffi::avformat_write_header(fmt_ctx_ptr, ptr::null_mut()) };
    if ret < 0 {
        bail!("Failed to write output header: {}", ff_err(ret));
    }

    // SAFETY: ownership of fmt_ctx_ptr transfers to ffmpeg-next Output wrapper.
    let octx = unsafe { ff::format::context::Output::wrap(fmt_ctx_ptr) };
    tracing::info!("Using software H.264 encoder: {codec_name}");
    Ok((enc_video, octx))
}

pub(super) fn create_software_h264_encoder(
    width: u32,
    height: u32,
    fps: u32,
    bitrate: u64,
    gop_size: u32,
) -> Result<ff::codec::encoder::video::Video> {
    let codec = ff::encoder::find_by_name("libx264")
        .or_else(|| ff::encoder::find_by_name("libopenh264"))
        .ok_or_else(|| anyhow::anyhow!("No H.264 software encoder found"))?;
    let codec_name = codec.name().to_string();

    let mut enc = {
        let ctx = ff::codec::Context::new_with_codec(codec);
        ctx.encoder().video()?
    };
    enc.set_width(width);
    enc.set_height(height);
    enc.set_format(ff::format::Pixel::YUV420P);
    enc.set_bit_rate(bitrate as usize);
    enc.set_gop(gop_size);
    // 90kHz media clock matches RTP directly. Eliminates 1/fps quantization
    // that previously caused sequential RTP timestamps during 60fps capture,
    // leading to 2x RTP time inflation and 10s+ browser jitter buffer growth.
    // See issue #25.
    enc.set_time_base(ff::Rational::new(1, 90_000));
    // Explicit framerate is REQUIRED when time_base is not 1/fps, otherwise
    // libx264 infers wrong fps from the 90kHz time_base and VBV rate control
    // breaks. Per Oracle review round for #25.
    enc.set_frame_rate(Some(ff::Rational::new(fps as i32, 1)));
    enc.set_max_b_frames(0);

    if codec_name == "libx264" {
        // SAFETY: priv_data and codec context belong to the unopened encoder;
        // each CString lives for the duration of its av_opt_set call.
        unsafe {
            let key = CString::new("preset").unwrap();
            let val = CString::new("veryfast").unwrap();
            ffi::av_opt_set((*enc.as_mut_ptr()).priv_data, key.as_ptr(), val.as_ptr(), 0);
            let key = CString::new("tune").unwrap();
            let val = CString::new("zerolatency").unwrap();
            ffi::av_opt_set((*enc.as_mut_ptr()).priv_data, key.as_ptr(), val.as_ptr(), 0);
            let key = CString::new("threads").unwrap();
            let val = CString::new("6").unwrap();
            ffi::av_opt_set((*enc.as_mut_ptr()).priv_data, key.as_ptr(), val.as_ptr(), 0);
            // High profile via AVCodecContext.profile (not x264opts - x264 rejects it there).
            // High enables CABAC + 8x8dct automatically.
            (*enc.as_mut_ptr()).profile = ffi::AV_PROFILE_H264_HIGH;
            // SAFETY: enc is a valid, initialized AVCodecContext from
            // avcodec_alloc_context3. Setting level is a simple i32 field
            // assignment on a properly aligned struct.
            (*enc.as_mut_ptr()).level = 42; // H.264 Level 4.2 (up to 1440p@30)
                                            // SAFETY: priv_data belongs to the unopened libx264 encoder context.
                                            // `forced-idr` is an FFmpeg-level private option (not x264-native),
                                            // so it must be set via av_opt_set, NOT via the x264opts string.
                                            // With forced-idr=1, setting AV_PICTURE_TYPE_I on an input frame
                                            // produces a true IDR NALU with inline SPS/PPS (repeat_headers=1).
            let key = CString::new("forced-idr").unwrap();
            let val = CString::new("1").unwrap();
            ffi::av_opt_set((*enc.as_mut_ptr()).priv_data, key.as_ptr(), val.as_ptr(), 0);
            let key = CString::new("x264opts").unwrap();
            // x264's vbv-maxrate unit is kbit/s and vbv-bufsize is kbit (NOT bps).
            // Confirmed via x264 source encoder/ratecontrol.c:658-661 which multiplies
            // these values by 1000 to convert kbit -> bit at use site. Passing bps makes
            // VBV effectively unbounded (5.5 Mbps becomes 5.5 Gbps, clipped to 2 Gbps).
            // See https://github.com/mirror/x264/blob/c24e06c2e184345ceb33eb20a15d1024d9fd3497/encoder/ratecontrol.c#L658-L661
            let vbv_maxrate_kbps = bitrate / 1000;
            let vbv_bufsize_kbps = (bitrate / 4) / 1000;
            let val = CString::new(format!(
                "repeat_headers=1:vbv-maxrate={vbv_maxrate_kbps}:vbv-bufsize={vbv_bufsize_kbps}"
            ))
            .unwrap();
            ffi::av_opt_set((*enc.as_mut_ptr()).priv_data, key.as_ptr(), val.as_ptr(), 0);
        }
    }

    let opened = enc
        .open()
        .map_err(|e| anyhow::anyhow!("Failed to open {codec_name} encoder: {e}"))?;
    tracing::info!("WebRTC encoder: {codec_name} {width}x{height} @ {fps}fps {bitrate}bps (profile High, preset veryfast)");
    Ok(opened.0)
}