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

This commit is contained in:
曾志威
2026-08-14 14:31:57 +08:00
parent 5db6b9ef68
commit 19a8e03a83
53 changed files with 9808 additions and 239 deletions
+4 -4
View File
@@ -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,
)
+3 -3
View File
@@ -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),
+22
View File
@@ -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"
+471
View File
@@ -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;
+166
View File
@@ -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
}
}
+302
View File
@@ -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")
}
}
+657
View File
@@ -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]
);
}
}
+186
View File
@@ -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();
}
}
File diff suppressed because it is too large Load Diff
+439
View File
@@ -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();
}
}
+172
View File
@@ -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()
}