use bevy::asset::io::Reader;
use bevy::asset::{
Asset, AssetLoader, AssetPath, ErasedLoadedAsset, Handle, LoadContext, LoadedUntypedAsset,
};
use bevy::ecs::reflect::AppTypeRegistry;
use bevy::prelude::*;
use bevy::reflect::TypePath;
use futures_lite::{AsyncReadExt, AsyncSeekExt};
use std::io;
use std::sync::Arc;
use thiserror::Error;
use crate::{DlcId, PackItem, Product};
use std::io::{Read, Seek, SeekFrom};
pub struct SyncReader<'a> {
inner: &'a mut dyn bevy::asset::io::Reader,
}
impl<'a> SyncReader<'a> {
pub fn new(inner: &'a mut dyn bevy::asset::io::Reader) -> Self {
SyncReader { inner }
}
}
impl<'a> Read for SyncReader<'a> {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
bevy::tasks::block_on(self.inner.read(buf))
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
}
}
impl<'a> Seek for SyncReader<'a> {
fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
match self.inner.seekable() {
Ok(seek) => bevy::tasks::block_on(seek.seek(pos))
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e)),
Err(_) => Err(std::io::Error::new(
std::io::ErrorKind::Other,
"reader not seekable",
)),
}
}
}
fn decompress_archive(
plaintext: &[u8],
) -> Result<std::collections::HashMap<String, Vec<u8>>, DlcLoaderError> {
use flate2::read::GzDecoder;
use std::io::Read;
use tar::Archive;
let mut archive = Archive::new(GzDecoder::new(std::io::Cursor::new(plaintext)));
let mut map = std::collections::HashMap::new();
for entry in archive
.entries()
.map_err(|e| DlcLoaderError::DecryptionFailed(format!("archive read failed: {}", e)))?
{
let mut file = entry.map_err(|e| {
DlcLoaderError::DecryptionFailed(format!("archive entry read failed: {}", e))
})?;
let path = file
.path()
.map_err(|e| DlcLoaderError::DecryptionFailed(format!("archive path error: {}", e)))?;
let path_str = path.to_string_lossy().replace("\\", "/");
let mut buf = Vec::new();
file.read_to_end(&mut buf).map_err(|e| {
DlcLoaderError::DecryptionFailed(format!("archive file read failed: {}", e))
})?;
map.insert(path_str, buf);
}
Ok(map)
}
pub(crate) fn decrypt_pack_entry_block_bytes<R: std::io::Read + std::io::Seek>(
reader: &mut R,
enc: &EncryptedAsset,
key: &crate::EncryptionKey,
full_path: &str,
) -> Result<Vec<u8>, DlcLoaderError> {
let _original_pos = reader
.stream_position()
.map_err(|e| DlcLoaderError::Io(e))?;
reader
.seek(std::io::SeekFrom::Start(0))
.map_err(|e| DlcLoaderError::Io(e))?;
let (_prod, _id, _ver, _entries, blocks) = crate::parse_encrypted_pack(&mut *reader)
.map_err(|e| DlcLoaderError::InvalidFormat(e.to_string()))?;
let block = blocks
.iter()
.find(|b| b.block_id == enc.block_id)
.ok_or_else(|| {
DlcLoaderError::DecryptionFailed(format!("block {} not found in pack", enc.block_id))
})?;
reader
.seek(std::io::SeekFrom::Start(block.file_offset))
.map_err(|e| DlcLoaderError::Io(e))?;
let mut limited = reader.take(block.encrypted_size as u64);
let pt_gz = crate::pack_format::decrypt_with_key(&key, &mut limited, &block.nonce)
.map_err(|e| DlcLoaderError::DecryptionFailed(e.to_string()))?;
let entries = decompress_archive(&pt_gz)?;
let label = match full_path.rsplit_once('#') {
Some((_, suffix)) => suffix,
None => full_path,
}
.replace("\\", "/");
entries.get(&label).cloned().ok_or_else(|| {
DlcLoaderError::DecryptionFailed(format!(
"entry '{}' not found in decrypted block {}",
label, enc.block_id
))
})
}
#[derive(Event, Clone)]
pub struct DlcPackLoaded {
dlc_id: DlcId,
pack: DlcPack,
}
impl DlcPackLoaded {
pub(crate) fn new(dlc_id: DlcId, pack: DlcPack) -> Self {
DlcPackLoaded { dlc_id, pack }
}
pub fn id(&self) -> &DlcId {
&self.dlc_id
}
pub fn pack(&self) -> &DlcPack {
&self.pack
}
}
pub(crate) fn fuzzy_type_path_match<'a>(stored: &'a str, expected: &'a str) -> bool {
let s = stored.trim_start_matches("::");
let e = expected.trim_start_matches("::");
if s == e {
return true;
}
if e.ends_with(s) && e.as_bytes().get(e.len() - s.len() - 1) == Some(&b':') {
return true;
}
if s.ends_with(e) && s.as_bytes().get(s.len() - e.len() - 1) == Some(&b':') {
return true;
}
false
}
pub trait ErasedSubAssetRegistrar: Send + Sync + 'static {
fn try_register(
&self,
label: String,
erased: ErasedLoadedAsset,
load_context: &mut LoadContext<'_>,
) -> Result<(), ErasedLoadedAsset>;
fn asset_type_path(&self) -> &'static str;
fn load_direct<'a>(
&'a self,
label: String,
fake_path: String,
reader: &'a mut dyn Reader,
load_context: &'a mut LoadContext<'_>,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), DlcLoaderError>> + Send + 'a>>;
}
pub struct TypedSubAssetRegistrar<A: Asset>(std::marker::PhantomData<A>);
impl<A: Asset> Default for TypedSubAssetRegistrar<A> {
fn default() -> Self {
Self(std::marker::PhantomData)
}
}
impl<A: Asset> ErasedSubAssetRegistrar for TypedSubAssetRegistrar<A> {
fn try_register(
&self,
label: String,
erased: ErasedLoadedAsset,
load_context: &mut LoadContext<'_>,
) -> Result<(), ErasedLoadedAsset> {
match erased.downcast::<A>() {
Ok(loaded) => {
load_context.add_loaded_labeled_asset(label, loaded);
Ok(())
}
Err(back) => Err(back),
}
}
fn asset_type_path(&self) -> &'static str {
A::type_path()
}
fn load_direct<'a>(
&'a self,
label: String,
fake_path: String,
reader: &'a mut dyn Reader,
load_context: &'a mut LoadContext<'_>,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), DlcLoaderError>> + Send + 'a>>
{
Box::pin(async move {
match load_context
.loader()
.with_static_type()
.immediate()
.with_reader(reader)
.load::<A>(fake_path)
.await
{
Ok(loaded) => {
load_context.add_loaded_labeled_asset(label, loaded);
Ok(())
}
Err(e) => Err(DlcLoaderError::DecryptionFailed(e.to_string())),
}
})
}
}
#[derive(Clone, Debug, Asset, TypePath)]
pub struct EncryptedAsset {
pub dlc_id: String,
pub original_extension: String,
pub type_path: Option<String>,
pub nonce: [u8; 12],
pub ciphertext: std::sync::Arc<[u8]>,
pub block_id: u32,
pub block_offset: u32,
pub size: u32,
}
impl EncryptedAsset {
pub(crate) fn decrypt_bytes(&self) -> Result<Vec<u8>, DlcLoaderError> {
let encrypt_key = crate::encrypt_key_registry::get(&self.dlc_id)
.ok_or_else(|| DlcLoaderError::DlcLocked(self.dlc_id.clone()))?;
crate::pack_format::decrypt_with_key(
&encrypt_key,
std::io::Cursor::new(&*self.ciphertext),
&self.nonce,
)
.map_err(|e| DlcLoaderError::DecryptionFailed(e.to_string()))
}
}
pub fn parse_encrypted(bytes: &[u8]) -> Result<EncryptedAsset, io::Error> {
if bytes.len() < 1 + 2 + 1 + 12 {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"encrypted file too small",
));
}
let version = bytes[0];
let mut offset = 1usize;
let dlc_len = u16::from_be_bytes([bytes[offset], bytes[offset + 1]]) as usize;
offset += 2;
if offset + dlc_len > bytes.len() {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"invalid dlc id length",
));
}
let dlc_id = String::from_utf8(bytes[offset..offset + dlc_len].to_vec())
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
offset += dlc_len;
let ext_len = bytes[offset] as usize;
offset += 1;
let original_extension = if ext_len == 0 {
"".to_string()
} else {
let s = String::from_utf8(bytes[offset..offset + ext_len].to_vec())
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
offset += ext_len;
s
};
let type_path = if version >= 1 {
if offset + 2 > bytes.len() {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"missing type_path length",
));
}
let tlen = u16::from_be_bytes([bytes[offset], bytes[offset + 1]]) as usize;
offset += 2;
if tlen == 0 {
None
} else {
if offset + tlen > bytes.len() {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"invalid type_path length",
));
}
let s = String::from_utf8(bytes[offset..offset + tlen].to_vec())
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
offset += tlen;
Some(s)
}
} else {
None
};
if offset + 12 > bytes.len() {
return Err(io::Error::new(io::ErrorKind::InvalidData, "missing nonce"));
}
let mut nonce = [0u8; 12];
nonce.copy_from_slice(&bytes[offset..offset + 12]);
offset += 12;
let ciphertext = bytes[offset..].into();
Ok(EncryptedAsset {
dlc_id,
original_extension,
type_path,
nonce,
ciphertext,
block_id: 0,
block_offset: 0,
size: 0,
})
}
#[derive(TypePath)]
pub struct DlcLoader<A: bevy::asset::Asset + 'static> {
#[allow(dead_code)]
type_registry: Arc<AppTypeRegistry>,
_marker: std::marker::PhantomData<A>,
}
impl<A> bevy::prelude::FromWorld for DlcLoader<A>
where
A: bevy::asset::Asset + 'static,
{
fn from_world(world: &mut bevy::prelude::World) -> Self {
let registry = world.resource::<AppTypeRegistry>().clone();
DlcLoader {
type_registry: Arc::new(registry),
_marker: std::marker::PhantomData,
}
}
}
#[derive(TypePath, Clone, Debug)]
pub struct DlcPackEntry {
path: String,
encrypted: EncryptedAsset,
}
impl DlcPackEntry {
pub fn new(path: String, encrypted: EncryptedAsset) -> Self {
DlcPackEntry { path, encrypted }
}
pub fn load_untyped(
&self,
asset_server: &bevy::prelude::AssetServer,
) -> Handle<LoadedUntypedAsset> {
asset_server.load_untyped(&self.path)
}
pub fn is_type<A: Asset>(&self) -> bool {
match self.encrypted.type_path.as_ref() {
Some(tp) => fuzzy_type_path_match(tp, A::type_path()),
None => false,
}
}
pub(crate) fn decrypt_bytes(&self) -> Result<Vec<u8>, DlcLoaderError> {
let entry_ek = crate::encrypt_key_registry::get_full(&self.encrypted.dlc_id)
.ok_or_else(|| DlcLoaderError::DlcLocked(self.encrypted.dlc_id.clone()))?;
let encrypt_key = entry_ek.key;
let path = entry_ek.path.ok_or_else(|| {
DlcLoaderError::DecryptionFailed(format!(
"no file path registered for DLC '{}', cannot decrypt",
self.encrypted.dlc_id
))
})?;
let mut file = std::fs::File::open(path).map_err(|e| {
DlcLoaderError::DecryptionFailed(format!("failed to open pack file: {}", e))
})?;
crate::asset_loader::decrypt_pack_entry_block_bytes(
&mut file,
&self.encrypted,
&encrypt_key,
&self.path,
)
}
pub fn path(&self) -> AssetPath<'_> {
bevy::asset::AssetPath::parse(&self.path)
}
pub fn entry_path(&self) -> &str {
&self.path
}
pub fn original_extension(&self) -> &String {
&self.encrypted.original_extension
}
pub fn type_path(&self) -> Option<&String> {
self.encrypted.type_path.as_ref()
}
}
impl From<(String, EncryptedAsset)> for DlcPackEntry {
fn from((path, encrypted): (String, EncryptedAsset)) -> Self {
DlcPackEntry { path, encrypted }
}
}
impl From<&(String, EncryptedAsset)> for DlcPackEntry {
fn from((path, encrypted): &(String, EncryptedAsset)) -> Self {
DlcPackEntry {
path: path.clone(),
encrypted: encrypted.clone(),
}
}
}
#[derive(Asset, TypePath, Clone, Debug)]
pub struct DlcPack {
dlc_id: DlcId,
product: Product,
version: u8,
metadata: crate::PackMetadata,
metadata_locked: bool,
entries: Vec<DlcPackEntry>,
pack_path: String,
}
impl DlcPack {
pub fn new(id: DlcId, product: Product, version: u8, entries: Vec<DlcPackEntry>) -> Self {
Self::new_with_metadata(id, product, version, crate::PackMetadata::new(), entries)
}
pub fn new_with_metadata(
id: DlcId,
product: Product,
version: u8,
metadata: crate::PackMetadata,
entries: Vec<DlcPackEntry>,
) -> Self {
Self::new_with_metadata_state(id, product, version, metadata, false, entries)
}
pub(crate) fn new_with_metadata_state(
id: DlcId,
product: Product,
version: u8,
metadata: crate::PackMetadata,
metadata_locked: bool,
entries: Vec<DlcPackEntry>,
) -> Self {
DlcPack {
dlc_id: id,
product,
version,
metadata,
metadata_locked,
entries,
pack_path: String::new(),
}
}
pub fn pack_path(&self) -> &str {
&self.pack_path
}
pub fn id(&self) -> &DlcId {
&self.dlc_id
}
pub fn product(&self) -> &str {
&self.product.0
}
pub fn version(&self) -> u8 {
self.version
}
pub fn metadata_locked(&self) -> bool {
self.metadata_locked
}
pub fn metadata_keys(&self) -> impl Iterator<Item = &str> + '_ {
self.metadata.keys().map(String::as_str)
}
pub fn has_metadata(&self, key: &str) -> bool {
self.metadata.contains_key(key)
}
pub fn get_metadata<T>(&self, key: &str) -> Result<Option<T>, crate::DlcPackMetadataError>
where
T: serde::de::DeserializeOwned,
{
if self.metadata_locked {
return Err(crate::DlcPackMetadataError::Locked);
}
match self.metadata.get(key) {
Some(value) => serde_json::from_value(value.clone())
.map(Some)
.map_err(|source| crate::DlcPackMetadataError::Deserialize {
key: key.to_string(),
source,
}),
None => Ok(None),
}
}
pub fn get_metadata_raw(
&self,
key: &str,
) -> Result<Option<&serde_json::Value>, crate::DlcPackMetadataError> {
if self.metadata_locked {
return Err(crate::DlcPackMetadataError::Locked);
}
Ok(self.metadata.get(key))
}
pub fn entries(&self) -> &[DlcPackEntry] {
&self.entries
}
pub fn find_entry(&self, path: &str) -> Option<&DlcPackEntry> {
self.entries
.iter()
.find(|e| e.path().to_string().ends_with(path) || e.path().path().ends_with(path))
}
pub fn find_by_type<A: Asset>(&self) -> Vec<&DlcPackEntry> {
self.entries
.iter()
.filter(|e| match e.type_path() {
Some(tp) => fuzzy_type_path_match(tp, A::type_path()),
None => false,
})
.collect()
}
pub fn decrypt_entry(
&self,
entry_path: &str,
) -> Result<Vec<u8>, crate::asset_loader::DlcLoaderError> {
let entry = self.find_entry(entry_path).ok_or_else(|| {
DlcLoaderError::DecryptionFailed(format!("entry not found: {}", entry_path))
})?;
entry.decrypt_bytes()
}
pub fn load<A: Asset>(
&self,
asset_server: &bevy::prelude::AssetServer,
entry_path: &str,
) -> Handle<A> {
if let Some(entry) = self.find_entry(entry_path) {
return asset_server.load::<A>(entry.path());
}
let normalized = entry_path.replace('\\', "/");
let requested_path = if !self.pack_path.is_empty() {
format!("{}#{}", self.pack_path, normalized)
} else if let Some((pack_path, _)) = self.entries.first().and_then(|entry| entry.path.split_once('#')) {
format!("{}#{}", pack_path, normalized)
} else {
normalized.clone()
};
warn!(
"DLC entry '{}' not found in pack '{}'; requesting '{}' so the asset server reports a normal load failure instead of panicking",
entry_path,
self.dlc_id,
requested_path,
);
asset_server.load::<A>(requested_path)
}
fn with_path(&self, path_string: String) -> DlcPack {
DlcPack {
dlc_id: self.dlc_id.clone(),
product: self.product.clone(),
version: self.version,
metadata: self.metadata.clone(),
metadata_locked: self.metadata_locked,
entries: self.entries.clone(),
pack_path: path_string,
}
}
}
#[derive(Clone, TypePath, serde::Serialize, serde::Deserialize)]
pub struct DlcPackLoaderSettings {}
impl Default for DlcPackLoaderSettings {
fn default() -> Self {
DlcPackLoaderSettings {}
}
}
#[derive(TypePath, Default)]
pub struct DlcPackLoader {
pub registrars: Vec<Box<dyn ErasedSubAssetRegistrar>>,
pub(crate) factories: Option<DlcPackRegistrarFactories>,
}
pub trait DlcPackRegistrarFactory: Send + Sync + 'static {
fn type_name(&self) -> &'static str;
fn create_registrar(&self) -> Box<dyn ErasedSubAssetRegistrar>;
}
pub struct TypedRegistrarFactory<T: Asset + 'static>(std::marker::PhantomData<T>);
impl<T: Asset + TypePath + 'static> DlcPackRegistrarFactory for TypedRegistrarFactory<T> {
fn type_name(&self) -> &'static str {
T::type_path()
}
fn create_registrar(&self) -> Box<dyn ErasedSubAssetRegistrar> {
Box::new(TypedSubAssetRegistrar::<T>::default())
}
}
impl<T: Asset + 'static> Default for TypedRegistrarFactory<T> {
fn default() -> Self {
TypedRegistrarFactory(std::marker::PhantomData)
}
}
use std::sync::RwLock;
#[derive(Clone, Resource)]
pub(crate) struct DlcPackRegistrarFactories(pub Arc<RwLock<Vec<Box<dyn DlcPackRegistrarFactory>>>>);
impl Default for DlcPackRegistrarFactories {
fn default() -> Self {
DlcPackRegistrarFactories(Arc::new(RwLock::new(Vec::new())))
}
}
pub(crate) fn default_pack_registrar_factories() -> Vec<Box<dyn DlcPackRegistrarFactory>> {
vec![
Box::new(TypedRegistrarFactory::<Image>::default()),
Box::new(TypedRegistrarFactory::<Scene>::default()),
Box::new(TypedRegistrarFactory::<bevy::mesh::Mesh>::default()),
Box::new(TypedRegistrarFactory::<Font>::default()),
Box::new(TypedRegistrarFactory::<AudioSource>::default()),
Box::new(TypedRegistrarFactory::<ColorMaterial>::default()),
Box::new(TypedRegistrarFactory::<bevy::pbr::StandardMaterial>::default()),
Box::new(TypedRegistrarFactory::<bevy::gltf::Gltf>::default()),
Box::new(TypedRegistrarFactory::<bevy::gltf::GltfMesh>::default()),
Box::new(TypedRegistrarFactory::<Shader>::default()),
Box::new(TypedRegistrarFactory::<DynamicScene>::default()),
Box::new(TypedRegistrarFactory::<AnimationClip>::default()),
Box::new(TypedRegistrarFactory::<AnimationGraph>::default()),
]
}
pub(crate) fn collect_pack_registrars(
factories: Option<&DlcPackRegistrarFactories>,
) -> Vec<Box<dyn ErasedSubAssetRegistrar>> {
use std::collections::HashSet;
let mut seen: HashSet<&'static str> = HashSet::new();
let mut out: Vec<Box<dyn ErasedSubAssetRegistrar>> = Vec::new();
if let Some(f) = factories {
let inner = f.0.read().unwrap();
for factory in inner.iter() {
out.push(factory.create_registrar());
seen.insert(factory.type_name());
}
}
for factory in default_pack_registrar_factories() {
if !seen.contains(factory.type_name()) {
out.push(factory.create_registrar());
seen.insert(factory.type_name());
}
}
out
}
impl AssetLoader for DlcPackLoader {
type Asset = DlcPack;
type Settings = DlcPackLoaderSettings;
type Error = DlcLoaderError;
fn extensions(&self) -> &[&str] {
&["dlcpack"]
}
async fn load(
&self,
reader: &mut dyn Reader,
_settings: &Self::Settings,
load_context: &mut LoadContext<'_>,
) -> Result<Self::Asset, Self::Error> {
let path_string = load_context.path().path().to_string_lossy().to_string();
let mut sync_reader = SyncReader::new(reader);
let mut parsed_pack = crate::parse_encrypted_pack_info(&mut sync_reader, None)
.map_err(|e| DlcLoaderError::InvalidFormat(e.to_string()))?;
if let Some(encrypt_key) = crate::encrypt_key_registry::get(parsed_pack.dlc_id.as_ref()) {
sync_reader.seek(SeekFrom::Start(0)).map_err(|_| {
DlcLoaderError::Io(io::Error::new(
io::ErrorKind::NotSeekable,
format!(
"reader not seekable, cannot decrypt metadata for pack '{}'",
path_string
),
))
})?;
parsed_pack = crate::parse_encrypted_pack_info(&mut sync_reader, Some(&encrypt_key))
.map_err(|e| DlcLoaderError::InvalidFormat(e.to_string()))?;
}
let crate::ParsedDlcPack {
product,
dlc_id,
version,
metadata,
metadata_locked,
entries: manifest_entries,
block_metadatas,
} = parsed_pack;
sync_reader.seek(SeekFrom::Start(0)).map_err(|_| {
DlcLoaderError::Io(io::Error::new(
io::ErrorKind::NotSeekable,
format!("reader not seekable, cannot load pack '{}'", path_string),
))
})?;
check_dlc_id_conflict(&dlc_id, &path_string)?;
if !crate::encrypt_key_registry::has(dlc_id.as_ref(), &path_string) {
crate::encrypt_key_registry::register_asset_path(dlc_id.as_ref(), &path_string);
}
let decrypted_items = {
match decrypt_pack_entries(&dlc_id, &manifest_entries, &block_metadatas, sync_reader) {
Ok(items) => Some(items),
Err(DlcLoaderError::DlcLocked(_)) => None,
Err(e) => return Err(e),
}
};
let mut out_entries = Vec::with_capacity(manifest_entries.len());
let mut unregistered_labels: Vec<String> = Vec::new();
let dynamic_regs = self
.factories
.as_ref()
.map(|f| crate::asset_loader::collect_pack_registrars(Some(f)));
let regs = dynamic_regs.unwrap_or_else(|| collect_pack_registrars(None));
for (path, enc) in manifest_entries.into_iter() {
let entry_label = path.replace('\\', "/");
let mut registered_as_labeled = false;
if let Some(ref items) = decrypted_items {
if let Some(item) = items.iter().find(|i| i.path() == path) {
let ext = item.ext().unwrap_or_default();
let type_path = item.type_path();
let plaintext = item.plaintext().to_vec();
let stem = std::path::Path::new(&entry_label)
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or("entry");
let fake_path = format!("{}.{}", stem, ext);
let mut vec_reader = bevy::asset::io::VecReader::new(plaintext.clone());
if let Some(tp) = type_path {
if let Some(registrar) = regs
.iter()
.find(|r| fuzzy_type_path_match(r.asset_type_path(), tp.as_str()))
{
match registrar
.load_direct(
entry_label.clone(),
fake_path.clone(),
&mut vec_reader,
load_context,
)
.await
{
Ok(()) => {
registered_as_labeled = true;
}
Err(e) => {
debug!(
"Static load for type '{}' failed: {}; falling back to extension dispatch",
tp, e
);
}
}
}
}
if !registered_as_labeled {
let mut vec_reader = bevy::asset::io::VecReader::new(plaintext.clone());
let result = load_context
.loader()
.immediate()
.with_reader(&mut vec_reader)
.with_unknown_type()
.load(fake_path.clone())
.await;
match result {
Ok(erased) => {
let mut remaining = Some(erased);
for registrar in regs.iter() {
let label = entry_label.clone();
let to_register = remaining.take().unwrap();
match registrar.try_register(label, to_register, load_context) {
Ok(()) => {
registered_as_labeled = true;
remaining = None;
break;
}
Err(back) => {
remaining = Some(back);
}
}
}
if let Some(_) = remaining {
warn!(
"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>()`",
entry_label, ext, path_string, entry_label
);
}
}
Err(e) => {
warn!(
"Failed to load entry '{}', extension='{}': {}",
entry_label, ext, e
);
}
}
}
}
}
let registered_path = format!("{}#{}", path_string, entry_label);
if !registered_as_labeled {
unregistered_labels.push(entry_label.clone());
}
out_entries.push(DlcPackEntry {
path: registered_path,
encrypted: enc,
});
}
if decrypted_items.is_some() && !unregistered_labels.is_empty() {
let example_label = &unregistered_labels[0];
let example_full = format!("{}#{}", path_string, example_label);
warn!(
"{} {} 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>()`.",
unregistered_labels.len(),
if unregistered_labels.len() == 1 {
"entry"
} else {
"entries"
},
path_string,
example_full,
);
}
Ok(DlcPack::new_with_metadata_state(
dlc_id.clone(),
product,
version as u8,
metadata,
metadata_locked,
out_entries,
)
.with_path(path_string))
}
}
fn check_dlc_id_conflict(dlc_id: &DlcId, path_string: &str) -> Result<(), DlcLoaderError> {
if let Some(existing_path) = crate::encrypt_key_registry::asset_path_for(dlc_id.as_ref()) {
if existing_path != path_string {
return Err(DlcLoaderError::DlcIdConflict(
dlc_id.to_string(),
existing_path,
path_string.to_string(),
));
}
}
Ok(())
}
fn decrypt_pack_entries<R: std::io::Read + std::io::Seek>(
dlc_id: &crate::DlcId,
entries: &[(String, EncryptedAsset)],
block_metadatas: &[crate::pack_format::BlockMetadata],
reader: R,
) -> Result<Vec<crate::PackItem>, DlcLoaderError> {
let encrypt_key = crate::encrypt_key_registry::get(dlc_id.as_ref())
.ok_or_else(|| DlcLoaderError::DlcLocked(dlc_id.to_string()))?;
let mut extracted_all = std::collections::HashMap::new();
let mut pr = crate::pack_format::PackReader::new(reader);
for block in block_metadatas {
pr.seek(std::io::SeekFrom::Start(block.file_offset))
.map_err(|e| DlcLoaderError::Io(e))?;
let pt = pr
.read_and_decrypt(&encrypt_key, block.encrypted_size as usize, &block.nonce)
.map_err(|e| {
let example = entries
.iter()
.find(|(_, enc)| enc.block_id == block.block_id)
.map(|(p, _)| p.as_str())
.unwrap_or("unknown");
DlcLoaderError::DecryptionFailed(format!(
"dlc='{}' entry='{}' (block {}) decryption failed: {}",
dlc_id, example, block.block_id, e
))
})?;
let extracted = decompress_archive(&pt)?;
extracted_all.extend(extracted);
}
let mut out = Vec::with_capacity(entries.len());
for (path, enc) in entries {
let normalized = path.replace("\\", "/");
let plaintext = extracted_all
.remove(&normalized)
.or_else(|| extracted_all.remove(path.as_str()))
.ok_or_else(|| {
DlcLoaderError::DecryptionFailed(format!("entry {} not found in any block", path))
})?;
let mut item = PackItem::new(path.clone(), plaintext)
.map_err(|e| DlcLoaderError::InvalidFormat(e.to_string()))?;
if !enc.original_extension.is_empty() {
item = item
.with_extension(enc.original_extension.clone())
.map_err(|e| DlcLoaderError::InvalidFormat(e.to_string()))?;
}
if let Some(tp) = &enc.type_path {
item = item.with_type_path(tp.clone());
}
out.push(item);
}
Ok(out)
}
#[derive(Error, Debug)]
pub enum DlcLoaderError {
#[error("IO error: {0}")]
Io(io::Error),
#[error("DLC locked: encrypt key not found for DLC id: {0}")]
DlcLocked(String),
#[error("Decryption failed: {0}")]
DecryptionFailed(String),
#[error("Invalid encrypted asset format: {0}")]
InvalidFormat(String),
#[error(
"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}"
)]
DlcIdConflict(String, String, String),
}
impl From<std::io::Error> for DlcLoaderError {
fn from(e: std::io::Error) -> Self {
DlcLoaderError::Io(e)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{EncryptionKey, PackItem};
#[allow(unused_imports)]
use secure_gate::{CloneableSecret, RevealSecret};
use serial_test::serial;
#[test]
#[serial]
fn encrypted_asset_decrypts_with_registry() {
let dlc_id = "standalone";
let key = EncryptionKey::new(rand::random());
crate::encrypt_key_registry::clear_all();
crate::encrypt_key_registry::insert(dlc_id, key.with_secret(|k| EncryptionKey::from(*k)));
let nonce = [0u8; 12];
let mut ciphertext = Vec::new();
{
let mut pw = crate::pack_format::PackWriter::new(&mut ciphertext);
pw.write_encrypted(&key, &nonce, b"hello").expect("encrypt");
}
let ct_len = ciphertext.len() as u32;
let enc = EncryptedAsset {
dlc_id: dlc_id.to_string(),
original_extension: "".to_string(),
type_path: None,
nonce,
ciphertext: ciphertext.into(),
block_id: 0,
block_offset: 0,
size: ct_len,
};
let plaintext = enc.decrypt_bytes().expect("decrypt");
assert_eq!(&plaintext, b"hello");
}
#[test]
#[serial]
fn dlcpack_accessors_work_and_fields_read() {
let entry = DlcPackEntry {
path: "a.txt".to_string(),
encrypted: EncryptedAsset {
dlc_id: "example_dlc".to_string(),
original_extension: "txt".to_string(),
type_path: None,
nonce: [0u8; 12],
ciphertext: vec![].into(),
block_id: 0,
block_offset: 0,
size: 0,
},
};
let pack = DlcPack::new(
DlcId::from("example_dlc"),
Product::from("test"),
4,
vec![entry.clone()],
);
assert_eq!(*pack.id(), DlcId::from("example_dlc"));
assert_eq!(pack.entries().len(), 1);
let found = pack.find_entry("a.txt").expect("entry present");
assert_eq!(found.path().path(), "a.txt");
assert_eq!(found.original_extension(), "txt");
assert!(found.type_path().is_none());
}
#[test]
#[serial]
fn decrypt_pack_entries_v4_without_key_returns_locked_error() {
crate::encrypt_key_registry::clear_all();
let dlc_id = crate::DlcId::from("locked_dlc");
let items = vec![PackItem::new("a.txt", b"hello".to_vec()).expect("pack item")];
let key = EncryptionKey::new(rand::random());
let _dlc_key = crate::DlcKey::generate_random();
let product = crate::Product::from("test");
let container = crate::pack_encrypted_pack(
&dlc_id,
&items,
&product,
&key,
crate::pack_format::DEFAULT_BLOCK_SIZE,
)
.expect("pack");
let mut cursor = std::io::Cursor::new(container);
let (_product, parsed_dlc_id, _version, parsed_entries, block_metadatas) =
crate::parse_encrypted_pack(&mut cursor).expect("parse");
let err = decrypt_pack_entries(&parsed_dlc_id, &parsed_entries, &block_metadatas, cursor)
.expect_err("should be locked");
match err {
DlcLoaderError::DlcLocked(id) => assert_eq!(id, "locked_dlc"),
_ => panic!("expected DlcLocked error, got {:?}", err),
}
}
#[test]
#[serial]
fn decrypt_pack_entries_v4_with_wrong_key_reports_entry_and_dlc() {
crate::encrypt_key_registry::clear_all();
let dlc_id = crate::DlcId::from("badkey_dlc");
let items = vec![PackItem::new("b.txt", b"world".to_vec()).expect("pack item")];
let real_key = EncryptionKey::new(rand::random());
let _dlc_key = crate::DlcKey::generate_random();
let product = crate::Product::from("test");
let container = crate::pack_encrypted_pack(
&dlc_id,
&items,
&product,
&real_key,
crate::pack_format::DEFAULT_BLOCK_SIZE,
)
.expect("pack");
let wrong_key: [u8; 32] = rand::random();
crate::encrypt_key_registry::insert(
&dlc_id.to_string(),
crate::EncryptionKey::from(wrong_key),
);
let mut cursor = std::io::Cursor::new(container);
let (_product, parsed_dlc_id, _version, parsed_entries, block_metadatas) =
crate::parse_encrypted_pack(&mut cursor).expect("parse");
let err = decrypt_pack_entries(&parsed_dlc_id, &parsed_entries, &block_metadatas, cursor)
.expect_err("should fail decryption");
match err {
DlcLoaderError::DecryptionFailed(msg) => {
assert!(msg.contains("dlc='badkey_dlc'"));
assert!(msg.contains("entry='b.txt'"));
assert!(msg.contains("authentication failed") || msg.contains("incorrect key"));
}
_ => panic!("expected DecryptionFailed, got {:?}", err),
}
}
#[test]
#[serial]
fn dlc_id_conflict_detection() {
crate::encrypt_key_registry::clear_all();
let dlc_id_str = "conflict_test_dlc";
let pack_path_1 = "existing_pack.dlcpack";
let pack_path_2 = "different_pack.dlcpack";
crate::encrypt_key_registry::register_asset_path(dlc_id_str, pack_path_1);
assert!(
!crate::encrypt_key_registry::check(dlc_id_str, pack_path_1),
"same pack path should NOT be a conflict"
);
let mut tries = 0;
while tries < 100 && !crate::encrypt_key_registry::check(dlc_id_str, pack_path_2) {
crate::encrypt_key_registry::register_asset_path(dlc_id_str, pack_path_1);
std::thread::sleep(std::time::Duration::from_millis(5));
tries += 1;
}
assert!(
crate::encrypt_key_registry::check(dlc_id_str, pack_path_2),
"different pack path SHOULD be detected as a conflict"
);
crate::encrypt_key_registry::clear_all();
}
#[test]
#[serial]
fn dlc_loader_conflict_helper_allows_same_path() {
crate::encrypt_key_registry::clear_all();
let dlc_id = crate::DlcId::from("foo");
let path = "same_pack.dlcpack";
crate::encrypt_key_registry::register_asset_path(dlc_id.as_ref(), path);
assert!(check_dlc_id_conflict(&dlc_id, path).is_ok());
}
#[test]
#[serial]
fn dlc_loader_conflict_helper_rejects_different_path() {
crate::encrypt_key_registry::clear_all();
let dlc_id = crate::DlcId::from("foo");
crate::encrypt_key_registry::register_asset_path(dlc_id.as_ref(), "other.dlcpack");
let err = check_dlc_id_conflict(&dlc_id, "new.dlcpack").expect_err("should conflict");
match err {
DlcLoaderError::DlcIdConflict(id, orig, newp) => {
assert_eq!(id, dlc_id.to_string());
assert_eq!(orig, "other.dlcpack");
assert_eq!(newp, "new.dlcpack");
}
_ => panic!("expected DlcIdConflict"),
}
}
#[test]
#[serial]
fn dlcpack_load_missing_entry_returns_handle_without_panicking() {
use bevy::asset::LoadState;
crate::encrypt_key_registry::clear_all();
let mut app = App::new();
app.add_plugins(MinimalPlugins);
app.add_plugins(bevy::asset::AssetPlugin::default());
app.init_asset::<crate::asset_loader::DlcPack>();
let asset_server = app.world().resource::<AssetServer>().clone();
let pack = DlcPack::new(
DlcId::from("missing_entry_dlc"),
Product::from("test"),
crate::DLC_PACK_VERSION_LATEST,
vec![DlcPackEntry {
path: "missing_entry.dlcpack#present.txt".to_string(),
encrypted: EncryptedAsset {
dlc_id: "missing_entry_dlc".to_string(),
original_extension: "txt".to_string(),
type_path: Some("examples::TextAsset".to_string()),
nonce: [0u8; 12],
ciphertext: Arc::new([]),
block_id: 0,
block_offset: 0,
size: 0,
},
}],
)
.with_path("missing_entry.dlcpack".to_string());
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
pack.load::<LoadedUntypedAsset>(&asset_server, "absent.txt")
}));
let handle = result.expect("missing entry load should not panic");
for _ in 0..10 {
app.update();
}
let state = asset_server.get_load_state(handle.id()).into();
assert!(
matches!(state, Some(LoadState::Loading) | Some(LoadState::Failed(_)) | None),
"unexpected load state: {:?}",
state
);
}
}
impl<A> AssetLoader for DlcLoader<A>
where
A: bevy::asset::Asset + TypePath + 'static,
{
type Asset = A;
type Settings = ();
type Error = DlcLoaderError;
async fn load(
&self,
reader: &mut dyn Reader,
_settings: &Self::Settings,
load_context: &mut LoadContext<'_>,
) -> Result<Self::Asset, Self::Error> {
let path_string = Some(load_context.path().path().to_string_lossy().to_string());
let mut bytes = Vec::new();
reader
.read_to_end(&mut bytes)
.await
.map_err(|e| DlcLoaderError::Io(e))?;
let enc =
parse_encrypted(&bytes).map_err(|e| DlcLoaderError::DecryptionFailed(e.to_string()))?;
if let Some(p) = &path_string {
crate::encrypt_key_registry::register_asset_path(&enc.dlc_id, p);
}
let plaintext = enc.decrypt_bytes().map_err(|e| {
match e {
DlcLoaderError::DecryptionFailed(msg) => DlcLoaderError::DecryptionFailed(format!(
"dlc='{}' path='{}' {}",
enc.dlc_id,
path_string
.clone()
.unwrap_or_else(|| "<unknown>".to_string()),
msg,
)),
other => other,
}
})?;
let ext = enc.original_extension;
let bytes_clone = plaintext.clone();
let stem = load_context
.path()
.path()
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or("dlc_decrypted");
let fake_path = format!("{}.{}", stem, ext);
{
let mut static_reader = bevy::asset::io::VecReader::new(bytes_clone.clone());
if let Ok(loaded) = load_context
.loader()
.with_static_type()
.immediate()
.with_reader(&mut static_reader)
.load::<A>(fake_path.clone())
.await
{
return Ok(loaded.take());
}
}
if !ext.is_empty() {
let mut ext_reader = bevy::asset::io::VecReader::new(bytes_clone.clone());
let attempt = load_context
.loader()
.immediate()
.with_reader(&mut ext_reader)
.with_unknown_type()
.load(fake_path.clone())
.await;
if let Ok(erased) = attempt {
match erased.downcast::<A>() {
Ok(loaded) => return Ok(loaded.take()),
Err(_) => {
return Err(DlcLoaderError::DecryptionFailed(format!(
"dlc loader: extension-based load succeeded but downcast to '{}' failed",
A::type_path(),
)));
}
}
} else if let Err(e) = attempt {
return Err(DlcLoaderError::DecryptionFailed(e.to_string()));
}
}
Err(DlcLoaderError::DecryptionFailed(format!(
"dlc loader: unable to load decrypted asset as {}{}",
A::type_path(),
if ext.is_empty() {
""
} else {
" (extension fallback also failed)"
}
)))
}
}