RemoteDesk/agent/agent-runtime/src/opus_capture.rs

use std::{fmt, time::Duration};

use crate::DESKTOP_MAX_OPUS_PACKET_BYTES;

pub const OPUS_SAMPLE_RATE: u32 = 48_000;
pub const OPUS_CHANNELS: u16 = 2;

#[derive(Clone, Debug)]
pub struct EncodedOpusPacket {
    pub data: Vec<u8>,
    pub duration: Duration,
    pub encode_latency: Duration,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum OpusCaptureError {
    RuntimeUnavailable,
    AudioMonitorUnavailable,
    PipelineFailure,
    OutputTimeout,
    InvalidOutput,
}

impl OpusCaptureError {
    #[must_use]
    pub const fn reason_code(self) -> &'static str {
        match self {
            Self::RuntimeUnavailable => "gstreamer_unavailable",
            Self::AudioMonitorUnavailable => "audio_monitor_unavailable",
            Self::PipelineFailure => "opus_pipeline_failed",
            Self::OutputTimeout => "opus_output_timeout",
            Self::InvalidOutput => "invalid_opus_output",
        }
    }
}

impl fmt::Display for OpusCaptureError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(self.reason_code())
    }
}

impl std::error::Error for OpusCaptureError {}

#[cfg(feature = "gstreamer-h264")]
mod implementation {
    use super::{
        DESKTOP_MAX_OPUS_PACKET_BYTES, EncodedOpusPacket, OPUS_CHANNELS, OPUS_SAMPLE_RATE,
        OpusCaptureError, valid_monitor_name,
    };
    use gstreamer as gst;
    use gstreamer::prelude::*;
    use gstreamer_app as gst_app;
    use gstreamer_app::prelude::*;
    use std::time::{Duration, Instant};

    const DEFAULT_MONITOR: &str = "@DEFAULT_MONITOR@";

    pub struct OpusCapture {
        pipeline: gst::Pipeline,
        sink: gst_app::AppSink,
    }

    impl OpusCapture {
        pub fn new() -> Result<Self, OpusCaptureError> {
            gst::init().map_err(|_| OpusCaptureError::RuntimeUnavailable)?;
            for factory in ["pulsesrc"] {
                if gst::ElementFactory::find(factory).is_none()
                    || gst::ElementFactory::find("opusenc").is_none()
                {
                    continue;
                }
                let description = format!(
                    "{factory} name=remotedesk_audio_source do-timestamp=true \
                     ! queue leaky=downstream max-size-buffers=8 max-size-bytes=0 max-size-time=0 \
                     ! audioconvert ! audioresample \
                     ! audio/x-raw,format=S16LE,rate={OPUS_SAMPLE_RATE},channels={OPUS_CHANNELS} \
                     ! opusenc bitrate=96000 frame-size=20 audio-type=restricted-lowdelay inband-fec=true packet-loss-percentage=5 dtx=true \
                     ! audio/x-opus,rate={OPUS_SAMPLE_RATE},channels={OPUS_CHANNELS},channel-mapping-family=0 \
                     ! appsink name=remotedesk_audio_sink sync=false max-buffers=8 drop=true"
                );
                let Ok(element) = gst::parse::launch(&description) else {
                    continue;
                };
                let Ok(pipeline) = element.downcast::<gst::Pipeline>() else {
                    continue;
                };
                let Some(source) = pipeline.by_name("remotedesk_audio_source") else {
                    continue;
                };
                let monitor = std::env::var("REMOTEDESK_AUDIO_MONITOR")
                    .ok()
                    .filter(|value| valid_monitor_name(value))
                    .unwrap_or_else(|| DEFAULT_MONITOR.to_owned());
                source.set_property("device", monitor);
                let Some(sink) = pipeline
                    .by_name("remotedesk_audio_sink")
                    .and_then(|element| element.downcast::<gst_app::AppSink>().ok())
                else {
                    continue;
                };
                if pipeline.set_state(gst::State::Playing).is_err() {
                    let _ = pipeline.set_state(gst::State::Null);
                    continue;
                }
                return Ok(Self { pipeline, sink });
            }
            Err(OpusCaptureError::AudioMonitorUnavailable)
        }

        pub fn next_packet(&mut self) -> Result<EncodedOpusPacket, OpusCaptureError> {
            let started = Instant::now();
            let sample = self
                .sink
                .try_pull_sample(gst::ClockTime::from_mseconds(250))
                .ok_or(OpusCaptureError::OutputTimeout)?;
            let caps = sample.caps().ok_or(OpusCaptureError::InvalidOutput)?;
            let structure = caps.structure(0).ok_or(OpusCaptureError::InvalidOutput)?;
            if structure.name() != "audio/x-opus"
                || structure.get::<i32>("rate").ok() != Some(OPUS_SAMPLE_RATE as i32)
                || structure.get::<i32>("channels").ok() != Some(i32::from(OPUS_CHANNELS))
            {
                return Err(OpusCaptureError::InvalidOutput);
            }
            let buffer = sample.buffer().ok_or(OpusCaptureError::InvalidOutput)?;
            let duration = buffer
                .duration()
                .map(|value| Duration::from_nanos(value.nseconds()))
                .filter(|value| {
                    *value >= Duration::from_micros(2_500) && *value <= Duration::from_millis(120)
                })
                .ok_or(OpusCaptureError::InvalidOutput)?;
            let mapped = buffer
                .map_readable()
                .map_err(|_| OpusCaptureError::InvalidOutput)?;
            let data = mapped.as_slice();
            if data.is_empty() || data.len() > DESKTOP_MAX_OPUS_PACKET_BYTES {
                return Err(OpusCaptureError::InvalidOutput);
            }
            Ok(EncodedOpusPacket {
                data: data.to_vec(),
                duration,
                encode_latency: started.elapsed(),
            })
        }
    }

    impl Drop for OpusCapture {
        fn drop(&mut self) {
            let _ = self.pipeline.set_state(gst::State::Null);
        }
    }
}

#[cfg(not(feature = "gstreamer-h264"))]
mod implementation {
    use super::{EncodedOpusPacket, OpusCaptureError};

    pub struct OpusCapture;

    impl OpusCapture {
        pub fn new() -> Result<Self, OpusCaptureError> {
            Err(OpusCaptureError::RuntimeUnavailable)
        }

        pub fn next_packet(&mut self) -> Result<EncodedOpusPacket, OpusCaptureError> {
            Err(OpusCaptureError::RuntimeUnavailable)
        }
    }
}

pub use implementation::OpusCapture;

fn valid_monitor_name(value: &str) -> bool {
    !value.is_empty()
        && value.len() <= 256
        && value
            .bytes()
            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'-' | b'_' | b'@'))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn monitor_name_is_bounded_and_not_pipeline_syntax() {
        assert!(valid_monitor_name("alsa_output.pci.monitor"));
        assert!(valid_monitor_name("@DEFAULT_MONITOR@"));
        assert!(!valid_monitor_name(""));
        assert!(!valid_monitor_name("monitor ! fakesink"));
        assert!(!valid_monitor_name("monitor\nnext"));
        assert!(!valid_monitor_name(&"x".repeat(257)));
    }
}