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
17pub 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
54fn 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
87pub(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 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 reader
119 .seek(std::io::SeekFrom::Start(block.file_offset))
120 .map_err(|e| DlcLoaderError::Io(e))?;
121
122 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 let entries = decompress_archive(&pt_gz)?;
130
131 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#[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 pub fn id(&self) -> &DlcId {
160 &self.dlc_id
161 }
162
163 pub fn pack(&self) -> &DlcPack {
165 &self.pack
166 }
167}
168
169pub(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 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
192pub 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 fn asset_type_path(&self) -> &'static str;
206
207 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
219pub 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#[derive(Clone, Debug, Asset, TypePath)]
277pub struct EncryptedAsset {
278 pub dlc_id: String,
279 pub original_extension: String,
280 pub type_path: Option<String>,
282 pub nonce: [u8; 12],
283 pub ciphertext: std::sync::Arc<[u8]>,
284 pub block_id: u32,
286 pub block_offset: u32,
287 pub size: u32,
288}
289
290impl EncryptedAsset {
291 pub(crate) fn decrypt_bytes(&self) -> Result<Vec<u8>, DlcLoaderError> {
295 let encrypt_key = crate::encrypt_key_registry::get(&self.dlc_id)
297 .ok_or_else(|| DlcLoaderError::DlcLocked(self.dlc_id.clone()))?;
298
299 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
309pub fn parse_encrypted(bytes: &[u8]) -> Result<EncryptedAsset, io::Error> {
311 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 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#[derive(TypePath)]
397pub struct DlcLoader<A: bevy::asset::Asset + 'static> {
398 #[allow(dead_code)]
401 type_registry: Arc<AppTypeRegistry>,
402 _marker: std::marker::PhantomData<A>,
403}
404
405impl<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 path: String,
424 encrypted: EncryptedAsset,
426}
427
428impl DlcPackEntry {
429 pub fn new(path: String, encrypted: EncryptedAsset) -> Self {
430 DlcPackEntry { path, encrypted }
431 }
432
433 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 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 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 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#[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 pub fn id(&self) -> &DlcId {
523 &self.dlc_id
524 }
525
526 pub fn product(&self) -> &str {
528 &self.product.0
529 }
530
531 pub fn version(&self) -> u8 {
533 self.version
534 }
535
536 pub fn entries(&self) -> &[DlcPackEntry] {
538 &self.entries
539 }
540
541 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 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 pub fn decrypt_entry(
562 &self,
563 entry_path: &str,
564 ) -> Result<Vec<u8>, crate::asset_loader::DlcLoaderError> {
565 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#[derive(TypePath, Default)]
603pub struct DlcPackLoader {
604 pub registrars: Vec<Box<dyn ErasedSubAssetRegistrar>>,
609 pub(crate) factories: Option<DlcPackRegistrarFactories>,
613}
614
615pub trait DlcPackRegistrarFactory: Send + Sync + 'static {
620 fn type_name(&self) -> &'static str;
621 fn create_registrar(&self) -> Box<dyn ErasedSubAssetRegistrar>;
622}
623
624pub 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#[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
660pub(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
682pub(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 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 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_dlc_id_conflict(&dlc_id, &path_string)?;
748
749 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 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 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 let mut registered_as_labeled = false;
788
789 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 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 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 debug!(
830 "Static load for type '{}' failed: {}; falling back to extension dispatch",
831 tp, e
832 );
833 }
834 }
835 }
836 }
837
838 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 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 decrypted_items.is_some() && !unregistered_labels.is_empty() {
907 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
933fn 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
954pub 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 let encrypt_key = crate::encrypt_key_registry::get(dlc_id.as_ref())
973 .ok_or_else(|| DlcLoaderError::DlcLocked(dlc_id.to_string()))?;
974
975 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 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 #[error("IO error: {0}")]
1036 Io(io::Error),
1037 #[error("DLC locked: encrypt key not found for DLC id: {0}")]
1039 DlcLocked(String),
1040 #[error("Decryption failed: {0}")]
1042 DecryptionFailed(String),
1043 #[error("Invalid encrypted asset format: {0}")]
1045 InvalidFormat(String),
1046 #[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 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 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 assert_eq!(*pack.id(), DlcId::from("example_dlc"));
1129 assert_eq!(pack.entries().len(), 1);
1130
1131 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 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 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 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 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 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 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 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 if let Some(p) = &path_string {
1285 crate::encrypt_key_registry::register_asset_path(&enc.dlc_id, p);
1286 }
1287
1288 let plaintext = enc.decrypt_bytes().map_err(|e| {
1291 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 let ext = enc.original_extension;
1310
1311 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 {
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 if !ext.is_empty() {
1344 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 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}