Document all-Rust migration and extend native media stack
ci / rust (push) Canceled after 0s
ci / web (push) Canceled after 0s
ci / package-preview (push) Canceled after 0s
ci / package-installer (push) Canceled after 0s
ci / linux-agent (push) Canceled after 0s
ci / edge-service (push) Canceled after 0s
ci / coturn-pop (push) Canceled after 0s
ci / package-windows-host (push) Canceled after 0s
ci / rust (push) Canceled after 0s
ci / web (push) Canceled after 0s
ci / package-preview (push) Canceled after 0s
ci / package-installer (push) Canceled after 0s
ci / linux-agent (push) Canceled after 0s
ci / edge-service (push) Canceled after 0s
ci / coturn-pop (push) Canceled after 0s
ci / package-windows-host (push) Canceled after 0s
This commit is contained in:
@@ -703,7 +703,7 @@ mod tests {
|
||||
|
||||
let selected = select_encoder(
|
||||
DesktopBackend::XorgShmCompatibility,
|
||||
ZeroCopyPolicy::Compatibility,
|
||||
ZeroCopyPolicy::RequiredEndToEnd,
|
||||
&capture,
|
||||
&candidates,
|
||||
)
|
||||
@@ -788,7 +788,7 @@ mod tests {
|
||||
|
||||
let selected = select_encoder(
|
||||
DesktopBackend::WaylandPipeWire,
|
||||
ZeroCopyPolicy::Compatibility,
|
||||
ZeroCopyPolicy::RequiredEndToEnd,
|
||||
&capture,
|
||||
&candidates,
|
||||
)
|
||||
@@ -828,7 +828,7 @@ mod tests {
|
||||
|
||||
let selected = select_encoder(
|
||||
DesktopBackend::WaylandPipeWire,
|
||||
ZeroCopyPolicy::Compatibility,
|
||||
ZeroCopyPolicy::RequiredEndToEnd,
|
||||
&capture,
|
||||
&candidates,
|
||||
)
|
||||
@@ -936,7 +936,7 @@ mod tests {
|
||||
transfer_mode: TransferMode::SameAdapterZeroCopy,
|
||||
};
|
||||
let mut pipeline = VideoPipeline::active(
|
||||
ZeroCopyPolicy::Compatibility,
|
||||
ZeroCopyPolicy::RequiredEndToEnd,
|
||||
Resolution::new(1920, 1080).unwrap(),
|
||||
zero_copy,
|
||||
)
|
||||
|
||||
@@ -989,7 +989,7 @@ mod tests {
|
||||
|
||||
let plan = MultiDisplayCapturePlan::create(
|
||||
DesktopBackend::WaylandPipeWire,
|
||||
ZeroCopyPolicy::Compatibility,
|
||||
ZeroCopyPolicy::RequiredEndToEnd,
|
||||
&topology,
|
||||
&sources,
|
||||
&DisplaySelection::all(7),
|
||||
@@ -1071,7 +1071,7 @@ mod tests {
|
||||
|
||||
let plan = MultiDisplayCapturePlan::create(
|
||||
DesktopBackend::WaylandPipeWire,
|
||||
ZeroCopyPolicy::Compatibility,
|
||||
ZeroCopyPolicy::RequiredEndToEnd,
|
||||
&topology,
|
||||
&sources,
|
||||
&DisplaySelection::all(7),
|
||||
@@ -1098,7 +1098,7 @@ mod tests {
|
||||
assert_eq!(
|
||||
MultiDisplayCapturePlan::create(
|
||||
DesktopBackend::WaylandPipeWire,
|
||||
ZeroCopyPolicy::Compatibility,
|
||||
ZeroCopyPolicy::RequiredEndToEnd,
|
||||
&topology,
|
||||
&sources,
|
||||
&DisplaySelection::all(7),
|
||||
|
||||
@@ -12,21 +12,43 @@ path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
anyhow = "1.0"
|
||||
base64 = "0.22"
|
||||
clap = { version = "4.5", features = ["derive"] }
|
||||
flate2 = "1.1"
|
||||
getrandom = { version = "0.3", features = ["std"] }
|
||||
hmac = "=0.13.0"
|
||||
memmap2 = "0.9"
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = "1.0"
|
||||
sha2 = "=0.11.0"
|
||||
tokio = { version = "1.47", features = ["macros", "net", "rt-multi-thread", "io-util", "sync", "time"] }
|
||||
zeroize = "1.8"
|
||||
windows = { version = "0.62.2", features = [
|
||||
"Win32_Foundation",
|
||||
"Win32_Graphics_Direct3D",
|
||||
"Win32_Graphics_Direct3D11",
|
||||
"Win32_Graphics_Dxgi",
|
||||
"Win32_Graphics_Dxgi_Common",
|
||||
"Win32_Graphics_Gdi",
|
||||
"Win32_Media_MediaFoundation",
|
||||
"Win32_System_Com",
|
||||
"Win32_System_Ole",
|
||||
"Win32_System_Variant",
|
||||
"Win32_UI_Input_KeyboardAndMouse",
|
||||
"Win32_UI_WindowsAndMessaging",
|
||||
] }
|
||||
|
||||
[target.'cfg(windows)'.dependencies]
|
||||
cpal = "0.17.3"
|
||||
opus2 = { version = "0.4.0", features = ["bundled"] }
|
||||
windows-sys = { version = "0.61.2", features = [
|
||||
"Win32_Foundation",
|
||||
"Win32_Security",
|
||||
"Win32_Security_Authorization",
|
||||
"Win32_Storage_FileSystem",
|
||||
"Win32_System_Threading",
|
||||
] }
|
||||
|
||||
[lints.rust]
|
||||
unsafe_code = "allow"
|
||||
|
||||
|
||||
@@ -0,0 +1,471 @@
|
||||
//! Windows WASAPI loopback capture and fixed-duration Opus packetization.
|
||||
//!
|
||||
//! The capture callback is deliberately tiny: it only copies normalized f32
|
||||
//! samples into a bounded queue. Resampling, Opus encoding, timestamps, and
|
||||
//! discontinuity marking stay on the media worker thread.
|
||||
|
||||
#[cfg(windows)]
|
||||
mod windows_audio {
|
||||
use cpal::traits::{DeviceTrait as _, HostTrait as _, StreamTrait as _};
|
||||
use opus2::{Application, Bitrate, Channels, Encoder};
|
||||
use std::io;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::mpsc::{Receiver, SyncSender, TryRecvError, TrySendError, sync_channel};
|
||||
use std::time::Instant;
|
||||
|
||||
const SAMPLE_RATE: u32 = 48_000;
|
||||
const CHANNELS: usize = 2;
|
||||
const FRAME_SAMPLES: usize = 960;
|
||||
const MAX_OPUS_BYTES: usize = 4 * 1024;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct CapturedAudioChunk {
|
||||
pub(crate) samples: Vec<f32>,
|
||||
pub(crate) capture_end_pts_100ns: u64,
|
||||
}
|
||||
|
||||
pub(crate) struct AudioCapture {
|
||||
receiver: Receiver<CapturedAudioChunk>,
|
||||
pub(crate) source_rate: u32,
|
||||
pub(crate) source_channels: usize,
|
||||
pub(crate) discontinuity: Arc<AtomicBool>,
|
||||
_stream: cpal::Stream,
|
||||
}
|
||||
|
||||
impl AudioCapture {
|
||||
pub(crate) fn new(session_started: Instant) -> io::Result<Self> {
|
||||
let host = cpal::default_host();
|
||||
let device = host.default_output_device().ok_or_else(|| {
|
||||
io::Error::new(
|
||||
io::ErrorKind::NotFound,
|
||||
"WASAPI default output device unavailable",
|
||||
)
|
||||
})?;
|
||||
let supported = device
|
||||
.default_output_config()
|
||||
.map_err(|error| io::Error::other(error.to_string()))?;
|
||||
let source_rate = supported.sample_rate();
|
||||
let source_channels = usize::from(supported.channels());
|
||||
if source_rate == 0 || source_channels == 0 || source_channels > 8 {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
"WASAPI loopback format is outside supported bounds",
|
||||
));
|
||||
}
|
||||
let config = supported.config();
|
||||
let (sender, receiver) = sync_channel(32);
|
||||
let discontinuity = Arc::new(AtomicBool::new(false));
|
||||
let callback_discontinuity = Arc::clone(&discontinuity);
|
||||
let error_discontinuity = Arc::clone(&discontinuity);
|
||||
let error_callback = move |_error: cpal::StreamError| {
|
||||
error_discontinuity.store(true, Ordering::Release);
|
||||
};
|
||||
let stream = match supported.sample_format() {
|
||||
cpal::SampleFormat::F32 => build_stream::<f32, _>(
|
||||
&device,
|
||||
&config,
|
||||
sender,
|
||||
callback_discontinuity,
|
||||
session_started,
|
||||
error_callback,
|
||||
)?,
|
||||
cpal::SampleFormat::I16 => build_stream::<i16, _>(
|
||||
&device,
|
||||
&config,
|
||||
sender,
|
||||
callback_discontinuity,
|
||||
session_started,
|
||||
error_callback,
|
||||
)?,
|
||||
cpal::SampleFormat::I32 => build_stream::<i32, _>(
|
||||
&device,
|
||||
&config,
|
||||
sender,
|
||||
callback_discontinuity,
|
||||
session_started,
|
||||
error_callback,
|
||||
)?,
|
||||
cpal::SampleFormat::U16 => build_stream::<u16, _>(
|
||||
&device,
|
||||
&config,
|
||||
sender,
|
||||
callback_discontinuity,
|
||||
session_started,
|
||||
error_callback,
|
||||
)?,
|
||||
cpal::SampleFormat::F64 => build_stream::<f64, _>(
|
||||
&device,
|
||||
&config,
|
||||
sender,
|
||||
callback_discontinuity,
|
||||
session_started,
|
||||
error_callback,
|
||||
)?,
|
||||
_ => {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::Unsupported,
|
||||
"WASAPI loopback sample format is unsupported",
|
||||
));
|
||||
}
|
||||
};
|
||||
stream
|
||||
.play()
|
||||
.map_err(|error| io::Error::other(error.to_string()))?;
|
||||
Ok(Self {
|
||||
receiver,
|
||||
source_rate,
|
||||
source_channels,
|
||||
discontinuity,
|
||||
_stream: stream,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn try_receive(&self) -> Result<Option<CapturedAudioChunk>, TryRecvError> {
|
||||
self.receiver.try_recv().map(Some)
|
||||
}
|
||||
|
||||
pub(crate) fn try_receive_latest(
|
||||
&self,
|
||||
) -> Result<Option<CapturedAudioChunk>, TryRecvError> {
|
||||
let mut latest = self.receiver.try_recv()?;
|
||||
loop {
|
||||
match self.receiver.try_recv() {
|
||||
Ok(chunk) => latest = chunk,
|
||||
Err(TryRecvError::Empty) => return Ok(Some(latest)),
|
||||
Err(TryRecvError::Disconnected) => return Err(TryRecvError::Disconnected),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn build_stream<T, E>(
|
||||
device: &cpal::Device,
|
||||
config: &cpal::StreamConfig,
|
||||
sender: SyncSender<CapturedAudioChunk>,
|
||||
discontinuity: Arc<AtomicBool>,
|
||||
session_started: Instant,
|
||||
error_callback: E,
|
||||
) -> io::Result<cpal::Stream>
|
||||
where
|
||||
T: cpal::SizedSample + cpal::Sample + Send + 'static,
|
||||
f32: cpal::FromSample<T>,
|
||||
E: FnMut(cpal::StreamError) + Send + 'static,
|
||||
{
|
||||
let callback = move |data: &[T], _info: &cpal::InputCallbackInfo| {
|
||||
let samples = data
|
||||
.iter()
|
||||
.map(|sample| sample.to_sample::<f32>())
|
||||
.collect();
|
||||
let capture_end_pts_100ns =
|
||||
duration_to_100ns(Instant::now().saturating_duration_since(session_started));
|
||||
match sender.try_send(CapturedAudioChunk {
|
||||
samples,
|
||||
capture_end_pts_100ns,
|
||||
}) {
|
||||
Ok(()) => {}
|
||||
Err(TrySendError::Full(_) | TrySendError::Disconnected(_)) => {
|
||||
discontinuity.store(true, Ordering::Release);
|
||||
}
|
||||
}
|
||||
};
|
||||
device
|
||||
.build_input_stream(config, callback, error_callback, None)
|
||||
.map_err(|error| io::Error::other(error.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) struct AudioPipeline {
|
||||
resampler: Resampler,
|
||||
encoder: Encoder,
|
||||
pending: Vec<f32>,
|
||||
sample_count: u64,
|
||||
sequence: u64,
|
||||
discontinuity: bool,
|
||||
timeline_base_pts_100ns: Option<u64>,
|
||||
needs_reanchor: bool,
|
||||
}
|
||||
|
||||
pub(crate) struct EncodedAudioPacket {
|
||||
pub(crate) sequence: u64,
|
||||
pub(crate) pts_100ns: u64,
|
||||
pub(crate) discontinuity: bool,
|
||||
pub(crate) data: Vec<u8>,
|
||||
}
|
||||
|
||||
impl AudioPipeline {
|
||||
pub(crate) fn new(source_rate: u32, source_channels: usize) -> io::Result<Self> {
|
||||
let mut encoder = Encoder::new(SAMPLE_RATE, Channels::Stereo, Application::Audio)
|
||||
.map_err(|error| io::Error::other(error.to_string()))?;
|
||||
encoder
|
||||
.set_bitrate(Bitrate::Bits(96_000))
|
||||
.map_err(|error| io::Error::other(error.to_string()))?;
|
||||
encoder
|
||||
.set_inband_fec(true)
|
||||
.map_err(|error| io::Error::other(error.to_string()))?;
|
||||
encoder
|
||||
.set_packet_loss_perc(5)
|
||||
.map_err(|error| io::Error::other(error.to_string()))?;
|
||||
encoder
|
||||
.set_dtx(true)
|
||||
.map_err(|error| io::Error::other(error.to_string()))?;
|
||||
Ok(Self {
|
||||
resampler: Resampler::new(source_rate, source_channels)?,
|
||||
encoder,
|
||||
pending: Vec::with_capacity(FRAME_SAMPLES * CHANNELS * 2),
|
||||
sample_count: 0,
|
||||
sequence: 1,
|
||||
discontinuity: false,
|
||||
timeline_base_pts_100ns: None,
|
||||
needs_reanchor: true,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn mark_discontinuity(&mut self) {
|
||||
self.discontinuity = true;
|
||||
self.needs_reanchor = true;
|
||||
self.pending.clear();
|
||||
self.resampler.reset();
|
||||
}
|
||||
|
||||
pub(crate) fn push_source(
|
||||
&mut self,
|
||||
samples: &[f32],
|
||||
capture_end_pts_100ns: u64,
|
||||
) -> io::Result<Vec<EncodedAudioPacket>> {
|
||||
if !samples.len().is_multiple_of(self.resampler.source_channels) {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
"audio callback frame is incomplete",
|
||||
));
|
||||
}
|
||||
if self.needs_reanchor {
|
||||
let source_frames = samples.len() / self.resampler.source_channels;
|
||||
let source_duration_100ns = u64::try_from(source_frames)
|
||||
.unwrap_or(u64::MAX)
|
||||
.saturating_mul(10_000_000)
|
||||
/ u64::from(self.resampler.source_rate);
|
||||
self.timeline_base_pts_100ns =
|
||||
Some(capture_end_pts_100ns.saturating_sub(source_duration_100ns));
|
||||
self.sample_count = 0;
|
||||
self.needs_reanchor = false;
|
||||
}
|
||||
self.pending.extend(self.resampler.push(samples)?);
|
||||
let frame_width = FRAME_SAMPLES * CHANNELS;
|
||||
let mut packets = Vec::new();
|
||||
while self.pending.len() >= frame_width {
|
||||
let frame = self.pending.drain(..frame_width).collect::<Vec<_>>();
|
||||
let data = self
|
||||
.encoder
|
||||
.encode_vec_float(&frame, MAX_OPUS_BYTES)
|
||||
.map_err(|error| io::Error::other(error.to_string()))?;
|
||||
let packet = EncodedAudioPacket {
|
||||
sequence: self.sequence,
|
||||
pts_100ns: self
|
||||
.timeline_base_pts_100ns
|
||||
.unwrap_or(capture_end_pts_100ns)
|
||||
.saturating_add(
|
||||
self.sample_count.saturating_mul(10_000_000) / u64::from(SAMPLE_RATE),
|
||||
),
|
||||
discontinuity: self.discontinuity,
|
||||
data,
|
||||
};
|
||||
self.sequence = self.sequence.saturating_add(1);
|
||||
self.sample_count = self.sample_count.saturating_add(FRAME_SAMPLES as u64);
|
||||
self.discontinuity = false;
|
||||
packets.push(packet);
|
||||
}
|
||||
Ok(packets)
|
||||
}
|
||||
}
|
||||
|
||||
struct Resampler {
|
||||
source_rate: u32,
|
||||
source_channels: usize,
|
||||
buffer: Vec<f32>,
|
||||
position_numerator: u64,
|
||||
}
|
||||
|
||||
impl Resampler {
|
||||
fn new(source_rate: u32, source_channels: usize) -> io::Result<Self> {
|
||||
if source_rate == 0 || source_channels == 0 || source_channels > 8 {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
"invalid source audio format",
|
||||
));
|
||||
}
|
||||
Ok(Self {
|
||||
source_rate,
|
||||
source_channels,
|
||||
buffer: Vec::new(),
|
||||
position_numerator: 0,
|
||||
})
|
||||
}
|
||||
|
||||
fn push(&mut self, samples: &[f32]) -> io::Result<Vec<f32>> {
|
||||
if !samples.len().is_multiple_of(self.source_channels) {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
"audio callback frame is incomplete",
|
||||
));
|
||||
}
|
||||
self.buffer.extend_from_slice(samples);
|
||||
let source_frames = self.buffer.len() / self.source_channels;
|
||||
let source_frames_u64 = u64::try_from(source_frames).map_err(|_| {
|
||||
io::Error::new(io::ErrorKind::InvalidData, "audio callback is too large")
|
||||
})?;
|
||||
let mut output = Vec::new();
|
||||
while self.position_numerator / u64::from(SAMPLE_RATE) < source_frames_u64 {
|
||||
let index_u64 = self.position_numerator / u64::from(SAMPLE_RATE);
|
||||
let index = usize::try_from(index_u64).map_err(|_| {
|
||||
io::Error::new(io::ErrorKind::InvalidData, "audio position overflow")
|
||||
})?;
|
||||
let remainder = self.position_numerator % u64::from(SAMPLE_RATE);
|
||||
if remainder != 0 && index + 1 >= source_frames {
|
||||
break;
|
||||
}
|
||||
let fraction = f32::from(u16::try_from(remainder).map_err(|_| {
|
||||
io::Error::new(io::ErrorKind::InvalidData, "audio fraction overflow")
|
||||
})?) / 48_000.0;
|
||||
for channel in 0..CHANNELS {
|
||||
let source_channel = if self.source_channels == 1 {
|
||||
0
|
||||
} else {
|
||||
channel.min(self.source_channels - 1)
|
||||
};
|
||||
let first = self.buffer[index * self.source_channels + source_channel];
|
||||
let second = if index + 1 < source_frames {
|
||||
self.buffer[(index + 1) * self.source_channels + source_channel]
|
||||
} else {
|
||||
first
|
||||
};
|
||||
output.push(first + (second - first) * fraction);
|
||||
}
|
||||
self.position_numerator = self
|
||||
.position_numerator
|
||||
.saturating_add(u64::from(self.source_rate));
|
||||
}
|
||||
let consumed = usize::try_from(self.position_numerator / u64::from(SAMPLE_RATE))
|
||||
.unwrap_or(source_frames);
|
||||
if consumed > 0 {
|
||||
let remove = consumed.min(source_frames);
|
||||
self.buffer.drain(..remove * self.source_channels);
|
||||
self.position_numerator = self.position_numerator.saturating_sub(
|
||||
u64::try_from(remove)
|
||||
.unwrap_or(u64::MAX)
|
||||
.saturating_mul(u64::from(SAMPLE_RATE)),
|
||||
);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn reset(&mut self) {
|
||||
self.buffer.clear();
|
||||
self.position_numerator = 0;
|
||||
}
|
||||
}
|
||||
|
||||
fn duration_to_100ns(duration: std::time::Duration) -> u64 {
|
||||
u64::try_from(duration.as_nanos() / 100).unwrap_or(u64::MAX)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn resampler_preserves_48khz_stereo_frames() {
|
||||
let input = vec![0.125_f32; FRAME_SAMPLES * 2];
|
||||
let output = Resampler::new(48_000, 2).unwrap().push(&input).unwrap();
|
||||
assert_eq!(output, input);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resampler_duplicates_mono_and_downsamples_96khz() {
|
||||
let mono = vec![0.25_f32; 960];
|
||||
let mono_output = Resampler::new(48_000, 1).unwrap().push(&mono).unwrap();
|
||||
assert_eq!(mono_output.len(), 960 * 2);
|
||||
assert!(
|
||||
mono_output
|
||||
.chunks_exact(2)
|
||||
.all(|frame| frame == [0.25, 0.25])
|
||||
);
|
||||
|
||||
let stereo_96khz = vec![0.5_f32; 960 * 2];
|
||||
let downsampled = Resampler::new(96_000, 2)
|
||||
.unwrap()
|
||||
.push(&stereo_96khz)
|
||||
.unwrap();
|
||||
assert_eq!(downsampled.len(), 480 * 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resampler_keeps_44100hz_ratio_across_callback_boundaries() {
|
||||
let mut resampler = Resampler::new(44_100, 2).unwrap();
|
||||
let mut output_frames = 0;
|
||||
for _ in 0..10 {
|
||||
output_frames += resampler.push(&vec![0.0; 441 * 2]).unwrap().len() / 2;
|
||||
}
|
||||
assert!((4_799..=4_800).contains(&output_frames));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn opus_pipeline_emits_bounded_monotonic_twenty_ms_packets() {
|
||||
let mut pipeline = AudioPipeline::new(48_000, 2).unwrap();
|
||||
let first = pipeline
|
||||
.push_source(&vec![0.0; FRAME_SAMPLES * 2], 1_000_000)
|
||||
.unwrap();
|
||||
let second = pipeline
|
||||
.push_source(&vec![0.0; FRAME_SAMPLES * 2], 1_200_000)
|
||||
.unwrap();
|
||||
assert_eq!(first.len(), 1);
|
||||
assert_eq!(second.len(), 1);
|
||||
assert_eq!(first[0].sequence, 1);
|
||||
assert_eq!(first[0].pts_100ns, 800_000);
|
||||
assert_eq!(second[0].sequence, 2);
|
||||
assert_eq!(second[0].pts_100ns, 1_000_000);
|
||||
assert!(!first[0].data.is_empty() && first[0].data.len() <= MAX_OPUS_BYTES);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discontinuity_is_applied_to_exactly_one_packet() {
|
||||
let mut pipeline = AudioPipeline::new(48_000, 2).unwrap();
|
||||
pipeline.mark_discontinuity();
|
||||
let first = pipeline
|
||||
.push_source(&vec![0.0; FRAME_SAMPLES * 2], 2_000_000)
|
||||
.unwrap();
|
||||
let second = pipeline
|
||||
.push_source(&vec![0.0; FRAME_SAMPLES * 2], 2_200_000)
|
||||
.unwrap();
|
||||
assert!(first[0].discontinuity);
|
||||
assert!(!second[0].discontinuity);
|
||||
assert_eq!(first[0].pts_100ns, 1_800_000);
|
||||
assert_eq!(second[0].pts_100ns, 2_000_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discontinuity_reanchors_pts_without_resetting_sequence() {
|
||||
let mut pipeline = AudioPipeline::new(48_000, 2).unwrap();
|
||||
let first = pipeline
|
||||
.push_source(&vec![0.0; FRAME_SAMPLES * 2], 200_000)
|
||||
.unwrap();
|
||||
pipeline.mark_discontinuity();
|
||||
let resumed = pipeline
|
||||
.push_source(&vec![0.0; FRAME_SAMPLES * 2], 5_200_000)
|
||||
.unwrap();
|
||||
assert_eq!(first[0].sequence, 1);
|
||||
assert_eq!(first[0].pts_100ns, 0);
|
||||
assert_eq!(resumed[0].sequence, 2);
|
||||
assert_eq!(resumed[0].pts_100ns, 5_000_000);
|
||||
assert!(resumed[0].discontinuity);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
pub(crate) use windows_audio::{AudioCapture, AudioPipeline, CapturedAudioChunk};
|
||||
|
||||
#[cfg(not(windows))]
|
||||
pub(crate) struct AudioCapture;
|
||||
|
||||
#[cfg(not(windows))]
|
||||
pub(crate) struct AudioPipeline;
|
||||
@@ -0,0 +1,166 @@
|
||||
//! Runtime capability probes. A registry/API presence check is not treated as
|
||||
//! an active zero-copy encoder; it only reports whether Windows exposes an
|
||||
//! approved hardware MFT that can be probed by the session worker.
|
||||
|
||||
use std::sync::OnceLock;
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub(crate) struct CapabilitySnapshot {
|
||||
pub(crate) desktop_duplication: bool,
|
||||
pub(crate) hardware_h264_mft: bool,
|
||||
pub(crate) hardware_h264_encoder_verified: bool,
|
||||
pub(crate) hardware_hevc_mft: bool,
|
||||
pub(crate) hardware_av1_mft: bool,
|
||||
pub(crate) hardware_pipeline_verified: bool,
|
||||
}
|
||||
|
||||
impl CapabilitySnapshot {
|
||||
#[must_use]
|
||||
pub(crate) const fn capture_backend(&self) -> &'static str {
|
||||
if self.desktop_duplication {
|
||||
"desktop-duplication-compatibility"
|
||||
} else {
|
||||
"gdi-bgra-zlib"
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub(crate) const fn pipeline_state(&self) -> &'static str {
|
||||
if self.hardware_pipeline_verified {
|
||||
"hardware_pipeline_verified"
|
||||
} else if self.desktop_duplication
|
||||
|| self.hardware_h264_mft
|
||||
|| self.hardware_hevc_mft
|
||||
|| self.hardware_av1_mft
|
||||
{
|
||||
"hardware_probe_partial"
|
||||
} else {
|
||||
"hardware_probe_failed"
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub(crate) fn encoder_available(&self, codec: &str) -> bool {
|
||||
match codec {
|
||||
"h264" => self.hardware_h264_mft,
|
||||
"hevc" => self.hardware_hevc_mft,
|
||||
"av1" => self.hardware_av1_mft,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static SNAPSHOT: OnceLock<CapabilitySnapshot> = OnceLock::new();
|
||||
|
||||
pub(crate) fn snapshot() -> &'static CapabilitySnapshot {
|
||||
SNAPSHOT.get_or_init(probe)
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
pub(crate) fn probe() -> CapabilitySnapshot {
|
||||
CapabilitySnapshot {
|
||||
desktop_duplication: crate::capture::DdaCapture::new().is_ok(),
|
||||
hardware_h264_mft: probe_hardware_encoder(
|
||||
windows::Win32::Media::MediaFoundation::MFVideoFormat_H264,
|
||||
),
|
||||
hardware_h264_encoder_verified: probe_hardware_h264_encoder(),
|
||||
hardware_hevc_mft: probe_hardware_encoder(
|
||||
windows::Win32::Media::MediaFoundation::MFVideoFormat_HEVC,
|
||||
),
|
||||
hardware_av1_mft: probe_hardware_encoder(
|
||||
windows::Win32::Media::MediaFoundation::MFVideoFormat_AV1,
|
||||
),
|
||||
hardware_pipeline_verified: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(windows))]
|
||||
pub(crate) const fn probe() -> CapabilitySnapshot {
|
||||
CapabilitySnapshot {
|
||||
desktop_duplication: false,
|
||||
hardware_h264_mft: false,
|
||||
hardware_h264_encoder_verified: false,
|
||||
hardware_hevc_mft: false,
|
||||
hardware_av1_mft: false,
|
||||
hardware_pipeline_verified: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
fn probe_hardware_h264_encoder() -> bool {
|
||||
let Ok(capture) = crate::capture::DdaCapture::new() else {
|
||||
return false;
|
||||
};
|
||||
let (width, height) = capture.dimensions();
|
||||
if width < 2 || height < 2 || width % 2 != 0 || height % 2 != 0 {
|
||||
return false;
|
||||
}
|
||||
let Ok(mut encoder) = crate::encoder::H264Encoder::new(width, height, 30, true) else {
|
||||
return false;
|
||||
};
|
||||
let Some(bytes) = usize::try_from(width)
|
||||
.ok()
|
||||
.and_then(|w| usize::try_from(height).ok().and_then(|h| w.checked_mul(h)))
|
||||
.and_then(|pixels| pixels.checked_mul(4))
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
let bgra = vec![0_u8; bytes];
|
||||
let Ok(nv12) = crate::encoder::bgra_to_nv12(width, height, &bgra) else {
|
||||
return false;
|
||||
};
|
||||
for _ in 0..3 {
|
||||
if encoder.encode(&nv12).is_err() {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
fn probe_hardware_encoder(codec: windows::core::GUID) -> bool {
|
||||
use windows::Win32::Media::MediaFoundation::{
|
||||
MF_VERSION, MFSTARTUP_FULL, MFStartup, MFT_CATEGORY_VIDEO_ENCODER, MFT_ENUM_FLAG_HARDWARE,
|
||||
MFT_ENUM_FLAG_SORTANDFILTER_APPROVED_ONLY, MFT_REGISTER_TYPE_INFO, MFTEnumEx,
|
||||
MFVideoFormat_NV12,
|
||||
};
|
||||
use windows::Win32::System::Com::{COINIT_MULTITHREADED, CoInitializeEx, CoTaskMemFree};
|
||||
|
||||
unsafe {
|
||||
if CoInitializeEx(None, COINIT_MULTITHREADED).is_err()
|
||||
|| MFStartup(MF_VERSION, MFSTARTUP_FULL).is_err()
|
||||
{
|
||||
return false;
|
||||
}
|
||||
let input = MFT_REGISTER_TYPE_INFO {
|
||||
guidMajorType: windows::Win32::Media::MediaFoundation::MFMediaType_Video,
|
||||
guidSubtype: MFVideoFormat_NV12,
|
||||
};
|
||||
let output = MFT_REGISTER_TYPE_INFO {
|
||||
guidMajorType: windows::Win32::Media::MediaFoundation::MFMediaType_Video,
|
||||
guidSubtype: codec,
|
||||
};
|
||||
let mut activations: *mut Option<windows::Win32::Media::MediaFoundation::IMFActivate> =
|
||||
std::ptr::null_mut();
|
||||
let mut count = 0_u32;
|
||||
let result = MFTEnumEx(
|
||||
MFT_CATEGORY_VIDEO_ENCODER,
|
||||
MFT_ENUM_FLAG_HARDWARE | MFT_ENUM_FLAG_SORTANDFILTER_APPROVED_ONLY,
|
||||
Some(&raw const input),
|
||||
Some(&raw const output),
|
||||
&raw mut activations,
|
||||
&raw mut count,
|
||||
)
|
||||
.is_ok()
|
||||
&& count > 0;
|
||||
if !activations.is_null() {
|
||||
let entries = std::slice::from_raw_parts_mut(activations, count as usize);
|
||||
for entry in entries {
|
||||
*entry = None;
|
||||
}
|
||||
CoTaskMemFree(Some(activations.cast()));
|
||||
}
|
||||
let _ = windows::Win32::Media::MediaFoundation::MFShutdown();
|
||||
result
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,302 @@
|
||||
//! Windows Desktop Duplication compatibility capture.
|
||||
//!
|
||||
//! The staging/readback path is deliberately reported as software/cpu_upload.
|
||||
//! It is a compatibility path, not an end-to-end GPU zero-copy pipeline. An
|
||||
//! IDD swap-chain consumer can replace this module without changing the agent
|
||||
//! control protocol or frame envelope.
|
||||
|
||||
#[cfg(windows)]
|
||||
mod windows_capture {
|
||||
use anyhow::Context as _;
|
||||
use flate2::{Compression, write::ZlibEncoder};
|
||||
use std::io::Write as _;
|
||||
use std::time::{Duration, Instant};
|
||||
use windows::Win32::Graphics::Direct3D::D3D_DRIVER_TYPE_UNKNOWN;
|
||||
use windows::Win32::Graphics::Direct3D11::{
|
||||
D3D11_BIND_FLAG, D3D11_CPU_ACCESS_READ, D3D11_CREATE_DEVICE_FLAG, D3D11_MAP_READ,
|
||||
D3D11_RESOURCE_MISC_FLAG, D3D11_TEXTURE2D_DESC, D3D11_USAGE_STAGING, D3D11CreateDevice,
|
||||
ID3D11DeviceContext, ID3D11Texture2D,
|
||||
};
|
||||
use windows::Win32::Graphics::Dxgi::Common::{DXGI_FORMAT_B8G8R8A8_UNORM, DXGI_SAMPLE_DESC};
|
||||
use windows::Win32::Graphics::Dxgi::{
|
||||
CreateDXGIFactory1, DXGI_ERROR_WAIT_TIMEOUT, DXGI_OUTDUPL_FRAME_INFO, IDXGIAdapter1,
|
||||
IDXGIOutput1, IDXGIOutputDuplication, IDXGIResource,
|
||||
};
|
||||
use windows::core::Interface as _;
|
||||
|
||||
const FRAME_MAGIC: &[u8; 4] = b"RDWF";
|
||||
const FRAME_VERSION: u8 = 2;
|
||||
const FRAME_HEADER_LEN: usize = 40;
|
||||
const CODEC_BGRA_ZLIB: u8 = 1;
|
||||
const MAX_DIMENSION: u32 = 16_384;
|
||||
const MAX_RAW_FRAME_BYTES: usize = 256 * 1024 * 1024;
|
||||
|
||||
pub struct EncodedFrame {
|
||||
pub bytes: Vec<u8>,
|
||||
pub capture_latency: Duration,
|
||||
pub encode_latency: Duration,
|
||||
}
|
||||
|
||||
pub struct CapturedFrame {
|
||||
pub width: u32,
|
||||
pub height: u32,
|
||||
pub pixels: Vec<u8>,
|
||||
pub capture_latency: Duration,
|
||||
}
|
||||
|
||||
pub struct DdaCapture {
|
||||
context: ID3D11DeviceContext,
|
||||
duplication: IDXGIOutputDuplication,
|
||||
staging: ID3D11Texture2D,
|
||||
width: u32,
|
||||
height: u32,
|
||||
pixels: Vec<u8>,
|
||||
}
|
||||
|
||||
impl DdaCapture {
|
||||
pub fn dimensions(&self) -> (u32, u32) {
|
||||
(self.width, self.height)
|
||||
}
|
||||
|
||||
pub fn new() -> anyhow::Result<Self> {
|
||||
unsafe {
|
||||
let factory = CreateDXGIFactory1::<windows::Win32::Graphics::Dxgi::IDXGIFactory1>()
|
||||
.context("CreateDXGIFactory1 failed")?;
|
||||
let mut selected: Option<(
|
||||
IDXGIAdapter1,
|
||||
windows::Win32::Graphics::Dxgi::IDXGIOutput,
|
||||
)> = None;
|
||||
for adapter_index in 0..32 {
|
||||
let Ok(adapter) = factory.EnumAdapters1(adapter_index) else {
|
||||
break;
|
||||
};
|
||||
for output_index in 0..16 {
|
||||
let Ok(output) = adapter.EnumOutputs(output_index) else {
|
||||
break;
|
||||
};
|
||||
let description =
|
||||
output.GetDesc().context("IDXGIOutput::GetDesc failed")?;
|
||||
if description.AttachedToDesktop.as_bool() {
|
||||
selected = Some((adapter, output));
|
||||
break;
|
||||
}
|
||||
}
|
||||
if selected.is_some() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let (adapter, output) =
|
||||
selected.ok_or_else(|| anyhow::anyhow!("no attached DXGI output"))?;
|
||||
let mut device = None;
|
||||
let mut context = None;
|
||||
D3D11CreateDevice(
|
||||
&adapter,
|
||||
D3D_DRIVER_TYPE_UNKNOWN,
|
||||
windows::Win32::Foundation::HMODULE::default(),
|
||||
D3D11_CREATE_DEVICE_FLAG(0),
|
||||
None,
|
||||
7,
|
||||
Some(&raw mut device),
|
||||
None,
|
||||
Some(&raw mut context),
|
||||
)
|
||||
.context("D3D11CreateDevice failed")?;
|
||||
let device = device.ok_or_else(|| anyhow::anyhow!("D3D11 returned no device"))?;
|
||||
let context =
|
||||
context.ok_or_else(|| anyhow::anyhow!("D3D11 returned no context"))?;
|
||||
let output1: IDXGIOutput1 =
|
||||
output.cast().context("DXGI output has no IDXGIOutput1")?;
|
||||
let duplication = output1
|
||||
.DuplicateOutput(&device)
|
||||
.context("DuplicateOutput failed")?;
|
||||
let desc = duplication.GetDesc();
|
||||
let width = desc.ModeDesc.Width;
|
||||
let height = desc.ModeDesc.Height;
|
||||
validate_geometry(width, height)?;
|
||||
let texture_desc = D3D11_TEXTURE2D_DESC {
|
||||
Width: width,
|
||||
Height: height,
|
||||
MipLevels: 1,
|
||||
ArraySize: 1,
|
||||
Format: DXGI_FORMAT_B8G8R8A8_UNORM,
|
||||
SampleDesc: DXGI_SAMPLE_DESC {
|
||||
Count: 1,
|
||||
Quality: 0,
|
||||
},
|
||||
Usage: D3D11_USAGE_STAGING,
|
||||
BindFlags: D3D11_BIND_FLAG(0).0 as u32,
|
||||
CPUAccessFlags: D3D11_CPU_ACCESS_READ.0 as u32,
|
||||
MiscFlags: D3D11_RESOURCE_MISC_FLAG(0).0 as u32,
|
||||
};
|
||||
let mut staging = None;
|
||||
device
|
||||
.CreateTexture2D(&texture_desc, None, Some(&raw mut staging))
|
||||
.context("CreateTexture2D staging resource failed")?;
|
||||
let staging =
|
||||
staging.ok_or_else(|| anyhow::anyhow!("D3D11 returned no staging texture"))?;
|
||||
let bytes = usize::try_from(width)
|
||||
.ok()
|
||||
.and_then(|w| usize::try_from(height).ok().and_then(|h| w.checked_mul(h)))
|
||||
.and_then(|pixels| pixels.checked_mul(4))
|
||||
.ok_or_else(|| anyhow::anyhow!("desktop frame size overflow"))?;
|
||||
Ok(Self {
|
||||
context,
|
||||
duplication,
|
||||
staging,
|
||||
width,
|
||||
height,
|
||||
pixels: vec![0; bytes],
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn capture_frame(&mut self) -> anyhow::Result<EncodedFrame> {
|
||||
let captured = self.capture_bgra_frame()?;
|
||||
let encode_started = Instant::now();
|
||||
let bytes = encode_frame(
|
||||
captured.width,
|
||||
captured.height,
|
||||
&captured.pixels,
|
||||
captured.capture_latency,
|
||||
Duration::ZERO,
|
||||
)?;
|
||||
Ok(EncodedFrame {
|
||||
bytes,
|
||||
capture_latency: captured.capture_latency,
|
||||
encode_latency: encode_started.elapsed(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn capture_bgra_frame(&mut self) -> anyhow::Result<CapturedFrame> {
|
||||
let capture_started = Instant::now();
|
||||
let mut info = DXGI_OUTDUPL_FRAME_INFO::default();
|
||||
let mut resource: Option<IDXGIResource> = None;
|
||||
unsafe {
|
||||
if let Err(error) =
|
||||
self.duplication
|
||||
.AcquireNextFrame(100, &raw mut info, &raw mut resource)
|
||||
{
|
||||
if error.code() == DXGI_ERROR_WAIT_TIMEOUT.into() {
|
||||
anyhow::bail!("dda_timeout");
|
||||
}
|
||||
return Err(error).context("AcquireNextFrame failed");
|
||||
}
|
||||
let mut mapped = false;
|
||||
let result = (|| {
|
||||
let resource = resource
|
||||
.ok_or_else(|| anyhow::anyhow!("DDA returned no frame resource"))?;
|
||||
let source: ID3D11Texture2D = resource
|
||||
.cast()
|
||||
.context("DDA resource is not a D3D11 texture")?;
|
||||
self.context.CopyResource(&self.staging, &source);
|
||||
let mut mapped_resource =
|
||||
windows::Win32::Graphics::Direct3D11::D3D11_MAPPED_SUBRESOURCE::default();
|
||||
self.context
|
||||
.Map(
|
||||
&self.staging,
|
||||
0,
|
||||
D3D11_MAP_READ,
|
||||
0,
|
||||
Some(&raw mut mapped_resource),
|
||||
)
|
||||
.context("Map staging texture failed")?;
|
||||
mapped = true;
|
||||
let row_bytes = usize::try_from(self.width).unwrap() * 4;
|
||||
let pitch = usize::try_from(mapped_resource.RowPitch)
|
||||
.context("invalid DDA row pitch")?;
|
||||
anyhow::ensure!(pitch >= row_bytes, "DDA row pitch is shorter than a row");
|
||||
anyhow::ensure!(
|
||||
!mapped_resource.pData.is_null(),
|
||||
"DDA map returned a null pointer"
|
||||
);
|
||||
let source_ptr = mapped_resource.pData.cast::<u8>();
|
||||
for row in 0..usize::try_from(self.height).unwrap() {
|
||||
let source_row =
|
||||
std::slice::from_raw_parts(source_ptr.add(row * pitch), row_bytes);
|
||||
self.pixels[row * row_bytes..(row + 1) * row_bytes]
|
||||
.copy_from_slice(source_row);
|
||||
}
|
||||
Ok::<(), anyhow::Error>(())
|
||||
})();
|
||||
if mapped {
|
||||
self.context.Unmap(&self.staging, 0);
|
||||
}
|
||||
let _ = self.duplication.ReleaseFrame();
|
||||
result?;
|
||||
}
|
||||
Ok(CapturedFrame {
|
||||
width: self.width,
|
||||
height: self.height,
|
||||
pixels: self.pixels.clone(),
|
||||
capture_latency: capture_started.elapsed(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_geometry(width: u32, height: u32) -> anyhow::Result<()> {
|
||||
anyhow::ensure!(
|
||||
width > 0 && height > 0 && width <= MAX_DIMENSION && height <= MAX_DIMENSION,
|
||||
"unsupported DDA output {width}x{height}"
|
||||
);
|
||||
let bytes = usize::try_from(width)
|
||||
.ok()
|
||||
.and_then(|w| usize::try_from(height).ok().and_then(|h| w.checked_mul(h)))
|
||||
.and_then(|pixels| pixels.checked_mul(4))
|
||||
.ok_or_else(|| anyhow::anyhow!("DDA frame size overflow"))?;
|
||||
anyhow::ensure!(bytes <= MAX_RAW_FRAME_BYTES, "DDA frame exceeds 256 MiB");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn encode_frame(
|
||||
width: u32,
|
||||
height: u32,
|
||||
pixels: &[u8],
|
||||
capture: Duration,
|
||||
encode: Duration,
|
||||
) -> anyhow::Result<Vec<u8>> {
|
||||
let mut compressor = ZlibEncoder::new(Vec::new(), Compression::fast());
|
||||
compressor.write_all(pixels)?;
|
||||
let compressed = compressor.finish()?;
|
||||
let raw_len = u32::try_from(pixels.len()).context("DDA raw frame is too large")?;
|
||||
let compressed_len =
|
||||
u32::try_from(compressed.len()).context("DDA compressed frame is too large")?;
|
||||
let mut frame = Vec::with_capacity(FRAME_HEADER_LEN + compressed.len());
|
||||
frame.extend_from_slice(FRAME_MAGIC);
|
||||
frame.push(FRAME_VERSION);
|
||||
frame.push(CODEC_BGRA_ZLIB);
|
||||
frame.extend_from_slice(&0_u16.to_le_bytes());
|
||||
frame.extend_from_slice(&width.to_le_bytes());
|
||||
frame.extend_from_slice(&height.to_le_bytes());
|
||||
frame.extend_from_slice(&raw_len.to_le_bytes());
|
||||
frame.extend_from_slice(&compressed_len.to_le_bytes());
|
||||
frame.extend_from_slice(
|
||||
&u64::try_from(capture.as_micros())
|
||||
.unwrap_or(u64::MAX)
|
||||
.to_le_bytes(),
|
||||
);
|
||||
frame.extend_from_slice(
|
||||
&u64::try_from(encode.as_micros())
|
||||
.unwrap_or(u64::MAX)
|
||||
.to_le_bytes(),
|
||||
);
|
||||
frame.extend_from_slice(&compressed);
|
||||
Ok(frame)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
pub(super) use windows_capture::DdaCapture;
|
||||
|
||||
#[cfg(not(windows))]
|
||||
pub(super) struct DdaCapture;
|
||||
|
||||
#[cfg(not(windows))]
|
||||
impl DdaCapture {
|
||||
pub(super) fn dimensions(&self) -> (u32, u32) {
|
||||
(0, 0)
|
||||
}
|
||||
|
||||
pub(super) fn new() -> anyhow::Result<Self> {
|
||||
anyhow::bail!("desktop_duplication_windows_only")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,657 @@
|
||||
//! Windows H.264 encoding and CPU-side compatibility conversion.
|
||||
//!
|
||||
//! The encoder accepts an NV12 frame and returns one bounded Annex-B access
|
||||
//! unit. The current Desktop Duplication path is a CPU readback, so callers
|
||||
//! must report that upload/copy boundary separately; this module does not
|
||||
//! claim end-to-end zero-copy merely because the MFT is hardware-backed.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
const MAX_DIMENSION: u32 = 16_384;
|
||||
const MAX_ACCESS_UNIT_BYTES: usize = 16 * 1024 * 1024;
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub(crate) struct EncodedAccessUnit {
|
||||
pub(crate) data: Vec<u8>,
|
||||
pub(crate) pts_100ns: i64,
|
||||
pub(crate) duration: Duration,
|
||||
pub(crate) encode_latency: Duration,
|
||||
pub(crate) keyframe: bool,
|
||||
pub(crate) encoder: &'static str,
|
||||
pub(crate) hardware: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
|
||||
pub(crate) enum EncoderError {
|
||||
InvalidConfiguration,
|
||||
InvalidFrame,
|
||||
RuntimeUnavailable,
|
||||
EncoderUnavailable,
|
||||
InputRejected,
|
||||
OutputUnavailable,
|
||||
InvalidOutput,
|
||||
KeyFrameUnavailable,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for EncoderError {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter.write_str(match self {
|
||||
Self::InvalidConfiguration => "invalid_encoder_configuration",
|
||||
Self::InvalidFrame => "invalid_nv12_frame",
|
||||
Self::RuntimeUnavailable => "media_foundation_unavailable",
|
||||
Self::EncoderUnavailable => "h264_encoder_unavailable",
|
||||
Self::InputRejected => "h264_encoder_input_rejected",
|
||||
Self::OutputUnavailable => "h264_encoder_output_unavailable",
|
||||
Self::InvalidOutput => "invalid_h264_access_unit",
|
||||
Self::KeyFrameUnavailable => "h264_keyframe_request_unavailable",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for EncoderError {}
|
||||
|
||||
pub(crate) fn validate_configuration(
|
||||
width: u32,
|
||||
height: u32,
|
||||
fps: u32,
|
||||
) -> Result<(), EncoderError> {
|
||||
if !(2..=MAX_DIMENSION).contains(&width)
|
||||
|| !(2..=MAX_DIMENSION).contains(&height)
|
||||
|| fps == 0
|
||||
|| fps > 240
|
||||
|| width % 2 != 0
|
||||
|| height % 2 != 0
|
||||
{
|
||||
return Err(EncoderError::InvalidConfiguration);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn expected_nv12_bytes(width: u32, height: u32) -> Option<usize> {
|
||||
usize::try_from(width)
|
||||
.ok()
|
||||
.and_then(|width| {
|
||||
usize::try_from(height)
|
||||
.ok()
|
||||
.and_then(|height| width.checked_mul(height))
|
||||
})
|
||||
.and_then(|y| y.checked_add(y / 2))
|
||||
}
|
||||
|
||||
/// Converts top-down BGRA8 into an 8-bit SDR NV12 frame.
|
||||
pub(crate) fn bgra_to_nv12(width: u32, height: u32, bgra: &[u8]) -> Result<Vec<u8>, EncoderError> {
|
||||
validate_configuration(width, height, 1)?;
|
||||
let expected = expected_nv12_bytes(width, height).ok_or(EncoderError::InvalidFrame)?;
|
||||
let bgra_len = usize::try_from(width)
|
||||
.ok()
|
||||
.and_then(|width| {
|
||||
usize::try_from(height)
|
||||
.ok()
|
||||
.and_then(|height| width.checked_mul(height))
|
||||
})
|
||||
.and_then(|pixels| pixels.checked_mul(4))
|
||||
.ok_or(EncoderError::InvalidFrame)?;
|
||||
if bgra.len() != bgra_len {
|
||||
return Err(EncoderError::InvalidFrame);
|
||||
}
|
||||
let width = usize::try_from(width).map_err(|_| EncoderError::InvalidFrame)?;
|
||||
let height = usize::try_from(height).map_err(|_| EncoderError::InvalidFrame)?;
|
||||
let y_plane_len = width
|
||||
.checked_mul(height)
|
||||
.ok_or(EncoderError::InvalidFrame)?;
|
||||
let mut output = vec![0_u8; expected];
|
||||
let (y_plane, uv_plane) = output.split_at_mut(y_plane_len);
|
||||
|
||||
for y in 0..height {
|
||||
for x in 0..width {
|
||||
let source = (y * width + x) * 4;
|
||||
let b = i32::from(bgra[source]);
|
||||
let g = i32::from(bgra[source + 1]);
|
||||
let r = i32::from(bgra[source + 2]);
|
||||
y_plane[y * width + x] = clamp_u8((66 * r + 129 * g + 25 * b + 128) / 256 + 16);
|
||||
}
|
||||
}
|
||||
for y in (0..height).step_by(2) {
|
||||
for x in (0..width).step_by(2) {
|
||||
let mut u = 0_i32;
|
||||
let mut v = 0_i32;
|
||||
for dy in 0..2 {
|
||||
for dx in 0..2 {
|
||||
let source = ((y + dy) * width + x + dx) * 4;
|
||||
let b = i32::from(bgra[source]);
|
||||
let g = i32::from(bgra[source + 1]);
|
||||
let r = i32::from(bgra[source + 2]);
|
||||
u += (-38 * r - 74 * g + 112 * b + 512) / 1024 + 128;
|
||||
v += (112 * r - 94 * g - 18 * b + 512) / 1024 + 128;
|
||||
}
|
||||
}
|
||||
let destination = (y / 2) * width + x;
|
||||
uv_plane[destination] = clamp_u8(u / 4);
|
||||
uv_plane[destination + 1] = clamp_u8(v / 4);
|
||||
}
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn clamp_u8(value: i32) -> u8 {
|
||||
value.clamp(0, 255) as u8
|
||||
}
|
||||
|
||||
pub(crate) fn annex_b_contains_idr(data: &[u8]) -> bool {
|
||||
let mut index = 0;
|
||||
while index + 4 <= data.len() {
|
||||
let start = if data[index..].starts_with(&[0, 0, 0, 1]) {
|
||||
index + 4
|
||||
} else if index + 3 <= data.len() && data[index..].starts_with(&[0, 0, 1]) {
|
||||
index + 3
|
||||
} else {
|
||||
index += 1;
|
||||
continue;
|
||||
};
|
||||
if start < data.len() && data[start] & 0x1f == 5 {
|
||||
return true;
|
||||
}
|
||||
index = start;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// Media Foundation may expose H.264 samples as byte-stream Annex-B or as
|
||||
/// AVC length-prefixed NAL units. Normalize both forms before the access unit
|
||||
/// crosses the process boundary.
|
||||
pub(crate) fn normalize_h264_access_unit(data: &[u8]) -> Result<Vec<u8>, EncoderError> {
|
||||
if data.is_empty() {
|
||||
return Err(EncoderError::InvalidOutput);
|
||||
}
|
||||
if data.starts_with(&[0, 0, 0, 1]) || data.starts_with(&[0, 0, 1]) {
|
||||
return Ok(data.to_vec());
|
||||
}
|
||||
let mut offset = 0_usize;
|
||||
let mut output = Vec::with_capacity(data.len() + 16);
|
||||
while offset + 4 <= data.len() {
|
||||
let length = u32::from_be_bytes(data[offset..offset + 4].try_into().unwrap()) as usize;
|
||||
offset += 4;
|
||||
if length == 0 || length > data.len().saturating_sub(offset) {
|
||||
return Err(EncoderError::InvalidOutput);
|
||||
}
|
||||
output.extend_from_slice(&[0, 0, 0, 1]);
|
||||
output.extend_from_slice(&data[offset..offset + length]);
|
||||
offset += length;
|
||||
}
|
||||
if offset != data.len() || output.is_empty() {
|
||||
return Err(EncoderError::InvalidOutput);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
mod windows_mft {
|
||||
use super::{
|
||||
EncodedAccessUnit, EncoderError, MAX_ACCESS_UNIT_BYTES, annex_b_contains_idr,
|
||||
validate_configuration,
|
||||
};
|
||||
use core::mem::ManuallyDrop;
|
||||
use std::collections::VecDeque;
|
||||
use std::time::{Duration, Instant};
|
||||
use windows::Win32::Media::MediaFoundation::{
|
||||
CLSID_MSH264EncoderMFT, CODECAPI_AVEncVideoForceKeyFrame, ICodecAPI, IMFActivate,
|
||||
IMFMediaType, IMFTransform, MF_E_NOTACCEPTING, MF_E_TRANSFORM_NEED_MORE_INPUT,
|
||||
MF_MT_AVG_BITRATE, MF_MT_FRAME_RATE, MF_MT_FRAME_SIZE, MF_MT_INTERLACE_MODE,
|
||||
MF_MT_MAJOR_TYPE, MF_MT_SUBTYPE, MF_VERSION, MFCreateMediaType, MFCreateMemoryBuffer,
|
||||
MFCreateSample, MFMediaType_Video, MFSTARTUP_FULL, MFShutdown, MFStartup,
|
||||
MFT_CATEGORY_VIDEO_ENCODER, MFT_ENUM_FLAG_HARDWARE,
|
||||
MFT_ENUM_FLAG_SORTANDFILTER_APPROVED_ONLY, MFT_MESSAGE_NOTIFY_BEGIN_STREAMING,
|
||||
MFT_MESSAGE_NOTIFY_START_OF_STREAM, MFT_OUTPUT_DATA_BUFFER, MFT_REGISTER_TYPE_INFO,
|
||||
MFVideoFormat_H264, MFVideoFormat_NV12, MFVideoInterlace_Progressive,
|
||||
};
|
||||
use windows::Win32::System::Com::{
|
||||
CLSCTX_INPROC_SERVER, COINIT_MULTITHREADED, CoCreateInstance, CoInitializeEx, CoTaskMemFree,
|
||||
};
|
||||
use windows::Win32::System::Variant::VARIANT;
|
||||
use windows::core::Interface as _;
|
||||
|
||||
const MAX_OUTPUTS_PER_INPUT: usize = 8;
|
||||
const MAX_PENDING_ACCESS_UNITS: usize = 16;
|
||||
|
||||
struct ComGuard;
|
||||
impl Drop for ComGuard {
|
||||
fn drop(&mut self) {
|
||||
unsafe { windows::Win32::System::Com::CoUninitialize() };
|
||||
}
|
||||
}
|
||||
|
||||
struct MfGuard;
|
||||
impl Drop for MfGuard {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
let _ = MFShutdown();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct H264Encoder {
|
||||
transform: IMFTransform,
|
||||
width: u32,
|
||||
height: u32,
|
||||
fps: u32,
|
||||
next_pts_100ns: i64,
|
||||
last_explicit_pts_100ns: Option<i64>,
|
||||
encoder: &'static str,
|
||||
hardware: bool,
|
||||
pending: VecDeque<EncodedAccessUnit>,
|
||||
_com: ComGuard,
|
||||
_mf: MfGuard,
|
||||
}
|
||||
|
||||
impl H264Encoder {
|
||||
pub(crate) fn new(
|
||||
width: u32,
|
||||
height: u32,
|
||||
fps: u32,
|
||||
hardware_only: bool,
|
||||
) -> Result<Self, EncoderError> {
|
||||
validate_configuration(width, height, fps)?;
|
||||
let com_result = unsafe { CoInitializeEx(None, COINIT_MULTITHREADED) };
|
||||
if com_result.is_err() && com_result != windows::Win32::Foundation::S_FALSE {
|
||||
return Err(EncoderError::RuntimeUnavailable);
|
||||
}
|
||||
let com = ComGuard;
|
||||
unsafe { MFStartup(MF_VERSION, MFSTARTUP_FULL) }
|
||||
.map_err(|_| EncoderError::RuntimeUnavailable)?;
|
||||
let mf = MfGuard;
|
||||
let (transform, encoder, hardware) = if hardware_only {
|
||||
activate_hardware_encoder()?.ok_or(EncoderError::EncoderUnavailable)?
|
||||
} else {
|
||||
let transform: IMFTransform = unsafe {
|
||||
CoCreateInstance(&CLSID_MSH264EncoderMFT, None, CLSCTX_INPROC_SERVER)
|
||||
}
|
||||
.map_err(|_| EncoderError::EncoderUnavailable)?;
|
||||
(transform, "mfh264enc", false)
|
||||
};
|
||||
configure_transform(&transform, width, height, fps)?;
|
||||
unsafe { transform.ProcessMessage(MFT_MESSAGE_NOTIFY_BEGIN_STREAMING, 0) }
|
||||
.map_err(|_| EncoderError::EncoderUnavailable)?;
|
||||
unsafe { transform.ProcessMessage(MFT_MESSAGE_NOTIFY_START_OF_STREAM, 0) }
|
||||
.map_err(|_| EncoderError::EncoderUnavailable)?;
|
||||
Ok(Self {
|
||||
transform,
|
||||
width,
|
||||
height,
|
||||
fps,
|
||||
next_pts_100ns: 0,
|
||||
last_explicit_pts_100ns: None,
|
||||
encoder,
|
||||
hardware,
|
||||
pending: VecDeque::new(),
|
||||
_com: com,
|
||||
_mf: mf,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn encode(
|
||||
&mut self,
|
||||
nv12: &[u8],
|
||||
) -> Result<Option<EncodedAccessUnit>, EncoderError> {
|
||||
let duration = Duration::from_nanos(1_000_000_000_u64 / u64::from(self.fps));
|
||||
let timestamp = self.next_pts_100ns;
|
||||
self.next_pts_100ns = self
|
||||
.next_pts_100ns
|
||||
.saturating_add(i64::try_from(duration.as_nanos() / 100).unwrap_or(i64::MAX));
|
||||
self.encode_timestamped(nv12, timestamp, duration)
|
||||
}
|
||||
|
||||
pub(crate) fn encode_at(
|
||||
&mut self,
|
||||
nv12: &[u8],
|
||||
requested_pts_100ns: i64,
|
||||
) -> Result<Option<EncodedAccessUnit>, EncoderError> {
|
||||
if requested_pts_100ns < 0 {
|
||||
return Err(EncoderError::InvalidFrame);
|
||||
}
|
||||
let timestamp = self
|
||||
.last_explicit_pts_100ns
|
||||
.map_or(requested_pts_100ns, |last| {
|
||||
requested_pts_100ns.max(last.saturating_add(1))
|
||||
});
|
||||
self.last_explicit_pts_100ns = Some(timestamp);
|
||||
let duration = Duration::from_nanos(1_000_000_000_u64 / u64::from(self.fps));
|
||||
self.encode_timestamped(nv12, timestamp, duration)
|
||||
}
|
||||
|
||||
fn encode_timestamped(
|
||||
&mut self,
|
||||
nv12: &[u8],
|
||||
timestamp: i64,
|
||||
duration: Duration,
|
||||
) -> Result<Option<EncodedAccessUnit>, EncoderError> {
|
||||
let started = Instant::now();
|
||||
let expected = super::expected_nv12_bytes(self.width, self.height)
|
||||
.ok_or(EncoderError::InvalidFrame)?;
|
||||
if nv12.len() != expected {
|
||||
return Err(EncoderError::InvalidFrame);
|
||||
}
|
||||
let sample = create_sample(nv12, timestamp, duration)?;
|
||||
let input = unsafe { self.transform.ProcessInput(0, &sample, 0) };
|
||||
if let Err(error) = input {
|
||||
if error.code() != MF_E_NOTACCEPTING {
|
||||
return Err(EncoderError::InputRejected);
|
||||
}
|
||||
let outputs = self.drain_outputs(started)?;
|
||||
self.queue_outputs(outputs)?;
|
||||
unsafe { self.transform.ProcessInput(0, &sample, 0) }
|
||||
.map_err(|_| EncoderError::InputRejected)?;
|
||||
}
|
||||
let outputs = self.drain_outputs(started)?;
|
||||
self.queue_outputs(outputs)?;
|
||||
Ok(self.pending.pop_front())
|
||||
}
|
||||
|
||||
pub(crate) fn request_keyframe(&self) -> Result<(), EncoderError> {
|
||||
let codec_api = self
|
||||
.transform
|
||||
.cast::<ICodecAPI>()
|
||||
.map_err(|_| EncoderError::KeyFrameUnavailable)?;
|
||||
let value = VARIANT::from(true);
|
||||
let force_keyframe = CODECAPI_AVEncVideoForceKeyFrame;
|
||||
unsafe { codec_api.SetValue(&raw const force_keyframe, &raw const value) }
|
||||
.map_err(|_| EncoderError::KeyFrameUnavailable)
|
||||
}
|
||||
|
||||
fn queue_outputs(&mut self, outputs: Vec<EncodedAccessUnit>) -> Result<(), EncoderError> {
|
||||
if self.pending.len().saturating_add(outputs.len()) > MAX_PENDING_ACCESS_UNITS {
|
||||
return Err(EncoderError::OutputUnavailable);
|
||||
}
|
||||
self.pending.extend(outputs);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn drain_outputs(
|
||||
&mut self,
|
||||
started: Instant,
|
||||
) -> Result<Vec<EncodedAccessUnit>, EncoderError> {
|
||||
let mut access_units = Vec::new();
|
||||
for _ in 0..MAX_OUTPUTS_PER_INPUT {
|
||||
let mut output = MFT_OUTPUT_DATA_BUFFER::default();
|
||||
output.pSample = ManuallyDrop::new(None);
|
||||
output.pEvents = ManuallyDrop::new(None);
|
||||
let mut status = 0_u32;
|
||||
let result = unsafe {
|
||||
self.transform.ProcessOutput(
|
||||
0,
|
||||
std::slice::from_mut(&mut output),
|
||||
&raw mut status,
|
||||
)
|
||||
};
|
||||
let sample = unsafe { ManuallyDrop::take(&mut output.pSample) };
|
||||
let _events = unsafe { ManuallyDrop::take(&mut output.pEvents) };
|
||||
match result {
|
||||
Ok(()) => {
|
||||
let sample = sample.ok_or(EncoderError::OutputUnavailable)?;
|
||||
access_units.push(self.access_unit_from_sample(sample, started.elapsed())?);
|
||||
}
|
||||
Err(error) if error.code() == MF_E_TRANSFORM_NEED_MORE_INPUT => break,
|
||||
Err(_) => return Err(EncoderError::OutputUnavailable),
|
||||
}
|
||||
}
|
||||
Ok(access_units)
|
||||
}
|
||||
|
||||
fn access_unit_from_sample(
|
||||
&self,
|
||||
sample: windows::Win32::Media::MediaFoundation::IMFSample,
|
||||
encode_latency: Duration,
|
||||
) -> Result<EncodedAccessUnit, EncoderError> {
|
||||
let buffer_count =
|
||||
unsafe { sample.GetBufferCount() }.map_err(|_| EncoderError::InvalidOutput)?;
|
||||
if buffer_count != 1 {
|
||||
return Err(EncoderError::InvalidOutput);
|
||||
}
|
||||
let buffer =
|
||||
unsafe { sample.GetBufferByIndex(0) }.map_err(|_| EncoderError::InvalidOutput)?;
|
||||
let current_length =
|
||||
unsafe { buffer.GetCurrentLength() }.map_err(|_| EncoderError::InvalidOutput)?;
|
||||
let mut data = vec![0_u8; MAX_ACCESS_UNIT_BYTES.min(current_length as usize)];
|
||||
let mut length = 0_u32;
|
||||
let mut locked = core::ptr::null_mut();
|
||||
unsafe { buffer.Lock(&raw mut locked, None, Some(&raw mut length)) }
|
||||
.map_err(|_| EncoderError::InvalidOutput)?;
|
||||
let result = if locked.is_null()
|
||||
|| usize::try_from(length).unwrap_or(usize::MAX) > MAX_ACCESS_UNIT_BYTES
|
||||
{
|
||||
Err(EncoderError::InvalidOutput)
|
||||
} else {
|
||||
data.resize(length as usize, 0);
|
||||
unsafe {
|
||||
core::ptr::copy_nonoverlapping(
|
||||
locked.cast::<u8>(),
|
||||
data.as_mut_ptr(),
|
||||
data.len(),
|
||||
)
|
||||
};
|
||||
Ok(())
|
||||
};
|
||||
let _ = unsafe { buffer.Unlock() };
|
||||
result?;
|
||||
if data.is_empty() {
|
||||
return Err(EncoderError::InvalidOutput);
|
||||
}
|
||||
let data = super::normalize_h264_access_unit(&data)?;
|
||||
let keyframe = annex_b_contains_idr(&data);
|
||||
let pts_100ns = unsafe { sample.GetSampleTime() }.unwrap_or(0);
|
||||
Ok(EncodedAccessUnit {
|
||||
data,
|
||||
pts_100ns,
|
||||
duration: Duration::from_nanos(1_000_000_000 / u64::from(self.fps)),
|
||||
encode_latency,
|
||||
keyframe,
|
||||
encoder: self.encoder,
|
||||
hardware: self.hardware,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn activate_hardware_encoder()
|
||||
-> Result<Option<(IMFTransform, &'static str, bool)>, EncoderError> {
|
||||
let input = MFT_REGISTER_TYPE_INFO {
|
||||
guidMajorType: MFMediaType_Video,
|
||||
guidSubtype: MFVideoFormat_NV12,
|
||||
};
|
||||
let output = MFT_REGISTER_TYPE_INFO {
|
||||
guidMajorType: MFMediaType_Video,
|
||||
guidSubtype: MFVideoFormat_H264,
|
||||
};
|
||||
let mut activations: *mut Option<IMFActivate> = core::ptr::null_mut();
|
||||
let mut count = 0_u32;
|
||||
let result = unsafe {
|
||||
windows::Win32::Media::MediaFoundation::MFTEnumEx(
|
||||
MFT_CATEGORY_VIDEO_ENCODER,
|
||||
MFT_ENUM_FLAG_HARDWARE | MFT_ENUM_FLAG_SORTANDFILTER_APPROVED_ONLY,
|
||||
Some(&raw const input),
|
||||
Some(&raw const output),
|
||||
&raw mut activations,
|
||||
&raw mut count,
|
||||
)
|
||||
};
|
||||
if result.is_err() || count == 0 {
|
||||
free_activations(activations, count);
|
||||
return Ok(None);
|
||||
}
|
||||
let transform = unsafe {
|
||||
(*activations)
|
||||
.as_ref()
|
||||
.ok_or(EncoderError::EncoderUnavailable)?
|
||||
.ActivateObject::<IMFTransform>()
|
||||
};
|
||||
free_activations(activations, count);
|
||||
Ok(Some((
|
||||
transform.map_err(|_| EncoderError::EncoderUnavailable)?,
|
||||
"hardware-mft",
|
||||
true,
|
||||
)))
|
||||
}
|
||||
|
||||
fn free_activations(pointer: *mut Option<IMFActivate>, count: u32) {
|
||||
if pointer.is_null() {
|
||||
return;
|
||||
}
|
||||
unsafe {
|
||||
for item in std::slice::from_raw_parts_mut(pointer, count as usize) {
|
||||
*item = None;
|
||||
}
|
||||
CoTaskMemFree(Some(pointer.cast()));
|
||||
}
|
||||
}
|
||||
|
||||
fn configure_transform(
|
||||
transform: &IMFTransform,
|
||||
width: u32,
|
||||
height: u32,
|
||||
fps: u32,
|
||||
) -> Result<(), EncoderError> {
|
||||
let input = media_type(MFVideoFormat_NV12, width, height, fps, None)?;
|
||||
let output = media_type(
|
||||
MFVideoFormat_H264,
|
||||
width,
|
||||
height,
|
||||
fps,
|
||||
Some(bitrate(width, height, fps)),
|
||||
)?;
|
||||
unsafe { transform.SetInputType(0, &input, 0) }
|
||||
.map_err(|_| EncoderError::EncoderUnavailable)?;
|
||||
unsafe { transform.SetOutputType(0, &output, 0) }
|
||||
.map_err(|_| EncoderError::EncoderUnavailable)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn media_type(
|
||||
subtype: windows::core::GUID,
|
||||
width: u32,
|
||||
height: u32,
|
||||
fps: u32,
|
||||
bitrate_value: Option<u32>,
|
||||
) -> Result<IMFMediaType, EncoderError> {
|
||||
let media_type =
|
||||
unsafe { MFCreateMediaType() }.map_err(|_| EncoderError::RuntimeUnavailable)?;
|
||||
unsafe { media_type.SetGUID(&MF_MT_MAJOR_TYPE, &MFMediaType_Video) }
|
||||
.map_err(|_| EncoderError::EncoderUnavailable)?;
|
||||
unsafe { media_type.SetGUID(&MF_MT_SUBTYPE, &subtype) }
|
||||
.map_err(|_| EncoderError::EncoderUnavailable)?;
|
||||
unsafe {
|
||||
media_type.SetUINT64(
|
||||
&MF_MT_FRAME_SIZE,
|
||||
(u64::from(width) << 32) | u64::from(height),
|
||||
)
|
||||
}
|
||||
.map_err(|_| EncoderError::EncoderUnavailable)?;
|
||||
unsafe { media_type.SetUINT64(&MF_MT_FRAME_RATE, (u64::from(fps) << 32) | 1) }
|
||||
.map_err(|_| EncoderError::EncoderUnavailable)?;
|
||||
unsafe {
|
||||
media_type.SetUINT32(&MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive.0 as u32)
|
||||
}
|
||||
.map_err(|_| EncoderError::EncoderUnavailable)?;
|
||||
if let Some(value) = bitrate_value {
|
||||
unsafe { media_type.SetUINT32(&MF_MT_AVG_BITRATE, value) }
|
||||
.map_err(|_| EncoderError::EncoderUnavailable)?;
|
||||
}
|
||||
Ok(media_type)
|
||||
}
|
||||
|
||||
fn bitrate(width: u32, height: u32, fps: u32) -> u32 {
|
||||
let value = u64::from(width)
|
||||
.saturating_mul(u64::from(height))
|
||||
.saturating_mul(u64::from(fps))
|
||||
.saturating_mul(2)
|
||||
.saturating_div(100);
|
||||
u32::try_from(value.clamp(2_000_000, 80_000_000)).unwrap_or(80_000_000)
|
||||
}
|
||||
|
||||
fn create_sample(
|
||||
data: &[u8],
|
||||
pts: i64,
|
||||
duration: Duration,
|
||||
) -> Result<windows::Win32::Media::MediaFoundation::IMFSample, EncoderError> {
|
||||
let buffer = unsafe {
|
||||
MFCreateMemoryBuffer(u32::try_from(data.len()).map_err(|_| EncoderError::InvalidFrame)?)
|
||||
}
|
||||
.map_err(|_| EncoderError::RuntimeUnavailable)?;
|
||||
let mut destination = core::ptr::null_mut();
|
||||
unsafe { buffer.Lock(&raw mut destination, None, None) }
|
||||
.map_err(|_| EncoderError::InvalidFrame)?;
|
||||
if destination.is_null() {
|
||||
let _ = unsafe { buffer.Unlock() };
|
||||
return Err(EncoderError::InvalidFrame);
|
||||
}
|
||||
unsafe {
|
||||
core::ptr::copy_nonoverlapping(data.as_ptr(), destination.cast::<u8>(), data.len());
|
||||
}
|
||||
unsafe { buffer.Unlock() }.map_err(|_| EncoderError::InvalidFrame)?;
|
||||
unsafe {
|
||||
buffer.SetCurrentLength(
|
||||
u32::try_from(data.len()).map_err(|_| EncoderError::InvalidFrame)?,
|
||||
)
|
||||
}
|
||||
.map_err(|_| EncoderError::InvalidFrame)?;
|
||||
let sample = unsafe { MFCreateSample() }.map_err(|_| EncoderError::RuntimeUnavailable)?;
|
||||
unsafe { sample.AddBuffer(&buffer) }.map_err(|_| EncoderError::InvalidFrame)?;
|
||||
unsafe { sample.SetSampleTime(pts) }.map_err(|_| EncoderError::InvalidFrame)?;
|
||||
unsafe {
|
||||
sample.SetSampleDuration(i64::try_from(duration.as_nanos() / 100).unwrap_or(i64::MAX))
|
||||
}
|
||||
.map_err(|_| EncoderError::InvalidFrame)?;
|
||||
Ok(sample)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
#[allow(unused_imports)]
|
||||
pub(crate) use windows_mft::H264Encoder;
|
||||
|
||||
#[cfg(not(windows))]
|
||||
pub(crate) struct H264Encoder;
|
||||
|
||||
#[cfg(not(windows))]
|
||||
impl H264Encoder {
|
||||
pub(crate) fn new(_: u32, _: u32, _: u32, _: bool) -> Result<Self, EncoderError> {
|
||||
Err(EncoderError::RuntimeUnavailable)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn bgra_to_nv12_has_expected_layout() {
|
||||
let bgra = vec![
|
||||
0_u8, 0, 255, 255, 0, 255, 0, 255, 255, 0, 0, 255, 255, 255, 255, 255,
|
||||
];
|
||||
let nv12 = bgra_to_nv12(2, 2, &bgra).unwrap();
|
||||
assert_eq!(nv12.len(), 6);
|
||||
assert!(nv12[..4].iter().all(|value| *value > 0));
|
||||
assert_eq!(nv12[4..].len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn odd_dimensions_and_wrong_lengths_fail_closed() {
|
||||
assert_eq!(
|
||||
validate_configuration(3, 2, 30),
|
||||
Err(EncoderError::InvalidConfiguration)
|
||||
);
|
||||
assert_eq!(bgra_to_nv12(2, 2, &[0; 3]), Err(EncoderError::InvalidFrame));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn idr_detection_accepts_annex_b_five() {
|
||||
assert!(annex_b_contains_idr(&[0, 0, 0, 1, 5, 1]));
|
||||
assert!(!annex_b_contains_idr(&[0, 0, 0, 1, 1, 1]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn avc_length_prefixed_output_is_normalized_to_annex_b() {
|
||||
let avc = [0_u8, 0, 0, 2, 5, 1, 0, 0, 0, 1, 1];
|
||||
assert_eq!(
|
||||
normalize_h264_access_unit(&avc).unwrap(),
|
||||
[0, 0, 0, 1, 5, 1, 0, 0, 0, 1, 1]
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,186 @@
|
||||
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
use hmac::{Hmac, KeyInit as _, Mac as _};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::Sha256;
|
||||
use std::io;
|
||||
use tokio::io::{
|
||||
AsyncBufReadExt as _, AsyncRead, AsyncReadExt as _, AsyncWrite, AsyncWriteExt as _, BufReader,
|
||||
};
|
||||
use zeroize::{Zeroize as _, Zeroizing};
|
||||
|
||||
pub const AGENT_PIPE_TOKEN_ENV: &str = "REMOTEDESK_AGENT_PIPE_TOKEN";
|
||||
const AUTH_PROTOCOL_VERSION: u8 = 1;
|
||||
const AUTH_KEY_BYTES: usize = 32;
|
||||
const AUTH_NONCE_BYTES: usize = 32;
|
||||
const AUTH_RESPONSE_MAX_BYTES: usize = 2 * 1024;
|
||||
const AUTH_DOMAIN: &[u8] = b"RemoteDesk Windows Agent pipe auth v1\0";
|
||||
|
||||
type HmacSha256 = Hmac<Sha256>;
|
||||
|
||||
pub struct PipeAuthenticator {
|
||||
key: Zeroizing<Vec<u8>>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct Challenge<'a> {
|
||||
kind: &'a str,
|
||||
protocol_version: u8,
|
||||
nonce: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct Authentication {
|
||||
kind: String,
|
||||
protocol_version: u8,
|
||||
mac: String,
|
||||
}
|
||||
|
||||
impl PipeAuthenticator {
|
||||
pub fn from_environment() -> io::Result<Self> {
|
||||
let encoded = std::env::var(AGENT_PIPE_TOKEN_ENV).map_err(|_| {
|
||||
io::Error::new(
|
||||
io::ErrorKind::PermissionDenied,
|
||||
format!("{AGENT_PIPE_TOKEN_ENV} is required for authenticated Named Pipe mode"),
|
||||
)
|
||||
})?;
|
||||
// The bootstrap is needed only while constructing the in-memory key.
|
||||
unsafe { std::env::remove_var(AGENT_PIPE_TOKEN_ENV) };
|
||||
let mut encoded = Zeroizing::new(encoded);
|
||||
let result = Self::from_encoded(&encoded);
|
||||
encoded.zeroize();
|
||||
result
|
||||
}
|
||||
|
||||
fn from_encoded(encoded: &str) -> io::Result<Self> {
|
||||
let key = Zeroizing::new(URL_SAFE_NO_PAD.decode(encoded).map_err(|_| {
|
||||
io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
"Agent Pipe token is not base64url",
|
||||
)
|
||||
})?);
|
||||
if key.len() != AUTH_KEY_BYTES {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
"Agent Pipe token must decode to exactly 32 bytes",
|
||||
));
|
||||
}
|
||||
Ok(Self { key })
|
||||
}
|
||||
|
||||
pub async fn authenticate<S>(&self, stream: &mut S) -> io::Result<()>
|
||||
where
|
||||
S: AsyncRead + AsyncWrite + Unpin,
|
||||
{
|
||||
let mut nonce = [0_u8; AUTH_NONCE_BYTES];
|
||||
getrandom::fill(&mut nonce).map_err(io::Error::other)?;
|
||||
let challenge = serde_json::to_vec(&Challenge {
|
||||
kind: "agent_auth_challenge",
|
||||
protocol_version: AUTH_PROTOCOL_VERSION,
|
||||
nonce: URL_SAFE_NO_PAD.encode(nonce),
|
||||
})
|
||||
.map_err(io::Error::other)?;
|
||||
stream.write_all(&challenge).await?;
|
||||
stream.write_all(b"\n").await?;
|
||||
stream.flush().await?;
|
||||
|
||||
let mut line = String::new();
|
||||
let mut reader = BufReader::new(stream).take((AUTH_RESPONSE_MAX_BYTES + 1) as u64);
|
||||
let bytes = reader.read_line(&mut line).await?;
|
||||
if bytes == 0
|
||||
|| bytes > AUTH_RESPONSE_MAX_BYTES
|
||||
|| !line.ends_with('\n')
|
||||
|| !line.is_ascii()
|
||||
{
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
"Agent Pipe authentication response is invalid",
|
||||
));
|
||||
}
|
||||
let authentication: Authentication = serde_json::from_str(&line)
|
||||
.map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "invalid Agent Pipe auth"))?;
|
||||
if authentication.kind != "agent_auth"
|
||||
|| authentication.protocol_version != AUTH_PROTOCOL_VERSION
|
||||
{
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::PermissionDenied,
|
||||
"Agent Pipe authentication protocol mismatch",
|
||||
));
|
||||
}
|
||||
let supplied = URL_SAFE_NO_PAD.decode(authentication.mac).map_err(|_| {
|
||||
io::Error::new(io::ErrorKind::InvalidData, "invalid Agent Pipe auth MAC")
|
||||
})?;
|
||||
let mut mac = HmacSha256::new_from_slice(&self.key)
|
||||
.map_err(|_| io::Error::other("invalid Agent Pipe HMAC key"))?;
|
||||
mac.update(AUTH_DOMAIN);
|
||||
mac.update(&nonce);
|
||||
mac.verify_slice(&supplied).map_err(|_| {
|
||||
io::Error::new(
|
||||
io::ErrorKind::PermissionDenied,
|
||||
"Agent Pipe authentication was rejected",
|
||||
)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn authentication_mac(key: &[u8], nonce: &[u8]) -> io::Result<Vec<u8>> {
|
||||
let mut mac = HmacSha256::new_from_slice(key)
|
||||
.map_err(|_| io::Error::other("invalid Agent Pipe HMAC key"))?;
|
||||
mac.update(AUTH_DOMAIN);
|
||||
mac.update(nonce);
|
||||
Ok(mac.finalize().into_bytes().to_vec())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::Value;
|
||||
|
||||
#[test]
|
||||
fn token_requires_exactly_32_base64url_bytes() {
|
||||
assert!(PipeAuthenticator::from_encoded("not-base64!").is_err());
|
||||
assert!(PipeAuthenticator::from_encoded(&URL_SAFE_NO_PAD.encode([7_u8; 31])).is_err());
|
||||
assert!(PipeAuthenticator::from_encoded(&URL_SAFE_NO_PAD.encode([7_u8; 32])).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn authentication_mac_matches_go_wire_vector() {
|
||||
let key = (0_u8..32).collect::<Vec<_>>();
|
||||
let nonce = (32_u8..64).collect::<Vec<_>>();
|
||||
let mac = authentication_mac(&key, &nonce).unwrap();
|
||||
assert_eq!(
|
||||
URL_SAFE_NO_PAD.encode(mac),
|
||||
"KinBDOBvVhz0eRv37iahq8BO0z8QdiI0Qnx7HZmcxIk"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn challenge_accepts_the_matching_response() {
|
||||
let key = [9_u8; 32];
|
||||
let authenticator = PipeAuthenticator::from_encoded(&URL_SAFE_NO_PAD.encode(key)).unwrap();
|
||||
let (mut server, client) = tokio::io::duplex(4 * 1024);
|
||||
let server_task =
|
||||
tokio::spawn(async move { authenticator.authenticate(&mut server).await });
|
||||
let (reader, mut writer) = tokio::io::split(client);
|
||||
let mut reader = BufReader::new(reader);
|
||||
let mut line = String::new();
|
||||
reader.read_line(&mut line).await.unwrap();
|
||||
let challenge: Value = serde_json::from_str(&line).unwrap();
|
||||
let nonce = URL_SAFE_NO_PAD
|
||||
.decode(challenge["nonce"].as_str().unwrap())
|
||||
.unwrap();
|
||||
let mac = authentication_mac(&key, &nonce).unwrap();
|
||||
let response = serde_json::json!({
|
||||
"kind": "agent_auth",
|
||||
"protocol_version": 1,
|
||||
"mac": URL_SAFE_NO_PAD.encode(mac)
|
||||
});
|
||||
writer
|
||||
.write_all(response.to_string().as_bytes())
|
||||
.await
|
||||
.unwrap();
|
||||
writer.write_all(b"\n").await.unwrap();
|
||||
server_task.await.unwrap().unwrap();
|
||||
}
|
||||
}
|
||||
+1041
-59
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,439 @@
|
||||
//! Bounded SPSC ring for encoded media IPC.
|
||||
//!
|
||||
//! The ring carries compressed access units only. It deliberately does not
|
||||
//! expose a GPU texture or a raw BGRA/NV12 surface to the Go service.
|
||||
|
||||
use memmap2::{MmapMut, MmapOptions};
|
||||
use serde::Serialize;
|
||||
use std::fs::{File, OpenOptions};
|
||||
use std::io;
|
||||
use std::mem::{align_of, size_of};
|
||||
use std::path::Path;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
const MAGIC: [u8; 8] = *b"RDMRING1";
|
||||
const VERSION: u32 = 1;
|
||||
const HEADER_BYTES: usize = size_of::<RingHeader>();
|
||||
const SLOT_HEADER_BYTES: usize = size_of::<SlotHeader>();
|
||||
const MAX_SLOTS: u32 = 256;
|
||||
const MAX_SLOT_BYTES: u32 = 16 * 1024 * 1024;
|
||||
|
||||
#[repr(C, align(8))]
|
||||
struct RingHeader {
|
||||
magic: [u8; 8],
|
||||
version: u32,
|
||||
slot_count: u32,
|
||||
slot_bytes: u32,
|
||||
reserved: u32,
|
||||
write_seq: AtomicU64,
|
||||
read_seq: AtomicU64,
|
||||
}
|
||||
|
||||
#[repr(C)]
|
||||
#[derive(Clone, Copy, Default)]
|
||||
struct SlotHeader {
|
||||
sequence: u64,
|
||||
pts: u64,
|
||||
kind: u32,
|
||||
length: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct MediaRecord {
|
||||
pub sequence: u64,
|
||||
pub pts: u64,
|
||||
pub kind: u32,
|
||||
pub payload: Vec<u8>,
|
||||
}
|
||||
|
||||
/// Control-plane description for a ring owned by one authenticated session.
|
||||
/// The ring contains encoded access units only; it is never a raw-surface
|
||||
/// transport. The creator must bind this descriptor to its session ACL before
|
||||
/// sending it to another process.
|
||||
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
|
||||
pub struct MediaRingDescriptor {
|
||||
pub schema_version: u16,
|
||||
pub session_id: String,
|
||||
pub generation: u64,
|
||||
pub owner_pid: u32,
|
||||
pub ring_path: String,
|
||||
pub slot_count: u32,
|
||||
pub slot_bytes: u32,
|
||||
pub video_transport: &'static str,
|
||||
pub audio_transport: &'static str,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub audio_ring: Option<AudioRingDescriptor>,
|
||||
}
|
||||
|
||||
/// Control-plane description for a separate Opus packet ring.
|
||||
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
|
||||
pub struct AudioRingDescriptor {
|
||||
pub schema_version: u16,
|
||||
pub session_id: String,
|
||||
pub generation: u64,
|
||||
pub owner_pid: u32,
|
||||
pub ring_path: String,
|
||||
pub slot_count: u32,
|
||||
pub slot_bytes: u32,
|
||||
pub audio_transport: &'static str,
|
||||
}
|
||||
|
||||
impl AudioRingDescriptor {
|
||||
pub fn validate(&self) -> io::Result<()> {
|
||||
if self.schema_version != VERSION as u16
|
||||
|| self.generation == 0
|
||||
|| self.owner_pid == 0
|
||||
|| self.session_id.is_empty()
|
||||
|| self.session_id.len() > 128
|
||||
|| !self
|
||||
.session_id
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_'))
|
||||
{
|
||||
return invalid_ring("audio ring descriptor identity is invalid");
|
||||
}
|
||||
if !Path::new(&self.ring_path).is_absolute() {
|
||||
return invalid_ring("audio ring path must be absolute");
|
||||
}
|
||||
if self.audio_transport != "shared_memory_opus_packets" {
|
||||
return invalid_ring("audio ring transport contract is invalid");
|
||||
}
|
||||
validate_geometry(self.slot_count, self.slot_bytes)
|
||||
}
|
||||
}
|
||||
|
||||
impl MediaRingDescriptor {
|
||||
pub fn validate(&self) -> io::Result<()> {
|
||||
if self.schema_version != VERSION as u16
|
||||
|| self.generation == 0
|
||||
|| self.owner_pid == 0
|
||||
|| self.session_id.is_empty()
|
||||
|| self.session_id.len() > 128
|
||||
|| !self
|
||||
.session_id
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_'))
|
||||
{
|
||||
return invalid_ring("media ring descriptor identity is invalid");
|
||||
}
|
||||
if !Path::new(&self.ring_path).is_absolute() {
|
||||
return invalid_ring("media ring path must be absolute");
|
||||
}
|
||||
if self.video_transport != "shared_memory_encoded_access_units"
|
||||
|| self.audio_transport != "separate_transport"
|
||||
{
|
||||
return invalid_ring("media ring transport contract is invalid");
|
||||
}
|
||||
validate_geometry(self.slot_count, self.slot_bytes)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct MediaRing {
|
||||
file: File,
|
||||
map: MmapMut,
|
||||
slot_count: u32,
|
||||
slot_bytes: u32,
|
||||
}
|
||||
|
||||
impl MediaRing {
|
||||
pub fn descriptor(
|
||||
&self,
|
||||
session_id: impl Into<String>,
|
||||
generation: u64,
|
||||
owner_pid: u32,
|
||||
ring_path: impl AsRef<Path>,
|
||||
) -> io::Result<MediaRingDescriptor> {
|
||||
let descriptor = MediaRingDescriptor {
|
||||
schema_version: VERSION as u16,
|
||||
session_id: session_id.into(),
|
||||
generation,
|
||||
owner_pid,
|
||||
ring_path: ring_path.as_ref().to_string_lossy().into_owned(),
|
||||
slot_count: self.slot_count,
|
||||
slot_bytes: self.slot_bytes,
|
||||
video_transport: "shared_memory_encoded_access_units",
|
||||
audio_transport: "separate_transport",
|
||||
audio_ring: None,
|
||||
};
|
||||
descriptor.validate()?;
|
||||
Ok(descriptor)
|
||||
}
|
||||
|
||||
pub fn audio_descriptor(
|
||||
&self,
|
||||
session_id: impl Into<String>,
|
||||
generation: u64,
|
||||
owner_pid: u32,
|
||||
ring_path: impl AsRef<Path>,
|
||||
) -> io::Result<AudioRingDescriptor> {
|
||||
let descriptor = AudioRingDescriptor {
|
||||
schema_version: VERSION as u16,
|
||||
session_id: session_id.into(),
|
||||
generation,
|
||||
owner_pid,
|
||||
ring_path: ring_path.as_ref().to_string_lossy().into_owned(),
|
||||
slot_count: self.slot_count,
|
||||
slot_bytes: self.slot_bytes,
|
||||
audio_transport: "shared_memory_opus_packets",
|
||||
};
|
||||
descriptor.validate()?;
|
||||
Ok(descriptor)
|
||||
}
|
||||
|
||||
pub fn create(path: impl AsRef<Path>, slot_count: u32, slot_bytes: u32) -> io::Result<Self> {
|
||||
validate_geometry(slot_count, slot_bytes)?;
|
||||
let path = path.as_ref();
|
||||
#[cfg(windows)]
|
||||
let file = crate::windows_security::create_current_user_file(path)?;
|
||||
#[cfg(not(windows))]
|
||||
let file = OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create_new(true)
|
||||
.open(path)?;
|
||||
let length = ring_bytes(slot_count, slot_bytes)?;
|
||||
file.set_len(length as u64)?;
|
||||
let mut map = unsafe { MmapOptions::new().len(length).map_mut(&file)? };
|
||||
unsafe {
|
||||
std::ptr::write(
|
||||
map.as_mut_ptr().cast::<RingHeader>(),
|
||||
RingHeader {
|
||||
magic: MAGIC,
|
||||
version: VERSION,
|
||||
slot_count,
|
||||
slot_bytes,
|
||||
reserved: 0,
|
||||
write_seq: AtomicU64::new(0),
|
||||
read_seq: AtomicU64::new(0),
|
||||
},
|
||||
);
|
||||
}
|
||||
Ok(Self {
|
||||
file,
|
||||
map,
|
||||
slot_count,
|
||||
slot_bytes,
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn open(path: impl AsRef<Path>) -> io::Result<Self> {
|
||||
let file = OpenOptions::new().read(true).write(true).open(path)?;
|
||||
let length = usize::try_from(file.metadata()?.len())
|
||||
.map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "ring file is too large"))?;
|
||||
if length < HEADER_BYTES {
|
||||
return invalid_ring("ring file is shorter than its header");
|
||||
}
|
||||
let map = unsafe { MmapOptions::new().len(length).map_mut(&file)? };
|
||||
let header = unsafe { &*map.as_ptr().cast::<RingHeader>() };
|
||||
if header.magic != MAGIC || header.version != VERSION {
|
||||
return invalid_ring("ring header magic or version is invalid");
|
||||
}
|
||||
validate_geometry(header.slot_count, header.slot_bytes)?;
|
||||
if ring_bytes(header.slot_count, header.slot_bytes)? != length {
|
||||
return invalid_ring("ring file size does not match its header");
|
||||
}
|
||||
Ok(Self {
|
||||
file,
|
||||
map,
|
||||
slot_count: header.slot_count,
|
||||
slot_bytes: header.slot_bytes,
|
||||
})
|
||||
}
|
||||
|
||||
/// Pushes one encoded access unit. A full ring returns false so the
|
||||
/// caller can drop this video packet without blocking input or audio.
|
||||
pub fn push(&mut self, record: &MediaRecord) -> io::Result<bool> {
|
||||
if record.payload.is_empty() || record.payload.len() > self.slot_bytes as usize {
|
||||
return invalid_ring("media record exceeds ring slot capacity");
|
||||
}
|
||||
let write = self.header().write_seq.load(Ordering::Relaxed);
|
||||
let read = self.header().read_seq.load(Ordering::Acquire);
|
||||
if write.wrapping_sub(read) >= u64::from(self.slot_count) {
|
||||
return Ok(false);
|
||||
}
|
||||
let slot = self.slot_mut(write);
|
||||
slot.payload[..record.payload.len()].copy_from_slice(&record.payload);
|
||||
*slot.header = SlotHeader {
|
||||
sequence: record.sequence,
|
||||
pts: record.pts,
|
||||
kind: record.kind,
|
||||
length: record.payload.len() as u32,
|
||||
};
|
||||
self.header()
|
||||
.write_seq
|
||||
.store(write.wrapping_add(1), Ordering::Release);
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn pop(&mut self) -> io::Result<Option<MediaRecord>> {
|
||||
let read = self.header().read_seq.load(Ordering::Relaxed);
|
||||
let write = self.header().write_seq.load(Ordering::Acquire);
|
||||
if read == write {
|
||||
return Ok(None);
|
||||
}
|
||||
let slot_bytes = self.slot_bytes as usize;
|
||||
let slot = self.slot_mut(read);
|
||||
let length = usize::try_from(slot.header.length).map_err(|_| {
|
||||
io::Error::new(io::ErrorKind::InvalidData, "ring record length overflow")
|
||||
})?;
|
||||
if length == 0 || length > slot_bytes {
|
||||
return invalid_ring("ring record length is invalid");
|
||||
}
|
||||
let record = MediaRecord {
|
||||
sequence: slot.header.sequence,
|
||||
pts: slot.header.pts,
|
||||
kind: slot.header.kind,
|
||||
payload: slot.payload[..length].to_vec(),
|
||||
};
|
||||
self.header()
|
||||
.read_seq
|
||||
.store(read.wrapping_add(1), Ordering::Release);
|
||||
Ok(Some(record))
|
||||
}
|
||||
|
||||
fn header(&self) -> &RingHeader {
|
||||
debug_assert_eq!(self.map.as_ptr().align_offset(align_of::<RingHeader>()), 0);
|
||||
unsafe { &*self.map.as_ptr().cast::<RingHeader>() }
|
||||
}
|
||||
|
||||
fn slot_mut(&mut self, sequence: u64) -> Slot<'_> {
|
||||
let offset =
|
||||
HEADER_BYTES + (sequence as usize % self.slot_count as usize) * self.slot_stride();
|
||||
let pointer = unsafe { self.map.as_mut_ptr().add(offset) };
|
||||
Slot {
|
||||
header: unsafe { &mut *pointer.cast::<SlotHeader>() },
|
||||
payload: unsafe {
|
||||
std::slice::from_raw_parts_mut(
|
||||
pointer.add(SLOT_HEADER_BYTES),
|
||||
self.slot_bytes as usize,
|
||||
)
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn slot_stride(&self) -> usize {
|
||||
SLOT_HEADER_BYTES + self.slot_bytes as usize
|
||||
}
|
||||
}
|
||||
|
||||
struct Slot<'a> {
|
||||
header: &'a mut SlotHeader,
|
||||
payload: &'a mut [u8],
|
||||
}
|
||||
|
||||
fn validate_geometry(slot_count: u32, slot_bytes: u32) -> io::Result<()> {
|
||||
if !(2..=MAX_SLOTS).contains(&slot_count) || !slot_count.is_power_of_two() {
|
||||
return invalid_ring("slot count must be a power of two between 2 and 256");
|
||||
}
|
||||
if slot_bytes == 0 || slot_bytes > MAX_SLOT_BYTES {
|
||||
return invalid_ring("slot capacity is invalid");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn ring_bytes(slot_count: u32, slot_bytes: u32) -> io::Result<usize> {
|
||||
let slots = usize::try_from(slot_count)
|
||||
.map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "slot count overflow"))?;
|
||||
let capacity = usize::try_from(slot_bytes)
|
||||
.map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "slot capacity overflow"))?;
|
||||
HEADER_BYTES
|
||||
.checked_add(
|
||||
slots
|
||||
.checked_mul(SLOT_HEADER_BYTES.checked_add(capacity).ok_or_else(|| {
|
||||
io::Error::new(io::ErrorKind::InvalidInput, "ring size overflow")
|
||||
})?)
|
||||
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "ring size overflow"))?,
|
||||
)
|
||||
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "ring size overflow"))
|
||||
}
|
||||
|
||||
fn invalid_ring<T>(message: &str) -> io::Result<T> {
|
||||
Err(io::Error::new(io::ErrorKind::InvalidData, message))
|
||||
}
|
||||
|
||||
impl Drop for MediaRing {
|
||||
fn drop(&mut self) {
|
||||
let _ = self.map.flush();
|
||||
let _ = self.file.sync_data();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
fn temp_path() -> std::path::PathBuf {
|
||||
let nonce = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos();
|
||||
std::env::temp_dir().join(format!("remotedesk-media-ring-{nonce}.bin"))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn producer_and_consumer_share_encoded_records() {
|
||||
let path = temp_path();
|
||||
let mut producer = MediaRing::create(&path, 2, 64).unwrap();
|
||||
let mut consumer = MediaRing::open(&path).unwrap();
|
||||
let record = MediaRecord {
|
||||
sequence: 7,
|
||||
pts: 11,
|
||||
kind: 1,
|
||||
payload: b"encoded-access-unit".to_vec(),
|
||||
};
|
||||
assert!(producer.push(&record).unwrap());
|
||||
assert_eq!(consumer.pop().unwrap(), Some(record));
|
||||
assert_eq!(consumer.pop().unwrap(), None);
|
||||
drop(consumer);
|
||||
drop(producer);
|
||||
std::fs::remove_file(path).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn full_ring_does_not_block() {
|
||||
let path = temp_path();
|
||||
let mut ring = MediaRing::create(&path, 2, 8).unwrap();
|
||||
let record = MediaRecord {
|
||||
sequence: 1,
|
||||
pts: 0,
|
||||
kind: 1,
|
||||
payload: vec![1],
|
||||
};
|
||||
assert!(ring.push(&record).unwrap());
|
||||
assert!(ring.push(&record).unwrap());
|
||||
assert!(!ring.push(&record).unwrap());
|
||||
drop(ring);
|
||||
std::fs::remove_file(path).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn descriptor_binds_session_and_encoded_only_transport() {
|
||||
let path = temp_path();
|
||||
let ring = MediaRing::create(&path, 2, 64).unwrap();
|
||||
let descriptor = ring.descriptor("session-1", 3, 42, &path).unwrap();
|
||||
assert_eq!(
|
||||
descriptor.video_transport,
|
||||
"shared_memory_encoded_access_units"
|
||||
);
|
||||
assert!(descriptor.validate().is_ok());
|
||||
drop(ring);
|
||||
std::fs::remove_file(path).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn descriptor_rejects_relative_or_unbound_identity() {
|
||||
let path = temp_path();
|
||||
let ring = MediaRing::create(&path, 2, 64).unwrap();
|
||||
assert!(ring.descriptor("bad session", 1, 42, &path).is_err());
|
||||
assert!(ring.descriptor("session", 0, 42, &path).is_err());
|
||||
assert!(
|
||||
ring.descriptor("session", 1, 42, "relative-ring.bin")
|
||||
.is_err()
|
||||
);
|
||||
drop(ring);
|
||||
std::fs::remove_file(path).unwrap();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,172 @@
|
||||
use std::ffi::{OsStr, c_void};
|
||||
use std::fs::File;
|
||||
use std::io;
|
||||
use std::mem;
|
||||
use std::os::windows::ffi::OsStrExt as _;
|
||||
use std::os::windows::io::{AsRawHandle as _, FromRawHandle as _, RawHandle};
|
||||
use std::path::Path;
|
||||
use std::ptr;
|
||||
use tokio::net::windows::named_pipe::{NamedPipeServer, PipeMode, ServerOptions};
|
||||
use windows_sys::Win32::Foundation::{
|
||||
ERROR_INSUFFICIENT_BUFFER, GENERIC_READ, GENERIC_WRITE, HANDLE, INVALID_HANDLE_VALUE, LocalFree,
|
||||
};
|
||||
use windows_sys::Win32::Security::Authorization::{
|
||||
ConvertSidToStringSidW, ConvertStringSecurityDescriptorToSecurityDescriptorW, SDDL_REVISION_1,
|
||||
};
|
||||
use windows_sys::Win32::Security::{
|
||||
GetTokenInformation, PSECURITY_DESCRIPTOR, SECURITY_ATTRIBUTES, TOKEN_QUERY, TOKEN_USER,
|
||||
TokenUser,
|
||||
};
|
||||
use windows_sys::Win32::Storage::FileSystem::{
|
||||
CREATE_NEW, CreateFileW, FILE_ATTRIBUTE_TEMPORARY, FILE_SHARE_DELETE, FILE_SHARE_READ,
|
||||
FILE_SHARE_WRITE,
|
||||
};
|
||||
use windows_sys::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken};
|
||||
|
||||
struct CurrentUserSecurity {
|
||||
descriptor: PSECURITY_DESCRIPTOR,
|
||||
}
|
||||
|
||||
impl CurrentUserSecurity {
|
||||
fn new() -> io::Result<Self> {
|
||||
let sid = current_user_sid()?;
|
||||
let sddl = wide(OsStr::new(&format!("D:P(A;;GA;;;{sid})")));
|
||||
let mut descriptor: PSECURITY_DESCRIPTOR = ptr::null_mut();
|
||||
// SAFETY: the SDDL is terminated and descriptor is writable output storage.
|
||||
if unsafe {
|
||||
ConvertStringSecurityDescriptorToSecurityDescriptorW(
|
||||
sddl.as_ptr(),
|
||||
SDDL_REVISION_1,
|
||||
&raw mut descriptor,
|
||||
ptr::null_mut(),
|
||||
)
|
||||
} == 0
|
||||
{
|
||||
return Err(io::Error::last_os_error());
|
||||
}
|
||||
Ok(Self { descriptor })
|
||||
}
|
||||
|
||||
fn attributes(&mut self) -> io::Result<SECURITY_ATTRIBUTES> {
|
||||
Ok(SECURITY_ATTRIBUTES {
|
||||
nLength: u32::try_from(mem::size_of::<SECURITY_ATTRIBUTES>())
|
||||
.map_err(|_| io::Error::other("SECURITY_ATTRIBUTES size is not representable"))?,
|
||||
lpSecurityDescriptor: self.descriptor.cast(),
|
||||
bInheritHandle: 0,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for CurrentUserSecurity {
|
||||
fn drop(&mut self) {
|
||||
if !self.descriptor.is_null() {
|
||||
// SAFETY: the descriptor was allocated by the SDDL conversion API.
|
||||
unsafe { LocalFree(self.descriptor.cast()) };
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn create_current_user_pipe(name: &str, first_instance: bool) -> io::Result<NamedPipeServer> {
|
||||
let mut security = CurrentUserSecurity::new()?;
|
||||
let mut attributes = security.attributes()?;
|
||||
let mut options = ServerOptions::new();
|
||||
options
|
||||
.pipe_mode(PipeMode::Byte)
|
||||
.max_instances(16)
|
||||
.reject_remote_clients(true)
|
||||
.first_pipe_instance(first_instance);
|
||||
// SAFETY: attributes and its descriptor remain alive for the complete creation call.
|
||||
unsafe {
|
||||
options.create_with_security_attributes_raw(name, (&raw mut attributes).cast::<c_void>())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn create_current_user_file(path: &Path) -> io::Result<File> {
|
||||
let mut security = CurrentUserSecurity::new()?;
|
||||
let mut attributes = security.attributes()?;
|
||||
let path = wide(path.as_os_str());
|
||||
// SAFETY: all pointers are valid for the duration of CreateFileW and the
|
||||
// security descriptor is owned by `security` until the call returns.
|
||||
let handle = unsafe {
|
||||
CreateFileW(
|
||||
path.as_ptr(),
|
||||
GENERIC_READ | GENERIC_WRITE,
|
||||
FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
|
||||
&raw mut attributes,
|
||||
CREATE_NEW,
|
||||
FILE_ATTRIBUTE_TEMPORARY,
|
||||
ptr::null_mut(),
|
||||
)
|
||||
};
|
||||
if handle == INVALID_HANDLE_VALUE {
|
||||
return Err(io::Error::last_os_error());
|
||||
}
|
||||
// SAFETY: CreateFileW returned a new owned kernel handle.
|
||||
Ok(unsafe { File::from_raw_handle(handle as RawHandle) })
|
||||
}
|
||||
|
||||
fn current_user_sid() -> io::Result<String> {
|
||||
let mut token: HANDLE = ptr::null_mut();
|
||||
// SAFETY: token points to writable storage and the process pseudo-handle is valid.
|
||||
if unsafe { OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &raw mut token) } == 0 {
|
||||
return Err(io::Error::last_os_error());
|
||||
}
|
||||
// SAFETY: OpenProcessToken returned an owned token handle.
|
||||
let token = unsafe { File::from_raw_handle(token as RawHandle) };
|
||||
let mut required = 0_u32;
|
||||
// SAFETY: a null buffer with zero length is the documented size query.
|
||||
unsafe {
|
||||
GetTokenInformation(
|
||||
token.as_raw_handle() as HANDLE,
|
||||
TokenUser,
|
||||
ptr::null_mut(),
|
||||
0,
|
||||
&raw mut required,
|
||||
)
|
||||
};
|
||||
let size_error = io::Error::last_os_error();
|
||||
if size_error.raw_os_error().map(i32::cast_unsigned) != Some(ERROR_INSUFFICIENT_BUFFER)
|
||||
|| required == 0
|
||||
{
|
||||
return Err(size_error);
|
||||
}
|
||||
let words = (required as usize).div_ceil(mem::size_of::<usize>());
|
||||
let mut buffer = vec![0_usize; words];
|
||||
// SAFETY: buffer contains at least `required` writable bytes.
|
||||
if unsafe {
|
||||
GetTokenInformation(
|
||||
token.as_raw_handle() as HANDLE,
|
||||
TokenUser,
|
||||
buffer.as_mut_ptr().cast::<c_void>(),
|
||||
required,
|
||||
&raw mut required,
|
||||
)
|
||||
} == 0
|
||||
{
|
||||
return Err(io::Error::last_os_error());
|
||||
}
|
||||
// SAFETY: TokenUser queries return a TOKEN_USER at the start of the buffer.
|
||||
let token_user = unsafe { &*buffer.as_ptr().cast::<TOKEN_USER>() };
|
||||
let mut sid_string = ptr::null_mut();
|
||||
// SAFETY: the SID is owned by the token information buffer and output is writable.
|
||||
if unsafe { ConvertSidToStringSidW(token_user.User.Sid, &raw mut sid_string) } == 0 {
|
||||
return Err(io::Error::last_os_error());
|
||||
}
|
||||
let mut length = 0_usize;
|
||||
// SAFETY: ConvertSidToStringSidW returns a terminated LocalAlloc string.
|
||||
unsafe {
|
||||
while *sid_string.add(length) != 0 {
|
||||
length += 1;
|
||||
}
|
||||
}
|
||||
// SAFETY: the preceding scan established the initialized UTF-16 range.
|
||||
let sid = String::from_utf16(unsafe { std::slice::from_raw_parts(sid_string, length) })
|
||||
.map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "current SID is invalid"));
|
||||
// SAFETY: the SID string was allocated by ConvertSidToStringSidW.
|
||||
unsafe { LocalFree(sid_string.cast()) };
|
||||
sid
|
||||
}
|
||||
|
||||
fn wide(value: &OsStr) -> Vec<u16> {
|
||||
value.encode_wide().chain([0]).collect()
|
||||
}
|
||||
Reference in New Issue
Block a user