use std::collections::HashMap;
use std::io::{Read, Seek, SeekFrom, Write};
use crate::{CodecResolver, Error, Packet, Result, StreamInfo};
pub trait Demuxer: Send {
fn format_name(&self) -> &str;
fn streams(&self) -> &[StreamInfo];
fn next_packet(&mut self) -> Result<Packet>;
fn set_active_streams(&mut self, _indices: &[u32]) {}
fn seek_to(&mut self, _stream_index: u32, _pts: i64) -> Result<i64> {
Err(Error::unsupported("this demuxer does not support seeking"))
}
fn metadata(&self) -> &[(String, String)] {
&[]
}
fn duration_micros(&self) -> Option<i64> {
None
}
fn attached_pictures(&self) -> &[crate::AttachedPicture] {
&[]
}
fn chapters(&self) -> &[crate::Chapter] {
&[]
}
fn attachments(&self) -> &[crate::Attachment] {
&[]
}
}
pub trait Muxer: Send {
fn format_name(&self) -> &str;
fn write_header(&mut self) -> Result<()>;
fn write_packet(&mut self, packet: &Packet) -> Result<()>;
fn write_trailer(&mut self) -> Result<()>;
}
pub type OpenDemuxerFn =
fn(input: Box<dyn ReadSeek>, codecs: &dyn CodecResolver) -> Result<Box<dyn Demuxer>>;
pub type OpenMuxerFn =
fn(output: Box<dyn WriteSeek>, streams: &[StreamInfo]) -> Result<Box<dyn Muxer>>;
pub struct ProbeData<'a> {
pub buf: &'a [u8],
pub ext: Option<&'a str>,
}
pub type ProbeScore = u8;
pub const MAX_PROBE_SCORE: ProbeScore = 100;
pub const PROBE_SCORE_EXTENSION: ProbeScore = 25;
pub type ContainerProbeFn = fn(probe: &ProbeData) -> ProbeScore;
#[non_exhaustive]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ProbeCandidate<'a> {
pub name: &'a str,
pub score: ProbeScore,
pub priority: i32,
pub order: usize,
}
struct ProbeEntry {
name: String,
probe: ContainerProbeFn,
priority: i32,
}
#[non_exhaustive]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ExtensionCandidate<'a> {
pub container: &'a str,
pub priority: i32,
pub order: usize,
}
struct ExtensionClaim {
container: String,
priority: i32,
order: usize,
}
pub trait ReadSeek: Read + Seek + Send {}
impl<T: Read + Seek + Send> ReadSeek for T {}
pub trait WriteSeek: Write + Seek + Send {}
impl<T: Write + Seek + Send> WriteSeek for T {}
#[derive(Default)]
pub struct ContainerRegistry {
demuxers: Vec<(String, OpenDemuxerFn)>,
muxers: Vec<(String, OpenMuxerFn)>,
extensions: HashMap<String, Vec<ExtensionClaim>>,
extension_seq: usize,
probes: Vec<ProbeEntry>,
}
impl ContainerRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn register_demuxer(&mut self, name: &str, open: OpenDemuxerFn) {
match self.demuxers.iter_mut().find(|(n, _)| n == name) {
Some(slot) => slot.1 = open,
None => self.demuxers.push((name.to_owned(), open)),
}
}
pub fn register_muxer(&mut self, name: &str, open: OpenMuxerFn) {
match self.muxers.iter_mut().find(|(n, _)| n == name) {
Some(slot) => slot.1 = open,
None => self.muxers.push((name.to_owned(), open)),
}
}
pub fn register_extension(&mut self, ext: &str, container_name: &str) {
self.register_extension_with_priority(ext, container_name, crate::DEFAULT_PRIORITY);
}
pub fn register_extension_with_priority(
&mut self,
ext: &str,
container_name: &str,
priority: i32,
) {
let order = self.extension_seq;
self.extension_seq += 1;
self.extensions
.entry(ext.to_lowercase())
.or_default()
.push(ExtensionClaim {
container: container_name.to_owned(),
priority,
order,
});
}
pub fn extension_candidates(&self, ext: &str) -> Vec<ExtensionCandidate<'_>> {
let Some(claims) = self.extensions.get(&ext.to_lowercase()) else {
return Vec::new();
};
let mut out: Vec<ExtensionCandidate<'_>> = claims
.iter()
.map(|c| ExtensionCandidate {
container: c.container.as_str(),
priority: c.priority,
order: c.order,
})
.collect();
out.sort_by(|a, b| a.priority.cmp(&b.priority).then(b.order.cmp(&a.order)));
out
}
pub fn register_probe(&mut self, container_name: &str, probe: ContainerProbeFn) {
self.register_probe_with_priority(container_name, probe, crate::DEFAULT_PRIORITY);
}
pub fn register_probe_with_priority(
&mut self,
container_name: &str,
probe: ContainerProbeFn,
priority: i32,
) {
if let Some(entry) = self.probes.iter_mut().find(|e| e.name == container_name) {
entry.probe = probe;
entry.priority = priority;
} else {
self.probes.push(ProbeEntry {
name: container_name.to_owned(),
probe,
priority,
});
}
}
pub fn probe_priority(&self, container_name: &str) -> Option<i32> {
self.probes
.iter()
.find(|e| e.name == container_name)
.map(|e| e.priority)
}
pub fn probe_candidates(&self, data: &ProbeData) -> Vec<ProbeCandidate<'_>> {
let mut out: Vec<ProbeCandidate<'_>> = self
.probes
.iter()
.enumerate()
.filter_map(|(order, e)| {
let score = (e.probe)(data);
(score != 0).then_some(ProbeCandidate {
name: e.name.as_str(),
score,
priority: e.priority,
order,
})
})
.collect();
out.sort_by(|a, b| {
b.score
.cmp(&a.score)
.then(a.priority.cmp(&b.priority))
.then(a.order.cmp(&b.order))
});
out
}
pub fn demuxer_names(&self) -> impl Iterator<Item = &str> {
self.demuxers.iter().map(|(n, _)| n.as_str())
}
pub fn muxer_names(&self) -> impl Iterator<Item = &str> {
self.muxers.iter().map(|(n, _)| n.as_str())
}
pub fn open_demuxer(
&self,
name: &str,
input: Box<dyn ReadSeek>,
codecs: &dyn CodecResolver,
) -> Result<Box<dyn Demuxer>> {
let open = self
.demuxers
.iter()
.find(|(n, _)| n == name)
.map(|(_, f)| *f)
.ok_or_else(|| Error::FormatNotFound(name.to_owned()))?;
open(input, codecs)
}
pub fn open_muxer(
&self,
name: &str,
output: Box<dyn WriteSeek>,
streams: &[StreamInfo],
) -> Result<Box<dyn Muxer>> {
let open = self
.muxers
.iter()
.find(|(n, _)| n == name)
.map(|(_, f)| *f)
.ok_or_else(|| Error::FormatNotFound(name.to_owned()))?;
open(output, streams)
}
pub fn container_for_extension(&self, ext: &str) -> Option<&str> {
let claims = self.extensions.get(&ext.to_lowercase())?;
let mut best: Option<&ExtensionClaim> = None;
for c in claims {
best = match best {
Some(b) if c.priority > b.priority => Some(b),
_ => Some(c),
};
}
best.map(|c| c.container.as_str())
}
pub fn probe_input(&self, input: &mut dyn ReadSeek, ext_hint: Option<&str>) -> Result<String> {
const PROBE_BUF_SIZE: usize = 256 * 1024;
let saved_pos = input.stream_position()?;
input.seek(SeekFrom::Start(0))?;
let mut buf = vec![0u8; PROBE_BUF_SIZE];
let mut got = 0;
while got < buf.len() {
match input.read(&mut buf[got..]) {
Ok(0) => break,
Ok(n) => got += n,
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
Err(e) => {
let _ = input.seek(SeekFrom::Start(saved_pos));
return Err(e.into());
}
}
}
buf.truncate(got);
input.seek(SeekFrom::Start(saved_pos))?;
let ext_lower = ext_hint.map(|s| s.to_ascii_lowercase());
let probe_data = ProbeData {
buf: &buf,
ext: ext_lower.as_deref(),
};
if let Some(best) = self.probe_candidates(&probe_data).first() {
return Ok(best.name.to_owned());
}
if let Some(ext) = ext_hint {
if let Some(name) = self.container_for_extension(ext) {
return Ok(name.to_owned());
}
}
Err(Error::FormatNotFound(
"no registered demuxer recognises this input".into(),
))
}
}
#[cfg(test)]
mod tests {
use super::*;
struct DummyDemuxer;
impl Demuxer for DummyDemuxer {
fn format_name(&self) -> &str {
"dummy"
}
fn streams(&self) -> &[StreamInfo] {
&[]
}
fn next_packet(&mut self) -> Result<Packet> {
Err(Error::Eof)
}
}
#[test]
fn default_seek_to_is_unsupported() {
let mut d = DummyDemuxer;
match d.seek_to(0, 0) {
Err(Error::Unsupported(_)) => {}
other => panic!(
"expected default seek_to to return Unsupported, got {:?}",
other
),
}
}
fn always(score: ProbeScore) -> ContainerProbeFn {
match score {
100 => |_: &ProbeData| 100,
90 => |_: &ProbeData| 90,
50 => |_: &ProbeData| 50,
_ => |_: &ProbeData| 0,
}
}
fn names<'a>(cands: &'a [ProbeCandidate<'a>]) -> Vec<&'a str> {
cands.iter().map(|c| c.name).collect()
}
fn probe_cursor(reg: &ContainerRegistry, bytes: &[u8]) -> String {
let mut cur = std::io::Cursor::new(bytes.to_vec());
reg.probe_input(&mut cur, None).expect("some probe matches")
}
#[test]
fn equal_scores_resolve_by_registration_order_across_fresh_registries() {
const ORDER: [&str; 6] = ["zeta", "alpha", "omicron", "beta", "kappa", "mu"];
for _ in 0..1000 {
let mut ctx = crate::RuntimeContext::new();
for name in ORDER {
ctx.containers.register_probe(name, always(100));
}
assert_eq!(probe_cursor(&ctx.containers, b"any bytes"), "zeta");
let cands = ctx.containers.probe_candidates(&ProbeData {
buf: b"any bytes",
ext: None,
});
assert_eq!(names(&cands), ORDER.to_vec());
}
}
#[test]
fn registration_order_is_the_only_remaining_discriminator() {
let mut a = ContainerRegistry::new();
a.register_probe("second-by-name", always(100));
a.register_probe("first-by-name", always(100));
assert_eq!(probe_cursor(&a, b"x"), "second-by-name");
let mut b = ContainerRegistry::new();
b.register_probe("first-by-name", always(100));
b.register_probe("second-by-name", always(100));
assert_eq!(probe_cursor(&b, b"x"), "first-by-name");
}
#[test]
fn lower_priority_number_beats_registration_order() {
let mut reg = ContainerRegistry::new();
reg.register_probe("earlier", always(100));
reg.register_probe_with_priority("later-but-preferred", always(100), 10);
assert_eq!(probe_cursor(®, b"x"), "later-but-preferred");
let cands = reg.probe_candidates(&ProbeData {
buf: b"x",
ext: None,
});
assert_eq!(names(&cands), vec!["later-but-preferred", "earlier"]);
assert_eq!(cands[0].priority, 10);
assert_eq!(cands[0].order, 1);
assert_eq!(cands[1].priority, crate::DEFAULT_PRIORITY);
assert_eq!(cands[1].order, 0);
}
#[test]
fn score_always_beats_priority() {
let mut reg = ContainerRegistry::new();
reg.register_probe_with_priority("confident-but-yielding", always(90), i32::MIN);
reg.register_probe("weakly-prioritised-but-certain", always(100));
assert_eq!(probe_cursor(®, b"x"), "weakly-prioritised-but-certain");
}
#[test]
fn probe_candidates_excludes_zero_scores_and_ranks_fully() {
let mut reg = ContainerRegistry::new();
reg.register_probe("silent", always(0));
reg.register_probe("half", always(50));
reg.register_probe_with_priority("full-late-preferred", always(100), 1);
reg.register_probe("full-early", always(100));
reg.register_probe("half-again", always(50));
let cands = reg.probe_candidates(&ProbeData {
buf: b"x",
ext: None,
});
assert_eq!(
names(&cands),
vec!["full-late-preferred", "full-early", "half", "half-again"]
);
assert_eq!(
cands.iter().map(|c| c.order).collect::<Vec<_>>(),
vec![2, 3, 1, 4]
);
assert!(cands.iter().all(|c| c.score != 0));
}
#[test]
fn re_registering_a_name_replaces_in_place() {
let mut reg = ContainerRegistry::new();
reg.register_probe_with_priority("first", always(100), 5);
reg.register_probe("second", always(100));
assert_eq!(reg.probe_priority("first"), Some(5));
assert_eq!(probe_cursor(®, b"x"), "first");
reg.register_probe("first", always(0));
assert_eq!(reg.probe_priority("first"), Some(crate::DEFAULT_PRIORITY));
assert_eq!(probe_cursor(®, b"x"), "second");
reg.register_probe("first", always(100));
let cands = reg.probe_candidates(&ProbeData {
buf: b"x",
ext: None,
});
assert_eq!(names(&cands), vec!["first", "second"]);
assert_eq!(cands[0].order, 0);
assert_eq!(reg.probe_priority("never-registered"), None);
}
#[test]
fn probe_input_falls_back_to_extension_when_nothing_scores() {
let mut reg = ContainerRegistry::new();
reg.register_probe("silent", always(0));
reg.register_extension("xyz", "by-extension");
let mut cur = std::io::Cursor::new(b"bytes".to_vec());
assert_eq!(
reg.probe_input(&mut cur, Some("XYZ")).unwrap(),
"by-extension"
);
assert!(matches!(
reg.probe_input(&mut cur, Some("abc")),
Err(Error::FormatNotFound(_))
));
}
#[test]
fn extension_equal_priority_most_recent_wins_deterministically() {
for _ in 0..1000 {
let mut ctx = crate::RuntimeContext::new();
ctx.containers.register_extension("SYN", "first");
ctx.containers.register_extension("syn", "second");
ctx.containers.register_extension("Syn", "third");
assert_eq!(ctx.containers.container_for_extension("syn"), Some("third"));
assert_eq!(ctx.containers.container_for_extension("SYN"), Some("third"));
let cands = ctx.containers.extension_candidates("syn");
assert_eq!(
cands.iter().map(|c| c.container).collect::<Vec<_>>(),
vec!["third", "second", "first"]
);
assert_eq!(
cands.iter().map(|c| c.order).collect::<Vec<_>>(),
vec![2, 1, 0]
);
}
}
#[test]
fn extension_priority_beats_registration_recency() {
let mut reg = ContainerRegistry::new();
reg.register_extension_with_priority("syn", "pinned", 10);
reg.register_extension("syn", "later-default");
reg.register_extension("syn", "latest-default");
assert_eq!(reg.container_for_extension("syn"), Some("pinned"));
let cands = reg.extension_candidates("syn");
assert_eq!(
cands.iter().map(|c| c.container).collect::<Vec<_>>(),
vec!["pinned", "latest-default", "later-default"]
);
assert_eq!(cands[0].priority, 10);
assert_eq!(cands[1].priority, crate::DEFAULT_PRIORITY);
reg.register_extension_with_priority("syn", "pinned-harder", 1);
assert_eq!(reg.container_for_extension("syn"), Some("pinned-harder"));
reg.register_extension("other", "elsewhere");
assert_eq!(reg.container_for_extension("other"), Some("elsewhere"));
assert!(reg.extension_candidates("nope").is_empty());
assert_eq!(reg.container_for_extension("nope"), None);
}
#[test]
fn extension_order_is_global_across_extensions() {
let mut reg = ContainerRegistry::new();
reg.register_extension("a", "x");
reg.register_extension("b", "y");
reg.register_extension("a", "z");
let a = reg.extension_candidates("a");
let b = reg.extension_candidates("b");
assert_eq!(a.iter().map(|c| c.order).collect::<Vec<_>>(), vec![2, 0]);
assert_eq!(b[0].order, 1);
}
fn open_dummy(
_input: Box<dyn ReadSeek>,
_codecs: &dyn CodecResolver,
) -> Result<Box<dyn Demuxer>> {
Ok(Box::new(DummyDemuxer))
}
fn open_dummy_mux(
_output: Box<dyn WriteSeek>,
_streams: &[StreamInfo],
) -> Result<Box<dyn Muxer>> {
Err(Error::unsupported("dummy muxer"))
}
#[test]
fn name_listings_follow_registration_order_and_replace_in_place() {
const ORDER: [&str; 5] = ["zeta", "alpha", "omicron", "beta", "kappa"];
for _ in 0..200 {
let mut reg = ContainerRegistry::new();
for n in ORDER {
reg.register_demuxer(n, open_dummy);
reg.register_muxer(n, open_dummy_mux);
}
reg.register_demuxer("alpha", open_dummy);
reg.register_muxer("omicron", open_dummy_mux);
assert_eq!(reg.demuxer_names().collect::<Vec<_>>(), ORDER.to_vec());
assert_eq!(reg.muxer_names().collect::<Vec<_>>(), ORDER.to_vec());
}
let mut reg = ContainerRegistry::new();
reg.register_demuxer("only", open_dummy);
let d = reg
.open_demuxer(
"only",
Box::new(std::io::Cursor::new(Vec::new())),
&crate::NullCodecResolver,
)
.unwrap();
assert_eq!(d.format_name(), "dummy");
assert!(matches!(
reg.open_demuxer(
"missing",
Box::new(std::io::Cursor::new(Vec::new())),
&crate::NullCodecResolver,
),
Err(Error::FormatNotFound(_))
));
assert!(matches!(
reg.open_muxer("missing", Box::new(std::io::Cursor::new(Vec::new())), &[]),
Err(Error::FormatNotFound(_))
));
}
#[test]
fn default_chapters_and_attachments_are_empty() {
let d = DummyDemuxer;
assert!(d.chapters().is_empty());
assert!(d.attachments().is_empty());
assert!(d.attached_pictures().is_empty());
assert!(d.metadata().is_empty());
assert_eq!(d.duration_micros(), None);
}
}