1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
use std::{
    cmp,
    collections::HashMap,
    convert::TryFrom,
    io::{self, Cursor, Read},
};

use bitstream_io::{BitRead, BitReader};
use byteorder::{LittleEndian, ReadBytesExt};
use flate2::read::ZlibDecoder;
use thiserror::Error;

use crate::{ext::read::*, save::*, MAGIC_BYTES};

lazy_static::lazy_static! {
    static ref DEFAULT_MATERIALS: Vec<String> = vec!["BMC_Hologram", "BMC_Plastic", "BMC_Glow", "BMC_Metallic", "BMC_Glass"].into_iter().map(|s| s.into()).collect();
}

/// A read error.
#[derive(Error, Debug)]
pub enum ReadError {
    #[error("generic io error")]
    IoError(#[from] io::Error),
    #[error("bad magic bytes (expected 'BRS')")]
    BadHeader,
    #[error("invalid data in header 1")]
    InvalidDataHeader1,
    #[error("invalid data in header 2")]
    InvalidDataHeader2,
    #[error("must read in sequence: header 1, header 2, [preview], bricks")]
    BadSectionReadOrder,
    #[error("invalid compressed section")]
    InvalidCompression,
}

/// A save reader, which reads data from its `reader` (a `Read + Seek`).
pub struct SaveReader<R: Read> {
    reader: R,
    pub version: u16,
    pub game_version: i32,

    header1_read: bool,
    header2_read: bool,
    preview_read: bool,
}

impl<R: Read> SaveReader<R> {
    /// Create a new save reader from an existing `reader`, a `Read + Seek`.
    pub fn new(mut reader: R) -> Result<Self, ReadError> {
        let mut magic = [0u8; 3];
        reader.read_exact(&mut magic)?;
        if magic != MAGIC_BYTES {
            return Err(ReadError::BadHeader);
        }

        let version = reader.read_u16::<LittleEndian>()?;
        let game_version = if version >= 8 {
            reader.read_i32::<LittleEndian>()?
        } else {
            0
        };

        Ok(SaveReader {
            version,
            game_version,
            reader,
            header1_read: false,
            header2_read: false,
            preview_read: version < 8,
        })
    }

    /// Skip the first header.
    pub fn skip_header1(&mut self) -> Result<(), ReadError> {
        skip_compressed(&mut self.reader)?;
        self.header1_read = true;
        Ok(())
    }

    /// Read the first header.
    pub fn read_header1(&mut self) -> Result<Header1, ReadError> {
        let (mut cursor, _) = read_compressed(&mut self.reader)?;

        // match map: a string
        let map = cursor.read_string()?;

        // match author name: a string
        let author_name = cursor.read_string()?;

        // match description: a string
        let description = cursor.read_string()?;

        // match author id: a uuid
        let author_uuid = cursor.read_uuid()?;

        // match host:
        // version >= 8: match a user (string followed by uuid)
        //         else: not provided
        let host = match self.version {
            _ if self.version >= 8 => {
                let name = cursor.read_string()?;
                let id = cursor.read_uuid()?;
                Some(User { name, id })
            }
            _ => None,
        };

        // match save time:
        // version >= 4: match 8 bytes
        //         else: not provided
        let save_time = match self.version {
            _ if self.version >= 4 => {
                let mut bytes = [0u8; 8]; // todo: figure out how to parse this
                cursor.read_exact(&mut bytes)?;
                Some(bytes)
            }
            _ => None,
        };

        // match brick count: an i32
        let brick_count = match cursor.read_i32::<LittleEndian>()? {
            count if count >= 0 => count,
            _ => return Err(ReadError::InvalidDataHeader1),
        } as u32;

        self.header1_read = true;
        Ok(Header1 {
            map,
            author: User {
                name: author_name,
                id: author_uuid,
            },
            description,
            host,
            save_time: save_time.unwrap_or([0u8; 8]),
            brick_count,
        })
    }

    /// Skip the second header.
    pub fn skip_header2(&mut self) -> Result<(), ReadError> {
        skip_compressed(&mut self.reader)?;
        self.header2_read = true;
        Ok(())
    }

