Skip to main content

bevy_dlc/
asset_loader.rs

1use bevy::asset::io::Reader;
2use bevy::asset::{
3    Asset, AssetLoader, AssetPath, ErasedLoadedAsset, Handle, LoadContext, LoadedUntypedAsset,
4};
5use bevy::ecs::reflect::AppTypeRegistry;
6use bevy::prelude::*;
7use bevy::reflect::TypePath;
8use futures_lite::{AsyncReadExt, AsyncSeekExt};
9use std::io;
10use std::sync::Arc;
11use thiserror::Error;
12
13use crate::{DlcId, PackItem, Product};
14
15use std::io::{Read, Seek, SeekFrom};
16
17/// Adapter that exposes a `std::io::Read + std::io::Seek` view over a
18/// `bevy::asset::io::Reader`.  This is used by pack-parsing routines so we can
19/// operate on the async reader without copying the entire file into memory.
20///
21/// The implementation simply blocks on the underlying async methods using
22/// [`pollster::block_on`].  Seeking works only when the wrapped reader is
23/// seekable; otherwise the `seek()` call returns an error.
24pub struct SyncReader<'a> {
25    inner: &'a mut dyn bevy::asset::io::Reader,
26}
27
28impl<'a> SyncReader<'a> {
29    pub fn new(inner: &'a mut dyn bevy::asset::io::Reader) -> Self {
30        SyncReader { inner }
31    }
32}
33
34impl<'a> Read for SyncReader<'a> {
35    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
36        bevy::tasks::block_on(self.inner.read(buf))
37            .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
38    }
39}
40
41impl<'a> Seek for SyncReader<'a> {
42    fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
43        match self.inner.seekable() {
44            Ok(seek) => bevy::tasks::block_on(seek.seek(pos))
45                .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e)),
46            Err(_) => Err(std::io::Error::new(
47                std::io::ErrorKind::Other,
48                "reader not seekable",
49            )),
50        }
51    }
52}
53
54/// Decompress a gzip‑compressed tar archive from `plaintext` and return a map from
55/// internal path -> contents.  Errors are mapped into `DlcLoaderError::DecryptionFailed`
56/// because the only callers are the pack loader and entry decrypt which already treat
57/// archive failures as decryption problems.
58fn decompress_archive(
59    plaintext: &[u8],
60) -> Result<std::collections::HashMap<String, Vec<u8>>, DlcLoaderError> {
61    use flate2::read::GzDecoder;
62    use std::io::Read;
63    use tar::Archive;
64
65    let mut archive = Archive::new(GzDecoder::new(std::io::Cursor::new(plaintext)));
66    let mut map = std::collections::HashMap::new();
67    for entry in archive
68        .entries()
69        .map_err(|e| DlcLoaderError::DecryptionFailed(format!("archive read failed: {}", e)))?
70    {
71        let mut file = entry.map_err(|e| {
72            DlcLoaderError::DecryptionFailed(format!("archive entry read failed: {}", e))
73        })?;
74        let path = file
75            .path()
76            .map_err(|e| DlcLoaderError::DecryptionFailed(format!("archive path error: {}", e)))?;
77        let path_str = path.to_string_lossy().replace("\\", "/");
78        let mut buf = Vec::new();
79        file.read_to_end(&mut buf).map_err(|e| {
80            DlcLoaderError::DecryptionFailed(format!("archive file read failed: {}", e))
81        })?;
82        map.insert(path_str, buf);
83    }
84    Ok(map)
85}
86
87/// Internal helper to decrypt a specific entry from a v4 pack by reading only
88/// the required data block from the file.
89pub(crate) fn decrypt_pack_entry_block_bytes<R: std::io::Read + std::io::Seek>(
90    reader: &mut R,
91    enc: &EncryptedAsset,
92    key: &crate::EncryptionKey,
93    full_path: &str,
94) -> Result<Vec<u8>, DlcLoaderError> {
95    // 1. Re-parse the pack to get the block metadata
96    let _original_pos = reader
97        .stream_position()
98        .map_err(|e| DlcLoaderError::Io(e))?;
99    reader
100        .seek(std::io::SeekFrom::Start(0))
101        .map_err(|e| DlcLoaderError::Io(e))?;
102
103    let (_prod, _id, _ver, _entries, blocks) = crate::parse_encrypted_pack(&mut *reader)
104        .map_err(|e| DlcLoaderError::InvalidFormat(e.to_string()))?;
105
106    let block = blocks
107        .iter()
108        .find(|b| b.block_id == enc.block_id)
109        .ok_or_else(|| {
110            DlcLoaderError::DecryptionFailed(format!("block {} not found in pack", enc.block_id))
111        })?;
112
113    // 2. Decrypt only the bytes corresponding to the desired block.  The
114    // pack format writes all blocks concatenated after the metadata, and
115    // `parse_encrypted_pack` leaves the reader positioned at the start of the
116    // ciphertext region.  We still seek explicitly to the recorded offset to
117    // be robust and to support callers that may have moved the reader.
118    reader
119        .seek(std::io::SeekFrom::Start(block.file_offset))
120        .map_err(|e| DlcLoaderError::Io(e))?;
121
122    // limit the reader to the block size so `decrypt_with_key` doesn't read
123    // past the boundary when multiple blocks exist.
124    let mut limited = reader.take(block.encrypted_size as u64);
125    let pt_gz = crate::pack_format::decrypt_with_key(&key, &mut limited, &block.nonce)
126        .map_err(|e| DlcLoaderError::DecryptionFailed(e.to_string()))?;
127
128    // 3. Decompress and find entry
129    let entries = decompress_archive(&pt_gz)?;
130
131    // Extract label from "pack.dlcpack#label"
132    let label = match full_path.rsplit_once('#') {
133        Some((_, suffix)) => suffix,
134        None => full_path,
135    }
136    .replace("\\", "/");
137
138    entries.get(&label).cloned().ok_or_else(|| {
139        DlcLoaderError::DecryptionFailed(format!(
140            "entry '{}' not found in decrypted block {}",
141            label, enc.block_id
142        ))
143    })
144}
145
146/// Event fired when a DLC pack is successfully loaded.
147#[derive(Event, Clone)]
148pub struct DlcPackLoaded {
149    dlc_id: DlcId,
150    pack: DlcPack,
151}
152
153impl DlcPackLoaded {
154    pub(crate) fn new(dlc_id: DlcId, pack: DlcPack) -> Self {
155        DlcPackLoaded { dlc_id, pack }
156    }
157
158    /// Return the DLC identifier for the loaded pack.
159    pub fn id(&self) -> &DlcId {
160        &self.dlc_id
161    }
162
163    /// Return a reference to the loaded `DlcPack`. The pack contains metadata about the DLC and provides methods to decrypt and load individual entries.
164    pub fn pack(&self) -> &DlcPack {
165        &self.pack
166    }
167}
168
169/// Fuzzy match for type paths, normalizing by trimming leading "::" to handle absolute vs relative paths.
170/// Also handles crate name differences by allowing suffix matches.
171pub(crate) fn fuzzy_type_path_match<'a>(stored: &'a str, expected: &'a str) -> bool {
172    let s = stored.trim_start_matches("::");
173    let e = expected.trim_start_matches("::");
174
175    if s == e {
176        return true;
177    }
178
179    // Allow suffix matching to handle differences in crate names (e.g. "my_crate::MyType" vs "MyType")
180    // or when one path is more specific than the other.
181    if e.ends_with(s) && e.as_bytes().get(e.len() - s.len() - 1) == Some(&b':') {
182        return true;
183    }
184
185    if s.ends_with(e) && s.as_bytes().get(s.len() - e.len() - 1) == Some(&b':') {
186        return true;
187    }
188
189    false
190}
191
192/// Attempts to downcast an `ErasedLoadedAsset` to `A` and, if successful,
193/// registers it as a labeled sub-asset in `load_context`.
194///
195/// Returns `true` when the asset was successfully registered.
196pub trait ErasedSubAssetRegistrar: Send + Sync + 'static {
197    fn try_register(
198        &self,
199        label: String,
200        erased: ErasedLoadedAsset,
201        load_context: &mut LoadContext<'_>,
202    ) -> Result<(), ErasedLoadedAsset>;
203
204    /// Return the `TypePath` of the asset type this registrar handles.
205    fn asset_type_path(&self) -> &'static str;
206
207    /// Attempt to load the asset directly using its static type, bypassing
208    /// extension dispatch. This is used when a `type_path` is provided by
209    /// the container.
210    fn load_direct<'a>(
211        &'a self,
212        label: String,
213        fake_path: String,
214        reader: &'a mut dyn Reader,
215        load_context: &'a mut LoadContext<'_>,
216    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), DlcLoaderError>> + Send + 'a>>;
217}
218
219/// Concrete implementation for asset type `A`.
220pub struct TypedSubAssetRegistrar<A: Asset>(std::marker::PhantomData<A>);
221
222impl<A: Asset> Default for TypedSubAssetRegistrar<A> {
223    fn default() -> Self {
224        Self(std::marker::PhantomData)
225    }
226}
227
228impl<A: Asset> ErasedSubAssetRegistrar for TypedSubAssetRegistrar<A> {
229    fn try_register(
230        &self,
231        label: String,
232        erased: ErasedLoadedAsset,
233        load_context: &mut LoadContext<'_>,
234    ) -> Result<(), ErasedLoadedAsset> {
235        match erased.downcast::<A>() {
236            Ok(loaded) => {
237                load_context.add_loaded_labeled_asset(label, loaded);
238                Ok(())
239            }
240            Err(back) => Err(back),
241        }
242    }
243
244    fn asset_type_path(&self) -> &'static str {
245        A::type_path()
246    }
247
248    fn load_direct<'a>(
249        &'a self,
250        label: String,
251        fake_path: String,
252        reader: &'a mut dyn Reader,
253        load_context: &'a mut LoadContext<'_>,
254    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), DlcLoaderError>> + Send + 'a>>
255    {
256        Box::pin(async move {
257            match load_context
258                .loader()
259                .with_static_type()
260                .immediate()
261                .with_reader(reader)
262                .load::<A>(fake_path)
263                .await
264            {
265                Ok(loaded) => {
266                    load_context.add_loaded_labeled_asset(label, loaded);
267                    Ok(())
268                }
269                Err(e) => Err(DlcLoaderError::DecryptionFailed(e.to_string())),
270            }
271        })
272    }
273}
274
275/// Represents a single encrypted file inside a `.dlcpack` container, along with its metadata (DLC ID, original extension, optional type path). The ciphertext is not decrypted at this stage; decryption is performed on demand by `DlcPackEntry::decrypt_bytes` using the global encrypt key registry.
276#[derive(Clone, Debug, Asset, TypePath)]
277pub struct EncryptedAsset {
278    pub dlc_id: String,
279    pub original_extension: String,
280    /// Optional serialized type identifier (e.g. `bevy::image::Image`)
281    pub type_path: Option<String>,
282    pub nonce: [u8; 12],
283    pub ciphertext: std::sync::Arc<[u8]>,
284    // --- v4 format extensions ---
285    pub block_id: u32,
286    pub block_offset: u32,
287    pub size: u32,
288}
289
290impl EncryptedAsset {
291    /// Decrypt the ciphertext contained in this `EncryptedAsset`, using the
292    /// global encryption-key registry to look up the correct key for
293    /// `self.dlc_id`.
294    pub(crate) fn decrypt_bytes(&self) -> Result<Vec<u8>, DlcLoaderError> {
295        // lookup key
296        let encrypt_key = crate::encrypt_key_registry::get(&self.dlc_id)
297            .ok_or_else(|| DlcLoaderError::DlcLocked(self.dlc_id.clone()))?;
298
299        // decrypt using the pack-format helper
300        crate::pack_format::decrypt_with_key(
301            &encrypt_key,
302            std::io::Cursor::new(&*self.ciphertext),
303            &self.nonce,
304        )
305        .map_err(|e| DlcLoaderError::DecryptionFailed(e.to_string()))
306    }
307}
308
309/// Parse the binary encrypted-asset format from a byte slice. This is used by the pack loader when parsing the pack metadata, and also by the `DlcLoader` when decrypting individual entries (since the entry metadata is stored in the same format as a standalone encrypted file).
310pub fn parse_encrypted(bytes: &[u8]) -> Result<EncryptedAsset, io::Error> {
311    // make sure we can read the fixed-size header fields without panicking:
312    // version (1 byte) + dlc_len (2 bytes) + ext_len (1 byte) + nonce (12 bytes)
313    // the remaining lengths (dlc_id, ext, type_path, ciphertext) are variable
314    // and validated later, so this check only guards the very earliest reads.
315    if bytes.len() < 1 + 2 + 1 + 12 {
316        return Err(io::Error::new(
317            io::ErrorKind::InvalidData,
318            "encrypted file too small",
319        ));
320    }
321    let version = bytes[0];
322    let mut offset = 1usize;
323
324    let dlc_len = u16::from_be_bytes([bytes[offset], bytes[offset + 1]]) as usize;
325    offset += 2;
326    if offset + dlc_len > bytes.len() {
327        return Err(io::Error::new(
328            io::ErrorKind::InvalidData,
329            "invalid dlc id length",
330        ));
331    }
332    let dlc_id = String::from_utf8(bytes[offset..offset + dlc_len].to_vec())
333        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
334    offset += dlc_len;
335
336    let ext_len = bytes[offset] as usize;
337    offset += 1;
338    let original_extension = if ext_len == 0 {
339        "".to_string()
340    } else {
341        let s = String::from_utf8(bytes[offset..offset + ext_len].to_vec())
342            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
343        offset += ext_len;
344        s
345    };
346
347    // version 1+ stores a serialized type identifier (u16 length + utf8 bytes)
348    let type_path = if version >= 1 {
349        if offset + 2 > bytes.len() {
350            return Err(io::Error::new(
351                io::ErrorKind::InvalidData,
352                "missing type_path length",
353            ));
354        }
355        let tlen = u16::from_be_bytes([bytes[offset], bytes[offset + 1]]) as usize;
356        offset += 2;
357        if tlen == 0 {
358            None
359        } else {
360            if offset + tlen > bytes.len() {
361                return Err(io::Error::new(
362                    io::ErrorKind::InvalidData,
363                    "invalid type_path length",
364                ));
365            }
366            let s = String::from_utf8(bytes[offset..offset + tlen].to_vec())
367                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
368            offset += tlen;
369            Some(s)
370        }
371    } else {
372        None
373    };
374
375    if offset + 12 > bytes.len() {
376        return Err(io::Error::new(io::ErrorKind::InvalidData, "missing nonce"));
377    }
378    let mut nonce = [0u8; 12];
379    nonce.copy_from_slice(&bytes[offset..offset + 12]);
380    offset += 12;
381    let ciphertext = bytes[offset..].into();
382
383    Ok(EncryptedAsset {
384        dlc_id,
385        original_extension,
386        type_path,
387        nonce,
388        ciphertext,
389        block_id: 0,
390        block_offset: 0,
391        size: 0,
392    })
393}
394
395/// A loader for individual encrypted files inside a `.dlcpack`.
396#[derive(TypePath)]
397pub struct DlcLoader<A: bevy::asset::Asset + 'static> {
398    /// Stored for potential future use (e.g., validating type_path matches `A`).
399    /// Currently unused because the generic `A` already specifies the target type.
400    #[allow(dead_code)]
401    type_registry: Arc<AppTypeRegistry>,
402    _marker: std::marker::PhantomData<A>,
403}
404
405// Provide FromWorld so the loader can be initialized by Bevy's App without
406// requiring `A: Default`.
407impl<A> bevy::prelude::FromWorld for DlcLoader<A>
408where
409    A: bevy::asset::Asset + 'static,
410{
411    fn from_world(world: &mut bevy::prelude::World) -> Self {
412        let registry = world.resource::<AppTypeRegistry>().clone();
413        DlcLoader {
414            type_registry: Arc::new(registry),
415            _marker: std::marker::PhantomData,
416        }
417    }
418}
419
420#[derive(TypePath, Clone, Debug)]
421pub struct DlcPackEntry {
422    /// Relative path inside the pack (as authored when packing)
423    path: String,
424    /// Encrypted asset metadata and ciphertext
425    encrypted: EncryptedAsset,
426}
427
428impl DlcPackEntry {
429    pub fn new(path: String, encrypted: EncryptedAsset) -> Self {
430        DlcPackEntry { path, encrypted }
431    }
432
433    /// Convenience: load this entry's registered path via `AssetServer::load`.
434    pub fn load_untyped(
435        &self,
436        asset_server: &bevy::prelude::AssetServer,
437    ) -> Handle<LoadedUntypedAsset> {
438        asset_server.load_untyped(&self.path)
439    }
440
441    /// Check if this entry is declared as type `A` (via the optional `type_path` in the container, which is independent of file extension). This is not used by the loader itself (which relies on extension-based dispatch) but can be used by user code to inspect entries or implement custom loading behavior.
442    pub fn is_type<A: Asset>(&self) -> bool {
443        match self.encrypted.type_path.as_ref() {
444            Some(tp) => fuzzy_type_path_match(tp, A::type_path()),
445            None => false,
446        }
447    }
448
449    /// Decrypt and return the plaintext bytes for this entry.
450    /// This consults the global encrypt-key registry and will return
451    /// `DlcLoaderError::DlcLocked` when the encrypt key is not present.
452    pub(crate) fn decrypt_bytes(&self) -> Result<Vec<u8>, DlcLoaderError> {
453        let entry_ek = crate::encrypt_key_registry::get_full(&self.encrypted.dlc_id)
454            .ok_or_else(|| DlcLoaderError::DlcLocked(self.encrypted.dlc_id.clone()))?;
455        let encrypt_key = entry_ek.key;
456
457        // v4 block‑based decryption only; older formats are no longer supported.
458        let path = entry_ek.path.ok_or_else(|| {
459            DlcLoaderError::DecryptionFailed(format!(
460                "no file path registered for DLC '{}', cannot decrypt",
461                self.encrypted.dlc_id
462            ))
463        })?;
464
465        let mut file = std::fs::File::open(path).map_err(|e| {
466            DlcLoaderError::DecryptionFailed(format!("failed to open pack file: {}", e))
467        })?;
468
469        crate::asset_loader::decrypt_pack_entry_block_bytes(
470            &mut file,
471            &self.encrypted,
472            &encrypt_key,
473            &self.path,
474        )
475    }
476
477    pub fn path(&self) -> AssetPath<'_> {
478        bevy::asset::AssetPath::parse(&self.path)
479    }
480
481    /// Get the raw entry path (relative path within the pack, without pack prefix)
482    pub fn entry_path(&self) -> &str {
483        &self.path
484    }
485
486    pub fn original_extension(&self) -> &String {
487        &self.encrypted.original_extension
488    }
489
490    pub fn type_path(&self) -> Option<&String> {
491        self.encrypted.type_path.as_ref()
492    }
493}
494
495impl From<(String, EncryptedAsset)> for DlcPackEntry {
496    fn from((path, encrypted): (String, EncryptedAsset)) -> Self {
497        DlcPackEntry { path, encrypted }
498    }
499}
500
501impl From<&(String, EncryptedAsset)> for DlcPackEntry {
502    fn from((path, encrypted): &(String, EncryptedAsset)) -> Self {
503        DlcPackEntry {
504            path: path.clone(),
505            encrypted: encrypted.clone(),
506        }
507    }
508}
509
510/// Represents a `.dlcpack` bundle (multiple encrypted entries).
511#[derive(Asset, TypePath, Clone, Debug)]
512pub struct DlcPack {
513    dlc_id: DlcId,
514    product: Product,
515    version: u8,
516    entries: Vec<DlcPackEntry>,
517
518    /// The file path from which this pack was loaded (for registry purposes)
519    pack_path: String,
520}
521
522impl DlcPack {
523    pub fn new(id: DlcId, product: Product, version: u8, entries: Vec<DlcPackEntry>) -> Self {
524        DlcPack {
525            dlc_id: id,
526            product,
527            version,
528            entries,
529            pack_path: String::new(),
530        }
531    }
532
533    /// Get the pack file path
534    pub fn pack_path(&self) -> &str {
535        &self.pack_path
536    }
537
538    /// Return the DLC identifier for this pack.
539    pub fn id(&self) -> &DlcId {
540        &self.dlc_id
541    }
542
543    /// Return the product name this pack belongs to.
544    pub fn product(&self) -> &str {
545        &self.product.0
546    }
547
548    /// Return the pack format version.
549    pub fn version(&self) -> u8 {
550        self.version
551    }
552
553    /// Return a slice of contained entries.
554    pub fn entries(&self) -> &[DlcPackEntry] {
555        &self.entries
556    }
557
558    /// Find an entry by its registered path
559    pub fn find_entry(&self, path: &str) -> Option<&DlcPackEntry> {
560        self.entries
561            .iter()
562            .find(|e| e.path().to_string().ends_with(path) || e.path().path().ends_with(path))
563    }
564
565    /// Find all entries that match the specified asset type `A`.
566    pub fn find_by_type<A: Asset>(&self) -> Vec<&DlcPackEntry> {
567        self.entries
568            .iter()
569            .filter(|e| match e.type_path() {
570                Some(tp) => fuzzy_type_path_match(tp, A::type_path()),
571                None => false,
572            })
573            .collect()
574    }
575
576    /// Decrypt an entry (accepts either `name` or `packfile.dlcpack#name`).
577    /// Returns plaintext or `DlcLocked`.
578    pub fn decrypt_entry(
579        &self,
580        entry_path: &str,
581    ) -> Result<Vec<u8>, crate::asset_loader::DlcLoaderError> {
582        // accept either "test.png" or "packfile.dlcpack#test.png" by
583        // checking both relative and absolute paths
584        let entry = self.find_entry(entry_path).ok_or_else(|| {
585            DlcLoaderError::DecryptionFailed(format!("entry not found: {}", entry_path))
586        })?;
587
588        entry.decrypt_bytes()
589    }
590
591    pub fn load<A: Asset>(
592        &self,
593        asset_server: &bevy::prelude::AssetServer,
594        entry_path: &str,
595    ) -> Handle<A> {
596        let entry = match self.find_entry(entry_path) {
597            Some(e) => e,
598            None => panic!("entry not found: {}", entry_path),
599        };
600        asset_server.load::<A>(entry.path())
601    }
602
603    fn with_path(&self, path_string: String) -> DlcPack {
604        DlcPack {
605            dlc_id: self.dlc_id.clone(),
606            product: self.product.clone(),
607            version: self.version,
608            entries: self.entries.clone(),
609            pack_path: path_string,
610        }
611    }
612}
613
614/// Settings for `DlcPackLoader` that control asset registration behavior.
615#[derive(Clone, TypePath, serde::Serialize, serde::Deserialize)]
616pub struct DlcPackLoaderSettings {}
617
618impl Default for DlcPackLoaderSettings {
619    fn default() -> Self {
620        DlcPackLoaderSettings {}
621    }
622}
623
624/// `AssetLoader` for `.dlcpack` bundles (contains multiple encrypted entries).
625///
626/// When the encrypt key is available in the registry at load time (i.e. the
627/// DLC is already unlocked), each entry is immediately decrypted and registered
628/// as a typed labeled sub-asset so `asset_server.load("pack.dlcpack#entry.png")`
629/// returns the correct `Handle<T>` for the extension's loader (e.g. `Handle<Image>`
630/// for `.png`). No asset type is hardcoded here — the correct type is determined
631/// purely by extension dispatch via Bevy's `immediate()` loader, and the result
632/// is downcast + registered using the list of `ErasedSubAssetRegistrar`s that
633/// `DlcPlugin::build` populates.
634///
635/// When the encrypt key is *not* yet available the pack is still loaded
636/// successfully (entries list is populated from the manifest) but the labeled
637/// sub-assets are not added — `reload_assets_on_unlock_system` will reload the
638/// pack once the key arrives, and the second load will succeed.
639#[derive(TypePath, Default)]
640pub struct DlcPackLoader {
641    /// Ordered list of per-type registrars. `DlcPlugin::build` pushes one
642    /// `TypedSubAssetRegistrar<A>` for every `A` it also registers via
643    /// `init_asset_loader::<DlcLoader<A>>()`. The loader tries each in turn
644    /// and uses the first successful downcast.
645    pub registrars: Vec<Box<dyn ErasedSubAssetRegistrar>>,
646    /// Optional shared reference to the `DlcPackRegistrarFactories` resource so
647    /// the loader can observe updates to the factory list at runtime without
648    /// requiring the asset loader to be re-registered.
649    pub(crate) factories: Option<DlcPackRegistrarFactories>,
650}
651
652/// Factory trait used to create `ErasedSubAssetRegistrar` instance.
653///
654/// Implement `TypedRegistrarFactory<T>` for asset types to produce a
655/// `TypedSubAssetRegistrar::<T>` at collection time.
656pub trait DlcPackRegistrarFactory: Send + Sync + 'static {
657    fn type_name(&self) -> &'static str;
658    fn create_registrar(&self) -> Box<dyn ErasedSubAssetRegistrar>;
659}
660
661/// Generic typed factory that constructs `TypedSubAssetRegistrar::<T>`.
662pub struct TypedRegistrarFactory<T: Asset + 'static>(std::marker::PhantomData<T>);
663
664impl<T: Asset + TypePath + 'static> DlcPackRegistrarFactory for TypedRegistrarFactory<T> {
665    fn type_name(&self) -> &'static str {
666        T::type_path()
667    }
668
669    fn create_registrar(&self) -> Box<dyn ErasedSubAssetRegistrar> {
670        Box::new(TypedSubAssetRegistrar::<T>::default())
671    }
672}
673
674impl<T: Asset + 'static> Default for TypedRegistrarFactory<T> {
675    fn default() -> Self {
676        TypedRegistrarFactory(std::marker::PhantomData)
677    }
678}
679
680use std::sync::RwLock;
681
682/// Internal factory resource used by `AppExt::register_dlc_type` so user code
683/// can request additional pack-registrars without pushing closures.
684///
685/// The resource wraps an `Arc<RwLock<_>>` so the registered `DlcPackLoader`
686/// instance can hold a cheap clone and observe updates made by
687/// `register_dlc_type(...)` without needing to re-register the loader.
688#[derive(Clone, Resource)]
689pub(crate) struct DlcPackRegistrarFactories(pub Arc<RwLock<Vec<Box<dyn DlcPackRegistrarFactory>>>>);
690
691impl Default for DlcPackRegistrarFactories {
692    fn default() -> Self {
693        DlcPackRegistrarFactories(Arc::new(RwLock::new(Vec::new())))
694    }
695}
696
697/// Return the default set of pack registrar factories used by `DlcPlugin`.
698///
699/// Using factory objects avoids closures and makes it trivial to add custom
700/// typed factories in user code (box a `TypedRegistrarFactory::<T>`).
701pub(crate) fn default_pack_registrar_factories() -> Vec<Box<dyn DlcPackRegistrarFactory>> {
702    vec![
703        Box::new(TypedRegistrarFactory::<Image>::default()),
704        Box::new(TypedRegistrarFactory::<Scene>::default()),
705        Box::new(TypedRegistrarFactory::<bevy::mesh::Mesh>::default()),
706        Box::new(TypedRegistrarFactory::<Font>::default()),
707        Box::new(TypedRegistrarFactory::<AudioSource>::default()),
708        Box::new(TypedRegistrarFactory::<ColorMaterial>::default()),
709        Box::new(TypedRegistrarFactory::<bevy::pbr::StandardMaterial>::default()),
710        Box::new(TypedRegistrarFactory::<bevy::gltf::Gltf>::default()),
711        Box::new(TypedRegistrarFactory::<bevy::gltf::GltfMesh>::default()),
712        Box::new(TypedRegistrarFactory::<Shader>::default()),
713        Box::new(TypedRegistrarFactory::<DynamicScene>::default()),
714        Box::new(TypedRegistrarFactory::<AnimationClip>::default()),
715        Box::new(TypedRegistrarFactory::<AnimationGraph>::default()),
716    ]
717}
718
719/// Build the final `registrars` vector by combining factory objects supplied via
720/// the `DlcPackRegistrarFactories` resource with the crate's default factories.
721pub(crate) fn collect_pack_registrars(
722    factories: Option<&DlcPackRegistrarFactories>,
723) -> Vec<Box<dyn ErasedSubAssetRegistrar>> {
724    use std::collections::HashSet;
725    let mut seen: HashSet<&'static str> = HashSet::new();
726    let mut out: Vec<Box<dyn ErasedSubAssetRegistrar>> = Vec::new();
727
728    if let Some(f) = factories {
729        let inner = f.0.read().unwrap();
730        for factory in inner.iter() {
731            out.push(factory.create_registrar());
732            seen.insert(factory.type_name());
733        }
734    }
735
736    for factory in default_pack_registrar_factories() {
737        if !seen.contains(factory.type_name()) {
738            out.push(factory.create_registrar());
739            seen.insert(factory.type_name());
740        }
741    }
742
743    out
744}
745
746impl AssetLoader for DlcPackLoader {
747    type Asset = DlcPack;
748    type Settings = DlcPackLoaderSettings;
749    type Error = DlcLoaderError;
750
751    fn extensions(&self) -> &[&str] {
752        &["dlcpack"]
753    }
754
755    async fn load(
756        &self,
757        reader: &mut dyn Reader,
758        _settings: &Self::Settings,
759        load_context: &mut LoadContext<'_>,
760    ) -> Result<Self::Asset, Self::Error> {
761        let path_string = load_context.path().path().to_string_lossy().to_string();
762
763        // Adapt the async `reader` to a synchronous `std::io::Read` so we can
764        // drive the existing pack‑parsing logic without buffering the whole file
765        // up front.  If the underlying reader is seekable we will also be able
766        // rewind it later in order to extract the raw bytes needed for
767        // decryption.
768        let mut sync_reader = SyncReader::new(reader);
769
770        let (product, dlc_id, version, manifest_entries, _block_metadatas) =
771            crate::parse_encrypted_pack(&mut sync_reader)
772                .map_err(|e| DlcLoaderError::InvalidFormat(e.to_string()))?;
773
774        // rewind the reader back to the start so decryption routines can re‑parse the file as needed.  If the reader is not seekable, error.
775        sync_reader.seek(SeekFrom::Start(0)).map_err(|_| {
776            DlcLoaderError::Io(io::Error::new(
777                io::ErrorKind::NotSeekable,
778                format!("reader not seekable, cannot load pack '{}'", path_string),
779            ))
780        })?;
781
782        // Check for DLC ID conflicts: reject if a DIFFERENT pack file is being loaded for the same DLC ID.
783        // Allow the same pack file to be loaded multiple times (e.g., when accessing labeled sub-assets).
784        check_dlc_id_conflict(&dlc_id, &path_string)?;
785
786        // Register this asset path for the dlc id so it can be reloaded on unlock.
787        // If the path already exists for this DLC ID, it's idempotent (same pack file).
788        if !crate::encrypt_key_registry::has(dlc_id.as_ref(), &path_string) {
789            crate::encrypt_key_registry::register_asset_path(dlc_id.as_ref(), &path_string);
790        }
791
792        // Try to decrypt all entries immediately when the encrypt key is present.
793        // Offload decryption + archive extraction to Bevy's compute thread pool so
794        // we don't block the asset loader threads on heavy CPU work. If the key
795        // is missing we still populate the manifest so callers can inspect
796        // entries; a reload after unlock will add the typed sub-assets.
797        let decrypted_items = {
798            match decrypt_pack_entries(&dlc_id, &manifest_entries, &_block_metadatas, sync_reader) {
799                Ok(items) => Some(items),
800                Err(DlcLoaderError::DlcLocked(_)) => None,
801                Err(e) => return Err(e),
802            }
803        };
804
805        let mut out_entries = Vec::with_capacity(manifest_entries.len());
806
807        let mut unregistered_labels: Vec<String> = Vec::new();
808
809        // Collect all available registrars once per pack load to avoid heavy
810        // overhead from shared resource locking/matching inside the loop.
811        let dynamic_regs = self
812            .factories
813            .as_ref()
814            .map(|f| crate::asset_loader::collect_pack_registrars(Some(f)));
815        let regs = dynamic_regs.unwrap_or_else(|| collect_pack_registrars(None));
816
817        for (path, enc) in manifest_entries.into_iter() {
818            let entry_label = path.replace('\\', "/");
819
820            // Track whether a typed labeled asset was successfully registered
821            // for this entry. If `false` after processing, the pack still
822            // contains the entry but no labeled asset will be available via
823            // `pack.dlcpack#entry` (AssetServer will report it as missing).
824            let mut registered_as_labeled = false;
825
826            // Try to load this entry as a typed sub-asset when plaintext is available.
827            if let Some(ref items) = decrypted_items {
828                if let Some(item) = items.iter().find(|i| i.path() == path) {
829                    let ext = item.ext().unwrap_or_default();
830                    let type_path = item.type_path();
831                    let plaintext = item.plaintext().to_vec();
832                    // Build a fake path with the correct extension so
833                    // `load_context.loader()` selects the right concrete loader
834                    // by extension (e.g. `.png` → ImageLoader, `.json` → JsonLoader).
835                    let stem = std::path::Path::new(&entry_label)
836                        .file_stem()
837                        .and_then(|s| s.to_str())
838                        .unwrap_or("entry");
839                    let fake_path = format!("{}.{}", stem, ext);
840
841                    let mut vec_reader = bevy::asset::io::VecReader::new(plaintext.clone());
842
843                    // 1. Guided load: if `type_path` is present in the container metadata,
844                    // attempt to find a matching registrar and load directly using
845                    // that type. This bypasses extension-based dispatch entirely.
846                    if let Some(tp) = type_path {
847                        if let Some(registrar) = regs
848                            .iter()
849                            .find(|r| fuzzy_type_path_match(r.asset_type_path(), tp.as_str()))
850                        {
851                            match registrar
852                                .load_direct(
853                                    entry_label.clone(),
854                                    fake_path.clone(),
855                                    &mut vec_reader,
856                                    load_context,
857                                )
858                                .await
859                            {
860                                Ok(()) => {
861                                    registered_as_labeled = true;
862                                }
863                                Err(e) => {
864                                    // if static load failed, we still have a chance with
865                                    // extension-based dispatch below (rare but possible).
866                                    debug!(
867                                        "Static load for type '{}' failed: {}; falling back to extension dispatch",
868                                        tp, e
869                                    );
870                                }
871                            }
872                        }
873                    }
874
875                    // 2. Extension dispatch: Bevy picks the right loader based on `fake_path`
876                    // extension. We then try to match the resulting erased asset against
877                    // all known registrars to register it as a labeled sub-asset.
878                    if !registered_as_labeled {
879                        let mut vec_reader = bevy::asset::io::VecReader::new(plaintext.clone());
880                        let result = load_context
881                            .loader()
882                            .immediate()
883                            .with_reader(&mut vec_reader)
884                            .with_unknown_type()
885                            .load(fake_path.clone())
886                            .await;
887
888                        match result {
889                            Ok(erased) => {
890                                let mut remaining = Some(erased);
891
892                                for registrar in regs.iter() {
893                                    let label = entry_label.clone();
894                                    let to_register = remaining.take().unwrap();
895                                    match registrar.try_register(label, to_register, load_context) {
896                                        Ok(()) => {
897                                            registered_as_labeled = true;
898                                            remaining = None;
899                                            break;
900                                        }
901                                        Err(back) => {
902                                            remaining = Some(back);
903                                        }
904                                    }
905                                }
906
907                                if let Some(_) = remaining {
908                                    warn!(
909                                        "DLC entry '{}' present in container but no registered asset type matched (extension='{}'); the asset will NOT be available as '{}#{}'. Register a loader with `app.register_dlc_type::<T>()`",
910                                        entry_label, ext, path_string, entry_label
911                                    );
912                                }
913                            }
914                            Err(e) => {
915                                warn!(
916                                    "Failed to load entry '{}', extension='{}': {}",
917                                    entry_label, ext, e
918                                );
919                            }
920                        }
921                    }
922                }
923            }
924
925            // Build the labeled registered path for the DlcPackEntry.
926            let registered_path = format!("{}#{}", path_string, entry_label);
927
928            if !registered_as_labeled {
929                unregistered_labels.push(entry_label.clone());
930            }
931
932            out_entries.push(DlcPackEntry {
933                path: registered_path,
934                encrypted: enc,
935            });
936        }
937
938        // If we actually had plaintext to attempt registration (i.e. the
939        // pack was unlocked at load time) then unregistered_labels indicates a
940        // genuine failure to match a loader.  When the pack is still locked we
941        // intentionally avoid logging anything here because a later reload when
942        // the key arrives will perform the real work and emit the warning.
943        if decrypted_items.is_some() && !unregistered_labels.is_empty() {
944            // provide a concrete example using the first unregistered label so
945            // the user can see how to reference it with the pack path.
946            let example_label = &unregistered_labels[0];
947            let example_full = format!("{}#{}", path_string, example_label);
948            warn!(
949                "{} {} in '{}' were not registered as labeled assets and will be inaccessible via '{}'. See earlier warnings for details or register the appropriate loader via `app.register_dlc_type::<T>()`.",
950                unregistered_labels.len(),
951                if unregistered_labels.len() == 1 {
952                    "entry"
953                } else {
954                    "entries"
955                },
956                path_string,
957                example_full,
958            );
959        }
960
961        Ok(
962            DlcPack::new(dlc_id.clone(), product, version as u8, out_entries)
963                .with_path(path_string),
964        )
965    }
966}
967
968/// Internal helper used by `DlcPackLoader` to determine whether a
969/// pack for `dlc_id` from `path_string` should be accepted or rejected.
970/// Registry holds at most one path per DLC ID; conflict only arises when
971/// that path is different.  Loading the same pack file twice (same path) is
972/// normal and should not trigger an error.
973fn check_dlc_id_conflict(dlc_id: &DlcId, path_string: &str) -> Result<(), DlcLoaderError> {
974    if let Some(existing_path) = crate::encrypt_key_registry::asset_path_for(dlc_id.as_ref()) {
975        if existing_path != path_string {
976            return Err(DlcLoaderError::DlcIdConflict(
977                dlc_id.to_string(),
978                existing_path,
979                path_string.to_string(),
980            ));
981        }
982    }
983    Ok(())
984}
985
986/// Decrypt all entries in a pack using the block-based decryption method. This is used by `DlcPackLoader` when the encrypt key is available at load time, and also by `DlcPackEntry::decrypt_bytes` for individual entries when accessed via `pack.dlcpack#entry`.
987fn decrypt_pack_entries<R: std::io::Read + std::io::Seek>(
988    dlc_id: &crate::DlcId,
989    entries: &[(String, EncryptedAsset)],
990    block_metadatas: &[crate::pack_format::BlockMetadata],
991    reader: R,
992) -> Result<Vec<crate::PackItem>, DlcLoaderError> {
993    // lookup encrypt key in global registry
994    let encrypt_key = crate::encrypt_key_registry::get(dlc_id.as_ref())
995        .ok_or_else(|| DlcLoaderError::DlcLocked(dlc_id.to_string()))?;
996
997    let mut extracted_all = std::collections::HashMap::new();
998    // use reader in a PackReader for convenience
999    let mut pr = crate::pack_format::PackReader::new(reader);
1000    for block in block_metadatas {
1001        pr.seek(std::io::SeekFrom::Start(block.file_offset))
1002            .map_err(|e| DlcLoaderError::Io(e))?;
1003        let pt = pr
1004            .read_and_decrypt(&encrypt_key, block.encrypted_size as usize, &block.nonce)
1005            .map_err(|e| {
1006                let example = entries
1007                    .iter()
1008                    .find(|(_, enc)| enc.block_id == block.block_id)
1009                    .map(|(p, _)| p.as_str())
1010                    .unwrap_or("unknown");
1011                DlcLoaderError::DecryptionFailed(format!(
1012                    "dlc='{}' entry='{}' (block {}) decryption failed: {}",
1013                    dlc_id, example, block.block_id, e
1014                ))
1015            })?;
1016        let extracted = decompress_archive(&pt)?;
1017        extracted_all.extend(extracted);
1018    }
1019
1020    let mut out = Vec::with_capacity(entries.len());
1021    for (path, enc) in entries {
1022        let normalized = path.replace("\\", "/");
1023        let plaintext = extracted_all
1024            .remove(&normalized)
1025            .or_else(|| extracted_all.remove(path.as_str()))
1026            .ok_or_else(|| {
1027                DlcLoaderError::DecryptionFailed(format!("entry {} not found in any block", path))
1028            })?;
1029
1030        let mut item = PackItem::new(path.clone(), plaintext)
1031            .map_err(|e| DlcLoaderError::InvalidFormat(e.to_string()))?;
1032        if !enc.original_extension.is_empty() {
1033            item = item
1034                .with_extension(enc.original_extension.clone())
1035                .map_err(|e| DlcLoaderError::InvalidFormat(e.to_string()))?;
1036        }
1037        if let Some(tp) = &enc.type_path {
1038            item = item.with_type_path(tp.clone());
1039        }
1040        out.push(item);
1041    }
1042
1043    Ok(out)
1044}
1045
1046#[derive(Error, Debug)]
1047pub enum DlcLoaderError {
1048    /// Used for any IO failure during pack loading (e.g. file not found, read error, etc).
1049    #[error("IO error: {0}")]
1050    Io(io::Error),
1051    /// Used when the encrypt key for the DLC ID is not found in the registry at load time, which means the DLC is still locked and entries cannot be decrypted yet. This is not a fatal error — the pack can still be loaded and inspected.
1052    #[error("DLC locked: encrypt key not found for DLC id: {0}")]
1053    DlcLocked(String),
1054    /// Used for any failure during decryption of an entry or archive, including authentication failures from incorrect keys and any errors from archive extraction or manifest-archive mismatches.
1055    #[error("Decryption failed: {0}")]
1056    DecryptionFailed(String),
1057    /// Used when the initial container-level decryption succeeds but the plaintext is malformed (e.g. gzip archive is corrupted, manifest metadata doesn't match archive contents, etc).
1058    #[error("Invalid encrypted asset format: {0}")]
1059    InvalidFormat(String),
1060    /// Used when a DLC ID conflict is detected: a different pack file is already registered for the same DLC ID. This likely indicates a configuration error (e.g. two different `.dlcpack` files with the same internal DLC ID, or the same `.dlcpack` being loaded from two different paths). The error includes both the original and new pack paths for debugging.
1061    #[error(
1062        "DLC ID conflict: a .dlcpack with DLC id '{0}' is already loaded; cannot load another pack with the same DLC id, original: {1}, new: {2}"
1063    )]
1064    DlcIdConflict(String, String, String),
1065}
1066
1067impl From<std::io::Error> for DlcLoaderError {
1068    fn from(e: std::io::Error) -> Self {
1069        DlcLoaderError::Io(e)
1070    }
1071}
1072
1073#[cfg(test)]
1074mod tests {
1075    use super::*;
1076    use crate::{EncryptionKey, PackItem};
1077    use secure_gate::ExposeSecret;
1078    use serial_test::serial;
1079
1080    #[test]
1081    #[serial]
1082    fn encrypted_asset_decrypts_with_registry() {
1083        // pick a DLC id and prepare a static random key material so we can
1084        // construct two distinct `EncryptionKey` instances that share the same
1085        // bytes (one goes in the registry, the other is used for encryption).
1086        let dlc_id = "standalone";
1087        let key = EncryptionKey::from_random();
1088
1089        crate::encrypt_key_registry::clear_all();
1090        crate::encrypt_key_registry::insert(dlc_id, key.with_secret(|k| EncryptionKey::from(*k)));
1091
1092        // build a small standalone encrypted blob using PackWriter
1093        let nonce = [0u8; 12];
1094        let mut ciphertext = Vec::new();
1095        {
1096            let mut pw = crate::pack_format::PackWriter::new(&mut ciphertext);
1097            pw.write_encrypted(&key, &nonce, b"hello").expect("encrypt");
1098        }
1099
1100        let ct_len = ciphertext.len() as u32;
1101        let enc = EncryptedAsset {
1102            dlc_id: dlc_id.to_string(),
1103            original_extension: "".to_string(),
1104            type_path: None,
1105            nonce,
1106            ciphertext: ciphertext.into(),
1107            block_id: 0,
1108            block_offset: 0,
1109            size: ct_len,
1110        };
1111
1112        let plaintext = enc.decrypt_bytes().expect("decrypt");
1113        assert_eq!(&plaintext, b"hello");
1114    }
1115
1116    #[test]
1117    #[serial]
1118    fn dlcpack_accessors_work_and_fields_read() {
1119        let entry = DlcPackEntry {
1120            path: "a.txt".to_string(),
1121            encrypted: EncryptedAsset {
1122                dlc_id: "example_dlc".to_string(),
1123                original_extension: "txt".to_string(),
1124                type_path: None,
1125                nonce: [0u8; 12],
1126                ciphertext: vec![].into(),
1127                block_id: 0,
1128                block_offset: 0,
1129                size: 0,
1130            },
1131        };
1132        let pack = DlcPack::new(
1133            DlcId::from("example_dlc"),
1134            Product::from("test"),
1135            4,
1136            vec![entry.clone()],
1137        );
1138
1139        // exercise getters (reads `dlc_id` + `entries` fields)
1140        assert_eq!(*pack.id(), DlcId::from("example_dlc"));
1141        assert_eq!(pack.entries().len(), 1);
1142
1143        // inspect an entry (reads `path`, `original_extension`)
1144        let found = pack.find_entry("a.txt").expect("entry present");
1145        assert_eq!(found.path().path(), "a.txt");
1146        assert_eq!(found.original_extension(), "txt");
1147        assert!(found.type_path().is_none());
1148    }
1149
1150    #[test]
1151    #[serial]
1152    fn decrypt_pack_entries_v4_without_key_returns_locked_error() {
1153        crate::encrypt_key_registry::clear_all();
1154        let dlc_id = crate::DlcId::from("locked_dlc");
1155        let items = vec![PackItem::new("a.txt", b"hello".to_vec()).expect("pack item")];
1156        let key = EncryptionKey::from_random();
1157        let _dlc_key = crate::DlcKey::generate_random();
1158        let product = crate::Product::from("test");
1159        let container = crate::pack_encrypted_pack(&dlc_id, &items, &product, &key, crate::pack_format::DEFAULT_BLOCK_SIZE).expect("pack");
1160
1161        let mut cursor = std::io::Cursor::new(container);
1162        let (_product, parsed_dlc_id, _version, parsed_entries, block_metadatas) =
1163            crate::parse_encrypted_pack(&mut cursor).expect("parse");
1164
1165        let err = decrypt_pack_entries(
1166            &parsed_dlc_id,
1167            &parsed_entries,
1168            &block_metadatas,
1169            cursor,
1170        )
1171        .expect_err("should be locked");
1172        match err {
1173            DlcLoaderError::DlcLocked(id) => assert_eq!(id, "locked_dlc"),
1174            _ => panic!("expected DlcLocked error, got {:?}", err),
1175        }
1176    }
1177
1178    #[test]
1179    #[serial]
1180    fn decrypt_pack_entries_v4_with_wrong_key_reports_entry_and_dlc() {
1181        crate::encrypt_key_registry::clear_all();
1182        let dlc_id = crate::DlcId::from("badkey_dlc");
1183        let items = vec![PackItem::new("b.txt", b"world".to_vec()).expect("pack item")];
1184        let real_key = EncryptionKey::from_random();
1185        let _dlc_key = crate::DlcKey::generate_random();
1186        let product = crate::Product::from("test");
1187        let container =
1188            crate::pack_encrypted_pack(&dlc_id, &items, &product, &real_key, crate::pack_format::DEFAULT_BLOCK_SIZE).expect("pack");
1189
1190        // insert an incorrect key for this DLC
1191        let wrong_key: [u8; 32] = rand::random();
1192        crate::encrypt_key_registry::insert(
1193            &dlc_id.to_string(),
1194            crate::EncryptionKey::from(wrong_key),
1195        );
1196
1197        let mut cursor = std::io::Cursor::new(container);
1198        let (_product, parsed_dlc_id, _version, parsed_entries, block_metadatas) =
1199            crate::parse_encrypted_pack(&mut cursor).expect("parse");
1200
1201        let err = decrypt_pack_entries(
1202            &parsed_dlc_id,
1203            &parsed_entries,
1204            &block_metadatas,
1205            cursor,
1206        )
1207        .expect_err("should fail decryption");
1208        match err {
1209            DlcLoaderError::DecryptionFailed(msg) => {
1210                assert!(msg.contains("dlc='badkey_dlc'"));
1211                assert!(msg.contains("entry='b.txt'"));
1212                // ensure inner cause is propagated (auth failed for wrong key)
1213                assert!(msg.contains("authentication failed") || msg.contains("incorrect key"));
1214            }
1215            _ => panic!("expected DecryptionFailed, got {:?}", err),
1216        }
1217    }
1218
1219    #[test]
1220    #[serial]
1221    fn dlc_id_conflict_detection() {
1222        // Verify conflict detection logic for a DLC ID.  We avoid checking the
1223        // registered path string directly because other tests may clear the
1224        // global registry concurrently; instead we rely solely on the `check`
1225        // helper which works atomically.
1226        crate::encrypt_key_registry::clear_all();
1227
1228        let dlc_id_str = "conflict_test_dlc";
1229        let pack_path_1 = "existing_pack.dlcpack";
1230        let pack_path_2 = "different_pack.dlcpack";
1231
1232        crate::encrypt_key_registry::register_asset_path(dlc_id_str, pack_path_1);
1233
1234        // same path never counts as a conflict
1235        assert!(
1236            !crate::encrypt_key_registry::check(dlc_id_str, pack_path_1),
1237            "same pack path should NOT be a conflict"
1238        );
1239
1240        // different path should trigger a conflict. the registry global is
1241        // shared across parallel test threads and other tests call
1242        // `clear_all()`, so the entry may be lost mid-check.  loop and
1243        // re-register until we observe the expected result or give up.
1244        let mut tries = 0;
1245        while tries < 100 && !crate::encrypt_key_registry::check(dlc_id_str, pack_path_2) {
1246            crate::encrypt_key_registry::register_asset_path(dlc_id_str, pack_path_1);
1247            std::thread::sleep(std::time::Duration::from_millis(5));
1248            tries += 1;
1249        }
1250        assert!(
1251            crate::encrypt_key_registry::check(dlc_id_str, pack_path_2),
1252            "different pack path SHOULD be detected as a conflict"
1253        );
1254
1255        crate::encrypt_key_registry::clear_all();
1256    }
1257
1258    #[test]
1259    #[serial]
1260    fn dlc_loader_conflict_helper_allows_same_path() {
1261        crate::encrypt_key_registry::clear_all();
1262        let dlc_id = crate::DlcId::from("foo");
1263        let path = "same_pack.dlcpack";
1264        crate::encrypt_key_registry::register_asset_path(dlc_id.as_ref(), path);
1265        // helper should consider loading the same path to be fine
1266        assert!(check_dlc_id_conflict(&dlc_id, path).is_ok());
1267    }
1268
1269    #[test]
1270    #[serial]
1271    fn dlc_loader_conflict_helper_rejects_different_path() {
1272        crate::encrypt_key_registry::clear_all();
1273        let dlc_id = crate::DlcId::from("foo");
1274        crate::encrypt_key_registry::register_asset_path(dlc_id.as_ref(), "other.dlcpack");
1275        let err = check_dlc_id_conflict(&dlc_id, "new.dlcpack").expect_err("should conflict");
1276        match err {
1277            DlcLoaderError::DlcIdConflict(id, orig, newp) => {
1278                assert_eq!(id, dlc_id.to_string());
1279                assert_eq!(orig, "other.dlcpack");
1280                assert_eq!(newp, "new.dlcpack");
1281            }
1282            _ => panic!("expected DlcIdConflict"),
1283        }
1284    }
1285}
1286
1287impl<A> AssetLoader for DlcLoader<A>
1288where
1289    A: bevy::asset::Asset + TypePath + 'static,
1290{
1291    type Asset = A;
1292    type Settings = ();
1293    type Error = DlcLoaderError;
1294
1295    async fn load(
1296        &self,
1297        reader: &mut dyn Reader,
1298        _settings: &Self::Settings,
1299        load_context: &mut LoadContext<'_>,
1300    ) -> Result<Self::Asset, Self::Error> {
1301        // capture the original requested path (for registry/bookkeeping)
1302        let path_string = Some(load_context.path().path().to_string_lossy().to_string());
1303
1304        let mut bytes = Vec::new();
1305        reader
1306            .read_to_end(&mut bytes)
1307            .await
1308            .map_err(|e| DlcLoaderError::Io(e))?;
1309
1310        let enc =
1311            parse_encrypted(&bytes).map_err(|e| DlcLoaderError::DecryptionFailed(e.to_string()))?;
1312
1313        // register this asset path for the dlc id so it can be reloaded on unlock
1314        if let Some(p) = &path_string {
1315            crate::encrypt_key_registry::register_asset_path(&enc.dlc_id, p);
1316        }
1317
1318        // decrypt using helper on `EncryptedAsset`; this hides the registry
1319        // lookup and error formatting so the loader remains lean.
1320        let plaintext = enc.decrypt_bytes().map_err(|e| {
1321            // augment the error message with the requested path for context
1322            match e {
1323                DlcLoaderError::DecryptionFailed(msg) => DlcLoaderError::DecryptionFailed(format!(
1324                    "dlc='{}' path='{}' {}",
1325                    enc.dlc_id,
1326                    path_string
1327                        .clone()
1328                        .unwrap_or_else(|| "<unknown>".to_string()),
1329                    msg,
1330                )),
1331                other => other,
1332            }
1333        })?;
1334
1335        // Choose an extension for the nested load so Bevy can pick a concrete
1336        // loader if one exists. We keep the extension around so we can retry
1337        // with it if the straightforward, static-type load fails. Prioritizing
1338        // a static-type request avoids the need to downcast an erased asset,
1339        // which is more efficient and sidesteps the edge cases where the
1340        // extension loader returns a different type.
1341        let ext = enc.original_extension;
1342
1343        // Keep plaintext bytes around so we can recreate readers as needed.
1344        let bytes_clone = plaintext.clone();
1345
1346        let stem = load_context
1347            .path()
1348            .path()
1349            .file_stem()
1350            .and_then(|s| s.to_str())
1351            .unwrap_or("dlc_decrypted");
1352        let fake_path = format!("{}.{}", stem, ext);
1353
1354        // First attempt a direct static-type load. This bypasses extension
1355        // dispatch entirely and returns a value of `A` if a loader exists for
1356        // that type. Only if this fails do we fall back to using the extension
1357        // and performing a downcast.
1358        {
1359            let mut static_reader = bevy::asset::io::VecReader::new(bytes_clone.clone());
1360            if let Ok(loaded) = load_context
1361                .loader()
1362                .with_static_type()
1363                .immediate()
1364                .with_reader(&mut static_reader)
1365                .load::<A>(fake_path.clone())
1366                .await
1367            {
1368                return Ok(loaded.take());
1369            }
1370        }
1371
1372        // Static load didn't succeed. Try using the extension to select a loader
1373        // and then downcast the result to `A`. This mirrors how the normal
1374        // AssetServer would work when loading a file from disk.
1375        if !ext.is_empty() {
1376            // rewind original reader and clone again
1377            let mut ext_reader = bevy::asset::io::VecReader::new(bytes_clone.clone());
1378            let attempt = load_context
1379                .loader()
1380                .immediate()
1381                .with_reader(&mut ext_reader)
1382                .with_unknown_type()
1383                .load(fake_path.clone())
1384                .await;
1385
1386            if let Ok(erased) = attempt {
1387                match erased.downcast::<A>() {
1388                    Ok(loaded) => return Ok(loaded.take()),
1389                    Err(_) => {
1390                        return Err(DlcLoaderError::DecryptionFailed(format!(
1391                            "dlc loader: extension-based load succeeded but downcast to '{}' failed",
1392                            A::type_path(),
1393                        )));
1394                    }
1395                }
1396            } else if let Err(e) = attempt {
1397                return Err(DlcLoaderError::DecryptionFailed(e.to_string()));
1398            }
1399        }
1400
1401        // If we reach here it means neither static nor extension-based loading
1402        // succeeded; return an appropriate error. The original static attempt
1403        // already logged a warning, so just surface a generic message.
1404        Err(DlcLoaderError::DecryptionFailed(format!(
1405            "dlc loader: unable to load decrypted asset as {}{}",
1406            A::type_path(),
1407            if ext.is_empty() {
1408                ""
1409            } else {
1410                " (extension fallback also failed)"
1411            }
1412        )))
1413    }
1414}