|
| 1 | +use crate::Error::InvalidAttributeData; |
| 2 | +use crate::byte_source::ByteSource; |
| 3 | +use crate::{Error, Result}; |
| 4 | + |
| 5 | +/// Represents the types of attributes in an Image file. |
| 6 | +#[derive(Clone, Debug, PartialEq, Eq)] |
| 7 | +enum AttributeType { |
| 8 | + End, |
| 9 | + Module, |
| 10 | + Parent, |
| 11 | + Base, |
| 12 | + Extension, |
| 13 | + Offset, |
| 14 | + Compressed, |
| 15 | + Uncompressed, |
| 16 | + Count, |
| 17 | +} |
| 18 | + |
| 19 | +impl TryFrom<u8> for AttributeType { |
| 20 | + type Error = Error; |
| 21 | + |
| 22 | + /// Converts a `u8` to an `AttributeType`. |
| 23 | + /// |
| 24 | + /// # Errors |
| 25 | + /// |
| 26 | + /// Returns an `Error::InvalidAttributeType` if the value does not correspond to any |
| 27 | + /// `AttributeType`. |
| 28 | + fn try_from(value: u8) -> Result<Self> { |
| 29 | + match value { |
| 30 | + 0 => Ok(AttributeType::End), |
| 31 | + 1 => Ok(AttributeType::Module), |
| 32 | + 2 => Ok(AttributeType::Parent), |
| 33 | + 3 => Ok(AttributeType::Base), |
| 34 | + 4 => Ok(AttributeType::Extension), |
| 35 | + 5 => Ok(AttributeType::Offset), |
| 36 | + 6 => Ok(AttributeType::Compressed), |
| 37 | + 7 => Ok(AttributeType::Uncompressed), |
| 38 | + 8 => Ok(AttributeType::Count), |
| 39 | + _ => Err(Error::InvalidAttributeType(value)), |
| 40 | + } |
| 41 | + } |
| 42 | +} |
| 43 | + |
| 44 | +/// Represents the attributes of a resource in an Image file. |
| 45 | +#[derive(Clone, Debug, PartialEq, Eq)] |
| 46 | +pub(crate) struct Attributes { |
| 47 | + /// The module offset. |
| 48 | + module_offset: usize, |
| 49 | + /// The parent offset. |
| 50 | + parent_offset: usize, |
| 51 | + /// The base offset. |
| 52 | + base_offset: usize, |
| 53 | + /// The extension offset. |
| 54 | + extension_offset: usize, |
| 55 | + /// The data offset. |
| 56 | + offset: usize, |
| 57 | + /// The compressed size. |
| 58 | + compressed_size: usize, |
| 59 | + /// The uncompressed size. |
| 60 | + uncompressed_size: usize, |
| 61 | +} |
| 62 | + |
| 63 | +impl Attributes { |
| 64 | + /// Parses attributes from a byte slice. |
| 65 | + pub fn from_bytes(byte_source: &ByteSource, mut offset: usize, limit: usize) -> Result<Self> { |
| 66 | + let mut attributes = Attributes { |
| 67 | + module_offset: 0, |
| 68 | + parent_offset: 0, |
| 69 | + base_offset: 0, |
| 70 | + extension_offset: 0, |
| 71 | + offset: 0, |
| 72 | + compressed_size: 0, |
| 73 | + uncompressed_size: 0, |
| 74 | + }; |
| 75 | + |
| 76 | + while offset < limit { |
| 77 | + let bytes = byte_source.get_bytes(offset..=offset)?; |
| 78 | + let data = bytes[0]; |
| 79 | + let attribute_type = AttributeType::try_from(data >> 3)?; |
| 80 | + if attribute_type == AttributeType::End { |
| 81 | + break; |
| 82 | + } |
| 83 | + |
| 84 | + let length = (data as usize & 0x7) + 1; |
| 85 | + offset += 1; |
| 86 | + let value = Self::read_value(byte_source, offset, length)?; |
| 87 | + |
| 88 | + match attribute_type { |
| 89 | + AttributeType::End => break, |
| 90 | + AttributeType::Module => { |
| 91 | + attributes.module_offset = value; |
| 92 | + } |
| 93 | + AttributeType::Parent => { |
| 94 | + attributes.parent_offset = value; |
| 95 | + } |
| 96 | + AttributeType::Base => { |
| 97 | + attributes.base_offset = value; |
| 98 | + } |
| 99 | + AttributeType::Extension => { |
| 100 | + attributes.extension_offset = value; |
| 101 | + } |
| 102 | + AttributeType::Offset => { |
| 103 | + attributes.offset = value; |
| 104 | + } |
| 105 | + AttributeType::Compressed => { |
| 106 | + attributes.compressed_size = value; |
| 107 | + } |
| 108 | + AttributeType::Uncompressed => { |
| 109 | + attributes.uncompressed_size = value; |
| 110 | + } |
| 111 | + AttributeType::Count => { |
| 112 | + // Count attribute is not used in this context |
| 113 | + } |
| 114 | + } |
| 115 | + |
| 116 | + offset += length; |
| 117 | + } |
| 118 | + |
| 119 | + Ok(attributes) |
| 120 | + } |
| 121 | + |
| 122 | + /// Reads a multiple byte value from the byte source. |
| 123 | + #[inline] |
| 124 | + fn read_value(byte_source: &ByteSource, offset: usize, length: usize) -> Result<usize> { |
| 125 | + let end = offset.checked_add(length).ok_or(InvalidAttributeData)?; |
| 126 | + let bytes = byte_source.get_bytes(offset..end)?; |
| 127 | + let mut buffer = [0u8; 8]; |
| 128 | + buffer[8 - bytes.len()..].copy_from_slice(&bytes); |
| 129 | + let value = u64::from_be_bytes(buffer); |
| 130 | + let value = usize::try_from(value)?; |
| 131 | + Ok(value) |
| 132 | + } |
| 133 | + |
| 134 | + /// Returns the module offset. |
| 135 | + pub fn module_offset(&self) -> usize { |
| 136 | + self.module_offset |
| 137 | + } |
| 138 | + |
| 139 | + /// Returns the parent offset. |
| 140 | + pub fn parent_offset(&self) -> usize { |
| 141 | + self.parent_offset |
| 142 | + } |
| 143 | + |
| 144 | + /// Returns the base offset. |
| 145 | + pub fn base_offset(&self) -> usize { |
| 146 | + self.base_offset |
| 147 | + } |
| 148 | + |
| 149 | + /// Returns the extension offset. |
| 150 | + pub fn extension_offset(&self) -> usize { |
| 151 | + self.extension_offset |
| 152 | + } |
| 153 | + |
| 154 | + /// Returns the data offset. |
| 155 | + pub fn offset(&self) -> usize { |
| 156 | + self.offset |
| 157 | + } |
| 158 | + |
| 159 | + /// Returns the compressed size. |
| 160 | + pub fn compressed_size(&self) -> usize { |
| 161 | + self.compressed_size |
| 162 | + } |
| 163 | + |
| 164 | + /// Returns the uncompressed size. |
| 165 | + pub fn uncompressed_size(&self) -> usize { |
| 166 | + self.uncompressed_size |
| 167 | + } |
| 168 | +} |
| 169 | + |
| 170 | +#[cfg(test)] |
| 171 | +mod tests { |
| 172 | + use super::*; |
| 173 | + |
| 174 | + #[test] |
| 175 | + fn test_attribute_type_conversion() -> Result<()> { |
| 176 | + assert_eq!(AttributeType::try_from(0)?, AttributeType::End); |
| 177 | + assert_eq!(AttributeType::try_from(1)?, AttributeType::Module); |
| 178 | + assert_eq!(AttributeType::try_from(2)?, AttributeType::Parent); |
| 179 | + assert_eq!(AttributeType::try_from(3)?, AttributeType::Base); |
| 180 | + assert_eq!(AttributeType::try_from(4)?, AttributeType::Extension); |
| 181 | + assert_eq!(AttributeType::try_from(5)?, AttributeType::Offset); |
| 182 | + assert_eq!(AttributeType::try_from(6)?, AttributeType::Compressed); |
| 183 | + assert_eq!(AttributeType::try_from(7)?, AttributeType::Uncompressed); |
| 184 | + assert_eq!(AttributeType::try_from(8)?, AttributeType::Count); |
| 185 | + assert!(matches!( |
| 186 | + AttributeType::try_from(9), |
| 187 | + Err(Error::InvalidAttributeType(9)) |
| 188 | + )); |
| 189 | + Ok(()) |
| 190 | + } |
| 191 | + |
| 192 | + #[test] |
| 193 | + fn test_attributes_from_bytes() -> Result<()> { |
| 194 | + let bytes = vec![ |
| 195 | + 9, 1, 191, 18, 7, 214, 127, 26, 7, 215, 162, 32, 1, 43, 6, 224, 205, 170, 57, 3, 160, 0, |
| 196 | + ]; |
| 197 | + let byte_source = ByteSource::Bytes(bytes); |
| 198 | + let attributes = Attributes::from_bytes(&byte_source, 0, byte_source.len()?)?; |
| 199 | + assert_eq!(attributes.module_offset(), 447); |
| 200 | + assert_eq!(attributes.parent_offset(), 513_663); |
| 201 | + assert_eq!(attributes.base_offset(), 513_954); |
| 202 | + assert_eq!(attributes.extension_offset(), 1); |
| 203 | + assert_eq!(attributes.offset(), 115_396_010); |
| 204 | + assert_eq!(attributes.compressed_size(), 0); |
| 205 | + assert_eq!(attributes.uncompressed_size(), 928); |
| 206 | + Ok(()) |
| 207 | + } |
| 208 | + |
| 209 | + #[test] |
| 210 | + fn test_read_value() -> Result<()> { |
| 211 | + let data = vec![0x01, 0x02, 0x03, 0x04]; |
| 212 | + let byte_source = ByteSource::Bytes(data); |
| 213 | + assert_eq!(Attributes::read_value(&byte_source, 0, 1)?, 0x01); |
| 214 | + assert_eq!(Attributes::read_value(&byte_source, 0, 2)?, 0x0102); |
| 215 | + assert_eq!(Attributes::read_value(&byte_source, 0, 3)?, 0x01_0203); |
| 216 | + assert_eq!(Attributes::read_value(&byte_source, 0, 4)?, 0x0102_0304); |
| 217 | + assert!(matches!( |
| 218 | + Attributes::read_value(&byte_source, 2, usize::MAX), |
| 219 | + Err(InvalidAttributeData) |
| 220 | + )); |
| 221 | + Ok(()) |
| 222 | + } |
| 223 | +} |
0 commit comments