    /// Read the second header.
    pub fn read_header2(&mut self) -> Result<Header2, ReadError> {
        if !self.header1_read {
            return Err(ReadError::BadSectionReadOrder);
        }

        let (mut cursor, _) = read_compressed(&mut self.reader)?;

        // match mods: an array of strings
        let mods = cursor.read_array(|r| r.read_string())?;

        // match brick assets: an array of strings
        let brick_assets = cursor.read_array(|r| r.read_string())?;

        // match colors: an array of 4 bytes each, BGRA
        let colors = cursor.read_array(|r| -> io::Result<Color> {
            let mut bytes = [0u8; 4];
            r.read_exact(&mut bytes)?;
            Ok(Color::from_bytes_bgra(bytes))
        })?;

        // match materials:
        // version >= 2: an array of strings
        //         else: a list of default materials (see top of file)
        let materials = match self.version {
            _ if self.version >= 2 => cursor.read_array(|r| r.read_string())?,
            _ => DEFAULT_MATERIALS.clone(),
        };

        // match brick owners:
        // version >= 3: match brick owner:
        //               version >= 8: a user (uuid followed by string), then an i32 for brick count
        //                       else: a user (uuid followed by string)
        let brick_owners = match self.version {
            _ if self.version >= 3 => cursor.read_array(|r| -> io::Result<BrickOwner> {
                match self.version {
                    _ if self.version >= 8 => {
                        let id = r.read_uuid()?;
                        let name = r.read_string()?;
                        let bricks = r.read_i32::<LittleEndian>()? as u32;
                        Ok(BrickOwner { name, id, bricks })
                    }
                    _ => {
                        let id = r.read_uuid()?;
                        let name = r.read_string()?;
                        Ok(BrickOwner::from(User { name, id }))
                    }
                }
            })?,
            _ => vec![],
        };

        // match physical materials
        // version >= 9: an array of strings
        //         else: not provided
        let physical_materials = match self.version {
            _ if self.version >= 9 => cursor.read_array(|r| r.read_string())?,
            _ => vec![],
        };

        self.header2_read = true;
        Ok(Header2 {
            mods,
            brick_assets,
            colors,
            materials,
            brick_owners,
            physical_materials,
        })
    }

    /// Read the preview in the save.
    ///
    /// The preview is an `Preview`, which might not exist (Preview::None).
    pub fn read_preview(&mut self) -> Result<Preview, ReadError> {
        if !self.header2_read {
            return Err(ReadError::BadSectionReadOrder);
        }

        if self.version < 8 {
            return Ok(Preview::None);
        }

        let preview = Preview::from_reader(&mut self.reader)?;
        self.preview_read = true;
        Ok(preview)
    }

    /// Skip over the preview section.
    pub fn skip_preview(&mut self) -> Result<(), ReadError> {
        if !self.header2_read {
            return Err(ReadError::BadSectionReadOrder);
        }

        if self.version < 8 {
            return Ok(());
        }

        if self.reader.read_u8()? != 0 {
            let len = self.reader.read_i32::<LittleEndian>()?;
            io::copy(&mut self.reader.by_ref().take(len as u64), &mut io::sink())?;
        }

        self.preview_read = true;
        Ok(())
    }

    /// Read the bricks and components from a save.
    pub fn read_bricks(
        &mut self,
        header1: &Header1,
        header2: &Header2,
    ) -> Result<(Vec<Brick>, HashMap<String, Component>), ReadError> {
        if !self.preview_read || !self.header2_read {
            return Err(ReadError::BadSectionReadOrder);
        }

        let (cursor, len) = read_compressed(&mut self.reader)?;
        let mut bits = BitReader::<_, bitstream_io::LittleEndian>::new(cursor);

        let brick_count = header1.brick_count as usize;
        let brick_asset_count = cmp::max(header2.brick_assets.len(), 2);
        let material_count = cmp::max(header2.materials.len(), 2);
        let physical_material_count = cmp::max(header2.physical_materials.len(), 2);

        let mut bricks = vec![];
        let mut components = HashMap::new();

        // loop over each brick
        loop {
            // align and break out of the loop if we've seeked far enough ahead
            bits.byte_align();
            if bricks.len() >= brick_count || bits.reader().unwrap().position() >= len as u64 {
                break;
            }

            let asset_name_index = bits.read_uint(brick_asset_count as u32)?;

            let size = match bits.read_bit()? {
                true => Size::Procedural(
                    bits.read_uint_packed()?,
                    bits.read_uint_packed()?,
                    bits.read_uint_packed()?,
                ),
                false => Size::Empty,
            };

            let position = (
                bits.read_int_packed()?,
                bits.read_int_packed()?,
                bits.read_int_packed()?,
            );

            let orientation = bits.read_uint(24)?;
            let direction = Direction::try_from(((orientation >> 2) % 6) as u8).unwrap();
            let rotation = Rotation::try_from((orientation & 3) as u8).unwrap();

            let collision = match self.version {
                _ if self.version >= 10 => Collision {
                    player: bits.read_bit()?,
                    weapon: bits.read_bit()?,
                    interaction: bits.read_bit()?,
                    tool: bits.read_bit()?,
                },
                _ => Collision::for_all(bits.read_bit()?),
            };

            let visibility = bits.read_bit()?;

            let material_index = match self.version {
                _ if self.version >= 8 => bits.read_uint(material_count as u32)?,
                _ => {
                    if bits.read_bit()? {
                        bits.read_uint_packed()?
                    } else {
                        1
                    }
                }
            };

            let physical_index = match self.version {
                _ if self.version >= 9 => bits.read_uint(physical_material_count as u32)?,
                _ => 0,
            };

            let material_intensity = match self.version {
                _ if self.version >= 9 => bits.read_uint(11)?,
                _ => 5,
            };

            let color = match bits.read_bit()? {
                true => match self.version {
                    _ if self.version >= 9 => {
                        let mut bytes = [0u8; 3];
                        bits.read_bytes(&mut bytes)?;
                        BrickColor::Unique(Color::from_bytes_rgb(bytes))
                    }
                    _ => {
                        let mut bytes = [0u8; 4];
                        bits.read_bytes(&mut bytes)?;
                        BrickColor::Unique(Color::from_bytes_bgra(bytes))
                    }
                },
                false => BrickColor::Index(bits.read_uint(header2.colors.len() as u32)?),
            };

            let owner_index = if self.version >= 3 {
                bits.read_uint_packed()?
            } else {
                0
            };

            let brick = Brick {
                asset_name_index,
                size,
                position,
                direction,
                rotation,
                collision,
                visibility,
                material_index,
                physical_index,
                material_intensity,
                color,
                owner_index,
                components: HashMap::new(),
            };

            bricks.push(brick);
        }

        // components
        if self.version >= 8 {
            let (mut cursor, _) = read_compressed(&mut self.reader)?;
            let len = cursor.read_i32::<LittleEndian>()?;

            for _ in 0..len {
                let name = cursor.read_string()?;

                let mut bit_bytes = vec![0u8; cursor.read_i32::<LittleEndian>()? as usize];
                cursor.read_exact(&mut bit_bytes)?;
                let mut bits =
                    BitReader::<_, bitstream_io::LittleEndian>::new(Cursor::new(bit_bytes));

                let version = bits.read_i32_le()?;
                let brick_indices = bits.read_array(|r| r.read_uint(brick_count as u32))?;

                let properties = bits
                    .read_array(|r| Ok((r.read_string()?, r.read_string()?)))?
                    .into_iter()
                    .collect::<HashMap<_, _>>();

                // components for each brick
                for &i in brick_indices.iter() {
                    let mut props = HashMap::new();
                    for (n, ty) in properties.iter() {
                        props.insert(n.clone(), bits.read_unreal_type(&ty)?);
                    }
                    bricks[i as usize].components.insert(name.clone(), props);
                }

                components.insert(
                    name,
                    Component {
                        version,
                        brick_indices,
                        properties,
                    },
                );
            }
        }

        Ok((bricks, components))
    }

