1use std::collections::HashMap;
8use std::sync::Arc;
9
10use crate::archive::{Codec, Container};
11use crate::error::{Error, Result};
12use crate::format::{FormatId, FormatKind, FormatMeta, MagicRule};
13
14#[derive(Clone, PartialEq, Eq, Debug)]
16pub struct FormatRow {
17 pub id: FormatId,
18 pub kind: FormatKind,
19 pub extensions: &'static [&'static str],
20 pub read: bool,
21 pub write: bool,
22 pub parallel: bool,
23}
24
25#[derive(Default)]
26pub struct Registry {
27 codecs: HashMap<FormatId, Arc<dyn Codec>>,
28 containers: HashMap<FormatId, Arc<dyn Container>>,
29 metas: HashMap<FormatId, FormatMeta>,
30 by_ext: HashMap<String, FormatId>,
31 magics: Vec<MagicRule>,
32}
33
34impl Registry {
35 pub fn new() -> Self {
36 Self::default()
37 }
38
39 fn index(&mut self, meta: FormatMeta) {
40 self.by_ext.retain(|_, v| *v != meta.id);
43 for ext in meta.extensions {
44 self.by_ext.insert(ext.to_ascii_lowercase(), meta.id);
45 }
46 self.magics.retain(|m| m.format != meta.id);
48 self.magics.extend_from_slice(meta.magics);
49 self.metas.insert(meta.id, meta);
50 }
51
52 pub fn register_codec(&mut self, codec: Arc<dyn Codec>, meta: FormatMeta) {
53 self.codecs.insert(meta.id, codec);
54 self.index(meta);
55 }
56
57 pub fn register_container(&mut self, container: Arc<dyn Container>, meta: FormatMeta) {
58 self.containers.insert(meta.id, container);
59 self.index(meta);
60 }
61
62 pub fn codec(&self, id: FormatId) -> Option<&Arc<dyn Codec>> {
63 self.codecs.get(&id)
64 }
65
66 pub fn container(&self, id: FormatId) -> Option<&Arc<dyn Container>> {
67 self.containers.get(&id)
68 }
69
70 pub fn require_codec(&self, id: FormatId) -> Result<&Arc<dyn Codec>> {
73 self.codec(id).ok_or(Error::FormatNotEnabled(id))
74 }
75
76 pub fn require_encoder(&self, id: FormatId) -> Result<&Arc<dyn Codec>> {
84 let codec = self.require_codec(id)?;
85 if !codec.caps().encode {
86 return Err(Error::CapabilityUnavailable {
87 format: id,
88 available: "read",
89 requested: "written",
90 });
91 }
92 Ok(codec)
93 }
94
95 pub fn require_decoder(&self, id: FormatId) -> Result<&Arc<dyn Codec>> {
97 let codec = self.require_codec(id)?;
98 if !codec.caps().decode {
99 return Err(Error::CapabilityUnavailable {
100 format: id,
101 available: "written",
102 requested: "read",
103 });
104 }
105 Ok(codec)
106 }
107
108 pub fn require_container(&self, id: FormatId) -> Result<&Arc<dyn Container>> {
109 self.container(id).ok_or(Error::FormatNotEnabled(id))
110 }
111
112 pub fn require_container_writer(&self, id: FormatId) -> Result<&Arc<dyn Container>> {
131 let container = self.require_container(id)?;
132 if !container.caps().write {
133 return Err(Error::CapabilityUnavailable {
134 format: id,
135 available: "read",
136 requested: "written",
137 });
138 }
139 Ok(container)
140 }
141
142 pub fn by_extension(&self, ext: &str) -> Option<FormatId> {
144 let key = ext.trim_start_matches('.').to_ascii_lowercase();
145 self.by_ext.get(&key).copied()
146 }
147
148 pub fn match_magic(&self, prefix: &[u8]) -> Vec<FormatId> {
150 let mut hits: Vec<(i16, FormatId)> = self
151 .magics
152 .iter()
153 .filter(|m| {
154 let end = match m.offset.checked_add(m.bytes.len()) {
155 Some(e) => e,
156 None => return false,
157 };
158 prefix.len() >= end && &prefix[m.offset..end] == m.bytes
159 })
160 .map(|m| {
161 (
162 self.metas.get(&m.format).map_or(0, |x| x.priority),
163 m.format,
164 )
165 })
166 .collect();
167 hits.sort_by(|a, b| b.0.cmp(&a.0).then_with(|| a.1.as_str().cmp(b.1.as_str())));
169 hits.dedup_by_key(|(_, id)| *id);
170 hits.into_iter().map(|(_, id)| id).collect()
171 }
172
173 pub fn has_magic(&self, id: FormatId) -> bool {
176 self.magics.iter().any(|m| m.format == id)
177 }
178
179 pub fn confirm(&self, id: FormatId, prefix: &[u8], whole: bool) -> Option<bool> {
197 let meta = self.metas.get(&id)?;
198 if !meta.magic_optional && self.has_magic(id) && !self.match_magic(prefix).contains(&id) {
199 return Some(false);
200 }
201 meta.confirm.map(|f| f(prefix, whole))
202 }
203
204 pub fn priority_of(&self, id: FormatId) -> i16 {
206 self.metas.get(&id).map_or(0, |m| m.priority)
207 }
208
209 pub fn matrix(&self) -> Vec<FormatRow> {
211 let mut rows: Vec<FormatRow> = self
212 .metas
213 .values()
214 .map(|meta| {
215 let (read, write, parallel) = match meta.kind {
216 FormatKind::Codec => match self.codecs.get(&meta.id) {
217 Some(c) => {
218 let caps = c.caps();
219 (
220 caps.decode,
221 caps.encode,
222 caps.parallel_encode || caps.parallel_decode,
223 )
224 }
225 None => (false, false, false),
226 },
227 FormatKind::Container => match self.containers.get(&meta.id) {
228 Some(c) => {
229 let caps = c.caps();
230 (caps.read, caps.write, false)
231 }
232 None => (false, false, false),
233 },
234 };
235 FormatRow {
236 id: meta.id,
237 kind: meta.kind,
238 extensions: meta.extensions,
239 read,
240 write,
241 parallel,
242 }
243 })
244 .collect();
245 rows.sort_by_key(|r| r.id.as_str());
246 rows
247 }
248}
249
250#[cfg(test)]
251mod tests {
252 use super::*;
253 use crate::archive::{DecodeOpts, EncodeOpts, Sink};
254 use crate::format::CodecCaps;
255 use crate::source::Source;
256 use crate::testing::{MOCK_CODEC, MOCK_CONTAINER, MockCodec, MockContainer};
257 use std::io::Write;
258
259 const CODEC_MAGIC: &[MagicRule] = &[MagicRule {
260 offset: 0,
261 bytes: b"MOCK",
262 format: MOCK_CODEC,
263 }];
264 const CONTAINER_MAGIC: &[MagicRule] = &[MagicRule {
265 offset: 0,
266 bytes: b"ME",
267 format: MOCK_CONTAINER,
268 }];
269
270 fn codec_meta() -> FormatMeta {
271 FormatMeta::codec(MOCK_CODEC, &["mk", "mock"], CODEC_MAGIC)
272 }
273
274 fn container_meta() -> FormatMeta {
275 FormatMeta::container(MOCK_CONTAINER, &["mar"], CONTAINER_MAGIC)
276 }
277
278 fn populated() -> Registry {
279 let mut r = Registry::new();
280 r.register_codec(Arc::new(MockCodec), codec_meta());
281 r.register_container(Arc::new(MockContainer), container_meta());
282 r
283 }
284
285 #[test]
286 fn an_empty_registry_reports_no_formats() {
287 assert!(Registry::new().matrix().is_empty());
288 }
289
290 #[test]
291 fn registered_formats_are_retrievable_by_id() {
292 let r = populated();
293 assert!(r.codec(MOCK_CODEC).is_some());
294 assert!(r.container(MOCK_CONTAINER).is_some());
295 assert!(r.container(MOCK_CODEC).is_none());
297 }
298
299 #[test]
300 fn a_missing_format_is_format_not_enabled_not_a_generic_error() {
301 let r = Registry::new();
308 let Err(err) = r.require_codec(MOCK_CODEC) else {
309 panic!("expected FormatNotEnabled for an unregistered codec");
310 };
311 assert!(matches!(err, crate::Error::FormatNotEnabled(id) if id == MOCK_CODEC));
312 assert_eq!(err.exit_code(), 3);
313 }
314
315 #[test]
316 fn extension_lookup_is_case_insensitive_and_dot_tolerant() {
317 let r = populated();
318 assert_eq!(r.by_extension("mk"), Some(MOCK_CODEC));
319 assert_eq!(r.by_extension("MOCK"), Some(MOCK_CODEC));
320 assert_eq!(r.by_extension(".mar"), Some(MOCK_CONTAINER));
321 assert_eq!(r.by_extension("nope"), None);
322 }
323
324 #[test]
325 fn magic_matching_finds_the_format_from_a_prefix() {
326 let r = populated();
327 assert_eq!(r.match_magic(b"MOCKpayload"), vec![MOCK_CODEC]);
328 assert_eq!(r.match_magic(b"ME\x05\x00hello"), vec![MOCK_CONTAINER]);
329 }
330
331 #[test]
332 fn magic_matching_returns_empty_rather_than_guessing() {
333 let r = populated();
334 assert!(r.match_magic(b"unrecognised bytes").is_empty());
335 assert!(
336 r.match_magic(b"").is_empty(),
337 "an empty prefix matches nothing"
338 );
339 assert!(
340 r.match_magic(b"M").is_empty(),
341 "a truncated prefix must not match"
342 );
343 }
344
345 #[test]
346 fn magic_respects_the_declared_offset() {
347 const OFFSET_MAGIC: &[MagicRule] = &[MagicRule {
348 offset: 257,
349 bytes: b"ustar",
350 format: FormatId::new("tar"),
351 }];
352 let mut r = Registry::new();
353 r.register_container(
354 Arc::new(MockContainer),
355 FormatMeta::container(FormatId::new("tar"), &["tar"], OFFSET_MAGIC),
356 );
357
358 let mut buf = vec![0u8; 512];
359 buf[257..262].copy_from_slice(b"ustar");
360 assert_eq!(r.match_magic(&buf), vec![FormatId::new("tar")]);
361
362 let mut wrong = vec![0u8; 512];
364 wrong[0..5].copy_from_slice(b"ustar");
365 assert!(r.match_magic(&wrong).is_empty());
366 }
367
368 #[test]
369 fn matrix_reports_capabilities_and_is_sorted() {
370 let r = populated();
371 let m = r.matrix();
372 assert_eq!(m.len(), 2);
373 assert!(m.windows(2).all(|w| w[0].id.as_str() <= w[1].id.as_str()));
374
375 let codec_row = m.iter().find(|row| row.id == MOCK_CODEC).unwrap();
376 assert_eq!(codec_row.kind, FormatKind::Codec);
377 assert!(codec_row.read && codec_row.write);
378 assert!(!codec_row.parallel, "the mock declares no parallelism");
379 assert_eq!(codec_row.extensions, &["mk", "mock"]);
380 }
381
382 #[test]
383 fn re_registering_an_id_replaces_it_rather_than_duplicating() {
384 let mut r = populated();
385 r.register_codec(Arc::new(MockCodec), codec_meta());
386 assert_eq!(r.matrix().len(), 2);
387 }
388
389 #[test]
390 fn re_registering_with_a_different_extension_list_drops_the_old_mapping() {
391 let mut r = populated();
392 r.register_codec(
393 Arc::new(MockCodec),
394 FormatMeta::codec(MOCK_CODEC, &["mk2"], CODEC_MAGIC),
395 );
396 assert_eq!(r.by_extension("mk2"), Some(MOCK_CODEC));
397 assert_eq!(
398 r.by_extension("mk"),
399 None,
400 "a stale extension must not survive re-registration"
401 );
402 assert_eq!(r.by_extension("mock"), None);
403 }
404
405 #[test]
406 fn has_magic_is_true_for_gzip_and_false_for_an_unregistered_magicless_format() {
407 let mut r = Registry::new();
408 r.register_codec(Arc::new(MockCodec), codec_meta());
409 const NONE: FormatId = FormatId::new("nomagic");
411 r.register_codec(Arc::new(MockCodec), FormatMeta::codec(NONE, &["nm"], &[]));
412 assert!(r.has_magic(MOCK_CODEC));
413 assert!(!r.has_magic(NONE));
414 assert!(!r.has_magic(FormatId::new("unregistered")));
415 }
416
417 #[test]
418 fn match_magic_orders_by_priority_before_id() {
419 const SHARED: &[MagicRule] = &[MagicRule {
421 offset: 0,
422 bytes: b"PK\x03\x04",
423 format: FormatId::new("zip"),
424 }];
425 const SHARED2: &[MagicRule] = &[MagicRule {
426 offset: 0,
427 bytes: b"PK\x03\x04",
428 format: FormatId::new("apk"),
429 }];
430
431 let mut r = Registry::new();
432 r.register_container(
434 Arc::new(MockContainer),
435 FormatMeta::container(FormatId::new("apk"), &["apk"], SHARED2).with_priority(-10),
436 );
437 r.register_container(
438 Arc::new(MockContainer),
439 FormatMeta::container(FormatId::new("zip"), &["zip"], SHARED),
440 );
441
442 let hits = r.match_magic(b"PK\x03\x04rest");
443 assert_eq!(
444 hits.first().copied(),
445 Some(FormatId::new("zip")),
446 "the base format must outrank its derivative"
447 );
448 assert_eq!(hits.len(), 2);
449 }
450
451 #[test]
452 fn equal_priority_still_orders_deterministically_by_id() {
453 const A: &[MagicRule] = &[MagicRule {
454 offset: 0,
455 bytes: b"XX",
456 format: FormatId::new("bbb"),
457 }];
458 const B: &[MagicRule] = &[MagicRule {
459 offset: 0,
460 bytes: b"XX",
461 format: FormatId::new("aaa"),
462 }];
463 let mut r = Registry::new();
464 r.register_codec(
465 Arc::new(MockCodec),
466 FormatMeta::codec(FormatId::new("bbb"), &["b"], A),
467 );
468 r.register_codec(
469 Arc::new(MockCodec),
470 FormatMeta::codec(FormatId::new("aaa"), &["a"], B),
471 );
472 assert_eq!(
473 r.match_magic(b"XXrest"),
474 vec![FormatId::new("aaa"), FormatId::new("bbb")]
475 );
476 }
477
478 #[test]
479 fn a_decode_only_codec_is_refused_an_encoder_rather_than_reaching_it() {
480 use crate::testing::{MOCK_CODEC, MockCodec};
481
482 struct ReadOnlyCodec;
485 impl Codec for ReadOnlyCodec {
486 fn id(&self) -> FormatId {
487 MOCK_CODEC
488 }
489 fn caps(&self) -> CodecCaps {
490 CodecCaps::decode_only()
491 }
492 fn decoder<'a>(
493 &self,
494 src: Box<dyn Source + 'a>,
495 o: &DecodeOpts,
496 ) -> Result<Box<dyn Source + 'a>> {
497 MockCodec.decoder(src, o)
498 }
499 fn encoder(&self, _d: Box<dyn Write + Send>, _o: &EncodeOpts) -> Result<Box<dyn Sink>> {
500 panic!("a capability check must stop the caller before it reaches here");
509 }
510 }
511
512 let mut reg = Registry::new();
513 reg.register_codec(
514 std::sync::Arc::new(ReadOnlyCodec),
515 FormatMeta::codec(MOCK_CODEC, &["mock"], &[]),
516 );
517
518 assert!(reg.require_decoder(MOCK_CODEC).is_ok());
520
521 let Err(err) = reg.require_encoder(MOCK_CODEC) else {
530 panic!("expected CapabilityUnavailable for a decode-only codec");
531 };
532 assert!(
533 matches!(err, Error::CapabilityUnavailable { .. }),
534 "got {err:?}"
535 );
536 assert!(
537 err.to_string().contains("mock-codec"),
538 "must name the format: {err}"
539 );
540 assert_eq!(
541 err.exit_code(),
542 3,
543 "a capability refusal must be distinguishable from other errors so a script can branch on it"
544 );
545 }
546
547 #[test]
552 fn a_read_only_container_is_refused_a_writer_rather_than_reaching_it() {
553 use crate::archive::{ArchiveRead, ArchiveWrite, Container, CreateOpts, OpenOpts};
554 use crate::format::ContainerCaps;
555 use crate::ladder::Resolved;
556
557 struct ReadOnlyContainer;
561 impl Container for ReadOnlyContainer {
562 fn id(&self) -> FormatId {
563 MOCK_CONTAINER
564 }
565 fn caps(&self) -> ContainerCaps {
566 ContainerCaps::read_only()
567 }
568 fn open(&self, resolved: Resolved, o: &OpenOpts) -> Result<Box<dyn ArchiveRead>> {
569 MockContainer.open(resolved, o)
570 }
571 fn create(
572 &self,
573 _dst: Box<dyn Sink>,
574 _o: &CreateOpts,
575 ) -> Result<Box<dyn ArchiveWrite>> {
576 panic!("a capability check must stop the caller before it reaches here");
581 }
582 }
583
584 let mut reg = Registry::new();
585 reg.register_container(
586 std::sync::Arc::new(ReadOnlyContainer),
587 FormatMeta::container(MOCK_CONTAINER, &["mar"], &[]),
588 );
589
590 assert!(reg.require_container(MOCK_CONTAINER).is_ok());
592
593 let Err(err) = reg.require_container_writer(MOCK_CONTAINER) else {
594 panic!("expected CapabilityUnavailable for a read-only container");
595 };
596 assert!(
597 matches!(err, Error::CapabilityUnavailable { .. }),
598 "got {err:?}"
599 );
600 assert!(
601 err.to_string().contains("mock-container"),
602 "must name the format: {err}"
603 );
604 assert_eq!(err.exit_code(), 3);
605
606 let mut ok = Registry::new();
609 ok.register_container(Arc::new(MockContainer), container_meta());
610 assert!(ok.require_container_writer(MOCK_CONTAINER).is_ok());
611 }
612
613 #[test]
614 fn confirm_answers_only_for_a_format_that_registered_a_function() {
615 fn starts_mock(p: &[u8], whole: bool) -> bool {
616 p.starts_with(b"MOCK") && !(whole && p.len() < 8)
617 }
618 let mut r = Registry::new();
619 r.register_codec(Arc::new(MockCodec), codec_meta().with_confirm(starts_mock));
620 r.register_container(Arc::new(MockContainer), container_meta());
621 assert_eq!(r.confirm(MOCK_CODEC, b"MOCK....", false), Some(true));
622 assert_eq!(r.confirm(MOCK_CODEC, b"MOCX....", false), Some(false));
623 assert_eq!(r.confirm(MOCK_CODEC, b"MOCK", false), Some(true));
625 assert_eq!(r.confirm(MOCK_CODEC, b"MOCK", true), Some(false));
626 assert_eq!(
627 r.confirm(MOCK_CONTAINER, b"ME", true),
628 None,
629 "no function registered: no verdict, not a refusal"
630 );
631 assert_eq!(r.confirm(FormatId::new("absent"), b"", false), None);
632
633 r.register_codec(Arc::new(MockCodec), codec_meta());
635 assert_eq!(r.confirm(MOCK_CODEC, b"MOCK", false), None);
636 }
637
638 #[test]
644 fn confirm_refuses_bytes_the_formats_own_magic_does_not_match() {
645 fn anything(_: &[u8], _: bool) -> bool {
646 true
647 }
648 let mut r = Registry::new();
649 r.register_codec(Arc::new(MockCodec), codec_meta().with_confirm(anything));
650 r.register_container(Arc::new(MockContainer), container_meta());
651 assert_eq!(r.confirm(MOCK_CODEC, b"MOCK....", false), Some(true));
652 assert_eq!(
653 r.confirm(MOCK_CODEC, b"\x05\x00\x00\x00", false),
654 Some(false)
655 );
656 assert_eq!(
657 r.confirm(MOCK_CONTAINER, b"xx", false),
658 Some(false),
659 "no function, but a magic that did not match"
660 );
661
662 r.register_codec(
663 Arc::new(MockCodec),
664 codec_meta().with_confirm(anything).with_magic_optional(),
665 );
666 assert_eq!(
667 r.confirm(MOCK_CODEC, b"\x05\x00\x00\x00", false),
668 Some(true)
669 );
670 }
671}