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 priority_of(&self, id: FormatId) -> i16 {
181 self.metas.get(&id).map_or(0, |m| m.priority)
182 }
183
184 pub fn matrix(&self) -> Vec<FormatRow> {
186 let mut rows: Vec<FormatRow> = self
187 .metas
188 .values()
189 .map(|meta| {
190 let (read, write, parallel) = match meta.kind {
191 FormatKind::Codec => match self.codecs.get(&meta.id) {
192 Some(c) => {
193 let caps = c.caps();
194 (
195 caps.decode,
196 caps.encode,
197 caps.parallel_encode || caps.parallel_decode,
198 )
199 }
200 None => (false, false, false),
201 },
202 FormatKind::Container => match self.containers.get(&meta.id) {
203 Some(c) => {
204 let caps = c.caps();
205 (caps.read, caps.write, false)
206 }
207 None => (false, false, false),
208 },
209 };
210 FormatRow {
211 id: meta.id,
212 kind: meta.kind,
213 extensions: meta.extensions,
214 read,
215 write,
216 parallel,
217 }
218 })
219 .collect();
220 rows.sort_by_key(|r| r.id.as_str());
221 rows
222 }
223}
224
225#[cfg(test)]
226mod tests {
227 use super::*;
228 use crate::archive::{DecodeOpts, EncodeOpts, Sink};
229 use crate::format::CodecCaps;
230 use crate::source::Source;
231 use crate::testing::{MOCK_CODEC, MOCK_CONTAINER, MockCodec, MockContainer};
232 use std::io::Write;
233
234 const CODEC_MAGIC: &[MagicRule] = &[MagicRule {
235 offset: 0,
236 bytes: b"MOCK",
237 format: MOCK_CODEC,
238 }];
239 const CONTAINER_MAGIC: &[MagicRule] = &[MagicRule {
240 offset: 0,
241 bytes: b"ME",
242 format: MOCK_CONTAINER,
243 }];
244
245 fn codec_meta() -> FormatMeta {
246 FormatMeta::codec(MOCK_CODEC, &["mk", "mock"], CODEC_MAGIC)
247 }
248
249 fn container_meta() -> FormatMeta {
250 FormatMeta::container(MOCK_CONTAINER, &["mar"], CONTAINER_MAGIC)
251 }
252
253 fn populated() -> Registry {
254 let mut r = Registry::new();
255 r.register_codec(Arc::new(MockCodec), codec_meta());
256 r.register_container(Arc::new(MockContainer), container_meta());
257 r
258 }
259
260 #[test]
261 fn an_empty_registry_reports_no_formats() {
262 assert!(Registry::new().matrix().is_empty());
263 }
264
265 #[test]
266 fn registered_formats_are_retrievable_by_id() {
267 let r = populated();
268 assert!(r.codec(MOCK_CODEC).is_some());
269 assert!(r.container(MOCK_CONTAINER).is_some());
270 assert!(r.container(MOCK_CODEC).is_none());
272 }
273
274 #[test]
275 fn a_missing_format_is_format_not_enabled_not_a_generic_error() {
276 let r = Registry::new();
283 let Err(err) = r.require_codec(MOCK_CODEC) else {
284 panic!("expected FormatNotEnabled for an unregistered codec");
285 };
286 assert!(matches!(err, crate::Error::FormatNotEnabled(id) if id == MOCK_CODEC));
287 assert_eq!(err.exit_code(), 3);
288 }
289
290 #[test]
291 fn extension_lookup_is_case_insensitive_and_dot_tolerant() {
292 let r = populated();
293 assert_eq!(r.by_extension("mk"), Some(MOCK_CODEC));
294 assert_eq!(r.by_extension("MOCK"), Some(MOCK_CODEC));
295 assert_eq!(r.by_extension(".mar"), Some(MOCK_CONTAINER));
296 assert_eq!(r.by_extension("nope"), None);
297 }
298
299 #[test]
300 fn magic_matching_finds_the_format_from_a_prefix() {
301 let r = populated();
302 assert_eq!(r.match_magic(b"MOCKpayload"), vec![MOCK_CODEC]);
303 assert_eq!(r.match_magic(b"ME\x05\x00hello"), vec![MOCK_CONTAINER]);
304 }
305
306 #[test]
307 fn magic_matching_returns_empty_rather_than_guessing() {
308 let r = populated();
309 assert!(r.match_magic(b"unrecognised bytes").is_empty());
310 assert!(
311 r.match_magic(b"").is_empty(),
312 "an empty prefix matches nothing"
313 );
314 assert!(
315 r.match_magic(b"M").is_empty(),
316 "a truncated prefix must not match"
317 );
318 }
319
320 #[test]
321 fn magic_respects_the_declared_offset() {
322 const OFFSET_MAGIC: &[MagicRule] = &[MagicRule {
323 offset: 257,
324 bytes: b"ustar",
325 format: FormatId::new("tar"),
326 }];
327 let mut r = Registry::new();
328 r.register_container(
329 Arc::new(MockContainer),
330 FormatMeta::container(FormatId::new("tar"), &["tar"], OFFSET_MAGIC),
331 );
332
333 let mut buf = vec![0u8; 512];
334 buf[257..262].copy_from_slice(b"ustar");
335 assert_eq!(r.match_magic(&buf), vec![FormatId::new("tar")]);
336
337 let mut wrong = vec![0u8; 512];
339 wrong[0..5].copy_from_slice(b"ustar");
340 assert!(r.match_magic(&wrong).is_empty());
341 }
342
343 #[test]
344 fn matrix_reports_capabilities_and_is_sorted() {
345 let r = populated();
346 let m = r.matrix();
347 assert_eq!(m.len(), 2);
348 assert!(m.windows(2).all(|w| w[0].id.as_str() <= w[1].id.as_str()));
349
350 let codec_row = m.iter().find(|row| row.id == MOCK_CODEC).unwrap();
351 assert_eq!(codec_row.kind, FormatKind::Codec);
352 assert!(codec_row.read && codec_row.write);
353 assert!(!codec_row.parallel, "the mock declares no parallelism");
354 assert_eq!(codec_row.extensions, &["mk", "mock"]);
355 }
356
357 #[test]
358 fn re_registering_an_id_replaces_it_rather_than_duplicating() {
359 let mut r = populated();
360 r.register_codec(Arc::new(MockCodec), codec_meta());
361 assert_eq!(r.matrix().len(), 2);
362 }
363
364 #[test]
365 fn re_registering_with_a_different_extension_list_drops_the_old_mapping() {
366 let mut r = populated();
367 r.register_codec(
368 Arc::new(MockCodec),
369 FormatMeta::codec(MOCK_CODEC, &["mk2"], CODEC_MAGIC),
370 );
371 assert_eq!(r.by_extension("mk2"), Some(MOCK_CODEC));
372 assert_eq!(
373 r.by_extension("mk"),
374 None,
375 "a stale extension must not survive re-registration"
376 );
377 assert_eq!(r.by_extension("mock"), None);
378 }
379
380 #[test]
381 fn has_magic_is_true_for_gzip_and_false_for_an_unregistered_magicless_format() {
382 let mut r = Registry::new();
383 r.register_codec(Arc::new(MockCodec), codec_meta());
384 const NONE: FormatId = FormatId::new("nomagic");
386 r.register_codec(Arc::new(MockCodec), FormatMeta::codec(NONE, &["nm"], &[]));
387 assert!(r.has_magic(MOCK_CODEC));
388 assert!(!r.has_magic(NONE));
389 assert!(!r.has_magic(FormatId::new("unregistered")));
390 }
391
392 #[test]
393 fn match_magic_orders_by_priority_before_id() {
394 const SHARED: &[MagicRule] = &[MagicRule {
396 offset: 0,
397 bytes: b"PK\x03\x04",
398 format: FormatId::new("zip"),
399 }];
400 const SHARED2: &[MagicRule] = &[MagicRule {
401 offset: 0,
402 bytes: b"PK\x03\x04",
403 format: FormatId::new("apk"),
404 }];
405
406 let mut r = Registry::new();
407 r.register_container(
409 Arc::new(MockContainer),
410 FormatMeta::container(FormatId::new("apk"), &["apk"], SHARED2).with_priority(-10),
411 );
412 r.register_container(
413 Arc::new(MockContainer),
414 FormatMeta::container(FormatId::new("zip"), &["zip"], SHARED),
415 );
416
417 let hits = r.match_magic(b"PK\x03\x04rest");
418 assert_eq!(
419 hits.first().copied(),
420 Some(FormatId::new("zip")),
421 "the base format must outrank its derivative"
422 );
423 assert_eq!(hits.len(), 2);
424 }
425
426 #[test]
427 fn equal_priority_still_orders_deterministically_by_id() {
428 const A: &[MagicRule] = &[MagicRule {
429 offset: 0,
430 bytes: b"XX",
431 format: FormatId::new("bbb"),
432 }];
433 const B: &[MagicRule] = &[MagicRule {
434 offset: 0,
435 bytes: b"XX",
436 format: FormatId::new("aaa"),
437 }];
438 let mut r = Registry::new();
439 r.register_codec(
440 Arc::new(MockCodec),
441 FormatMeta::codec(FormatId::new("bbb"), &["b"], A),
442 );
443 r.register_codec(
444 Arc::new(MockCodec),
445 FormatMeta::codec(FormatId::new("aaa"), &["a"], B),
446 );
447 assert_eq!(
448 r.match_magic(b"XXrest"),
449 vec![FormatId::new("aaa"), FormatId::new("bbb")]
450 );
451 }
452
453 #[test]
454 fn a_decode_only_codec_is_refused_an_encoder_rather_than_reaching_it() {
455 use crate::testing::{MOCK_CODEC, MockCodec};
456
457 struct ReadOnlyCodec;
460 impl Codec for ReadOnlyCodec {
461 fn id(&self) -> FormatId {
462 MOCK_CODEC
463 }
464 fn caps(&self) -> CodecCaps {
465 CodecCaps::decode_only()
466 }
467 fn decoder(&self, src: Box<dyn Source>, o: &DecodeOpts) -> Result<Box<dyn Source>> {
468 MockCodec.decoder(src, o)
469 }
470 fn encoder(&self, _d: Box<dyn Write + Send>, _o: &EncodeOpts) -> Result<Box<dyn Sink>> {
471 panic!("a capability check must stop the caller before it reaches here");
480 }
481 }
482
483 let mut reg = Registry::new();
484 reg.register_codec(
485 std::sync::Arc::new(ReadOnlyCodec),
486 FormatMeta::codec(MOCK_CODEC, &["mock"], &[]),
487 );
488
489 assert!(reg.require_decoder(MOCK_CODEC).is_ok());
491
492 let Err(err) = reg.require_encoder(MOCK_CODEC) else {
501 panic!("expected CapabilityUnavailable for a decode-only codec");
502 };
503 assert!(
504 matches!(err, Error::CapabilityUnavailable { .. }),
505 "got {err:?}"
506 );
507 assert!(
508 err.to_string().contains("mock-codec"),
509 "must name the format: {err}"
510 );
511 assert_eq!(
512 err.exit_code(),
513 3,
514 "a capability refusal must be distinguishable from other errors so a script can branch on it"
515 );
516 }
517
518 #[test]
523 fn a_read_only_container_is_refused_a_writer_rather_than_reaching_it() {
524 use crate::archive::{ArchiveRead, ArchiveWrite, Container, CreateOpts, OpenOpts};
525 use crate::format::ContainerCaps;
526 use crate::ladder::Resolved;
527
528 struct ReadOnlyContainer;
532 impl Container for ReadOnlyContainer {
533 fn id(&self) -> FormatId {
534 MOCK_CONTAINER
535 }
536 fn caps(&self) -> ContainerCaps {
537 ContainerCaps::read_only()
538 }
539 fn open(&self, resolved: Resolved, o: &OpenOpts) -> Result<Box<dyn ArchiveRead>> {
540 MockContainer.open(resolved, o)
541 }
542 fn create(
543 &self,
544 _dst: Box<dyn Sink>,
545 _o: &CreateOpts,
546 ) -> Result<Box<dyn ArchiveWrite>> {
547 panic!("a capability check must stop the caller before it reaches here");
552 }
553 }
554
555 let mut reg = Registry::new();
556 reg.register_container(
557 std::sync::Arc::new(ReadOnlyContainer),
558 FormatMeta::container(MOCK_CONTAINER, &["mar"], &[]),
559 );
560
561 assert!(reg.require_container(MOCK_CONTAINER).is_ok());
563
564 let Err(err) = reg.require_container_writer(MOCK_CONTAINER) else {
565 panic!("expected CapabilityUnavailable for a read-only container");
566 };
567 assert!(
568 matches!(err, Error::CapabilityUnavailable { .. }),
569 "got {err:?}"
570 );
571 assert!(
572 err.to_string().contains("mock-container"),
573 "must name the format: {err}"
574 );
575 assert_eq!(err.exit_code(), 3);
576
577 let mut ok = Registry::new();
580 ok.register_container(Arc::new(MockContainer), container_meta());
581 assert!(ok.require_container_writer(MOCK_CONTAINER).is_ok());
582 }
583}