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