kornia-io 0.2.0

/// A module for capturing video streams from v4l2 cameras.
pub mod camera;

/// A module for capturing video streams from different sources.
pub mod capture;

/// Error types for the stream module.
pub mod error;

/// A module for capturing video streams from rtsp sources.
pub mod rtsp;

/// A module for capturing video streams from v4l cameras.
pub mod v4l2;

/// A module for capturing video streams from video files.
pub mod video;

pub use crate::stream::camera::{CameraCapture, CameraCaptureConfig};
pub use crate::stream::capture::StreamCapture;
pub use crate::stream::error::StreamCaptureError;
pub use crate::stream::rtsp::RTSPCameraConfig;
pub use crate::stream::v4l2::V4L2CameraConfig;
pub use crate::stream::video::VideoWriter;

use std::any::Any;
use std::sync::Arc;

/// Quotes a user-supplied value for interpolation into a `gst_parse_launch` pipeline string.
///
/// Values are wrapped in double quotes so whitespace, `!` and `=` cannot introduce new elements or
/// properties. Characters that would terminate or escape the quoted string (`"`, `\`) and control
/// characters are rejected outright.
///
/// # Errors
///
/// Returns [`StreamCaptureError::InvalidConfig`] if `value` contains `"`, `\` or a control
/// character.
pub(crate) fn quote_pipeline_value(value: &str) -> Result<String, StreamCaptureError> {
    if value
        .chars()
        .any(|c| c == '"' || c == '\\' || c.is_control())
    {
        return Err(StreamCaptureError::InvalidConfig(format!(
            "value {value:?} contains characters not allowed in a pipeline description"
        )));
    }
    Ok(format!("\"{value}\""))
}

/// Sets the `location` property of the element called `name` in `pipeline` to `path`.
///
/// File paths are set as a property after parsing instead of being interpolated into the
/// `gst_parse_launch` description, so any path (spaces, `!`, `"`, Windows `\`
/// separators, ...) works verbatim and can never inject pipeline syntax.
///
/// # Errors
///
/// Returns [`StreamCaptureError::GetElementByNameError`] if the pipeline has no element
/// called `name`, or [`StreamCaptureError::InvalidConfig`] if that element has no
/// writable string `location` property.
pub(crate) fn set_location_property(
    pipeline: &gstreamer::Pipeline,
    name: &str,
    path: &std::path::Path,
) -> Result<(), StreamCaptureError> {
    use gstreamer::prelude::*;
    let element = pipeline
        .by_name(name)
        .ok_or(StreamCaptureError::GetElementByNameError)?;
    // `set_property` panics on a missing, read-only or non-string property.
    let writable_string = element.find_property("location").is_some_and(|p| {
        p.value_type() == String::static_type()
            && p.flags().contains(gstreamer::glib::ParamFlags::WRITABLE)
    });
    if !writable_string {
        return Err(StreamCaptureError::InvalidConfig(format!(
            "element {name:?} has no writable string `location` property"
        )));
    }
    element.set_property("location", path.to_string_lossy().as_ref());
    Ok(())
}

use kornia_image::Image;
use kornia_tensor::resource::{MemoryDomain, MemoryResource};

/// A proper [`MemoryResource`] for a GStreamer-mapped buffer (sysmem).
///
/// Holds a `MappedBuffer<Readable>` which:
/// - keeps the GStreamer buffer's reference count alive, and
/// - keeps the map handle active so the host pointer remains valid.
///
/// `Drop` is implicit: when `GstResource` is dropped, `_map` drops first, which
/// unmaps the buffer and releases the GStreamer buffer reference — exactly once.
pub struct GstResource {
    /// The mapped, readable GStreamer buffer.
    ///
    /// Keeping this field alive keeps both the memory map and the buffer ref-count
    /// alive. `Drop` on this field unmaps and releases automatically.
    pub _map: gstreamer::buffer::MappedBuffer<gstreamer::buffer::Readable>,
}

// SAFETY: gstreamer Buffers are ref-counted and thread-safe; the MappedBuffer holds
// a read-only map. Once mapped, the pointer is valid until the map is released on Drop.
unsafe impl Send for GstResource {}
unsafe impl Sync for GstResource {}

impl MemoryResource for GstResource {
    /// Returns the host pointer to the mapped GStreamer buffer data.
    fn as_ptr(&self) -> *mut u8 {
        // MappedBuffer<Readable>::as_ptr returns *const u8; we cast to *mut u8 as the
        // MemoryResource trait requires *mut u8.  The Image built from this is read-only
        // in practice (the buffer is only mapped for reading), so callers must not write.
        self._map.as_ptr() as *mut u8
    }

    /// Returns the size in bytes of the mapped region.
    fn len_bytes(&self) -> usize {
        self._map.len()
    }

