Mercurial > touhou

comparison src/th06/pbg3.rs @ 637:afa012bb8021

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Hello Rust!
author Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
date Wed, 03 Jul 2019 16:27:12 +0200
parents
children 7bde50132735
comparison
equal deleted inserted replaced
636:4fa0a8e7d941 637:afa012bb8021
1 //! PBG3 archive files handling.
2 //!
3 //! This module provides classes for handling the PBG3 file format.
4 //! The PBG3 format is the archive format used by Touhou 6: EoSD.
5 //!
6 //! PBG3 files are merely a bitstream composed of a header, a file
7 //! table, and LZSS-compressed files.
8
9 use crate::util::bitstream::BitStream;
10 use crate::util::lzss;
11 use std::io;
12 use std::collections::hash_map::{self, HashMap};
13
14 /// Helper struct to handle strings and integers in PBG3 bitstreams.
15 pub struct PBG3BitStream<R: io::Read + io::Seek> {
16 bitstream: BitStream<R>,
17 }
18
19 impl<R: io::Read + io::Seek> PBG3BitStream<R> {
20 /// Create a bitstream capable of reading u32 and strings.
21 pub fn new(bitstream: BitStream<R>) -> PBG3BitStream<R> {
22 PBG3BitStream {
23 bitstream,
24 }
25 }
26
27 /// Seek inside the bitstream, ditching any unused data read.
28 pub fn seek(&mut self, seek_from: io::SeekFrom) -> io::Result<u64> {
29 self.bitstream.seek(seek_from)
30 }
31
32 /// Return the current position in the stream.
33 pub fn tell(&mut self) -> io::Result<u64> {
34 self.bitstream.seek(io::SeekFrom::Current(0))
35 }
36
37 /// Read a given amount of bits.
38 pub fn read(&mut self, nb_bits: usize) -> io::Result<usize> {
39 self.bitstream.read(nb_bits)
40 }
41
42 /// Read a given amount of bytes.
43 pub fn read_bytes(&mut self, nb_bytes: usize) -> io::Result<Vec<u8>> {
44 self.bitstream.read_bytes(nb_bytes)
45 }
46
47 /// Read an integer from the bitstream.
48 ///
49 /// Integers have variable sizes. They begin with a two-bit value indicating
50 /// the number of (non-aligned) bytes to read.
51 pub fn read_u32(&mut self) -> io::Result<u32> {
52 let size = self.read(2)?;
53 Ok(self.read((size + 1) * 8)? as u32)
54 }
55
56 /// Read a string from the bitstream.
57 ///
58 /// Strings are stored as NULL-terminated sequences of bytes.
59 /// The only catch is that they are not byte-aligned.
60 pub fn read_string(&mut self, mut max_size: usize) -> io::Result<Vec<u8>> {
61 let mut buf = Vec::new();
62 while max_size > 0 {
63 let byte = self.read(8)? as u8;
64 if byte == 0 {
65 break;
66 }
67 buf.push(byte);
68 max_size -= 1;
69 }
70 Ok(buf)
71 }
72 }
73
74 type Entry = (u32, u32, u32, u32, u32);
75
76 /// Handle PBG3 archive files.
77 ///
78 /// PBG3 is a file archive format used in Touhou 6: EoSD.
79 /// This class provides a representation of such files, as well as functions to
80 /// read and extract files from a PBG3 archive.
81 pub struct PBG3<R: io::Read + io::Seek> {
82 /// List of PBG3Entry objects describing files present in the archive.
83 entries: HashMap<String, Entry>,
84
85 /// PBG3BitStream struct.
86 bitstream: PBG3BitStream<R>,
87 }
88
89 impl<R: io::Read + io::Seek> PBG3<R> {
90 /// Create a PBG3 archive.
91 fn new(entries: HashMap<String, Entry>, bitstream: PBG3BitStream<R>) -> PBG3<R> {
92 PBG3 {
93 entries,
94 bitstream,
95 }
96 }
97
98 /// Open a PBG3 archive.
99 pub fn from_file(mut file: R) -> io::Result<PBG3<R>> {
100 let mut magic = [0; 4];
101 file.read(&mut magic)?;
102 if &magic != b"PBG3" {
103 return Err(io::Error::new(io::ErrorKind::Other, "Wrong magic!"));
104 }
105
106 let bitstream = BitStream::new(file);
107 let mut bitstream = PBG3BitStream::new(bitstream);
108 let mut entries = HashMap::new();
109
110 let nb_entries = bitstream.read_u32()?;
111 let offset = bitstream.read_u32()?;
112 bitstream.seek(io::SeekFrom::Start(offset as u64))?;
113
114 for _ in 0..nb_entries {
115 let unknown_1 = bitstream.read_u32()?;
116 let unknown_2 = bitstream.read_u32()?;
117 let checksum = bitstream.read_u32()?; // Checksum of *compressed data*
118 let offset = bitstream.read_u32()?;
119 let size = bitstream.read_u32()?;
120 let name = bitstream.read_string(255)?;
121 // XXX: no unwrap!
122 let name = String::from_utf8(name).unwrap();
123 entries.insert(name, (unknown_1, unknown_2, checksum, offset, size));
124 }
125
126 Ok(PBG3::new(entries, bitstream))
127 }
128
129 /// List all file entries in this PBG3 archive.
130 pub fn list_files(&self) -> hash_map::Keys<String, Entry> {
131 self.entries.keys()
132 }
133
134 /// Read a single file from this PBG3 archive.
135 pub fn get_file(&mut self, filename: String, check: bool) -> io::Result<Vec<u8>> {
136 // XXX: no unwrap!
137 let (_unknown_1, _unknown_2, checksum, offset, size) = self.entries.get(&filename).unwrap();
138 self.bitstream.seek(io::SeekFrom::Start(*offset as u64))?;
139 let data = lzss::decompress(&mut self.bitstream.bitstream, *size as usize, 0x2000, 13, 4, 3)?;
140 if check {
141 // Verify the checksum.
142 let compressed_size = self.bitstream.tell()? as u32 - *offset;
143 self.bitstream.seek(io::SeekFrom::Start(*offset as u64))?;
144 let mut value: u32 = 0;
145 for c in self.bitstream.read_bytes(compressed_size as usize)? {
146 value += c as u32;
147 value &= 0xffffffff;
148 }
149 if value != *checksum {
150 return Err(io::Error::new(io::ErrorKind::Other, "Corrupted data!"));
151 }
152 }
153 Ok(data)
154 }
155 }
156
157 #[cfg(test)]
158 mod tests {
159 use super::*;
160 use crate::util::SeekableSlice;
161 use std::fs::File;
162
163 #[test]
164 fn bitstream() {
165 let data = SeekableSlice::new(b"Hello world!\0");
166 let bitstream = BitStream::new(data);
167 let mut pbg3 = PBG3BitStream::new(bitstream);
168 assert_eq!(pbg3.read_string(42).unwrap(), b"Hello world!");
169 }
170
171 #[test]
172 fn file_present() {
173 let file = File::open("/home/linkmauve/games/pc/東方/TH06 ~ The Embodiment of Scarlet Devil/MD.DAT").unwrap();
174 let file = io::BufReader::new(file);
175 let pbg3 = PBG3::from_file(file).unwrap();
176 let files = pbg3.list_files().cloned().collect::<Vec<String>>();
177 assert!(files.contains(&String::from("th06_01.pos")));
178 }
179
180 #[test]
181 fn check_all_files() {
182 let file = File::open("/home/linkmauve/games/pc/東方/TH06 ~ The Embodiment of Scarlet Devil/MD.DAT").unwrap();
183 let file = io::BufReader::new(file);
184 let mut pbg3 = PBG3::from_file(file).unwrap();
185 let files = pbg3.list_files().cloned().collect::<Vec<String>>();
186 for filename in files {
187 pbg3.get_file(filename, true).unwrap();
188 }
189 }
190 }