criware_utf_core/
packet.rs

1use std::{
2    alloc::{self, Layout},
3    io::{Cursor, Read, Write},
4    ops::{Deref, DerefMut},
5};
6
7use crate::{Error, IOErrorHelper, Result, Table};
8
9mod cri_encryption;
10
11fn aligned_vec(initial_size: usize, size: usize) -> Vec<u8> {
12    let minimum_size = size.div_ceil(64) << 6;
13    let layout = Layout::from_size_align(minimum_size, 64).expect("Invalid layout");
14    unsafe {
15        let ptr = alloc::alloc(layout);
16        if ptr.is_null() {
17            alloc::handle_alloc_error(layout);
18        }
19        Vec::from_raw_parts(ptr, initial_size, minimum_size)
20    }
21}
22fn aligned_vec_empty() -> Vec<u8> {
23    aligned_vec(0, 256)
24}
25fn aligned_vec_full(size: usize) -> Vec<u8> {
26    aligned_vec(size, size)
27}
28
29/**
30Packed, encryptable UTF table
31 */
32pub struct Packet<T: Table> {
33    prefix: &'static [u8; 4],
34    encrypted: bool,
35    unknown_value: u32,
36    table: T,
37}
38
39impl<T: Table> Packet<T> {
40    /**
41    Creates a new UTF table packet with the given prefix
42
43    The table itself is initialized with `T::new()`
44     */
45    pub fn new(prefix: &'static [u8; 4]) -> Self {
46        Self::from_table(T::new(), prefix)
47    }
48
49    /**
50    Creates a new UTF table packet with the given prefix
51     */
52    pub fn from_table(table: T, prefix: &'static [u8; 4]) -> Self {
53        Packet {
54            prefix,
55            encrypted: false,
56            unknown_value: 0,
57            table,
58        }
59    }
60
61    /**
62    Reads a UTF table packet from the given stream, verifying that it has
63    the given 4-byte prefix.
64     */
65    pub fn read_packet(reader: &mut dyn Read, prefix: &'static [u8; 4]) -> Result<Self> {
66        let mut header = [0u8; 16];
67        reader.read_exact(&mut header).io("UTF packet header")?;
68        if prefix != &header[0..4] {
69            return Err(Error::WrongTableSchema);
70        }
71        let unknown_value = u32::from_le_bytes(header[4..8].try_into().unwrap());
72        let table_size = u64::from_le_bytes(header[8..16].try_into().unwrap());
73        if table_size < 32 {
74            return Err(Error::MalformedHeader);
75        }
76        let mut table_data = aligned_vec_full(table_size as usize);
77        let mut decrypted_table_data = aligned_vec_full(table_size as usize);
78        reader
79            .read_exact(table_data.as_mut_slice())
80            .io("UTF table")?;
81        if &table_data[0..4] == b"@UTF" {
82            return Ok(Packet {
83                prefix,
84                encrypted: false,
85                unknown_value,
86                table: T::read(&mut Cursor::new(table_data))?,
87            });
88        }
89        if !cri_encryption::can_decrypt(table_data.as_slice()) {
90            return Err(Error::DecryptionError);
91        }
92        cri_encryption::decrypt(table_data.as_slice(), decrypted_table_data.as_mut_slice());
93        if &decrypted_table_data[0..4] == b"@UTF" {
94            return Ok(Packet {
95                prefix,
96                encrypted: true,
97                unknown_value,
98                table: T::read(&mut Cursor::new(decrypted_table_data))?,
99            });
100        }
101        return Err(Error::DecryptionError);
102    }
103
104    /**
105    Writes a UTF table packet to the given stream.
106     */
107    pub fn write_packet(&self, writer: &mut dyn Write) -> Result<()> {
108        let mut table_buffer = Cursor::new(aligned_vec_empty());
109        self.table.write(&mut table_buffer)?;
110        let table_buffer = {
111            let mut buffer = table_buffer.into_inner();
112            if buffer.capacity() < buffer.len().div_ceil(64) << 6 {
113                buffer.reserve_exact(64 - (buffer.len() & 63));
114            }
115            if self.encrypted {
116                let mut new_buffer = aligned_vec_full(buffer.len());
117                cri_encryption::encrypt(buffer.as_slice(), new_buffer.as_mut_slice());
118                new_buffer
119            } else {
120                buffer
121            }
122        };
123        writer.write_all(self.prefix).io("UTF packet header")?;
124        writer
125            .write_all(&u32::to_le_bytes(self.unknown_value))
126            .io("UTF packet header")?;
127        writer
128            .write_all(&u64::to_le_bytes(table_buffer.len() as u64))
129            .io("UTF packet header")?;
130        writer
131            .write_all(table_buffer.as_slice())
132            .io("UTF packet table")?;
133        Ok(())
134    }
135
136    /**
137    Returns whether or not the table is encrypted
138     */
139    pub fn is_encrypted(&self) -> bool {
140        self.encrypted
141    }
142
143    /**
144    Disables encryption for this packet
145     */
146    pub fn disable_encryption(&mut self) {
147        self.encrypted = false;
148    }
149
150    /**
151    Enables encryption for this packet
152     */
153    pub fn enable_encryption(&mut self) {
154        self.encrypted = true;
155    }
156}
157
158impl<T: Table> Deref for Packet<T> {
159    type Target = T;
160
161    fn deref(&self) -> &Self::Target {
162        &self.table
163    }
164}
165
166impl<T: Table> DerefMut for Packet<T> {
167    fn deref_mut(&mut self) -> &mut Self::Target {
168        &mut self.table
169    }
170}
criware_utf_core/packet.rs

criware_utf_core/
packet.rs
