Skip to main content

rpkg_rs/resource/
package_builder.rs

1use std::borrow::Borrow;
2use std::collections::HashMap;
3use std::fs::File;
4use std::io;
5use std::io::{BufWriter, Read, Seek, Write};
6use std::path::{Path, PathBuf};
7
8use crate::resource::pdefs::{PartitionId, PartitionType};
9use crate::resource::resource_package::{
10    ChunkType, PackageHeader, PackageMetadata, PackageOffsetFlags, PackageOffsetInfo,
11    PackageVersion, ResourceHeader, ResourcePackage, ResourcePackageSource,
12    ResourceReferenceCountAndFlags, ResourceReferenceFlags,
13};
14use crate::resource::resource_partition::PatchId;
15use crate::resource::runtime_resource_id::RuntimeResourceID;
16use crate::{GlacierResource, GlacierResourceError, WoaVersion};
17use binrw::BinWrite;
18use binrw::__private::Required;
19use binrw::io::Cursor;
20use binrw::meta::WriteEndian;
21use indexmap::{IndexMap, IndexSet};
22use lzzzz::{lz4, lz4_hc};
23use thiserror::Error;
24use crate::resource::resource_info::ResourceInfo;
25
26/// `PackageResourceBlob` is an enum representing various types of package resource stores, which can
27/// include files, file sections, and memory buffers, optionally compressed or scrambled.
28enum PackageResourceBlob {
29    File {
30        path: PathBuf,
31        size: u32,
32        compression_level: Option<i32>,
33        should_scramble: bool,
34    },
35    FileAtOffset {
36        path: PathBuf,
37        offset: u64,
38        size: u32,
39        compressed_size: Option<u32>,
40        is_scrambled: bool,
41    },
42    Memory {
43        data: Vec<u8>,
44        compression_level: Option<i32>,
45        should_scramble: bool,
46    },
47    CompressedMemory {
48        data: Vec<u8>,
49        decompressed_size: Option<u32>,
50        is_scrambled: bool,
51    },
52}
53
54impl PackageResourceBlob {
55    /// The (uncompressed) size of the resource blob in bytes.
56    pub fn size(&self) -> u32 {
57        match self {
58            PackageResourceBlob::File { size, .. } => *size,
59            PackageResourceBlob::FileAtOffset { size, .. } => *size,
60            PackageResourceBlob::Memory { data, .. } => data.len() as u32,
61            PackageResourceBlob::CompressedMemory {
62                data,
63                decompressed_size,
64                ..
65            } => match decompressed_size {
66                Some(size) => *size,
67                None => data.len() as u32,
68            },
69        }
70    }
71}
72
73/// A builder for creating a resource within a ResourcePackage
74pub struct PackageResourceBuilder {
75    rrid: RuntimeResourceID,
76    blob: PackageResourceBlob,
77    resource_type: [u8; 4],
78    system_memory_requirement: u32,
79    video_memory_requirement: u32,
80    // We store references in a vector because their order is important and there can be duplicates.
81    references: Vec<(RuntimeResourceID, ResourceReferenceFlags)>,
82}
83
84#[derive(Debug, Error)]
85pub enum PackageResourceBuilderError {
86    #[error("Error reading the file: {0}")]
87    IoError(#[from] io::Error),
88
89    #[error("File is too large")]
90    FileTooLarge,
91
92    #[error("The offset you provided is after the end of the file")]
93    InvalidFileOffset,
94
95    #[error("The size you provided extends beyond the end of the file")]
96    InvalidFileBlobSize,
97
98    #[error("Resource types must be exactly 4 characters")]
99    InvalidResourceType,
100
101    #[error("Internal Glacier resource error")]
102    GlacierResourceError(#[from] GlacierResourceError),
103}
104
105/// A builder for creating a resource within a ResourcePackage.
106impl PackageResourceBuilder {
107    /// Converts a resource type string to a byte array.
108    /// Characters are reversed since everything is little endian.
109    fn resource_type_to_bytes(resource_type: &str) -> Result<[u8; 4], PackageResourceBuilderError> {
110        resource_type
111            .chars()
112            .rev()
113            .collect::<String>()
114            .as_bytes()
115            .try_into()
116            .map_err(|_| PackageResourceBuilderError::InvalidResourceType)
117    }
118
119    /// Create a new resource builder from a file on disk.
120    ///
121    /// # Arguments
122    /// * `rrid` - The resource ID of the resource.
123    /// * `resource_type` - The type of the resource.
124    /// * `path` - The path to the file.
125    /// * `compression_level` - The compression level to use for the file, or None for no compression.
126    /// * `should_scramble` - Whether the file data should be scrambled.
127    pub fn from_file(
128        rrid: RuntimeResourceID,
129        resource_type: &str,
130        path: &Path,
131        compression_level: Option<i32>,
132        should_scramble: bool,
133    ) -> Result<Self, PackageResourceBuilderError> {
134        let file_size = path
135            .metadata()
136            .map_err(PackageResourceBuilderError::IoError)?
137            .len();
138
139        if file_size >= u32::MAX as u64 {
140            return Err(PackageResourceBuilderError::FileTooLarge);
141        }
142
143        Ok(Self {
144            rrid,
145            resource_type: Self::resource_type_to_bytes(resource_type)?,
146            system_memory_requirement: file_size as u32,
147            video_memory_requirement: u32::MAX,
148            references: vec![],
149            blob: PackageResourceBlob::File {
150                path: path.to_path_buf(),
151                size: file_size as u32,
152                compression_level,
153                should_scramble,
154            },
155        })
156    }
157
158    /// Create a new resource builder from a file on disk, but only reading a part of it.
159    ///
160    /// # Arguments
161    /// * `rrid` - The resource ID of the resource.
162    /// * `resource_type` - The type of the resource.
163    /// * `path` - The path to the file.
164    /// * `offset` - The offset of the file to start reading from.
165    /// * `size` - The size of the data.
166    /// * `compressed_size` - The compressed size of the data, if the resource is compressed.
167    /// * `is_scrambled` - Whether the data is scrambled.
168    fn from_file_at_offset(
169        rrid: RuntimeResourceID,
170        resource_type: &str,
171        path: &Path,
172        offset: u64,
173        size: u32,
174        compressed_size: Option<u32>,
175        is_scrambled: bool,
176    ) -> Result<Self, PackageResourceBuilderError> {
177        let file_size = path
178            .metadata()
179            .map_err(PackageResourceBuilderError::IoError)?
180            .len();
181
182        if offset >= file_size {
183            return Err(PackageResourceBuilderError::InvalidFileOffset);
184        }
185
186        let read_size = compressed_size.unwrap_or(size);
187
188        if offset + read_size as u64 > file_size {
189            return Err(PackageResourceBuilderError::InvalidFileBlobSize);
190        }
191
192        Ok(Self {
193            rrid,
194            resource_type: Self::resource_type_to_bytes(resource_type)?,
195            system_memory_requirement: size,
196            video_memory_requirement: u32::MAX,
197            references: vec![],
198            blob: PackageResourceBlob::FileAtOffset {
199                path: path.to_path_buf(),
200                offset,
201                size,
202                compressed_size,
203                is_scrambled,
204            },
205        })
206    }
207
208    /// Create a new resource builder from a (possibly compressed) in-memory blob.
209    ///
210    /// # Arguments
211    /// * `rrid` - The resource ID of the resource.
212    /// * `resource_type` - The type of the resource.
213    /// * `data` - The data of the resource.
214    /// * `decompressed_size` - The decompressed size of the data, if the resource is compressed.
215    /// * `is_scrambled` - Whether the data is scrambled.
216    fn from_compressed_memory(
217        rrid: RuntimeResourceID,
218        resource_type: &str,
219        data: Vec<u8>,
220        decompressed_size: Option<u32>,
221        is_scrambled: bool,
222    ) -> Result<Self, PackageResourceBuilderError> {
223        if data.len() > u32::MAX as usize {
224            return Err(PackageResourceBuilderError::FileTooLarge);
225        }
226
227        let real_size = decompressed_size.unwrap_or(data.len() as u32);
228
229        Ok(Self {
230            rrid,
231            resource_type: Self::resource_type_to_bytes(resource_type)?,
232            system_memory_requirement: real_size,
233            video_memory_requirement: u32::MAX,
234            references: vec![],
235            blob: PackageResourceBlob::CompressedMemory {
236                data,
237                decompressed_size,
238                is_scrambled,
239            },
240        })
241    }
242
243    /// Create a new resource builder from an in-memory blob.
244    ///
245    /// This is similar to `from_compressed_memory`, but it expects the data to be uncompressed and
246    /// can optionally compress and scramble it.
247    ///
248    /// # Arguments
249    /// * `rrid` - The resource ID of the resource.
250    /// * `resource_type` - The type of the resource.
251    /// * `data` - The data of the resource.
252    /// * `compression_level` - The compression level to use for the data, or None for no compression.
253    /// * `should_scramble` - Whether the data should be scrambled.
254    pub fn from_memory(
255        rrid: RuntimeResourceID,
256        resource_type: &str,
257        data: Vec<u8>,
258        compression_level: Option<i32>,
259        should_scramble: bool,
260    ) -> Result<Self, PackageResourceBuilderError> {
261        if data.len() > u32::MAX as usize {
262            return Err(PackageResourceBuilderError::FileTooLarge);
263        }
264
265        let real_size = data.len() as u32;
266
267        Ok(Self {
268            rrid,
269            resource_type: Self::resource_type_to_bytes(resource_type)?,
270            system_memory_requirement: real_size,
271            video_memory_requirement: u32::MAX,
272            references: vec![],
273            blob: PackageResourceBlob::Memory {
274                data,
275                compression_level,
276                should_scramble,
277            },
278        })
279    }
280
281    /// Create a new resource builder from a `ResourceInfo` struct.
282    ///
283    /// This is similar to `from_compressed_memory`, but it expects the data to be uncompressed and
284    /// can optionally compress and scramble it.
285    ///
286    /// # Arguments
287    /// * `resource_info` - The resource info of the resource.
288    /// * `data` - The data of the resource.
289    /// * `is_packaged_data` - Should be enabled when the data is packaged (compressed and/or scrambled)
290    pub fn from_resource_info(
291        resource_info: ResourceInfo,
292        data: Vec<u8>,
293        is_packaged_data: bool,
294    ) -> Result<Self, PackageResourceBuilderError> {
295        if data.len() > u32::MAX as usize {
296            return Err(PackageResourceBuilderError::FileTooLarge);
297        }
298
299        let blob = if is_packaged_data {
300            PackageResourceBlob::CompressedMemory {
301                data,
302                decompressed_size: Some(resource_info.size()),
303                is_scrambled: resource_info.is_scrambled(),
304            }
305        } else {
306            PackageResourceBlob::Memory {
307                data,
308                compression_level: if resource_info.is_compressed() {Some(12)} else {None},
309                should_scramble: resource_info.is_scrambled(),
310            }
311        };
312
313        Ok(Self {
314            rrid: resource_info.entry.runtime_resource_id,
315            resource_type: resource_info.header.resource_type,
316            system_memory_requirement: resource_info.header.system_memory_requirement,
317            video_memory_requirement: resource_info.header.video_memory_requirement,
318            references: resource_info.references().to_vec(),
319            blob,
320        })
321    }
322
323    /// Create a new resource builder from a GlacierResource.
324    ///
325    /// # Arguments
326    /// * `rrid` - The resource ID of the resource.
327    /// * `glacier_resource` - A reference to an object implementing the `GlacierResource` trait.
328    /// * `woa_version` - The HITMAN game version you want to construct the GlacierResource for
329    pub fn from_glacier_resource<G: GlacierResource>(
330        rrid: RuntimeResourceID,
331        glacier_resource: &G,
332        woa_version: WoaVersion
333    ) -> Result<Self, PackageResourceBuilderError> {
334        let system_memory_requirement = glacier_resource.system_memory_requirement();
335        let video_memory_requirement = glacier_resource.video_memory_requirement();
336        let data = glacier_resource
337            .serialize(woa_version)
338            .map_err(PackageResourceBuilderError::GlacierResourceError)?;
339
340        Ok(Self {
341            rrid,
342            resource_type: G::resource_type().into_iter().rev().collect::<Vec<_>>().try_into().unwrap(),
343            system_memory_requirement: u32::try_from(system_memory_requirement).unwrap_or(u32::MAX),
344            video_memory_requirement: u32::try_from(video_memory_requirement).unwrap_or(u32::MAX),
345            references: vec![],
346            blob: PackageResourceBlob::Memory {
347                data,
348                compression_level: if glacier_resource.should_compress() {Some(12)} else {None},
349                should_scramble: glacier_resource.should_scramble(),
350            },
351        })
352    }
353
354    /// Adds a reference to the resource.
355    ///
356    /// This specifies that this resource depends on / references another resource.
357    ///
358    /// # Arguments
359    /// * `rrid` - The resource ID of the reference.
360    /// * `flags` - The flags of the reference.
361    pub fn with_reference(
362        &mut self,
363        rrid: RuntimeResourceID,
364        flags: ResourceReferenceFlags,
365    ) -> &mut Self {
366        self.references.push((rrid, flags));
367        self
368    }
369
370    /// Adds multiple references to the resource at once.
371    ///
372    /// Accepts any iterable of `(RuntimeResourceID, ResourceReferenceFlags)` pairs.
373    ///
374    /// # Examples
375    /// builder.with_references(vec![(id1, flags1), (id2, flags2)]);
376    pub fn with_references<I, P>(&mut self, refs: I) -> &mut Self
377    where
378        I: IntoIterator<Item = P>,
379        P: Borrow<(RuntimeResourceID, ResourceReferenceFlags)>,
380        (RuntimeResourceID, ResourceReferenceFlags): Copy,
381    {
382        self.references.extend(refs.into_iter().map(|p| *p.borrow()));
383        self
384    }
385    /// Sets the memory requirements of the resource.
386    ///
387    /// # Arguments
388    /// * `system_memory_requirement` - The system memory requirement of the resource.
389    /// * `video_memory_requirement` - The video memory requirement of the resource.
390    pub fn with_memory_requirements(
391        &mut self,
392        system_memory_requirement: u32,
393        video_memory_requirement: u32,
394    ) -> &mut Self {
395        self.system_memory_requirement = system_memory_requirement;
396        self.video_memory_requirement = video_memory_requirement;
397        self
398    }
399}
400
401/// A builder for creating a ResourcePackage.
402/// ```
403/// # use rpkg_rs::resource::package_builder::{PackageBuilderError, PackageResourceBuilder};
404/// # use rpkg_rs::resource::pdefs::{PartitionId, PartitionType};
405/// # use rpkg_rs::resource::resource_package::PackageVersion;
406/// # use rpkg_rs::resource::resource_partition::PatchId;
407/// # use rpkg_rs::resource::runtime_resource_id::RuntimeResourceID;
408/// # use rpkg_rs::resource::package_builder::PackageBuilder;
409/// # use std::error::Error;
410/// # use std::fs;
411/// # fn main() -> Result<(), Box<dyn Error>>{
412/// #   let temp_dir = tempfile::tempdir()?;
413/// #   let output_path = temp_dir.path();
414///
415///     let mut builder = PackageBuilder::new_with_patch_id(PartitionId::default(), PatchId::Base);
416///     builder.with_resource(PackageResourceBuilder::from_memory(RuntimeResourceID::default(), "TYPE", vec![0,1,2,3,4,5], None, false).unwrap());
417///     builder.build(PackageVersion::RPKGv2, output_path)?;
418///
419///     assert!(temp_dir.path().join("chunk0.rpkg").exists());
420/// #   Ok(())
421/// # }
422pub struct PackageBuilder {
423    partition_id: PartitionId,
424    patch_id: PatchId,
425    use_legacy_references: bool,
426    resources: IndexMap<RuntimeResourceID, PackageResourceBuilder>,
427    unneeded_resources: IndexSet<RuntimeResourceID>,
428}
429
430#[derive(Debug, Error)]
431pub enum PackageBuilderError {
432    #[error("Error writing the file: {0}")]
433    IoError(#[from] io::Error),
434
435    #[error("Error serializing the package: {0}")]
436    SerializationError(#[from] binrw::Error),
437
438    #[error("Unneeded resources are only supported when building a patch package")]
439    UnneededResourcesNotSupported,
440
441    #[error("Building patch but no patch ID was provided")]
442    NoPatchId,
443
444    #[error("Too many resources in the package")]
445    TooManyResources,
446
447    #[error("A resource has too many references")]
448    TooManyReferences,
449
450    #[error("Resource type is not valid")]
451    InvalidResourceType,
452
453    #[error("Cannot build from a resource package without a source")]
454    NoSource,
455
456    #[error("Could not duplicate resource {0} from the source package: {1}")]
457    CannotDuplicateResource(RuntimeResourceID, PackageResourceBuilderError),
458
459    #[error("LZ4 compression error: {0}")]
460    Lz4CompressionError(#[from] lzzzz::Error),
461
462    #[error("Invalid partition id index cannot be greater than 255")]
463    InvalidPartitionIdIndex,
464
465    #[error("Patch id cannot be greater than 255")]
466    InvalidPatchId,
467}
468
469struct OffsetTableResult {
470    offset_table_size: u32,
471    resource_entry_offsets: HashMap<RuntimeResourceID, u64>,
472}
473
474struct MetadataTableResult {
475    metadata_table_size: u32,
476}
477
478/// A writer that xors the data with a predefined key.
479struct XorWriter<'a, W: Write + Seek> {
480    writer: &'a mut W,
481}
482
483impl<W: Write + Seek> Write for XorWriter<'_, W> {
484    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
485        let str_xor = [0xdc, 0x45, 0xa6, 0x9c, 0xd3, 0x72, 0x4c, 0xab];
486
487        for (index, byte) in buf.iter().enumerate() {
488            let xored_byte = *byte ^ str_xor[index % str_xor.len()];
489            self.writer.write_all(&[xored_byte])?;
490        }
491
492        Ok(buf.len())
493    }
494
495    fn flush(&mut self) -> Result<(), io::Error> {
496        self.writer.flush()
497    }
498}
499
500impl PackageBuilder {
501    /// Creates a new package builder.
502    ///
503    /// # Arguments
504    /// * `chunk_id` - The chunk ID of the package. e.g. chunk0
505    /// * `chunk_type` - The chunk type of the package.
506    pub fn new(chunk_id: u8, chunk_type: ChunkType) -> Self {
507        Self {
508            partition_id: PartitionId {
509                part_type: match chunk_type {
510                    ChunkType::Standard => PartitionType::Standard,
511                    ChunkType::Addon => PartitionType::Addon,
512                },
513                index: chunk_id as usize,
514            },
515            use_legacy_references: false,
516            patch_id: PatchId::Base,
517            resources: IndexMap::new(),
518            unneeded_resources: IndexSet::new(),
519        }
520    }
521
522    /// Creates a new package builder using the given partition id and patch id.
523    ///
524    /// # Arguments
525    /// * `partition_id` - The partition id of the package.
526    /// * `patch_id` - The patch id of the package.
527    pub fn new_with_patch_id(partition_id: PartitionId, patch_id: PatchId) -> Self {
528        Self {
529            partition_id,
530            patch_id,
531            use_legacy_references: false,
532            resources: IndexMap::new(),
533            unneeded_resources: IndexSet::new(),
534        }
535    }
536
537    /// Creates a new package builder by duplicating an existing ResourcePackage.
538    ///
539    /// # Arguments
540    /// * `resource_package` - The ResourcePackage to duplicate.
541    pub fn from_resource_package(
542        resource_package: &ResourcePackage,
543    ) -> Result<Self, PackageBuilderError> {
544        let source = resource_package
545            .source
546            .as_ref()
547            .ok_or(PackageBuilderError::NoSource)?;
548
549        let mut package = Self {
550            partition_id: PartitionId {
551                part_type: match resource_package
552                    .metadata
553                    .as_ref()
554                    .map(|m| m.chunk_type)
555                    .unwrap_or_default()
556                {
557                    ChunkType::Standard => PartitionType::Standard,
558                    ChunkType::Addon => PartitionType::Addon,
559                },
560                index: resource_package
561                    .metadata
562                    .as_ref()
563                    .map(|m| m.chunk_id)
564                    .unwrap_or_default() as usize,
565            },
566            patch_id: match resource_package
567                .metadata
568                .as_ref()
569                .map(|m| m.patch_id)
570                .unwrap_or_default()
571            {
572                0 => PatchId::Base,
573                x => PatchId::Patch(x as usize),
574            },
575            use_legacy_references: false,
576            resources: IndexMap::new(),
577            unneeded_resources: IndexSet::new(),
578        };
579
580        for (rrid, resource) in &resource_package.resources {
581            let mut builder = match source {
582                ResourcePackageSource::File(source_path) => {
583                    PackageResourceBuilder::from_file_at_offset(
584                        *rrid,
585                        &resource.data_type(),
586                        source_path,
587                        resource.entry.data_offset,
588                        resource.header.data_size,
589                        resource.compressed_size(),
590                        resource.is_scrambled(),
591                    )
592                    .map_err(|e| PackageBuilderError::CannotDuplicateResource(*rrid, e))?
593                }
594
595                ResourcePackageSource::Memory(source_data) => {
596                    let read_size = resource
597                        .compressed_size()
598                        .unwrap_or(resource.header.data_size);
599
600                    let start_offset = resource.entry.data_offset as usize;
601                    let end_offset = start_offset + read_size as usize;
602
603                    let decompressed_size = if resource.is_compressed() {
604                        Some(resource.header.data_size)
605                    } else {
606                        None
607                    };
608
609                    PackageResourceBuilder::from_compressed_memory(
610                        *rrid,
611                        &resource.data_type(),
612                        source_data[start_offset..end_offset].to_vec(),
613                        decompressed_size,
614                        resource.is_scrambled(),
615                    )
616                    .map_err(|e| PackageBuilderError::CannotDuplicateResource(*rrid, e))?
617                }
618            };
619
620            builder.with_memory_requirements(
621                resource.system_memory_requirement(),
622                resource.video_memory_requirement(),
623            );
624
625            for (rrid, flags) in resource.references() {
626                builder.with_reference(*rrid, *flags);
627            }
628
629            package.with_resource(builder);
630        }
631
632        for rrid in resource_package.unneeded_resource_ids() {
633            package.with_unneeded_resource(*rrid);
634        }
635
636        Ok(package)
637    }
638
639    /// Sets the partition ID of the package.
640    pub fn with_partition_id(&mut self, partition_id: &PartitionId) -> &mut Self {
641        self.partition_id = partition_id.clone();
642        self
643    }
644
645    /// Sets the patch ID of the package.
646    pub fn with_patch_id(&mut self, patch_id: &PatchId) -> &mut Self {
647        self.patch_id = *patch_id;
648        self
649    }
650
651    /// When this flag is set it will build the reference flags with the legacy format
652    pub fn use_legacy_references(&mut self) -> &mut Self {
653        self.use_legacy_references = true;
654        self
655    }
656
657    /// Adds a resource to the package.
658    ///
659    /// If a resource with the same resource ID already exists, it will be overwritten.
660    ///
661    /// # Arguments
662    /// * `resource` - The resource to add to the package.
663    pub fn with_resource(&mut self, resource: PackageResourceBuilder) -> &mut Self {
664        self.resources.insert(resource.rrid, resource);
665        self
666    }
667
668    /// Adds multiple resources to the package at once.
669    ///
670    /// If a resource with the same resource ID already exists, it will be overwritten.
671    ///
672    /// # Arguments
673    /// * `resources` - An iterable of `PackageResourceBuilder`s to add.
674    pub fn with_resources<I>(&mut self, resources: I) -> &mut Self
675    where
676        I: IntoIterator<Item = PackageResourceBuilder>,
677    {
678        self.resources
679            .extend(resources.into_iter().map(|r| (r.rrid, r)));
680        self
681    }
682
683    /// Adds an unneeded resource to the package.
684    ///
685    /// # Arguments
686    /// * `rrid` - The resource ID of the resource.
687    pub fn with_unneeded_resource(&mut self, rrid: RuntimeResourceID) -> &mut Self {
688        self.unneeded_resources.insert(rrid);
689        self
690    }
691
692
693    /// Marks multiple resources as unneeded.
694    ///
695    /// # Arguments
696    /// * `rrids` - An iterable of resource IDs to mark as unneeded.
697    pub fn with_unneeded_resources<I, P>(&mut self, rrids: I) -> &mut Self
698    where
699        I: IntoIterator<Item = P>,
700        P: Borrow<RuntimeResourceID>, 
701        RuntimeResourceID: Copy,
702    {
703        self.unneeded_resources.extend(rrids.into_iter().map(|p| *p.borrow()));
704        self
705    }
706
707    /// Patches data at a given offset and returns to the previous position.
708    fn backpatch<W: Write + Seek, T: BinWrite + WriteEndian>(
709        writer: &mut W,
710        patch_offset: u64,
711        data: &T,
712    ) -> Result<(), PackageBuilderError>
713    where
714        for<'a> T::Args<'a>: Required,
715    {
716        let current_offset = writer
717            .stream_position()
718            .map_err(PackageBuilderError::IoError)?;
719        writer
720            .seek(io::SeekFrom::Start(patch_offset))
721            .map_err(PackageBuilderError::IoError)?;
722        data.write(writer)
723            .map_err(PackageBuilderError::SerializationError)?;
724        writer
725            .seek(io::SeekFrom::Start(current_offset))
726            .map_err(PackageBuilderError::IoError)?;
727        Ok(())
728    }
729
730    /// Writes the offset table to the given writer.
731    fn write_offset_table<W: Write + Seek>(
732        &self,
733        writer: &mut W,
734    ) -> Result<OffsetTableResult, PackageBuilderError> {
735        // We need to keep a map of rrid => offset to patch the data offsets later.
736        let mut resource_entry_offsets = HashMap::new();
737        let offset_table_start = writer
738            .stream_position()
739            .map_err(PackageBuilderError::IoError)?;
740
741        for (rrid, _) in &self.resources {
742            let current_offset = writer
743                .stream_position()
744                .map_err(PackageBuilderError::IoError)?;
745
746            let resource_entry = PackageOffsetInfo {
747                runtime_resource_id: *rrid,
748                data_offset: 0,
749                flags: PackageOffsetFlags::new(),
750            };
751
752            resource_entry
753                .write(writer)
754                .map_err(PackageBuilderError::SerializationError)?;
755            resource_entry_offsets.insert(*rrid, current_offset);
756        }
757
758        // Write the offset table size.
759        let offset_table_end = writer
760            .stream_position()
761            .map_err(PackageBuilderError::IoError)?;
762        let offset_table_size = offset_table_end - offset_table_start;
763
764        if offset_table_size > u32::MAX as u64 {
765            return Err(PackageBuilderError::TooManyResources);
766        }
767
768        Ok(OffsetTableResult {
769            offset_table_size: offset_table_size as u32,
770            resource_entry_offsets,
771        })
772    }
773
774    /// Writes the metadata table to the given writer.
775    fn write_metadata_table<W: Write + Seek>(
776        &self,
777        writer: &mut W,
778        legacy_references: bool,
779    ) -> Result<MetadataTableResult, PackageBuilderError> {
780        let metadata_table_start = writer
781            .stream_position()
782            .map_err(PackageBuilderError::IoError)?;
783
784        for (_, resource) in &self.resources {
785            let metadata_offset = writer
786                .stream_position()
787                .map_err(PackageBuilderError::IoError)?;
788
789            // Write the resource metadata followed by the references table if there are any.
790            // We set the references chunk size to 0, and we'll patch it later.
791            let mut resource_metadata = ResourceHeader {
792                resource_type: resource.resource_type,
793                references_chunk_size: 0,
794                states_chunk_size: 0,
795                data_size: resource.blob.size(),
796                system_memory_requirement: resource.system_memory_requirement,
797                video_memory_requirement: resource.video_memory_requirement,
798                references: Vec::new(),
799            };
800
801            resource_metadata
802                .write(writer)
803                .map_err(PackageBuilderError::SerializationError)?;
804
805            // Write the references table if there are any.
806            if !resource.references.is_empty() {
807                let reference_table_start = writer
808                    .stream_position()
809                    .map_err(PackageBuilderError::IoError)?;
810
811                let reference_count_and_flags = ResourceReferenceCountAndFlags::new()
812                    .with_reference_count(resource.references.len() as u32)
813                    .with_is_new_format(!legacy_references)
814                    .with_always_true(true);
815
816                reference_count_and_flags
817                    .write(writer)
818                    .map_err(PackageBuilderError::SerializationError)?;
819
820                // In legacy mode, we write resource ids first, then flags.
821                // In new mode, we do the opposite. We also use the appropriate version of the flags.
822                if legacy_references {
823                    for (rrid, _) in &resource.references {
824                        rrid.write(writer)
825                            .map_err(PackageBuilderError::SerializationError)?;
826                    }
827
828                    for (_, flags) in &resource.references {
829                        flags
830                            .to_legacy()
831                            .write(writer)
832                            .map_err(PackageBuilderError::SerializationError)?;
833                    }
834                } else {
835                    for (_, flags) in &resource.references {
836                        flags
837                            .to_standard()
838                            .write(writer)
839                            .map_err(PackageBuilderError::SerializationError)?;
840                    }
841
842                    for (rrid, _) in &resource.references {
843                        rrid.write(writer)
844                            .map_err(PackageBuilderError::SerializationError)?;
845                    }
846                }
847
848                let reference_table_end = writer
849                    .stream_position()
850                    .map_err(PackageBuilderError::IoError)?;
851                let reference_table_size = reference_table_end - reference_table_start;
852
853                if reference_table_size > u32::MAX as u64 {
854                    return Err(PackageBuilderError::TooManyReferences);
855                }
856
857                // Calculate the size and patch the metadata.
858                resource_metadata.references_chunk_size = reference_table_size as u32;
859                PackageBuilder::backpatch(writer, metadata_offset, &resource_metadata)?;
860            }
861        }
862
863        // Write the metadata table size.
864        let metadata_table_end = writer
865            .stream_position()
866            .map_err(PackageBuilderError::IoError)?;
867        let metadata_table_size = metadata_table_end - metadata_table_start;
868
869        if metadata_table_size > u32::MAX as u64 {
870            return Err(PackageBuilderError::TooManyResources);
871        }
872
873        Ok(MetadataTableResult {
874            metadata_table_size: metadata_table_size as u32,
875        })
876    }
877
878    /// Builds the package, writing it to the given writer.
879    fn build_internal<W: Write + Seek>(
880        &self,
881        version: PackageVersion,
882        writer: &mut W,
883    ) -> Result<(), PackageBuilderError> {
884        // Perform some basic validation.
885        if !self.unneeded_resources.is_empty() && self.patch_id.is_base() {
886            return Err(PackageBuilderError::UnneededResourcesNotSupported);
887        }
888
889        // First create a base header. We'll fill it and patch it later.
890        let mut header = ResourcePackage {
891            source: None,
892            magic: match version {
893                PackageVersion::RPKGv1 => *b"GKPR",
894                PackageVersion::RPKGv2 => *b"2KPR",
895            },
896            metadata: match version {
897                PackageVersion::RPKGv1 => None,
898                PackageVersion::RPKGv2 => Some(PackageMetadata {
899                    unknown: 1,
900                    chunk_id: self.partition_id.index as u8,
901                    chunk_type: match self.partition_id.part_type {
902                        PartitionType::Addon => ChunkType::Addon,
903                        _ => ChunkType::Standard,
904                    },
905                    patch_id: match self.patch_id {
906                        PatchId::Base => 0,
907                        PatchId::Patch(x) => x as u8,
908                    },
909                    language_tag: *b"xx",
910                }),
911            },
912            header: PackageHeader {
913                file_count: self.resources.len() as u32,
914                offset_table_size: 0,
915                metadata_table_size: 0,
916            },
917            unneeded_resource_count: self.unneeded_resources.len() as u32,
918            unneeded_resources: Some(self.unneeded_resources.iter().copied().collect()),
919            resources: IndexMap::new(),
920        };
921
922        // Write the header and the tables.
923        header
924            .write_args(writer, (self.patch_id.is_patch(),))
925            .map_err(PackageBuilderError::SerializationError)?;
926
927        let offset_table_result = self.write_offset_table(writer)?;
928        let metadata_table_result =
929            self.write_metadata_table(writer, self.use_legacy_references)?;
930
931        // Now that we're done writing the tables, let's patch the header.
932        header.header.offset_table_size = offset_table_result.offset_table_size;
933        header.header.metadata_table_size = metadata_table_result.metadata_table_size;
934        PackageBuilder::backpatch(writer, 0, &header)?;
935
936        // Write the resource data.
937        for (rrid, resource) in &self.resources {
938            let data_offset = writer
939                .stream_position()
940                .map_err(PackageBuilderError::IoError)?;
941
942            let (compressed_size, is_scrambled) = match &resource.blob {
943                PackageResourceBlob::File {
944                    path,
945                    size,
946                    compression_level,
947                    should_scramble,
948                } => {
949                    let mut file = File::open(path).map_err(PackageBuilderError::IoError)?;
950
951                    // Wrap our writer in a XorWriter if we should scramble.
952                    let mut data_writer: Box<dyn Write> = match should_scramble {
953                        true => Box::new(XorWriter { writer }),
954                        false => Box::new(&mut *writer),
955                    };
956
957                    let compressed_size = match compression_level {
958                        Some(level) => {
959                            // TODO: Switch to streaming API.
960                            let mut compressed_buffer =
961                                vec![0; lz4::max_compressed_size(*size as usize)];
962                            let mut decompressed_data = vec![0; *size as usize];
963                            file.read_exact(&mut decompressed_data)
964                                .map_err(PackageBuilderError::IoError)?;
965
966                            let compressed_size = match version {
967                                PackageVersion::RPKGv1 => lz4::compress(
968                                    &decompressed_data,
969                                    &mut compressed_buffer,
970                                    *level,
971                                )?,
972                                PackageVersion::RPKGv2 => lz4_hc::compress(
973                                    &decompressed_data,
974                                    &mut compressed_buffer,
975                                    *level,
976                                )?,
977                            };
978
979                            // Write the compressed data.
980                            data_writer
981                                .write_all(&compressed_buffer[..compressed_size])
982                                .map_err(PackageBuilderError::IoError)?;
983
984                            Some(compressed_size as u32)
985                        }
986
987                        None => {
988                            io::copy(&mut file, &mut data_writer)
989                                .map_err(PackageBuilderError::IoError)?;
990                            None
991                        }
992                    };
993
994                    (compressed_size, *should_scramble)
995                }
996
997                PackageResourceBlob::FileAtOffset {
998                    path,
999                    offset,
1000                    size,
1001                    compressed_size,
1002                    is_scrambled,
1003                } => {
1004                    let size_to_copy = compressed_size.unwrap_or_else(|| *size);
1005
1006                    let mut file = File::open(path).map_err(PackageBuilderError::IoError)?;
1007                    file.seek(io::SeekFrom::Start(*offset))
1008                        .map_err(PackageBuilderError::IoError)?;
1009                    io::copy(&mut file.take(size_to_copy as u64), writer)
1010                        .map_err(PackageBuilderError::IoError)?;
1011
1012                    (*compressed_size, *is_scrambled)
1013                }
1014
1015                PackageResourceBlob::CompressedMemory {
1016                    data,
1017                    decompressed_size,
1018                    is_scrambled,
1019                } => {
1020                    writer
1021                        .write_all(data)
1022                        .map_err(PackageBuilderError::IoError)?;
1023                    let compressed_size = decompressed_size.map(|_| data.len() as u32);
1024                    (compressed_size, *is_scrambled)
1025                }
1026
1027                PackageResourceBlob::Memory {
1028                    data,
1029                    compression_level,
1030                    should_scramble,
1031                } => {
1032                    // Wrap our writer in a XorWriter if we should scramble.
1033                    let mut data_writer: Box<dyn Write> = match should_scramble {
1034                        true => Box::new(XorWriter { writer }),
1035                        false => Box::new(&mut *writer),
1036                    };
1037
1038                    let compressed_size = match compression_level {
1039                        Some(level) => {
1040                            // TODO: Switch to streaming API.
1041                            let mut compressed_buffer =
1042                                vec![0; lz4::max_compressed_size(data.len())];
1043                            let compressed_size = match version {
1044                                PackageVersion::RPKGv1 => {
1045                                    lz4::compress(data, &mut compressed_buffer, *level)?
1046                                }
1047                                PackageVersion::RPKGv2 => {
1048                                    lz4_hc::compress(data, &mut compressed_buffer, *level)?
1049                                }
1050                            };
1051
1052                            // Write the compressed data.
1053                            data_writer
1054                                .write_all(&compressed_buffer[..compressed_size])
1055                                .map_err(PackageBuilderError::IoError)?;
1056
1057                            Some(compressed_size as u32)
1058                        }
1059
1060                        None => {
1061                            data_writer
1062                                .write_all(data)
1063                                .map_err(PackageBuilderError::IoError)?;
1064                            None
1065                        }
1066                    };
1067
1068                    (compressed_size, *should_scramble)
1069                }
1070            };
1071
1072            // Patch the offset info.
1073            // If the resource is not compressed, we set the compressed size to 0.
1074            let final_compressed_size = compressed_size.unwrap_or(0);
1075
1076            let offset_info = PackageOffsetInfo {
1077                runtime_resource_id: *rrid,
1078                data_offset,
1079                flags: PackageOffsetFlags::new()
1080                    .with_compressed_size(final_compressed_size)
1081                    .with_is_scrambled(is_scrambled),
1082            };
1083
1084            let patch_offset = offset_table_result.resource_entry_offsets[rrid];
1085            PackageBuilder::backpatch(writer, patch_offset, &offset_info)?;
1086        }
1087
1088        Ok(())
1089    }
1090
1091    #[deprecated(since = "1.1.1", note = "use `build_to_file` instead")]
1092    pub fn build(
1093        self,
1094        version: PackageVersion,
1095        output_path: &Path,
1096    ) -> Result<(), PackageBuilderError> {
1097        self.build_to_file(version, output_path)
1098    }
1099
1100    /// Builds the package for the given version and writes it to the given writer.
1101    ///
1102    /// # Arguments
1103    /// * `version` - The version of the package to build.
1104    /// * `writer` - The struct implementing the Write and Seek traits.
1105    pub fn build_to_writer<W: Write + Seek>(self, version: PackageVersion, writer: &mut W) -> Result<(), PackageBuilderError>{
1106        self.build_internal(version, writer)
1107    }
1108    
1109    /// Builds the package for the given version and writes it to the given path.
1110    ///
1111    /// # Arguments
1112    /// * `version` - The version of the package to build.
1113    /// * `output_path` - The path to the output file.
1114    pub fn build_to_file<P: AsRef<Path>>(self, version: PackageVersion, output_path: P) -> Result<(), PackageBuilderError> {
1115        let output_path: &Path = output_path.as_ref();
1116        let output_file = match output_path.is_dir() {
1117            true => output_path.join(self.partition_id.to_filename(self.patch_id)),
1118            false => output_path.to_path_buf(),
1119        };
1120
1121        let file = File::create(output_file).map_err(PackageBuilderError::IoError)?;
1122        let mut writer = BufWriter::new(file);
1123        let result = self.build_internal(version, &mut writer);
1124        writer.flush()?;
1125        result
1126    }
1127
1128    #[deprecated(since = "1.1.1", note = "use `build_to_vec` instead")]
1129    pub fn build_in_memory(self, version: PackageVersion) -> Result<Vec<u8>, PackageBuilderError> {
1130        self.build_to_vec(version)
1131    }
1132
1133    /// Builds the package for the given version and returns it as a byte vector.
1134    ///
1135    /// # Arguments
1136    /// * `version` - The version of the package to build.
1137    pub fn build_to_vec(self, version: PackageVersion) -> Result<Vec<u8>, PackageBuilderError> {
1138        let mut writer = Cursor::new(vec![]);
1139
1140        self.build_internal(version, &mut writer)?;
1141        Ok(writer.into_inner())
1142    }
1143}