    /// GStreamer system memory is host-accessible.
    fn domain(&self) -> MemoryDomain {
        MemoryDomain::Host
    }

    /// Downcast hook.
    fn as_any(&self) -> &dyn Any {
        self
    }

    /// Mutable downcast hook.
    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }
}

/// Construct a borrowed [`Image`] backed by a GStreamer [`MappedBuffer`].
///
/// # Arguments
///
/// * `size` - image dimensions.
/// * `mapped_buffer` - the read-mapped GStreamer buffer; ownership is transferred
///   into a [`GstResource`] keepalive that is Arc-shared with the tensor's `ForeignResource`.
///
/// # Returns
///
/// An `Image<u8, 3>` whose memory is the GStreamer buffer.
/// The buffer remains live (and the pointer valid) for exactly the lifetime of the
/// returned `Image`; dropping the `Image` releases the buffer ref exactly once.
///
/// # Safety
///
/// The caller must ensure the pointer has not been aliased as mutable elsewhere.
pub(crate) fn image_from_gst_buffer(
    size: kornia_image::ImageSize,
    mapped_buffer: gstreamer::buffer::MappedBuffer<gstreamer::buffer::Readable>,
) -> Result<kornia_image::Image<u8, 3>, crate::stream::error::StreamCaptureError> {
    // Capture pointer and length BEFORE moving mapped_buffer into GstResource.
    let data_ptr: *const u8 = mapped_buffer.as_ptr();
    let data_len: usize = mapped_buffer.len();

    // Defense-in-depth: verify the buffer is large enough for an RGB24 frame.
    // The pipeline normally enforces RGB caps, but a misconfigured or non-RGB
    // pipeline would silently produce out-of-bounds stride-based access otherwise.
    let expected_len = size
        .width
        .checked_mul(size.height)
        .and_then(|n| n.checked_mul(3))
        .ok_or_else(|| {
            crate::stream::error::StreamCaptureError::InvalidImageFormat(format!(
                "frame dimensions overflow: {}x{}",
                size.width, size.height
            ))
        })?;
    if data_len < expected_len {
        return Err(
            crate::stream::error::StreamCaptureError::BufferSizeMismatch {
                expected: expected_len,
                got: data_len,
            },
        );
    }

    // Move the MappedBuffer into a GstResource; its Drop releases the buffer.
    let resource = GstResource {
        _map: mapped_buffer,
    };
    let keepalive: Arc<dyn Any + Send + Sync> = Arc::new(resource);

    // SAFETY:
    // - `data_ptr` is non-null as GStreamer sysmem buffers are always non-null.
    // - We verified `data_len >= expected_len` above, preventing out-of-bounds reads.
    // - `keepalive` (GstResource) holds the map alive for the lifetime of the Image.
    let image = unsafe {
        Image::<u8, 3>::from_borrowed_host_readonly(size, data_ptr, keepalive)
            .map_err(crate::stream::error::StreamCaptureError::ImageError)?
    };

    Ok(image)
}

#[cfg(test)]
mod tests {
    use crate::stream::StreamCapture;

    /// A quoted value containing pipeline syntax must be parsed as a single property value,
    /// not as additional elements.
    #[test]
    fn quoted_value_cannot_inject_elements() -> Result<(), Box<dyn std::error::Error>> {
        use gstreamer::prelude::*;
        gstreamer::init()?;
        let evil = "/tmp/a b ! fakesink name=injected location=x";
        let desc = format!(
            "filesrc name=src location={}",
            super::quote_pipeline_value(evil)?
        );
        let bin = gstreamer::parse::launch(&desc)?
            .dynamic_cast::<gstreamer::Bin>()
            .ok();
        // A single element is returned as-is rather than wrapped in a bin.
        assert!(bin.is_none(), "value was split into multiple elements");
        let src = gstreamer::parse::launch(&desc)?;
        assert_eq!(
            src.property::<Option<String>>("location").as_deref(),
            Some(evil)
        );

        assert!(super::quote_pipeline_value("a\" ! fakesink").is_err());
        assert!(super::quote_pipeline_value("a\\").is_err());
        assert!(super::quote_pipeline_value("a\nb").is_err());
        Ok(())
    }

