dvb_si/descriptors/extension/
registry.rs1use alloc::boxed::Box;
22use alloc::collections::BTreeMap;
23
24use super::{ExtensionBodyDef, ExtensionDescriptor, validate_and_split};
25use crate::error::Result;
26
27broadcast_common::declare_object_trait! {
28 pub trait ExtensionObject, "extension descriptor"
36}
37
38pub(crate) type CustomParse =
46 Box<dyn for<'a> Fn(&'a [u8]) -> Result<Box<dyn ExtensionObject>> + Send + Sync>;
47
48#[derive(Default)]
63pub struct ExtensionRegistry {
64 custom: BTreeMap<u8, CustomParse>,
65}
66
67impl ExtensionRegistry {
68 #[must_use]
70 pub fn new() -> Self {
71 Self::default()
72 }
73
74 #[must_use]
77 pub fn has_custom(&self, tag_extension: u8) -> bool {
78 self.custom.contains_key(&tag_extension)
79 }
80
81 pub fn register<T>(&mut self) -> &mut Self
93 where
94 T: for<'a> ExtensionBodyDef<'a> + ExtensionObject + 'static,
95 {
96 let tag_ext = T::TAG_EXTENSION;
97 self.custom.insert(
98 tag_ext,
99 Box::new(|sel| {
100 Ok(Box::new(<T as broadcast_common::Parse>::parse(sel)?)
101 as Box<dyn ExtensionObject>)
102 }),
103 );
104 self
105 }
106
107 pub fn parse_body<'a>(
111 &self,
112 tag_extension: u8,
113 selector: &'a [u8],
114 ) -> Result<RegisteredExtension<'a>> {
115 if let Some(parse_fn) = self.custom.get(&tag_extension) {
116 let value = parse_fn(selector)?;
117 Ok(RegisteredExtension::Custom {
118 tag_extension,
119 value,
120 })
121 } else {
122 let body = super::parse_body(tag_extension, selector)?;
123 Ok(RegisteredExtension::Builtin(ExtensionDescriptor {
124 tag_extension,
125 body,
126 }))
127 }
128 }
129
130 pub fn parse<'a>(&self, bytes: &'a [u8]) -> Result<RegisteredExtension<'a>> {
138 let (tag_extension, sel) = validate_and_split(bytes)?;
139 self.parse_body(tag_extension, sel)
140 }
141}
142
143#[derive(Debug)]
149#[cfg_attr(feature = "serde", derive(serde::Serialize))]
150#[cfg_attr(feature = "serde", serde(rename_all = "camelCase"))]
151#[non_exhaustive]
152pub enum RegisteredExtension<'a> {
153 Builtin(super::ExtensionDescriptor<'a>),
156 Custom {
158 tag_extension: u8,
160 #[cfg_attr(feature = "serde", serde(serialize_with = "serialize_erased"))]
163 value: Box<dyn ExtensionObject>,
164 },
165}
166
167#[cfg(test)]
168mod tests {
169 use super::*;
170 use crate::descriptors::extension::{ExtensionBodyDef, TAG, TAG_EXTENSION_LEN};
171 use crate::error::Error;
172
173 const TEST_TAG_EXTENSION: u8 = 0x40;
174
175 #[derive(Debug, PartialEq, Eq)]
176 #[cfg_attr(feature = "serde", derive(serde::Serialize))]
177 struct MyExtBody {
178 payload: Vec<u8>,
179 }
180
181 impl ExtensionBodyDef<'_> for MyExtBody {
182 const TAG_EXTENSION: u8 = TEST_TAG_EXTENSION;
183 const NAME: &'static str = "MY_EXT_BODY";
184 }
185
186 impl<'a> broadcast_common::Parse<'a> for MyExtBody {
187 type Error = crate::error::Error;
188 fn parse(sel: &'a [u8]) -> Result<Self> {
189 Ok(Self {
190 payload: sel.to_vec(),
191 })
192 }
193 }
194
195 fn wrap_ext(tag_ext: u8, sel: &[u8]) -> Vec<u8> {
196 let mut v = vec![TAG, (sel.len() + TAG_EXTENSION_LEN) as u8, tag_ext];
197 v.extend_from_slice(sel);
198 v
199 }
200
201 #[test]
202 fn custom_extension_parsed_and_downcastable() {
203 let mut reg = ExtensionRegistry::new();
204 reg.register::<MyExtBody>();
205
206 let sel = [0xDE, 0xAD, 0xBE];
207 let bytes = wrap_ext(TEST_TAG_EXTENSION, &sel);
208 let re = reg.parse(&bytes).unwrap();
209 match re {
210 RegisteredExtension::Custom {
211 tag_extension,
212 value,
213 } => {
214 assert_eq!(tag_extension, TEST_TAG_EXTENSION);
215 let concrete = value
216 .downcast_ref::<MyExtBody>()
217 .expect("downcast should succeed");
218 assert_eq!(concrete.payload, sel);
219 }
220 other => panic!("expected Custom, got {other:?}"),
221 }
222 }
223
224 #[test]
225 fn unregistered_tag_extension_yields_builtin() {
226 use crate::descriptors::extension::ExtensionBody;
227 let reg = ExtensionRegistry::new();
228 let d = crate::descriptors::extension::ExtensionDescriptor {
230 tag_extension: 0x0B,
231 body: ExtensionBody::ServiceRelocated(
232 crate::descriptors::extension::ServiceRelocated {
233 old_original_network_id: 1,
234 old_transport_stream_id: 2,
235 old_service_id: 3,
236 },
237 ),
238 };
239 let mut buf = vec![0u8; d.serialized_len()];
240 use broadcast_common::Serialize;
241 d.serialize_into(&mut buf).unwrap();
242
243 let re = reg.parse(&buf).unwrap();
244 match re {
245 RegisteredExtension::Builtin(d) => {
246 assert_eq!(d.tag_extension, 0x0B);
247 assert!(matches!(d.body, ExtensionBody::ServiceRelocated(_)));
248 }
249 other => panic!("expected Builtin, got {other:?}"),
250 }
251 }
252
253 #[test]
254 fn unknown_tag_extension_yields_builtin_raw() {
255 use crate::descriptors::extension::ExtensionBody;
256 let reg = ExtensionRegistry::new();
257 let sel = [0xAA, 0xBB];
258 let bytes = wrap_ext(0xFE, &sel);
259 let re = reg.parse(&bytes).unwrap();
260 match re {
261 RegisteredExtension::Builtin(d) => {
262 assert_eq!(d.tag_extension, 0xFE);
263 assert!(matches!(d.body, ExtensionBody::Raw(b) if b == sel));
264 }
265 other => panic!("expected Builtin, got {other:?}"),
266 }
267 }
268
269 #[test]
270 fn parse_rejects_wrong_tag() {
271 let reg = ExtensionRegistry::new();
272 let raw = [0x43, 1, 0x04];
273 assert!(matches!(
274 reg.parse(&raw).unwrap_err(),
275 Error::InvalidDescriptor { tag: 0x43, .. }
276 ));
277 }
278
279 #[test]
280 fn parse_rejects_short_buffer() {
281 let reg = ExtensionRegistry::new();
282 let raw = [TAG];
283 assert!(matches!(
284 reg.parse(&raw).unwrap_err(),
285 Error::BufferTooShort(..)
286 ));
287 }
288
289 #[test]
290 fn parse_rejects_empty_body() {
291 let reg = ExtensionRegistry::new();
292 let raw = [TAG, 0];
293 assert!(matches!(
294 reg.parse(&raw).unwrap_err(),
295 Error::InvalidDescriptor { tag: TAG, .. }
296 ));
297 }
298
299 #[cfg(feature = "serde")]
300 #[test]
301 fn custom_variant_serializes_via_erased_serde() {
302 let mut reg = ExtensionRegistry::new();
303 reg.register::<MyExtBody>();
304
305 let bytes = wrap_ext(TEST_TAG_EXTENSION, &[0x01, 0x02]);
306 let re = reg.parse(&bytes).unwrap();
307 let json = serde_json::to_value(&re).unwrap();
308 let custom = json.get("custom").expect("expected 'custom' key");
309 assert_eq!(custom["tag_extension"], TEST_TAG_EXTENSION as u64);
310 assert_eq!(custom["value"]["payload"], serde_json::json!([1, 2]));
311 }
312}