use crate::big_data::{BigDataSlices, BIG_DATA_SEGMENT_SIZE};
use crate::error::{NtHiveError, Result};
use crate::helpers::byte_subrange;
use crate::hive::Hive;
use crate::string::NtHiveNameString;
use ::byteorder::{BigEndian, ByteOrder, LittleEndian};
use bitflags::bitflags;
use core::convert::TryInto;
use core::mem;
use core::ops::{Deref, Range};
use core::ptr;
use enumn::N;
use memoffset::offset_of;
use zerocopy::*;
#[cfg(feature = "alloc")]
use {alloc::string::String, alloc::vec::Vec, core::char, core::iter};
const DATA_STORED_IN_DATA_OFFSET: u32 = 0x8000_0000;
bitflags! {
struct KeyValueFlags: u16 {
const VALUE_COMP_NAME = 0x0001;
}
}
#[derive(Clone)]
pub enum KeyValueData<'a, B: ByteSlice> {
Small(&'a [u8]),
Big(BigDataSlices<'a, B>),
}
impl<'a, B> KeyValueData<'a, B>
where
B: ByteSlice,
{
#[cfg(feature = "alloc")]
pub fn into_vec(self) -> Result<Vec<u8>> {
match self {
KeyValueData::Small(data) => Ok(data.to_vec()),
KeyValueData::Big(iter) => {
let mut data = Vec::new();
for slice_data in iter {
let slice_data = slice_data?;
data.extend_from_slice(slice_data);
}
Ok(data)
}
}
}
}
#[derive(Clone, Copy, Debug, Eq, N, PartialEq)]
#[repr(u32)]
pub enum KeyValueDataType {
RegNone = 0x0000_0000,
RegSZ = 0x0000_0001,
RegExpandSZ = 0x0000_0002,
RegBinary = 0x0000_0003,
RegDWord = 0x0000_0004,
RegDWordBigEndian = 0x0000_0005,
RegLink = 0x0000_0006,
RegMultiSZ = 0x0000_0007,
RegResourceList = 0x0000_0008,
RegFullResourceDescriptor = 0x0000_0009,
RegResourceRequirementsList = 0x0000_000a,
RegQWord = 0x0000_000b,
}
#[allow(dead_code)]
#[derive(AsBytes, FromBytes, Unaligned)]
#[repr(packed)]
struct KeyValueHeader {
signature: [u8; 2],
name_length: U16<LittleEndian>,
data_size: U32<LittleEndian>,
data_offset: U32<LittleEndian>,
data_type: U32<LittleEndian>,
flags: U16<LittleEndian>,
spare: U16<LittleEndian>,
}
#[derive(Clone)]
pub struct KeyValue<H: Deref<Target = Hive<B>>, B: ByteSlice> {
hive: H,
header_range: Range<usize>,
data_range: Range<usize>,
}
impl<H, B> KeyValue<H, B>
where
H: Deref<Target = Hive<B>>,
B: ByteSlice,
{
pub(crate) fn new(hive: H, cell_range: Range<usize>) -> Result<Self> {
let header_range = byte_subrange(&cell_range, mem::size_of::<KeyValueHeader>())
.ok_or_else(|| NtHiveError::InvalidHeaderSize {
offset: hive.offset_of_data_offset(cell_range.start),
expected: mem::size_of::<KeyValueHeader>(),
actual: cell_range.len(),
})?;
let data_range = header_range.end..cell_range.end;
let key_value = Self {
hive,
header_range,
data_range,
};
key_value.validate_signature()?;
Ok(key_value)
}
fn header(&self) -> LayoutVerified<&[u8], KeyValueHeader> {
LayoutVerified::new(&self.hive.data[self.header_range.clone()]).unwrap()
}
pub fn data(&self) -> Result<KeyValueData<B>> {
let header = self.header();
let data_size = header.data_size.get();
let data_stored_in_data_offset = data_size & DATA_STORED_IN_DATA_OFFSET > 0;
let data_size = (data_size & !DATA_STORED_IN_DATA_OFFSET) as usize;
if data_stored_in_data_offset {
if data_size > mem::size_of::<u32>() {
return Err(NtHiveError::InvalidSizeField {
offset: self.hive.offset_of_field(&header.data_size),
expected: mem::size_of::<u32>(),
actual: data_size,
});
}
let data_start = self.header_range.start + offset_of!(KeyValueHeader, data_offset);
let data_end = data_start + data_size;
Ok(KeyValueData::Small(&self.hive.data[data_start..data_end]))
} else if data_size <= BIG_DATA_SEGMENT_SIZE {
let cell_range = self
.hive
.cell_range_from_data_offset(header.data_offset.get())?;
if cell_range.len() < data_size {
return Err(NtHiveError::InvalidDataSize {
offset: self.hive.offset_of_data_offset(cell_range.start),
expected: data_size,
actual: cell_range.len(),
});
}
let data_start = cell_range.start;
let data_end = data_start + data_size;
Ok(KeyValueData::Small(&self.hive.data[data_start..data_end]))
} else {
let cell_range = self
.hive
.cell_range_from_data_offset(header.data_offset.get())?;
let iter = BigDataSlices::new(
&self.hive,
data_size as u32,
self.hive.offset_of_field(&header.data_size),
cell_range,
)?;
Ok(KeyValueData::Big(iter))
}
}
#[cfg(feature = "alloc")]
fn utf16le_to_string_lossy<'a, I>(iter: I) -> Result<String>
where
I: Iterator<Item = Result<&'a [u8]>>,
{
let mut string = String::new();
for slice_data in iter {
let slice_data = slice_data?;
let u16_iter = slice_data
.chunks_exact(2)
.map(|two_bytes| u16::from_le_bytes(two_bytes.try_into().unwrap()));
string.reserve(u16_iter.len());
let char_iter =
char::decode_utf16(u16_iter).map(|x| x.unwrap_or(char::REPLACEMENT_CHARACTER));
for c in char_iter {
if c == '\0' {
return Ok(string);
} else {
string.push(c);
}
}
}
Ok(string)
}
#[cfg(feature = "alloc")]
pub fn string_data(&self) -> Result<String> {
match self.data_type()? {
KeyValueDataType::RegSZ | KeyValueDataType::RegExpandSZ => (),
data_type => {
return Err(NtHiveError::InvalidKeyValueDataType {
expected: &[KeyValueDataType::RegSZ, KeyValueDataType::RegExpandSZ],
actual: data_type,
});
}
}
match self.data()? {
KeyValueData::Small(data) => Self::utf16le_to_string_lossy(iter::once(Ok(data))),
KeyValueData::Big(iter) => Self::utf16le_to_string_lossy(iter),
}
}
pub fn dword_data(&self) -> Result<u32> {
if let KeyValueData::Small(data) = self.data()? {
if data.len() != mem::size_of::<u32>() {
return Err(NtHiveError::InvalidDataSize {
offset: self.hive.offset_of_field(&data),
expected: mem::size_of::<u32>(),
actual: data.len(),
});
}
match self.data_type()? {
KeyValueDataType::RegDWord => Ok(LittleEndian::read_u32(data)),
KeyValueDataType::RegDWordBigEndian => Ok(BigEndian::read_u32(data)),
data_type => Err(NtHiveError::InvalidKeyValueDataType {
expected: &[
KeyValueDataType::RegDWord,
KeyValueDataType::RegDWordBigEndian,
],
actual: data_type,
}),
}
} else {
Err(NtHiveError::InvalidDataSize {
offset: self
.hive
.offset_of_data_offset(self.header().data_offset.get() as usize),
expected: mem::size_of::<u32>(),
actual: self.data_size() as usize,
})
}
}
#[cfg(feature = "alloc")]
fn multi_utf16le_to_string_lossy<'a, I>(iter: I) -> Result<Vec<String>>
where
I: Iterator<Item = Result<&'a [u8]>>,
{
let mut strings = Vec::new();
let mut string = String::new();
for slice_data in iter {
let slice_data = slice_data?;
let u16_iter = slice_data
.chunks_exact(2)
.map(|two_bytes| u16::from_le_bytes(two_bytes.try_into().unwrap()));
let char_iter =
char::decode_utf16(u16_iter).map(|x| x.unwrap_or(char::REPLACEMENT_CHARACTER));
for c in char_iter {
if c == '\0' {
if string.is_empty() {
return Ok(strings);
}
strings.push(string);
string = String::new();
} else {
string.push(c);
}
}
}
Ok(strings)
}
#[cfg(feature = "alloc")]
pub fn multi_string_data(&self) -> Result<Vec<String>> {
match self.data_type()? {
KeyValueDataType::RegMultiSZ => (),
data_type => {
return Err(NtHiveError::InvalidKeyValueDataType {
expected: &[KeyValueDataType::RegMultiSZ],
actual: data_type,
});
}
}
match self.data()? {
KeyValueData::Small(data) => Self::multi_utf16le_to_string_lossy(iter::once(Ok(data))),
KeyValueData::Big(iter) => Self::multi_utf16le_to_string_lossy(iter),
}
}
pub fn qword_data(&self) -> Result<u64> {
if let KeyValueData::Small(data) = self.data()? {
if data.len() != mem::size_of::<u64>() {
return Err(NtHiveError::InvalidDataSize {
offset: self.hive.offset_of_field(&data),
expected: mem::size_of::<u64>(),
actual: data.len(),
});
}
match self.data_type()? {
KeyValueDataType::RegQWord => Ok(LittleEndian::read_u64(data)),
data_type => Err(NtHiveError::InvalidKeyValueDataType {
expected: &[KeyValueDataType::RegQWord],
actual: data_type,
}),
}
} else {
Err(NtHiveError::InvalidDataSize {
offset: self
.hive
.offset_of_data_offset(self.header().data_offset.get() as usize),
expected: mem::size_of::<u64>(),
actual: self.data_size() as usize,
})
}
}
pub fn data_size(&self) -> u32 {
let header = self.header();
header.data_size.get() & !DATA_STORED_IN_DATA_OFFSET
}
pub fn data_type(&self) -> Result<KeyValueDataType> {
let header = self.header();
let data_type_code = header.data_type.get();
KeyValueDataType::n(data_type_code).ok_or_else(|| {
NtHiveError::UnsupportedKeyValueDataType {
offset: self.hive.offset_of_field(&header.data_type),
actual: data_type_code,
}
})
}
pub fn name(&self) -> Result<NtHiveNameString> {
let header = self.header();
let flags = KeyValueFlags::from_bits_truncate(header.flags.get());
let name_length = header.name_length.get() as usize;
let name_range = byte_subrange(&self.data_range, name_length).ok_or_else(|| {
NtHiveError::InvalidSizeField {
offset: self.hive.offset_of_field(&header.name_length),
expected: name_length as usize,
actual: self.data_range.len(),
}
})?;
let name_bytes = &self.hive.data[name_range];
if flags.contains(KeyValueFlags::VALUE_COMP_NAME) {
Ok(NtHiveNameString::Latin1(name_bytes))
} else {
Ok(NtHiveNameString::Utf16LE(name_bytes))
}
}
fn validate_signature(&self) -> Result<()> {
let header = self.header();
let signature = &header.signature;
let expected_signature = b"vk";
if signature == expected_signature {
Ok(())
} else {
Err(NtHiveError::InvalidTwoByteSignature {
offset: self.hive.offset_of_field(signature),
expected: expected_signature,
actual: *signature,
})
}
}
}
impl<B> PartialEq for KeyValue<&Hive<B>, B>
where
B: ByteSlice,
{
fn eq(&self, other: &Self) -> bool {
ptr::eq(self.hive, other.hive)
&& self.header_range == other.header_range
&& self.data_range == other.data_range
}
}
impl<B> Eq for KeyValue<&Hive<B>, B> where B: ByteSlice {}
#[cfg(test)]
mod tests {
use crate::*;
#[test]
fn test_data() {
let testhive = crate::helpers::tests::testhive_vec();
let hive = Hive::new(testhive.as_ref()).unwrap();
let root_key_node = hive.root_key_node().unwrap();
let key_node = root_key_node.subkey("data-test").unwrap().unwrap();
let key_value = key_node.value("reg-sz").unwrap().unwrap();
assert_eq!(key_value.data_type().unwrap(), KeyValueDataType::RegSZ);
assert_eq!(key_value.string_data().unwrap(), "sz-test");
let key_value = key_node
.value("reg-sz-with-terminating-nul")
.unwrap()
.unwrap();
assert_eq!(key_value.data_type().unwrap(), KeyValueDataType::RegSZ);
assert_eq!(key_value.string_data().unwrap(), "sz-test");
let key_value = key_node.value("reg-expand-sz").unwrap().unwrap();
assert_eq!(
key_value.data_type().unwrap(),
KeyValueDataType::RegExpandSZ
);
assert_eq!(key_value.string_data().unwrap(), "sz-test");
let key_value = key_node.value("reg-multi-sz").unwrap().unwrap();
assert_eq!(key_value.data_type().unwrap(), KeyValueDataType::RegMultiSZ);
assert_eq!(
key_value.multi_string_data().unwrap(),
vec!["multi-sz-test", "line2"]
);
let key_value = key_node.value("dword").unwrap().unwrap();
assert_eq!(key_value.data_type().unwrap(), KeyValueDataType::RegDWord);
assert_eq!(key_value.dword_data().unwrap(), 42);
let key_value = key_node.value("dword-big-endian").unwrap().unwrap();
assert_eq!(
key_value.data_type().unwrap(),
KeyValueDataType::RegDWordBigEndian
);
assert_eq!(key_value.dword_data().unwrap(), 42 << 24);
let key_value = key_node.value("qword").unwrap().unwrap();
assert_eq!(key_value.data_type().unwrap(), KeyValueDataType::RegQWord);
assert_eq!(key_value.qword_data().unwrap(), u64::MAX);
let key_value = key_node.value("binary").unwrap().unwrap();
let key_value_data = key_value.data().unwrap();
assert_eq!(key_value.data_type().unwrap(), KeyValueDataType::RegBinary);
assert!(matches!(key_value_data, KeyValueData::Small(_)));
assert_eq!(key_value_data.into_vec().unwrap(), vec![1, 2, 3, 4, 5]);
}
}