    /// File paths are set as a property, so characters that `quote_pipeline_value`
    /// rejects (e.g. Windows `\` separators) or that are pipeline syntax round-trip.
    #[test]
    fn location_property_accepts_any_path() -> Result<(), Box<dyn std::error::Error>> {
        use gstreamer::prelude::*;
        gstreamer::init()?;
        let path = std::path::Path::new(r#"C:\videos\my clip ! fakesink name=x "q".mp4"#);
        let pipeline = gstreamer::parse::launch("filesrc name=src ! fakesink")?
            .dynamic_cast::<gstreamer::Pipeline>()
            .map_err(|_| "not a pipeline")?;
        super::set_location_property(&pipeline, "src", path)?;
        let src = pipeline.by_name("src").ok_or("missing src")?;
        assert_eq!(
            src.property::<Option<String>>("location").as_deref(),
            path.to_str()
        );
        // Only the two parsed elements exist; nothing was injected.
        assert_eq!(pipeline.children().len(), 2);
        assert!(super::set_location_property(&pipeline, "missing", path).is_err());
        // An element without a `location` property is an error, not a panic.
        let fakesink = pipeline
            .children()
            .into_iter()
            .find(|e| e.name() != "src")
            .ok_or("missing fakesink")?;
        assert!(super::set_location_property(&pipeline, &fakesink.name(), path).is_err());
        Ok(())
    }

    /// Verifies that capturing N frames with `videotestsrc` succeeds, that the pixel
    /// data is readable through the Image slice (proving the GstResource keepalive is
    /// active), and that dropping each Image releases the buffer exactly once (no
    /// crash / no double-unmap — validated by the clean exit without sanitizer errors).
    ///
    /// Uses `videotestsrc` (no camera or display required).
    #[test]
    fn gst_resource_capture_n_frames_and_drop() -> Result<(), Box<dyn std::error::Error>> {
        const N_FRAMES: usize = 5;
        const WIDTH: usize = 8;
        const HEIGHT: usize = 4;

        if !gstreamer::INITIALIZED.load(std::sync::atomic::Ordering::Relaxed) {
            gstreamer::init()?;
        }

        let pipeline_desc = format!(
            "videotestsrc num-buffers={n} ! \
             video/x-raw,format=RGB,width={w},height={h},framerate=30/1 ! \
             appsink name=sink sync=false",
            n = N_FRAMES,
            w = WIDTH,
            h = HEIGHT,
        );

        let mut capture = StreamCapture::new(&pipeline_desc)?;
        capture.start()?;

        let mut frames_received = 0usize;
        // Poll until we have all N frames (videotestsrc with num-buffers is bounded).
        // We attempt up to 5×N polls to avoid an infinite loop.
        let max_polls = N_FRAMES * 5;
        for _ in 0..max_polls {
            if let Some(image) = capture.grab_rgb8()? {
                // 1. Verify dimensions.
                assert_eq!(image.width(), WIDTH, "frame width mismatch");
                assert_eq!(image.height(), HEIGHT, "frame height mismatch");
                assert_eq!(image.num_channels(), 3, "frame channels mismatch");

                // 2. Read pixel data — proves the GstResource keepalive is active and
                //    the underlying mapped buffer is still valid.
                let slice = image.as_slice();
                assert_eq!(
                    slice.len(),
                    WIDTH * HEIGHT * 3,
                    "frame pixel count mismatch"
                );
                // Access first and last byte to ensure the mapping is live.
                let _ = slice[0];
                let _ = slice[slice.len() - 1];

                frames_received += 1;

                // 3. `image` drops here — GstResource::Drop unmaps and releases the
                //    GStreamer buffer ref exactly once.  A double-free or use-after-free
                //    would crash here (or be caught by valgrind/asan in CI).
            }
            if frames_received >= N_FRAMES {
                break;
            }
            std::thread::sleep(std::time::Duration::from_millis(10));
        }

        // Close pipeline before asserting so Close errors don't shadow the count.
        capture.close()?;

        assert_eq!(
            frames_received, N_FRAMES,
            "expected {N_FRAMES} frames but received {frames_received}"
        );

        Ok(())
    }

    /// Validates the buffer-size guard arithmetic used in `image_from_gst_buffer`.
    ///
    /// A real `MappedBuffer<Readable>` requires a live GStreamer pipeline and cannot
    /// be constructed in a pure unit test, so this test asserts the validation formula
    /// `expected = width * height * 3` for representative boundary values.
    #[test]
    fn gst_buffer_size_validation_arithmetic() {
        // expected bytes for an RGB24 frame of various sizes
        assert_eq!(8usize * 4 * 3, 96, "8x4 RGB24 = 96 bytes");
        assert_eq!(640usize * 480 * 3, 921_600, "640x480 RGB24 = 921600 bytes");
        assert_eq!(
            1920usize * 1080 * 3,
            6_220_800,
            "1920x1080 RGB24 = 6220800 bytes"
        );
        // A buffer smaller than the expected size must be rejected.
        // The actual rejection is inside image_from_gst_buffer; this test
        // documents the check boundary: expected-1 < expected → rejected.
        let width = 8usize;
        let height = 4usize;
        let expected = width * height * 3;
        assert!(
            (expected - 1) < expected,
            "a buffer of size expected-1 is strictly smaller than expected"
        );
    }
}