    /// Read all parts of a save into a `SaveData`.
    pub fn read_all(&mut self) -> Result<SaveData, ReadError> {
        let header1 = self.read_header1()?;
        let header2 = self.read_header2()?;
        let preview = self.read_preview()?;
        let (bricks, components) = self.read_bricks(&header1, &header2)?;

        Ok(SaveData {
            version: self.version,
            game_version: self.game_version,
            header1,
            header2,
            preview,
            bricks,
            components,
        })
    }

    /// Read all parts of a save (except the preview) into a `SaveData`.
    pub fn read_all_skip_preview(&mut self) -> Result<SaveData, ReadError> {
        let header1 = self.read_header1()?;
        let header2 = self.read_header2()?;
        self.skip_preview()?;
        let (bricks, components) = self.read_bricks(&header1, &header2)?;

        Ok(SaveData {
            version: self.version,
            game_version: self.game_version,
            header1,
            header2,
            preview: Preview::None,
            bricks,
            components,
        })
    }
}

/// Read a compressed section from a `Read`, following the BRS spec for compressed sections.
fn read_compressed(reader: &mut impl Read) -> Result<(Cursor<Vec<u8>>, i32), ReadError> {
    let (uncompressed_size, compressed_size) = (
        reader.read_i32::<LittleEndian>()?,
        reader.read_i32::<LittleEndian>()?,
    );
    if uncompressed_size < 0 || compressed_size < 0 || compressed_size >= uncompressed_size {
        return Err(ReadError::InvalidCompression);
    }

    let mut bytes = vec![0u8; uncompressed_size as usize];

    if compressed_size == 0 {
        // no need to decompress first
        reader.read_exact(&mut bytes)?;
    } else {
        // decompress first, then read
        let mut compressed = vec![0u8; compressed_size as usize];
        reader.read_exact(&mut compressed)?;
        ZlibDecoder::new(&compressed[..]).read_exact(&mut bytes)?;
    }

    Ok((Cursor::new(bytes), uncompressed_size))
}

/// Read a compressed section from a `Read`, discarding its contents.
fn skip_compressed(reader: &mut impl Read) -> Result<(), ReadError> {
    let (uncompressed_size, compressed_size) = (
        reader.read_i32::<LittleEndian>()?,
        reader.read_i32::<LittleEndian>()?,
    );
    if uncompressed_size < 0 || compressed_size < 0 || compressed_size >= uncompressed_size {
        return Err(ReadError::InvalidCompression);
    }

    io::copy(
        &mut reader.take(if compressed_size == 0 {
            uncompressed_size as u64
        } else {
            compressed_size as u64
        }),
        &mut io::sink(),
    )?;

    Ok(())
}