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    pub fn original_extension(&self) -> &String {
482        &self.encrypted.original_extension
483    }
484
485    pub fn type_path(&self) -> Option<&String> {
486        self.encrypted.type_path.as_ref()
487    }
488}
489
490impl From<(String, EncryptedAsset)> for DlcPackEntry {
491    fn from((path, encrypted): (String, EncryptedAsset)) -> Self {
492        DlcPackEntry { path, encrypted }
493    }
494}
495
496impl From<&(String, EncryptedAsset)> for DlcPackEntry {
497    fn from((path, encrypted): &(String, EncryptedAsset)) -> Self {
498        DlcPackEntry { path: path.clone(), encrypted: encrypted.clone() }
499    }
500}
501
502/// Represents a `.dlcpack` bundle (multiple encrypted entries).
503#[derive(Asset, TypePath, Clone, Debug)]
504pub struct DlcPack {
505    dlc_id: DlcId,
506    product: Product,
507    version: u8,
508    entries: Vec<DlcPackEntry>,
509}
510
511impl DlcPack {
512    pub fn new(id: DlcId, product: Product, version: u8, entries: Vec<DlcPackEntry>) -> Self {
513        DlcPack {
514            dlc_id: id,
515            product,
516            version,
517            entries,
518        }
519    }
520
521    /// Return the DLC identifier for this pack.
522    pub fn id(&self) -> &DlcId {
523        &self.dlc_id
524    }
525
526    /// Return the product name this pack belongs to.
527    pub fn product(&self) -> &str {
528        &self.product.0
529    }
530
531    /// Return the pack format version.
532    pub fn version(&self) -> u8 {
533        self.version
534    }
535
536    /// Return a slice of contained entries.
537    pub fn entries(&self) -> &[DlcPackEntry] {
538        &self.entries
539    }
540
541    /// Find an entry by its registered path
542    pub fn find_entry(&self, path: &str) -> Option<&DlcPackEntry> {
543        self.entries
544            .iter()
545            .find(|e| e.path().to_string().ends_with(path) || e.path().path().ends_with(path))
546    }
547
548    /// Find all entries that match the specified asset type `A`.
549    pub fn find_by_type<A: Asset>(&self) -> Vec<&DlcPackEntry> {
550        self.entries
551            .iter()
552            .filter(|e| match e.type_path() {
553                Some(tp) => fuzzy_type_path_match(tp, A::type_path()),
554                None => false,
555            })
556            .collect()
557    }
558
559    /// Decrypt an entry (accepts either `name` or `packfile.dlcpack#name`).
560    /// Returns plaintext or `DlcLocked`.
561    pub fn decrypt_entry(
562        &self,
563        entry_path: &str,
564    ) -> Result<Vec<u8>, crate::asset_loader::DlcLoaderError> {
565        // accept either "test.png" or "packfile.dlcpack#test.png" by
566        // checking both relative and absolute paths
567        let entry = self.find_entry(entry_path).ok_or_else(|| {
568            DlcLoaderError::DecryptionFailed(format!("entry not found: {}", entry_path))
569        })?;
570
571        entry.decrypt_bytes()
572    }
573
574    pub fn load<A: Asset>(
575        &self,
576        asset_server: &bevy::prelude::AssetServer,
577        entry_path: &str,
578    ) -> Option<Handle<A>> {
579        let entry = match self.find_entry(entry_path) {
580            Some(e) => e,
581            None => return None,
582        };
583        Some(asset_server.load::<A>(entry.path()))
584    }
585}
586
587/// `AssetLoader` for `.dlcpack` bundles (contains multiple encrypted entries).
588///
589/// When the encrypt key is available in the registry at load time (i.e. the
590/// DLC is already unlocked), each entry is immediately decrypted and registered
591/// as a typed labeled sub-asset so `asset_server.load("pack.dlcpack#entry.png")`
592/// returns the correct `Handle<T>` for the extension's loader (e.g. `Handle<Image>`
593/// for `.png`). No asset type is hardcoded here — the correct type is determined
594/// purely by extension dispatch via Bevy's `immediate()` loader, and the result
595/// is downcast + registered using the list of `ErasedSubAssetRegistrar`s that
596/// `DlcPlugin::build` populates.
597///
598/// When the encrypt key is *not* yet available the pack is still loaded
599/// successfully (entries list is populated from the manifest) but the labeled
600/// sub-assets are not added — `reload_assets_on_unlock_system` will reload the
601/// pack once the key arrives, and the second load will succeed.
602#[derive(TypePath, Default)]
603pub struct DlcPackLoader {
604    /// Ordered list of per-type registrars. `DlcPlugin::build` pushes one
605    /// `TypedSubAssetRegistrar<A>` for every `A` it also registers via
606    /// `init_asset_loader::<DlcLoader<A>>()`. The loader tries each in turn
607    /// and uses the first successful downcast.
608    pub registrars: Vec<Box<dyn ErasedSubAssetRegistrar>>,
609    /// Optional shared reference to the `DlcPackRegistrarFactories` resource so
610    /// the loader can observe updates to the factory list at runtime without
611    /// requiring the asset loader to be re-registered.
612    pub(crate) factories: Option<DlcPackRegistrarFactories>,
613}
614
615/// Factory trait used to create `ErasedSubAssetRegistrar` instance.
616///
617/// Implement `TypedRegistrarFactory<T>` for asset types to produce a
618/// `TypedSubAssetRegistrar::<T>` at collection time.
619pub trait DlcPackRegistrarFactory: Send + Sync + 'static {
620    fn type_name(&self) -> &'static str;
621    fn create_registrar(&self) -> Box<dyn ErasedSubAssetRegistrar>;
622}
623
624/// Generic typed factory that constructs `TypedSubAssetRegistrar::<T>`.
625pub struct TypedRegistrarFactory<T: Asset + 'static>(std::marker::PhantomData<T>);
626
627impl<T: Asset + TypePath + 'static> DlcPackRegistrarFactory for TypedRegistrarFactory<T> {
628    fn type_name(&self) -> &'static str {
629        T::type_path()
630    }
631
632    fn create_registrar(&self) -> Box<dyn ErasedSubAssetRegistrar> {
633        Box::new(TypedSubAssetRegistrar::<T>::default())
634    }
635}
636
637impl<T: Asset + 'static> Default for TypedRegistrarFactory<T> {
638    fn default() -> Self {
639        TypedRegistrarFactory(std::marker::PhantomData)
640    }
641}
642
643use std::sync::RwLock;
644
645/// Internal factory resource used by `AppExt::register_dlc_type` so user code
646/// can request additional pack-registrars without pushing closures.
647///
648/// The resource wraps an `Arc<RwLock<_>>` so the registered `DlcPackLoader`
649/// instance can hold a cheap clone and observe updates made by
650/// `register_dlc_type(...)` without needing to re-register the loader.
651#[derive(Clone, Resource)]
652pub(crate) struct DlcPackRegistrarFactories(pub Arc<RwLock<Vec<Box<dyn DlcPackRegistrarFactory>>>>);
653
654impl Default for DlcPackRegistrarFactories {
655    fn default() -> Self {
656        DlcPackRegistrarFactories(Arc::new(RwLock::new(Vec::new())))
657    }
658}
659
660/// Return the default set of pack registrar factories used by `DlcPlugin`.
661///
662/// Using factory objects avoids closures and makes it trivial to add custom
663/// typed factories in user code (box a `TypedRegistrarFactory::<T>`).
664pub(crate) fn default_pack_registrar_factories() -> Vec<Box<dyn DlcPackRegistrarFactory>> {
665    vec![
666        Box::new(TypedRegistrarFactory::<Image>::default()),
667        Box::new(TypedRegistrarFactory::<Scene>::default()),
668        Box::new(TypedRegistrarFactory::<bevy::mesh::Mesh>::default()),
669        Box::new(TypedRegistrarFactory::<Font>::default()),
670        Box::new(TypedRegistrarFactory::<AudioSource>::default()),
671        Box::new(TypedRegistrarFactory::<ColorMaterial>::default()),
672        Box::new(TypedRegistrarFactory::<bevy::pbr::StandardMaterial>::default()),
673        Box::new(TypedRegistrarFactory::<bevy::gltf::Gltf>::default()),
674        Box::new(TypedRegistrarFactory::<bevy::gltf::GltfMesh>::default()),
675        Box::new(TypedRegistrarFactory::<Shader>::default()),
676        Box::new(TypedRegistrarFactory::<DynamicScene>::default()),
677        Box::new(TypedRegistrarFactory::<AnimationClip>::default()),
678        Box::new(TypedRegistrarFactory::<AnimationGraph>::default()),
679    ]
680}
681
682/// Build the final `registrars` vector by combining factory objects supplied via
683/// the `DlcPackRegistrarFactories` resource with the crate's default factories.
684pub(crate) fn collect_pack_registrars(
685    factories: Option<&DlcPackRegistrarFactories>,
686) -> Vec<Box<dyn ErasedSubAssetRegistrar>> {
687    use std::collections::HashSet;
688    let mut seen: HashSet<&'static str> = HashSet::new();
689    let mut out: Vec<Box<dyn ErasedSubAssetRegistrar>> = Vec::new();
690
691    if let Some(f) = factories {
692        let inner = f.0.read().unwrap();
693        for factory in inner.iter() {
694            out.push(factory.create_registrar());
695            seen.insert(factory.type_name());
696        }
697    }
698
699    for factory in default_pack_registrar_factories() {
700        if !seen.contains(factory.type_name()) {
701            out.push(factory.create_registrar());
702            seen.insert(factory.type_name());
703        }
704    }
705
706    out
707}
708
709impl AssetLoader for DlcPackLoader {
710    type Asset = DlcPack;
711    type Settings = ();
712    type Error = DlcLoaderError;
713
714    fn extensions(&self) -> &[&str] {
715        &["dlcpack"]
716    }
717
718    async fn load(
719        &self,
720        reader: &mut dyn Reader,
721        _settings: &Self::Settings,
722        load_context: &mut LoadContext<'_>,
723    ) -> Result<Self::Asset, Self::Error> {
724        let path_string = load_context.path().path().to_string_lossy().to_string();
725
726        // Adapt the async `reader` to a synchronous `std::io::Read` so we can
727        // drive the existing pack‑parsing logic without buffering the whole file
728        // up front.  If the underlying reader is seekable we will also be able
729        // rewind it later in order to extract the raw bytes needed for
730        // decryption.
731        let mut sync_reader = SyncReader::new(reader);
732
733        let (product, dlc_id, version, manifest_entries, _block_metadatas) =
734            crate::parse_encrypted_pack(&mut sync_reader)
735                .map_err(|e| DlcLoaderError::InvalidFormat(e.to_string()))?;
736
737        // rewind the reader back to the start so decryption routines can re‑parse the file as needed.  If the reader is not seekable, error.
738        sync_reader.seek(SeekFrom::Start(0)).map_err(|_| {
739            DlcLoaderError::Io(io::Error::new(
740                io::ErrorKind::NotSeekable,
741                format!("reader not seekable, cannot load pack '{}'", path_string),
742            ))
743        })?;
744
745        // Check for DLC ID conflicts: reject if a DIFFERENT pack file is being loaded for the same DLC ID.
746        // Allow the same pack file to be loaded multiple times (e.g., when accessing labeled sub-assets).
747        check_dlc_id_conflict(&dlc_id, &path_string)?;
748
749        // Register this asset path for the dlc id so it can be reloaded on unlock.
750        // If the path already exists for this DLC ID, it's idempotent (same pack file).
751        if !crate::encrypt_key_registry::has(dlc_id.as_ref(), &path_string) {
752            crate::encrypt_key_registry::register_asset_path(dlc_id.as_ref(), &path_string);
753        }
754
755        // Try to decrypt all entries immediately when the encrypt key is present.
756        // Offload decryption + archive extraction to Bevy's compute thread pool so
757        // we don't block the asset loader threads on heavy CPU work. If the key
758        // is missing we still populate the manifest so callers can inspect
759        // entries; a reload after unlock will add the typed sub-assets.
760        let decrypted_items = {
761            match decrypt_pack_entries(sync_reader) {
762                Ok::<(DlcId, Vec<PackItem>), DlcLoaderError>((_id, items)) => Some(items),
763                Err(DlcLoaderError::DlcLocked(_)) => None,
764                Err(e) => return Err(e),
765            }
766        };
767
768        let mut out_entries = Vec::with_capacity(manifest_entries.len());
769
770        let mut unregistered_labels: Vec<String> = Vec::new();
771
772        // Collect all available registrars once per pack load to avoid heavy
773        // overhead from shared resource locking/matching inside the loop.
774        let dynamic_regs = self
775            .factories
776            .as_ref()
777            .map(|f| crate::asset_loader::collect_pack_registrars(Some(f)));
778        let regs = dynamic_regs.unwrap_or_else(|| collect_pack_registrars(None));
779
780        for (path, enc) in manifest_entries.into_iter() {
781            let entry_label = path.replace('\\', "/");
782
783            // Track whether a typed labeled asset was successfully registered
784            // for this entry. If `false` after processing, the pack still
785            // contains the entry but no labeled asset will be available via
786            // `pack.dlcpack#entry` (AssetServer will report it as missing).
787            let mut registered_as_labeled = false;
788
789            // Try to load this entry as a typed sub-asset when plaintext is available.
790            if let Some(ref items) = decrypted_items {
791                if let Some(item) = items.iter().find(|i| i.path() == path) {
792                    let ext = item.ext().unwrap_or_default();
793                    let type_path = item.type_path();
794                    let plaintext = item.plaintext().to_vec();
795                    // Build a fake path with the correct extension so
796                    // `load_context.loader()` selects the right concrete loader
797                    // by extension (e.g. `.png` → ImageLoader, `.json` → JsonLoader).
798                    let stem = std::path::Path::new(&entry_label)
799                        .file_stem()
800                        .and_then(|s| s.to_str())
801                        .unwrap_or("entry");
802                    let fake_path = format!("{}.{}", stem, ext);
803
804                    let mut vec_reader = bevy::asset::io::VecReader::new(plaintext.clone());
805
806                    // 1. Guided load: if `type_path` is present in the container metadata,
807                    // attempt to find a matching registrar and load directly using
808                    // that type. This bypasses extension-based dispatch entirely.
809                    if let Some(tp) = type_path {
810                        if let Some(registrar) = regs
811                            .iter()
812                            .find(|r| fuzzy_type_path_match(r.asset_type_path(), tp.as_str()))
813                        {
814                            match registrar
815                                .load_direct(
816                                    entry_label.clone(),
817                                    fake_path.clone(),
818                                    &mut vec_reader,
819                                    load_context,
820                                )
821                                .await
822                            {
823                                Ok(()) => {
824                                    registered_as_labeled = true;
825                                }
826                                Err(e) => {
827                                    // if static load failed, we still have a chance with
828                                    // extension-based dispatch below (rare but possible).
829                                    debug!(
830                                        "Static load for type '{}' failed: {}; falling back to extension dispatch",
831                                        tp, e
832                                    );
833                                }
834                            }
835                        }
836                    }
837
838                    // 2. Extension dispatch: Bevy picks the right loader based on `fake_path`
839                    // extension. We then try to match the resulting erased asset against
840                    // all known registrars to register it as a labeled sub-asset.
841                    if !registered_as_labeled {
842                        let mut vec_reader = bevy::asset::io::VecReader::new(plaintext.clone());
843                        let result = load_context
844                            .loader()
845                            .immediate()
846                            .with_reader(&mut vec_reader)
847                            .with_unknown_type()
848                            .load(fake_path.clone())
849                            .await;
850
851                        match result {
852                            Ok(erased) => {
853                                let mut remaining = Some(erased);
854
855                                for registrar in regs.iter() {
856                                    let label = entry_label.clone();
857                                    let to_register = remaining.take().unwrap();
858                                    match registrar.try_register(label, to_register, load_context) {
859                                        Ok(()) => {
860                                            registered_as_labeled = true;
861                                            remaining = None;
862                                            break;
863                                        }
864                                        Err(back) => {
865                                            remaining = Some(back);
866                                        }
867                                    }
868                                }
869
870                                if let Some(_) = remaining {
871                                    warn!(
872                                        "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>()`",
873                                        entry_label, ext, path_string, entry_label
874                                    );
875                                }
876                            }
877                            Err(e) => {
878                                warn!(
879                                    "Failed to load entry '{}', extension='{}': {}",
880                                    entry_label, ext, e
881                                );
882                            }
883                        }
884                    }
885                }
886            }
887
888            // Build the labeled registered path for the DlcPackEntry.
889            let registered_path = format!("{}#{}", path_string, entry_label);
890
891            if !registered_as_labeled {
892                unregistered_labels.push(entry_label.clone());
893            }
894
895            out_entries.push(DlcPackEntry {
896                path: registered_path,
897                encrypted: enc,
898            });
899        }
900
901        // If we actually had plaintext to attempt registration (i.e. the
902        // pack was unlocked at load time) then unregistered_labels indicates a
903        // genuine failure to match a loader.  When the pack is still locked we
904        // intentionally avoid logging anything here because a later reload when
905        // the key arrives will perform the real work and emit the warning.
906        if decrypted_items.is_some() && !unregistered_labels.is_empty() {
907            // provide a concrete example using the first unregistered label so
908            // the user can see how to reference it with the pack path.
909            let example_label = &unregistered_labels[0];
910            let example_full = format!("{}#{}", path_string, example_label);
911            warn!(
912                "{} {} 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>()`.",
913                unregistered_labels.len(),
914                if unregistered_labels.len() == 1 {
915                    "entry"
916                } else {
917                    "entries"
918                },
919                path_string,
920                example_full,
921            );
922        }
923
924        Ok(DlcPack::new(
925            dlc_id.clone(),
926            product,
927            version as u8,
928            out_entries,
929        ))
930    }
931}
932
933/// Internal helper used by `DlcPackLoader` to determine whether a
934/// pack for `dlc_id` from `path_string` should be accepted or rejected.
935/// Registry holds at most one path per DLC ID; conflict only arises when
936/// that path is different.  Loading the same pack file twice (same path) is
937/// normal and should not trigger an error.
938fn check_dlc_id_conflict(
939    dlc_id: &DlcId,
940    path_string: &str,
941) -> Result<(), DlcLoaderError> {
942    if let Some(existing_path) = crate::encrypt_key_registry::asset_path_for(dlc_id.as_ref()) {
943        if existing_path != path_string {
944            return Err(DlcLoaderError::DlcIdConflict(
945                dlc_id.to_string(),
946                existing_path,
947                path_string.to_string(),
948            ));
949        }
950    }
951    Ok(())
952}
953
954/// Decrypt entries out of a `.dlcpack` container.
955/// Returns `(dlc_id, Vec<(path, original_extension, plaintext)>)`.
956/// Decrypt entries out of a `.dlcpack` container.
957///
958/// This no-longer requires keeping the entire file in memory; the caller
959/// provides a `Read+Seek` reader.  Only the current v4 hybrid format is
960/// supported, earlier versions have been removed along with their legacy
961/// semantics.
962///
963/// Returns `(dlc_id, Vec<(path, original_extension, plaintext)>)`.
964pub fn decrypt_pack_entries<R: std::io::Read + std::io::Seek>(
965    mut reader: R,
966) -> Result<(crate::DlcId, Vec<crate::PackItem>), DlcLoaderError> {
967    let (_product, dlc_id, version, entries, block_metadatas) =
968        crate::parse_encrypted_pack(&mut reader)
969            .map_err(|e| DlcLoaderError::InvalidFormat(e.to_string()))?;
970
971    // lookup encrypt key in global registry
972    let encrypt_key = crate::encrypt_key_registry::get(dlc_id.as_ref())
973        .ok_or_else(|| DlcLoaderError::DlcLocked(dlc_id.to_string()))?;
974
975    // only v4 is supported anymore
976    if version != 4 {
977        return Err(DlcLoaderError::InvalidFormat(format!(
978            "unsupported pack version: {}",
979            version
980        )));
981    }
982
983    let mut extracted_all = std::collections::HashMap::new();
984    // use reader in a PackReader for convenience
985    let mut pr = crate::pack_format::PackReader::new(reader);
986    for block in block_metadatas {
987        pr.seek(std::io::SeekFrom::Start(block.file_offset))
988            .map_err(|e| DlcLoaderError::Io(e))?;
989        let pt = pr
990            .read_and_decrypt(&encrypt_key, block.encrypted_size as usize, &block.nonce)
991            .map_err(|e| {
992                let example = entries
993                    .iter()
994                    .find(|(_, enc)| enc.block_id == block.block_id)
995                    .map(|(p, _)| p.as_str())
996                    .unwrap_or("unknown");
997                DlcLoaderError::DecryptionFailed(format!(
998                    "dlc='{}' entry='{}' (block {}) decryption failed: {}",
999                    dlc_id, example, block.block_id, e
1000                ))
1001            })?;
1002        let extracted = decompress_archive(&pt)?;
1003        extracted_all.extend(extracted);
1004    }
1005
1006    let mut out = Vec::with_capacity(entries.len());
1007    for (path, enc) in entries {
1008        let normalized = path.replace("\\", "/");
1009        let plaintext = extracted_all
1010            .remove(&normalized)
1011            .or_else(|| extracted_all.remove(&path))
1012            .ok_or_else(|| {
1013                DlcLoaderError::DecryptionFailed(format!("entry {} not found in any block", path))
1014            })?;
1015
1016        let mut item = PackItem::new(path.clone(), plaintext)
1017            .map_err(|e| DlcLoaderError::InvalidFormat(e.to_string()))?;
1018        if !enc.original_extension.is_empty() {
1019            item = item
1020                .with_extension(enc.original_extension)
1021                .map_err(|e| DlcLoaderError::InvalidFormat(e.to_string()))?;
1022        }
1023        if let Some(tp) = enc.type_path {
1024            item = item.with_type_path(tp);
1025        }
1026        out.push(item);
1027    }
1028
1029    Ok((dlc_id, out))
1030}
1031
1032#[derive(Error, Debug)]
1033pub enum DlcLoaderError {
1034    /// Used for any IO failure during pack loading (e.g. file not found, read error, etc).
1035    #[error("IO error: {0}")]
1036    Io(io::Error),
1037    /// 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.
1038    #[error("DLC locked: encrypt key not found for DLC id: {0}")]
1039    DlcLocked(String),
1040    /// 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.
1041    #[error("Decryption failed: {0}")]
1042    DecryptionFailed(String),
1043    /// 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).
1044    #[error("Invalid encrypted asset format: {0}")]
1045    InvalidFormat(String),
1046    /// 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.
1047    #[error(
1048        "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}"
1049    )]
1050    DlcIdConflict(String, String, String),
1051}
1052
1053impl From<std::io::Error> for DlcLoaderError {
1054    fn from(e: std::io::Error) -> Self {
1055        DlcLoaderError::Io(e)
1056    }
1057}
1058
1059#[cfg(test)]
1060mod tests {
1061    use super::*;
1062    use crate::{EncryptionKey, PackItem};
1063    use secure_gate::ExposeSecret;
1064    use serial_test::serial;
1065
1066    #[test]
1067    #[serial]
1068    fn encrypted_asset_decrypts_with_registry() {
1069        // pick a DLC id and prepare a static random key material so we can
1070        // construct two distinct `EncryptionKey` instances that share the same
1071        // bytes (one goes in the registry, the other is used for encryption).
1072        let dlc_id = "standalone";
1073        let key = EncryptionKey::from_random();
1074
1075        crate::encrypt_key_registry::clear_all();
1076        crate::encrypt_key_registry::insert(dlc_id, key.with_secret(|k| {
1077            EncryptionKey::from(*k)
1078        }));
1079
1080        // build a small standalone encrypted blob using PackWriter
1081        let nonce = [0u8; 12];
1082        let mut ciphertext = Vec::new();
1083        {
1084            let mut pw = crate::pack_format::PackWriter::new(&mut ciphertext);
1085            pw.write_encrypted(&key, &nonce, b"hello").expect("encrypt");
1086        }
1087
1088        let ct_len = ciphertext.len() as u32;
1089        let enc = EncryptedAsset {
1090            dlc_id: dlc_id.to_string(),
1091            original_extension: "".to_string(),
1092            type_path: None,
1093            nonce,
1094            ciphertext: ciphertext.into(),
1095            block_id: 0,
1096            block_offset: 0,
1097            size: ct_len,
1098        };
1099
1100        let plaintext = enc.decrypt_bytes().expect("decrypt");
1101        assert_eq!(&plaintext, b"hello");
1102    }
1103
1104    #[test]
1105    #[serial]
1106    fn dlcpack_accessors_work_and_fields_read() {
1107        let entry = DlcPackEntry {
1108            path: "a.txt".to_string(),
1109            encrypted: EncryptedAsset {
1110                dlc_id: "example_dlc".to_string(),
1111                original_extension: "txt".to_string(),
1112                type_path: None,
1113                nonce: [0u8; 12],
1114                ciphertext: vec![].into(),
1115                block_id: 0,
1116                block_offset: 0,
1117                size: 0,
1118            },
1119        };
1120        let pack = DlcPack::new(
1121            DlcId::from("example_dlc"),
1122            Product::from("test"),
1123            4,
1124            vec![entry.clone()],
1125        );
1126
1127        // exercise getters (reads `dlc_id` + `entries` fields)
1128        assert_eq!(*pack.id(), DlcId::from("example_dlc"));
1129        assert_eq!(pack.entries().len(), 1);
1130
1131        // inspect an entry (reads `path`, `original_extension`)
1132        let found = pack.find_entry("a.txt").expect("entry present");
1133        assert_eq!(found.path().path(), "a.txt");
1134        assert_eq!(found.original_extension(), "txt");
1135        assert!(found.type_path().is_none());
1136    }
1137
1138    #[test]
1139    #[serial]
1140    fn decrypt_pack_entries_without_key_returns_locked_error() {
1141        crate::encrypt_key_registry::clear_all();
1142        let dlc_id = crate::DlcId::from("locked_dlc");
1143        let items = vec![PackItem::new("a.txt", b"hello".to_vec()).expect("pack item")];
1144        let key = EncryptionKey::from_random();
1145        let _dlc_key = crate::DlcKey::generate_random();
1146        let product = crate::Product::from("test");
1147        let container = crate::pack_encrypted_pack(&dlc_id, &items, &product, &key).expect("pack");
1148
1149        let err =
1150            decrypt_pack_entries(std::io::Cursor::new(container)).expect_err("should be locked");
1151        match err {
1152            DlcLoaderError::DlcLocked(id) => assert_eq!(id, "locked_dlc"),
1153            _ => panic!("expected DlcLocked error, got {:?}", err),
1154        }
1155    }
1156
1157    #[test]
1158    #[serial]
1159    fn decrypt_pack_entries_with_wrong_key_reports_entry_and_dlc() {
1160        crate::encrypt_key_registry::clear_all();
1161        let dlc_id = crate::DlcId::from("badkey_dlc");
1162        let items = vec![PackItem::new("b.txt", b"world".to_vec()).expect("pack item")];
1163        let real_key = EncryptionKey::from_random();
1164        let _dlc_key = crate::DlcKey::generate_random();
1165        let product = crate::Product::from("test");
1166        let container =
1167            crate::pack_encrypted_pack(&dlc_id, &items, &product, &real_key).expect("pack");
1168
1169        // insert an incorrect key for this DLC
1170        let wrong_key: [u8; 32] = rand::random();
1171        crate::encrypt_key_registry::insert(
1172            &dlc_id.to_string(),
1173            crate::EncryptionKey::from(wrong_key),
1174        );
1175
1176        let err = decrypt_pack_entries(std::io::Cursor::new(container))
1177            .expect_err("should fail decryption");
1178        match err {
1179            DlcLoaderError::DecryptionFailed(msg) => {
1180                assert!(msg.contains("dlc='badkey_dlc'"));
1181                assert!(msg.contains("entry='b.txt'"));
1182                // ensure inner cause is propagated (auth failed for wrong key)
1183                assert!(msg.contains("authentication failed") || msg.contains("incorrect key"));
1184            }
1185            _ => panic!("expected DecryptionFailed, got {:?}", err),
1186        }
1187    }
1188
1189    #[test]
1190    #[serial]
1191    fn dlc_id_conflict_detection() {
1192        // Verify conflict detection logic for a DLC ID.  We avoid checking the
1193        // registered path string directly because other tests may clear the
1194        // global registry concurrently; instead we rely solely on the `check`
1195        // helper which works atomically.
1196        crate::encrypt_key_registry::clear_all();
1197
1198        let dlc_id_str = "conflict_test_dlc";
1199        let pack_path_1 = "existing_pack.dlcpack";
1200        let pack_path_2 = "different_pack.dlcpack";
1201
1202        crate::encrypt_key_registry::register_asset_path(dlc_id_str, pack_path_1);
1203
1204        // same path never counts as a conflict
1205        assert!(
1206            !crate::encrypt_key_registry::check(dlc_id_str, pack_path_1),
1207            "same pack path should NOT be a conflict"
1208        );
1209
1210        // different path should trigger a conflict. the registry global is
1211        // shared across parallel test threads and other tests call
1212        // `clear_all()`, so the entry may be lost mid-check.  loop and
1213        // re-register until we observe the expected result or give up.
1214        let mut tries = 0;
1215        while tries < 100 && !crate::encrypt_key_registry::check(dlc_id_str, pack_path_2) {
1216            crate::encrypt_key_registry::register_asset_path(dlc_id_str, pack_path_1);
1217            std::thread::sleep(std::time::Duration::from_millis(5));
1218            tries += 1;
1219        }
1220        assert!(
1221            crate::encrypt_key_registry::check(dlc_id_str, pack_path_2),
1222            "different pack path SHOULD be detected as a conflict"
1223        );
1224
1225        crate::encrypt_key_registry::clear_all();
1226    }
1227
1228    #[test]
1229    #[serial]
1230    fn dlc_loader_conflict_helper_allows_same_path() {
1231        crate::encrypt_key_registry::clear_all();
1232        let dlc_id = crate::DlcId::from("foo");
1233        let path = "same_pack.dlcpack";
1234        crate::encrypt_key_registry::register_asset_path(dlc_id.as_ref(), path);
1235        // helper should consider loading the same path to be fine
1236        assert!(check_dlc_id_conflict(&dlc_id, path).is_ok());
1237    }
1238
1239    #[test]
1240    #[serial]
1241    fn dlc_loader_conflict_helper_rejects_different_path() {
1242        crate::encrypt_key_registry::clear_all();
1243        let dlc_id = crate::DlcId::from("foo");
1244        crate::encrypt_key_registry::register_asset_path(dlc_id.as_ref(), "other.dlcpack");
1245        let err = check_dlc_id_conflict(&dlc_id, "new.dlcpack").expect_err("should conflict");
1246        match err {
1247            DlcLoaderError::DlcIdConflict(id, orig, newp) => {
1248                assert_eq!(id, dlc_id.to_string());
1249                assert_eq!(orig, "other.dlcpack");
1250                assert_eq!(newp, "new.dlcpack");
1251            }
1252            _ => panic!("expected DlcIdConflict"),
1253        }
1254    }
1255}
1256
1257impl<A> AssetLoader for DlcLoader<A>
1258where
1259    A: bevy::asset::Asset + TypePath + 'static,
1260{
1261    type Asset = A;
1262    type Settings = ();
1263    type Error = DlcLoaderError;
1264
1265    async fn load(
1266        &self,
1267        reader: &mut dyn Reader,
1268        _settings: &Self::Settings,
1269        load_context: &mut LoadContext<'_>,
1270    ) -> Result<Self::Asset, Self::Error> {
1271        // capture the original requested path (for registry/bookkeeping)
1272        let path_string = Some(load_context.path().path().to_string_lossy().to_string());
1273
1274        let mut bytes = Vec::new();
1275        reader
1276            .read_to_end(&mut bytes)
1277            .await
1278            .map_err(|e| DlcLoaderError::Io(e))?;
1279
1280        let enc =
1281            parse_encrypted(&bytes).map_err(|e| DlcLoaderError::DecryptionFailed(e.to_string()))?;
1282
1283        // register this asset path for the dlc id so it can be reloaded on unlock
1284        if let Some(p) = &path_string {
1285            crate::encrypt_key_registry::register_asset_path(&enc.dlc_id, p);
1286        }
1287
1288        // decrypt using helper on `EncryptedAsset`; this hides the registry
1289        // lookup and error formatting so the loader remains lean.
1290        let plaintext = enc.decrypt_bytes().map_err(|e| {
1291            // augment the error message with the requested path for context
1292            match e {
1293                DlcLoaderError::DecryptionFailed(msg) => DlcLoaderError::DecryptionFailed(format!(
1294                    "dlc='{}' path='{}' {}",
1295                    enc.dlc_id,
1296                    path_string.clone().unwrap_or_else(|| "<unknown>".to_string()),
1297                    msg,
1298                )),
1299                other => other,
1300            }
1301        })?;
1302
1303        // Choose an extension for the nested load so Bevy can pick a concrete
1304        // loader if one exists. We keep the extension around so we can retry
1305        // with it if the straightforward, static-type load fails. Prioritizing
1306        // a static-type request avoids the need to downcast an erased asset,
1307        // which is more efficient and sidesteps the edge cases where the
1308        // extension loader returns a different type.
1309        let ext = enc.original_extension;
1310
1311        // Keep plaintext bytes around so we can recreate readers as needed.
1312        let bytes_clone = plaintext.clone();
1313
1314        let stem = load_context
1315            .path()
1316            .path()
1317            .file_stem()
1318            .and_then(|s| s.to_str())
1319            .unwrap_or("dlc_decrypted");
1320        let fake_path = format!("{}.{}", stem, ext);
1321
1322        // First attempt a direct static-type load. This bypasses extension
1323        // dispatch entirely and returns a value of `A` if a loader exists for
1324        // that type. Only if this fails do we fall back to using the extension
1325        // and performing a downcast.
1326        {
1327            let mut static_reader = bevy::asset::io::VecReader::new(bytes_clone.clone());
1328            if let Ok(loaded) = load_context
1329                .loader()
1330                .with_static_type()
1331                .immediate()
1332                .with_reader(&mut static_reader)
1333                .load::<A>(fake_path.clone())
1334                .await
1335            {
1336                return Ok(loaded.take());
1337            }
1338        }
1339
1340        // Static load didn't succeed. Try using the extension to select a loader
1341        // and then downcast the result to `A`. This mirrors how the normal
1342        // AssetServer would work when loading a file from disk.
1343        if !ext.is_empty() {
1344            // rewind original reader and clone again
1345            let mut ext_reader = bevy::asset::io::VecReader::new(bytes_clone.clone());
1346            let attempt = load_context
1347                .loader()
1348                .immediate()
1349                .with_reader(&mut ext_reader)
1350                .with_unknown_type()
1351                .load(fake_path.clone())
1352                .await;
1353
1354            if let Ok(erased) = attempt {
1355                match erased.downcast::<A>() {
1356                    Ok(loaded) => return Ok(loaded.take()),
1357                    Err(_) => {
1358                        return Err(DlcLoaderError::DecryptionFailed(format!(
1359                            "dlc loader: extension-based load succeeded but downcast to '{}' failed",
1360                            A::type_path(),
1361                        )));
1362                    }
1363                }
1364            } else if let Err(e) = attempt {
1365                return Err(DlcLoaderError::DecryptionFailed(e.to_string()));
1366            }
1367        }
1368
1369        // If we reach here it means neither static nor extension-based loading
1370        // succeeded; return an appropriate error. The original static attempt
1371        // already logged a warning, so just surface a generic message.
1372        Err(DlcLoaderError::DecryptionFailed(format!(
1373            "dlc loader: unable to load decrypted asset as {}{}",
1374            A::type_path(),
1375            if ext.is_empty() {
1376                ""
1377            } else {
1378                " (extension fallback also failed)"
1379            }
1380        )))
1381    }
1382}