use chrono::{DateTime, TimeZone, Utc};
use std::collections::HashMap;
use std::hash::{Hash, Hasher};
const TIMESTAMP_EXTENDED_TYPE: u8 = 121;
fn try_parse_iso8601(s: &str) -> Option<i64> {
DateTime::parse_from_rfc3339(s)
.ok()
.map(|dt| dt.timestamp())
}
fn format_iso8601(epoch: i64) -> String {
Utc.timestamp_opt(epoch, 0)
.single()
.map(|dt| dt.to_rfc3339_opts(chrono::SecondsFormat::Secs, true))
.unwrap_or_else(|| format!("{epoch}"))
}
mod validation;
pub use validation::{
validate_data_section, validate_data_value_pointers, validate_data_value_utf8,
validate_value_strings_utf8, DataFormatStats, DataFormatValidationResult,
PointerValidationError, PointerValidationResult, PointerValidationStats, MAX_POINTER_DEPTH,
MAX_TOTAL_DEPTH,
};
#[derive(Debug, Clone, PartialEq)]
pub enum DataValue {
#[allow(dead_code)]
Pointer(u32),
String(String),
Double(f64),
Bytes(Vec<u8>),
Uint16(u16),
Uint32(u32),
Map(HashMap<String, Self>),
Int32(i32),
Uint64(u64),
Uint128(u128),
Array(Vec<Self>),
Bool(bool),
Float(f32),
Timestamp(i64),
}
impl serde::Serialize for DataValue {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
match self {
Self::Pointer(_) => Err(serde::ser::Error::custom(
"Pointer is an internal type and cannot be serialized to JSON",
)),
Self::String(s) => serializer.serialize_str(s),
Self::Double(d) => serializer.serialize_f64(*d),
Self::Bytes(b) => serializer.serialize_bytes(b),
Self::Uint16(n) => serializer.serialize_u16(*n),
Self::Uint32(n) => serializer.serialize_u32(*n),
Self::Map(m) => m.serialize(serializer),
Self::Int32(n) => serializer.serialize_i32(*n),
Self::Uint64(n) => serializer.serialize_u64(*n),
Self::Uint128(n) => serializer.serialize_u128(*n),
Self::Array(a) => a.serialize(serializer),
Self::Bool(b) => serializer.serialize_bool(*b),
Self::Float(f) => serializer.serialize_f32(*f),
Self::Timestamp(epoch) => serializer.serialize_str(&format_iso8601(*epoch)),
}
}
}
impl<'de> serde::Deserialize<'de> for DataValue {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
struct DataValueVisitor;
impl<'de> serde::de::Visitor<'de> for DataValueVisitor {
type Value = DataValue;
fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
formatter.write_str("a valid MMDB data value")
}
fn visit_bool<E>(self, v: bool) -> Result<DataValue, E> {
Ok(DataValue::Bool(v))
}
fn visit_i32<E>(self, v: i32) -> Result<DataValue, E> {
Ok(DataValue::Int32(v))
}
fn visit_i64<E>(self, v: i64) -> Result<DataValue, E>
where
E: serde::de::Error,
{
if v >= 0 {
if v <= i64::from(u16::MAX) {
Ok(DataValue::Uint16(u16::try_from(v).unwrap()))
} else if v <= i64::from(u32::MAX) {
Ok(DataValue::Uint32(u32::try_from(v).unwrap()))
} else {
Ok(DataValue::Uint64(u64::try_from(v).unwrap()))
}
} else if v >= i64::from(i32::MIN) {
Ok(DataValue::Int32(i32::try_from(v).unwrap()))
} else {
Err(serde::de::Error::custom(format!(
"value {v} is outside the supported signed integer range \
({} to {}). MMDB format only supports Int32. \
Consider using a string or unsigned integer instead.",
i32::MIN,
i32::MAX
)))
}
}
fn visit_u64<E>(self, v: u64) -> Result<DataValue, E> {
if v <= u64::from(u16::MAX) {
Ok(DataValue::Uint16(u16::try_from(v).unwrap()))
} else if v <= u64::from(u32::MAX) {
Ok(DataValue::Uint32(u32::try_from(v).unwrap()))
} else {
Ok(DataValue::Uint64(v))
}
}
fn visit_f32<E>(self, v: f32) -> Result<DataValue, E> {
Ok(DataValue::Float(v))
}
fn visit_f64<E>(self, v: f64) -> Result<DataValue, E> {
Ok(DataValue::Double(v))
}
fn visit_str<E>(self, v: &str) -> Result<DataValue, E> {
if let Some(epoch) = try_parse_iso8601(v) {
return Ok(DataValue::Timestamp(epoch));
}
Ok(DataValue::String(v.to_string()))
}
fn visit_string<E>(self, v: String) -> Result<DataValue, E> {
if let Some(epoch) = try_parse_iso8601(&v) {
return Ok(DataValue::Timestamp(epoch));
}
Ok(DataValue::String(v))
}
fn visit_bytes<E>(self, v: &[u8]) -> Result<DataValue, E> {
Ok(DataValue::Bytes(v.to_vec()))
}
fn visit_byte_buf<E>(self, v: Vec<u8>) -> Result<DataValue, E> {
Ok(DataValue::Bytes(v))
}
fn visit_seq<A>(self, mut seq: A) -> Result<DataValue, A::Error>
where
A: serde::de::SeqAccess<'de>,
{
let mut array = Vec::new();
while let Some(value) = seq.next_element()? {
array.push(value);
}
Ok(DataValue::Array(array))
}
fn visit_map<A>(self, mut map: A) -> Result<DataValue, A::Error>
where
A: serde::de::MapAccess<'de>,
{
let mut hash_map = HashMap::new();
while let Some((key, value)) = map.next_entry()? {
hash_map.insert(key, value);
}
Ok(DataValue::Map(hash_map))
}
}
deserializer.deserialize_any(DataValueVisitor)
}
}
impl Hash for DataValue {
fn hash<H: Hasher>(&self, state: &mut H) {
std::mem::discriminant(self).hash(state);
match self {
Self::Pointer(v) | Self::Uint32(v) => v.hash(state),
Self::String(v) => v.hash(state),
Self::Double(v) => {
v.to_bits().hash(state);
}
Self::Bytes(v) => v.hash(state),
Self::Uint16(v) => v.hash(state),
Self::Map(m) => {
let mut keys: Vec<&String> = m.keys().collect();
keys.sort_unstable();
keys.len().hash(state);
for key in keys {
key.hash(state);
m[key].hash(state);
}
}
Self::Int32(v) => v.hash(state),
Self::Uint64(v) => v.hash(state),
Self::Uint128(v) => v.hash(state),
Self::Array(v) => {
v.len().hash(state);
for item in v {
item.hash(state);
}
}
Self::Bool(v) => v.hash(state),
Self::Float(v) => {
v.to_bits().hash(state);
}
Self::Timestamp(v) => v.hash(state),
}
}
}
pub struct DataEncoder {
buffer: Vec<u8>,
dedup_map: HashMap<Vec<u8>, u32>,
string_cache: HashMap<String, u32>,
intern_strings: bool,
}
impl DataEncoder {
#[must_use]
pub fn new() -> Self {
Self {
buffer: Vec::new(),
dedup_map: HashMap::new(),
string_cache: HashMap::new(),
intern_strings: true,
}
}
#[must_use]
pub fn new_without_interning() -> Self {
Self {
buffer: Vec::new(),
dedup_map: HashMap::new(),
string_cache: HashMap::new(),
intern_strings: false,
}
}
pub fn encode(&mut self, value: &DataValue) -> u32 {
let saved_intern = self.intern_strings;
self.intern_strings = false;
let mut temp = Vec::new();
Self::encode_to_buffer(value, &mut temp);
self.intern_strings = saved_intern;
if let Some(&offset) = self.dedup_map.get(&temp) {
return offset;
}
let offset = u32::try_from(self.buffer.len()).expect("Data section exceeds u32::MAX bytes");
self.encode_value_interned(value);
self.dedup_map.insert(temp, offset);
offset
}
#[must_use]
pub fn into_bytes(self) -> Vec<u8> {
self.buffer
}
#[must_use]
pub fn size(&self) -> usize {
self.buffer.len()
}
fn encode_value_interned(&mut self, value: &DataValue) {
match value {
DataValue::String(s) if self.intern_strings => {
if let Some(&existing_offset) = self.string_cache.get(s) {
Self::encode_pointer(existing_offset, &mut self.buffer);
} else {
let offset = u32::try_from(self.buffer.len())
.expect("Data section exceeds u32::MAX bytes");
Self::encode_string(s, &mut self.buffer);
self.string_cache.insert(s.clone(), offset);
}
}
DataValue::Map(m) => self.encode_map_interned(m),
DataValue::Array(a) => self.encode_array_interned(a),
_ => Self::encode_to_buffer(value, &mut self.buffer),
}
}
fn encode_to_buffer(value: &DataValue, buffer: &mut Vec<u8>) {
match value {
DataValue::Pointer(offset) => Self::encode_pointer(*offset, buffer),
DataValue::String(s) => Self::encode_string(s, buffer),
DataValue::Double(d) => Self::encode_double(*d, buffer),
DataValue::Bytes(b) => Self::encode_bytes(b, buffer),
DataValue::Uint16(n) => Self::encode_uint16(*n, buffer),
DataValue::Uint32(n) => Self::encode_uint32(*n, buffer),
DataValue::Map(m) => Self::encode_map(m, buffer),
DataValue::Int32(n) => Self::encode_int32(*n, buffer),
DataValue::Uint64(n) => Self::encode_uint64(*n, buffer),
DataValue::Uint128(n) => Self::encode_uint128(*n, buffer),
DataValue::Array(a) => Self::encode_array(a, buffer),
DataValue::Bool(b) => Self::encode_bool(*b, buffer),
DataValue::Float(f) => Self::encode_float(*f, buffer),
DataValue::Timestamp(t) => Self::encode_timestamp(*t, buffer),
}
}
fn encode_pointer(offset: u32, buffer: &mut Vec<u8>) {
let size = if offset < 2048 {
0 } else if offset < 2048 + 524288 {
1 } else if offset < 2048 + 524288 + 134217728 {
2 } else {
3 };
match size {
0 => {
let high_3_bits = ((offset >> 8) & 0x7) as u8;
let low_8_bits = (offset & 0xFF) as u8;
let ctrl = 0x20 | high_3_bits; buffer.push(ctrl);
buffer.push(low_8_bits);
}
1 => {
let adjusted = offset - 2048;
let high_3_bits = ((adjusted >> 16) & 0x7) as u8;
let mid_8_bits = ((adjusted >> 8) & 0xFF) as u8;
let low_8_bits = (adjusted & 0xFF) as u8;
let ctrl = 0x20 | (1 << 3) | high_3_bits; buffer.push(ctrl);
buffer.push(mid_8_bits);
buffer.push(low_8_bits);
}
2 => {
let adjusted = offset - 526336;
let high_3_bits = ((adjusted >> 24) & 0x7) as u8;
let b0 = ((adjusted >> 16) & 0xFF) as u8;
let b1 = ((adjusted >> 8) & 0xFF) as u8;
let b2 = (adjusted & 0xFF) as u8;
let ctrl = 0x20 | (2 << 3) | high_3_bits; buffer.push(ctrl);
buffer.push(b0);
buffer.push(b1);
buffer.push(b2);
}
_ => {
let ctrl = 0x20 | (3 << 3); buffer.push(ctrl);
buffer.extend_from_slice(&offset.to_be_bytes());
}
}
}
fn encode_string(s: &str, buffer: &mut Vec<u8>) {
let bytes = s.as_bytes();
Self::encode_with_size(2, bytes.len(), buffer);
buffer.extend_from_slice(bytes);
}
fn encode_double(d: f64, buffer: &mut Vec<u8>) {
buffer.push(0x68); buffer.extend_from_slice(&d.to_be_bytes());
}
fn encode_bytes(b: &[u8], buffer: &mut Vec<u8>) {
Self::encode_with_size(4, b.len(), buffer);
buffer.extend_from_slice(b);
}
fn encode_uint16(n: u16, buffer: &mut Vec<u8>) {
buffer.push(0xA2); buffer.extend_from_slice(&n.to_be_bytes());
}
fn encode_uint32(n: u32, buffer: &mut Vec<u8>) {
buffer.push(0xC4); buffer.extend_from_slice(&n.to_be_bytes());
}
fn encode_map_interned(&mut self, m: &HashMap<String, DataValue>) {
Self::encode_with_size(7, m.len(), &mut self.buffer);
let mut pairs: Vec<_> = m.iter().collect();
pairs.sort_by_key(|(k, _)| *k);
for (key, value) in pairs {
if self.intern_strings {
if let Some(&existing_offset) = self.string_cache.get(key) {
Self::encode_pointer(existing_offset, &mut self.buffer);
} else {
let offset = u32::try_from(self.buffer.len())
.expect("Data section exceeds u32::MAX bytes");
Self::encode_string(key, &mut self.buffer);
self.string_cache.insert(key.clone(), offset);
}
} else {
Self::encode_string(key, &mut self.buffer);
}
self.encode_value_interned(value);
}
}
fn encode_map(m: &HashMap<String, DataValue>, buffer: &mut Vec<u8>) {
Self::encode_with_size(7, m.len(), buffer);
let mut pairs: Vec<_> = m.iter().collect();
pairs.sort_by_key(|(k, _)| *k);
for (key, value) in pairs {
Self::encode_string(key, buffer);
Self::encode_to_buffer(value, buffer);
}
}
fn encode_int32(n: i32, buffer: &mut Vec<u8>) {
buffer.push(0x04); buffer.push(0x01); buffer.extend_from_slice(&n.to_be_bytes());
}
fn encode_uint64(n: u64, buffer: &mut Vec<u8>) {
buffer.push(0x08); buffer.push(0x02); buffer.extend_from_slice(&n.to_be_bytes());
}
fn encode_uint128(n: u128, buffer: &mut Vec<u8>) {
buffer.push(0x10); buffer.push(0x03); buffer.extend_from_slice(&n.to_be_bytes());
}
fn encode_array_interned(&mut self, a: &[DataValue]) {
let size = a.len();
if size < 29 {
self.buffer.push(u8::try_from(size).unwrap());
} else if size < 29 + 256 {
self.buffer.push(29);
self.buffer.push(u8::try_from(size - 29).unwrap());
} else if size < 29 + 256 + 65536 {
self.buffer.push(30);
let adjusted = size - 29 - 256;
self.buffer
.extend_from_slice(&u16::try_from(adjusted).unwrap().to_be_bytes());
} else {
self.buffer.push(31);
let adjusted = size - 29 - 256 - 65536;
self.buffer
.extend_from_slice(&u32::try_from(adjusted).unwrap().to_be_bytes()[1..]);
}
self.buffer.push(0x04);
for value in a {
self.encode_value_interned(value);
}
}
fn encode_array(a: &[DataValue], buffer: &mut Vec<u8>) {
let size = a.len();
if size < 29 {
buffer.push(u8::try_from(size).unwrap());
} else if size < 29 + 256 {
buffer.push(29);
buffer.push(u8::try_from(size - 29).unwrap());
} else if size < 29 + 256 + 65536 {
buffer.push(30);
let adjusted = size - 29 - 256;
buffer.extend_from_slice(&u16::try_from(adjusted).unwrap().to_be_bytes());
} else {
buffer.push(31);
let adjusted = size - 29 - 256 - 65536;
buffer.extend_from_slice(&u32::try_from(adjusted).unwrap().to_be_bytes()[1..]);
}
buffer.push(0x04);
for value in a {
Self::encode_to_buffer(value, buffer);
}
}
fn encode_bool(b: bool, buffer: &mut Vec<u8>) {
if b {
buffer.push(0x01); } else {
buffer.push(0x00); }
buffer.push(0x07); }
fn encode_float(f: f32, buffer: &mut Vec<u8>) {
buffer.push(0x04); buffer.push(0x08); buffer.extend_from_slice(&f.to_be_bytes());
}
fn encode_timestamp(epoch: i64, buffer: &mut Vec<u8>) {
buffer.push(0x08); buffer.push(TIMESTAMP_EXTENDED_TYPE);
buffer.extend_from_slice(&epoch.to_be_bytes());
}
fn encode_with_size(type_id: u8, size: usize, buffer: &mut Vec<u8>) {
let type_bits = type_id << 5;
if size < 29 {
buffer.push(type_bits | u8::try_from(size).unwrap());
} else if size < 29 + 256 {
buffer.push(type_bits | 29);
buffer.push(u8::try_from(size - 29).unwrap());
} else if size < 29 + 256 + 65536 {
buffer.push(type_bits | 30);
let adjusted = size - 29 - 256;
buffer.extend_from_slice(&u16::try_from(adjusted).unwrap().to_be_bytes());
} else {
buffer.push(type_bits | 31);
let adjusted = size - 29 - 256 - 65536;
buffer.extend_from_slice(&u32::try_from(adjusted).unwrap().to_be_bytes()[1..]);
}
}
}
impl Default for DataEncoder {
fn default() -> Self {
Self::new()
}
}
pub struct DataDecoder<'a> {
buffer: &'a [u8],
base_offset: usize,
}
impl<'a> DataDecoder<'a> {
#[must_use]
pub fn new(buffer: &'a [u8], base_offset: usize) -> Self {
Self {
buffer,
base_offset,
}
}
pub fn decode(&self, offset: u32) -> Result<DataValue, &'static str> {
let mut cursor = offset as usize;
if cursor < self.base_offset {
return Err("Offset before base");
}
cursor -= self.base_offset;
let value = self.decode_at(&mut cursor)?;
self.resolve_pointers(value)
}
fn decode_at(&self, cursor: &mut usize) -> Result<DataValue, &'static str> {
if *cursor >= self.buffer.len() {
return Err("Cursor out of bounds");
}
let ctrl = self.buffer[*cursor];
*cursor += 1;
let type_id = ctrl >> 5;
let payload = ctrl & 0x1F;
match type_id {
0 => self.decode_extended(cursor, payload),
1 => self.decode_pointer(cursor, payload),
2 => self.decode_string(cursor, payload),
3 => self.decode_double(cursor),
4 => self.decode_bytes(cursor, payload),
5 => self.decode_uint16(cursor, payload),
6 => self.decode_uint32(cursor, payload),
7 => self.decode_map(cursor, payload),
_ => Err("Invalid type"),
}
}
fn decode_extended(
&self,
cursor: &mut usize,
size_from_ctrl: u8,
) -> Result<DataValue, &'static str> {
if *cursor >= self.buffer.len() {
return Err("Extended type truncated");
}
let raw_ext_type = self.buffer[*cursor];
let type_id = 7 + raw_ext_type;
*cursor += 1;
match type_id {
8 => self.decode_int32(cursor, size_from_ctrl), 9 => self.decode_uint64(cursor, size_from_ctrl), 10 => self.decode_uint128(cursor, size_from_ctrl), 11 => self.decode_array(cursor, size_from_ctrl), 14 => Ok(DataValue::Bool(size_from_ctrl != 0)), 15 => self.decode_float(cursor, size_from_ctrl), 128 => self.decode_timestamp(cursor, size_from_ctrl), _ => {
eprintln!(
"Unknown extended type: raw_ext_type={}, type_id={}, size_from_ctrl={}, offset={}",
raw_ext_type, type_id, size_from_ctrl, *cursor - 1
);
Err("Unknown extended type")
}
}
}
fn decode_pointer(&self, cursor: &mut usize, payload: u8) -> Result<DataValue, &'static str> {
let size_bits = (payload >> 3) & 0x3; let offset = match size_bits {
0 => {
if *cursor >= self.buffer.len() {
return Err("Pointer data truncated");
}
let low_3_bits = u32::from(payload & 0x7);
let next_byte = u32::from(self.buffer[*cursor]);
*cursor += 1;
(low_3_bits << 8) | next_byte
}
1 => {
if *cursor + 1 >= self.buffer.len() {
return Err("Pointer data truncated");
}
let low_3_bits = u32::from(payload & 0x7);
let b0 = u32::from(self.buffer[*cursor]);
let b1 = u32::from(self.buffer[*cursor + 1]);
*cursor += 2;
2048 + ((low_3_bits << 16) | (b0 << 8) | b1)
}
2 => {
if *cursor + 2 >= self.buffer.len() {
return Err("Pointer data truncated");
}
let low_3_bits = u32::from(payload & 0x7);
let b0 = u32::from(self.buffer[*cursor]);
let b1 = u32::from(self.buffer[*cursor + 1]);
let b2 = u32::from(self.buffer[*cursor + 2]);
*cursor += 3;
526336 + ((low_3_bits << 24) | (b0 << 16) | (b1 << 8) | b2)
}
3 => {
if *cursor + 3 >= self.buffer.len() {
return Err("Pointer data truncated");
}
let mut bytes = [0u8; 4];
bytes.copy_from_slice(&self.buffer[*cursor..*cursor + 4]);
*cursor += 4;
u32::from_be_bytes(bytes)
}
_ => return Err("Invalid pointer size"),
};
Ok(DataValue::Pointer(offset))
}
fn decode_string(&self, cursor: &mut usize, size_bits: u8) -> Result<DataValue, &'static str> {
let len = self.decode_size(cursor, size_bits)?;
if *cursor + len > self.buffer.len() {
return Err("String data out of bounds");
}
let s = std::str::from_utf8(&self.buffer[*cursor..*cursor + len])
.map_err(|_| "Invalid UTF-8")?;
*cursor += len;
Ok(DataValue::String(s.to_string()))
}
fn decode_double(&self, cursor: &mut usize) -> Result<DataValue, &'static str> {
if *cursor + 8 > self.buffer.len() {
return Err("Double data out of bounds");
}
let mut bytes = [0u8; 8];
bytes.copy_from_slice(&self.buffer[*cursor..*cursor + 8]);
*cursor += 8;
Ok(DataValue::Double(f64::from_be_bytes(bytes)))
}
fn decode_bytes(&self, cursor: &mut usize, size_bits: u8) -> Result<DataValue, &'static str> {
let len = self.decode_size(cursor, size_bits)?;
if *cursor + len > self.buffer.len() {
return Err("Bytes data out of bounds");
}
let bytes = self.buffer[*cursor..*cursor + len].to_vec();
*cursor += len;
Ok(DataValue::Bytes(bytes))
}
fn decode_uint16(&self, cursor: &mut usize, size_bits: u8) -> Result<DataValue, &'static str> {
let size = self.decode_size(cursor, size_bits)?;
if size > 2 {
return Err("Uint16 size too large");
}
if *cursor + size > self.buffer.len() {
return Err("Uint16 data out of bounds");
}
let mut value = 0u16;
for i in 0..size {
value = (value << 8) | u16::from(self.buffer[*cursor + i]);
}
*cursor += size;
Ok(DataValue::Uint16(value))
}
fn decode_uint32(&self, cursor: &mut usize, size_bits: u8) -> Result<DataValue, &'static str> {
let size = self.decode_size(cursor, size_bits)?;
if size > 4 {
return Err("Uint32 size too large");
}
if *cursor + size > self.buffer.len() {
return Err("Uint32 data out of bounds");
}
let mut value = 0u32;
for i in 0..size {
value = (value << 8) | u32::from(self.buffer[*cursor + i]);
}
*cursor += size;
Ok(DataValue::Uint32(value))
}
fn decode_map(&self, cursor: &mut usize, size_bits: u8) -> Result<DataValue, &'static str> {
let count = self.decode_size(cursor, size_bits)?;
let mut map = HashMap::new();
for _ in 0..count {
let key_value = self.decode_at(cursor)?;
let key = match key_value {
DataValue::String(s) => s,
DataValue::Pointer(offset) => {
match self.decode(offset)? {
DataValue::String(s) => s,
_ => return Err("Pointer in map key must point to string"),
}
}
_ => return Err("Map key must be string or pointer to string"),
};
let value = self.decode_at(cursor)?;
map.insert(key, value);
}
Ok(DataValue::Map(map))
}
fn decode_int32(&self, cursor: &mut usize, size_bits: u8) -> Result<DataValue, &'static str> {
let size = self.decode_size(cursor, size_bits)?;
if size > 4 {
return Err("Int32 size too large");
}
if *cursor + size > self.buffer.len() {
return Err("Int32 data out of bounds");
}
let mut value = 0i32;
if size > 0 {
let is_negative = (self.buffer[*cursor] & 0x80) != 0;
if is_negative {
value = -1;
}
for i in 0..size {
value = (value << 8) | i32::from(self.buffer[*cursor + i]);
}
}
*cursor += size;
Ok(DataValue::Int32(value))
}
fn decode_uint64(&self, cursor: &mut usize, size_bits: u8) -> Result<DataValue, &'static str> {
let size = self.decode_size(cursor, size_bits)?;
if size > 8 {
return Err("Uint64 size too large");
}
if *cursor + size > self.buffer.len() {
return Err("Uint64 data out of bounds");
}
let mut value = 0u64;
for i in 0..size {
value = (value << 8) | u64::from(self.buffer[*cursor + i]);
}
*cursor += size;
Ok(DataValue::Uint64(value))
}
fn decode_uint128(&self, cursor: &mut usize, size_bits: u8) -> Result<DataValue, &'static str> {
let size = self.decode_size(cursor, size_bits)?;
if size > 16 {
return Err("Uint128 size too large");
}
if *cursor + size > self.buffer.len() {
return Err("Uint128 data out of bounds");
}
let mut value = 0u128;
for i in 0..size {
value = (value << 8) | u128::from(self.buffer[*cursor + i]);
}
*cursor += size;
Ok(DataValue::Uint128(value))
}
fn decode_array(&self, cursor: &mut usize, size_bits: u8) -> Result<DataValue, &'static str> {
let count = self.decode_size(cursor, size_bits)?;
let mut array = Vec::with_capacity(count);
for _ in 0..count {
array.push(self.decode_at(cursor)?);
}
Ok(DataValue::Array(array))
}
fn decode_float(&self, cursor: &mut usize, size_bits: u8) -> Result<DataValue, &'static str> {
if size_bits != 4 {
return Err("Float must be 4 bytes");
}
if *cursor + 4 > self.buffer.len() {
return Err("Float data out of bounds");
}
let mut bytes = [0u8; 4];
bytes.copy_from_slice(&self.buffer[*cursor..*cursor + 4]);
*cursor += 4;
Ok(DataValue::Float(f32::from_be_bytes(bytes)))
}
fn decode_timestamp(
&self,
cursor: &mut usize,
size_bits: u8,
) -> Result<DataValue, &'static str> {
if size_bits != 8 {
return Err("Timestamp must be 8 bytes");
}
if *cursor + 8 > self.buffer.len() {
return Err("Timestamp data out of bounds");
}
let mut bytes = [0u8; 8];
bytes.copy_from_slice(&self.buffer[*cursor..*cursor + 8]);
*cursor += 8;
Ok(DataValue::Timestamp(i64::from_be_bytes(bytes)))
}
fn decode_size(&self, cursor: &mut usize, size_bits: u8) -> Result<usize, &'static str> {
match size_bits {
0..=28 => Ok(size_bits as usize),
29 => {
if *cursor >= self.buffer.len() {
return Err("Size byte out of bounds");
}
let size = self.buffer[*cursor] as usize;
*cursor += 1;
Ok(29 + size)
}
30 => {
if *cursor + 2 > self.buffer.len() {
return Err("Size bytes out of bounds");
}
let mut bytes = [0u8; 2];
bytes.copy_from_slice(&self.buffer[*cursor..*cursor + 2]);
*cursor += 2;
Ok(29 + 256 + u16::from_be_bytes(bytes) as usize)
}
31 => {
if *cursor + 3 > self.buffer.len() {
return Err("Size bytes out of bounds");
}
let b0 = self.buffer[*cursor] as usize;
let b1 = self.buffer[*cursor + 1] as usize;
let b2 = self.buffer[*cursor + 2] as usize;
*cursor += 3;
Ok(29 + 256 + 65536 + ((b0 << 16) | (b1 << 8) | b2))
}
_ => Err("Invalid size encoding"),
}
}
fn resolve_pointers(&self, value: DataValue) -> Result<DataValue, &'static str> {
match value {
DataValue::Pointer(offset) => {
let mut cursor = offset as usize;
if cursor < self.base_offset {
return Err("Pointer offset before base");
}
cursor -= self.base_offset;
let pointed_value = self.decode_at(&mut cursor)?;
self.resolve_pointers(pointed_value)
}
DataValue::Map(entries) => {
let mut resolved_map = HashMap::new();
for (key, val) in entries {
resolved_map.insert(key, self.resolve_pointers(val)?);
}
Ok(DataValue::Map(resolved_map))
}
DataValue::Array(items) => {
let mut resolved_array = Vec::new();
for item in items {
resolved_array.push(self.resolve_pointers(item)?);
}
Ok(DataValue::Array(resolved_array))
}
other => Ok(other),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_encode_decode_all_types() {
let mut encoder = DataEncoder::new();
let string_val = DataValue::String("hello".to_string());
let uint16_val = DataValue::Uint16(12345);
let uint32_val = DataValue::Uint32(0xDEADBEEF);
let uint64_val = DataValue::Uint64(0x123456789ABCDEF0);
let uint128_val = DataValue::Uint128(0x0123456789ABCDEF0123456789ABCDEF);
let int32_val = DataValue::Int32(-42);
let double_val = DataValue::Double(std::f64::consts::PI);
let float_val = DataValue::Float(std::f32::consts::E);
let bool_val = DataValue::Bool(true);
let bytes_val = DataValue::Bytes(vec![0xDE, 0xAD, 0xBE, 0xEF]);
let offsets = [
encoder.encode(&string_val),
encoder.encode(&uint16_val),
encoder.encode(&uint32_val),
encoder.encode(&uint64_val),
encoder.encode(&uint128_val),
encoder.encode(&int32_val),
encoder.encode(&double_val),
encoder.encode(&float_val),
encoder.encode(&bool_val),
encoder.encode(&bytes_val),
];
let bytes = encoder.into_bytes();
let decoder = DataDecoder::new(&bytes, 0);
let values = vec![
string_val,
uint16_val,
uint32_val,
uint64_val,
uint128_val,
int32_val,
double_val,
float_val,
bool_val,
bytes_val,
];
for (offset, expected) in offsets.iter().zip(values.iter()) {
let decoded = decoder.decode(*offset).unwrap();
assert_eq!(&decoded, expected);
}
}
#[test]
fn test_encode_decode_map() {
let mut encoder = DataEncoder::new();
let mut map = HashMap::new();
map.insert("country".to_string(), DataValue::String("US".to_string()));
map.insert("asn".to_string(), DataValue::Uint32(13335));
map.insert("score".to_string(), DataValue::Double(0.95));
let value = DataValue::Map(map.clone());
let offset = encoder.encode(&value);
let bytes = encoder.into_bytes();
let decoder = DataDecoder::new(&bytes, 0);
let decoded = decoder.decode(offset).unwrap();
assert_eq!(decoded, value);
}
#[test]
fn test_encode_decode_array() {
let mut encoder = DataEncoder::new();
let value = DataValue::Array(vec![
DataValue::String("tag1".to_string()),
DataValue::String("tag2".to_string()),
DataValue::Uint32(123),
DataValue::Bool(false),
]);
let offset = encoder.encode(&value);
let bytes = encoder.into_bytes();
let decoder = DataDecoder::new(&bytes, 0);
let decoded = decoder.decode(offset).unwrap();
assert_eq!(decoded, value);
}
#[test]
fn test_deduplication() {
let mut encoder = DataEncoder::new();
let value = DataValue::String("test".to_string());
let offset1 = encoder.encode(&value);
let offset2 = encoder.encode(&value);
let offset3 = encoder.encode(&value);
assert_eq!(offset1, offset2);
assert_eq!(offset2, offset3);
let value2 = DataValue::String("different".to_string());
let offset4 = encoder.encode(&value2);
assert_ne!(offset1, offset4);
}
#[test]
fn test_complex_nested_structure() {
let mut encoder = DataEncoder::new();
let mut threat_data = HashMap::new();
threat_data.insert(
"threat_level".to_string(),
DataValue::String("high".to_string()),
);
threat_data.insert(
"category".to_string(),
DataValue::String("malware".to_string()),
);
threat_data.insert("confidence".to_string(), DataValue::Float(0.98));
threat_data.insert("first_seen".to_string(), DataValue::Uint64(1704067200));
let mut indicators = HashMap::new();
indicators.insert("ip_count".to_string(), DataValue::Uint32(42));
indicators.insert("domain_count".to_string(), DataValue::Uint32(15));
threat_data.insert("indicators".to_string(), DataValue::Map(indicators));
threat_data.insert(
"tags".to_string(),
DataValue::Array(vec![
DataValue::String("botnet".to_string()),
DataValue::String("c2".to_string()),
]),
);
threat_data.insert("active".to_string(), DataValue::Bool(true));
let value = DataValue::Map(threat_data);
let offset = encoder.encode(&value);
let bytes = encoder.into_bytes();
let decoder = DataDecoder::new(&bytes, 0);
let decoded = decoder.decode(offset).unwrap();
assert_eq!(decoded, value);
}
#[test]
fn test_large_strings() {
let mut encoder = DataEncoder::new();
let short = "x".repeat(28); let medium = "x".repeat(100); let long = "x".repeat(1000);
let offset1 = encoder.encode(&DataValue::String(short.clone()));
let offset2 = encoder.encode(&DataValue::String(medium.clone()));
let offset3 = encoder.encode(&DataValue::String(long.clone()));
let bytes = encoder.into_bytes();
let decoder = DataDecoder::new(&bytes, 0);
assert_eq!(decoder.decode(offset1).unwrap(), DataValue::String(short));
assert_eq!(decoder.decode(offset2).unwrap(), DataValue::String(medium));
assert_eq!(decoder.decode(offset3).unwrap(), DataValue::String(long));
}
#[test]
fn test_string_interning() {
let mut encoder = DataEncoder::new();
let mut map1 = HashMap::new();
map1.insert(
"threat_level".to_string(),
DataValue::String("high".to_string()),
);
map1.insert(
"category".to_string(),
DataValue::String("malware".to_string()),
);
map1.insert("score".to_string(), DataValue::Uint32(95));
let mut map2 = HashMap::new();
map2.insert(
"threat_level".to_string(),
DataValue::String("high".to_string()),
); map2.insert(
"category".to_string(),
DataValue::String("phishing".to_string()),
);
map2.insert("score".to_string(), DataValue::Uint32(88));
let mut map3 = HashMap::new();
map3.insert(
"threat_level".to_string(),
DataValue::String("high".to_string()),
); map3.insert(
"category".to_string(),
DataValue::String("malware".to_string()),
); map3.insert("score".to_string(), DataValue::Uint32(92));
let offset1 = encoder.encode(&DataValue::Map(map1.clone()));
let offset2 = encoder.encode(&DataValue::Map(map2.clone()));
let offset3 = encoder.encode(&DataValue::Map(map3.clone()));
let bytes_with_interning = encoder.into_bytes();
let mut encoder_no_intern = DataEncoder::new_without_interning();
encoder_no_intern.encode(&DataValue::Map(map1.clone()));
encoder_no_intern.encode(&DataValue::Map(map2.clone()));
encoder_no_intern.encode(&DataValue::Map(map3.clone()));
let bytes_no_interning = encoder_no_intern.into_bytes();
println!("With interning: {} bytes", bytes_with_interning.len());
println!("Without interning: {} bytes", bytes_no_interning.len());
println!(
"Savings: {} bytes ({:.1}%)",
bytes_no_interning.len() - bytes_with_interning.len(),
100.0 * (bytes_no_interning.len() - bytes_with_interning.len()) as f64
/ bytes_no_interning.len() as f64
);
assert!(bytes_with_interning.len() < bytes_no_interning.len());
let decoder = DataDecoder::new(&bytes_with_interning, 0);
let decoded1 = decoder.decode(offset1).unwrap();
let decoded2 = decoder.decode(offset2).unwrap();
let decoded3 = decoder.decode(offset3).unwrap();
assert_eq!(decoded1, DataValue::Map(map1));
assert_eq!(decoded2, DataValue::Map(map2));
assert_eq!(decoded3, DataValue::Map(map3));
}
#[test]
fn test_string_interning_in_arrays() {
let mut encoder = DataEncoder::new();
let array = DataValue::Array(vec![
DataValue::String("botnet".to_string()),
DataValue::String("c2".to_string()),
DataValue::String("botnet".to_string()), DataValue::String("malware".to_string()),
DataValue::String("c2".to_string()), ]);
let offset = encoder.encode(&array);
let bytes = encoder.into_bytes();
let decoder = DataDecoder::new(&bytes, 0);
let decoded = decoder.decode(offset).unwrap();
assert_eq!(decoded, array);
}
#[test]
fn test_pointer_encoding() {
let mut encoder = DataEncoder::new();
let target_data = DataValue::String("shared_value".to_string());
let target_offset = encoder.encode(&target_data);
let mut map = HashMap::new();
map.insert(
"direct".to_string(),
DataValue::String("direct_value".to_string()),
);
map.insert("ptr_ref".to_string(), DataValue::Pointer(target_offset));
let map_offset = encoder.encode(&DataValue::Map(map));
let bytes = encoder.into_bytes();
let decoder = DataDecoder::new(&bytes, 0);
let decoded = decoder.decode(map_offset).unwrap();
if let DataValue::Map(decoded_map) = decoded {
assert_eq!(
decoded_map.get("direct"),
Some(&DataValue::String("direct_value".to_string()))
);
assert_eq!(
decoded_map.get("ptr_ref"),
Some(&DataValue::String("shared_value".to_string()))
);
} else {
panic!("Expected Map, got {decoded:?}");
}
}
#[test]
fn test_large_negative_integer_rejected() {
let json = format!("{}", i64::MIN);
let result: Result<DataValue, _> = serde_json::from_str(&json);
assert!(result.is_err());
let err = result.unwrap_err().to_string();
assert!(err.contains("outside the supported signed integer range"));
}
#[test]
fn test_i32_min_accepted() {
let json = format!("{}", i32::MIN);
let result: Result<DataValue, _> = serde_json::from_str(&json);
assert!(result.is_ok());
assert_eq!(result.unwrap(), DataValue::Int32(i32::MIN));
}
#[test]
fn test_timestamp_binary_roundtrip() {
let mut encoder = DataEncoder::new();
let epoch = 1727894671i64; let value = DataValue::Timestamp(epoch);
let offset = encoder.encode(&value);
let bytes = encoder.into_bytes();
let decoder = DataDecoder::new(&bytes, 0);
let decoded = decoder.decode(offset).unwrap();
assert_eq!(decoded, DataValue::Timestamp(epoch));
}
#[test]
fn test_timestamp_json_serialize() {
let value = DataValue::Timestamp(1727894671);
let json = serde_json::to_string(&value).unwrap();
assert_eq!(json, "\"2024-10-02T18:44:31Z\"");
}
#[test]
fn test_timestamp_json_deserialize() {
let json = "\"2024-10-02T18:44:31Z\"";
let value: DataValue = serde_json::from_str(json).unwrap();
assert_eq!(value, DataValue::Timestamp(1727894671));
}
#[test]
fn test_timestamp_with_fractional_seconds() {
let json = "\"2024-10-02T18:44:31.123456Z\"";
let value: DataValue = serde_json::from_str(json).unwrap();
if let DataValue::Timestamp(epoch) = value {
assert_eq!(epoch, 1727894671);
} else {
panic!("Expected Timestamp, got {value:?}");
}
}
#[test]
fn test_non_timestamp_string_stays_string() {
let json = "\"hello world\"";
let value: DataValue = serde_json::from_str(json).unwrap();
assert_eq!(value, DataValue::String("hello world".to_string()));
}
#[test]
fn test_timestamp_negative_epoch() {
let mut encoder = DataEncoder::new();
let epoch = -86400i64; let value = DataValue::Timestamp(epoch);
let offset = encoder.encode(&value);
let bytes = encoder.into_bytes();
let decoder = DataDecoder::new(&bytes, 0);
let decoded = decoder.decode(offset).unwrap();
assert_eq!(decoded, DataValue::Timestamp(epoch));
}
#[test]
fn test_timestamp_in_map() {
let mut encoder = DataEncoder::new();
let mut map = HashMap::new();
map.insert("first_seen".to_string(), DataValue::Timestamp(1727894671));
map.insert("last_seen".to_string(), DataValue::Timestamp(1727981071));
map.insert("name".to_string(), DataValue::String("test".to_string()));
let offset = encoder.encode(&DataValue::Map(map.clone()));
let bytes = encoder.into_bytes();
let decoder = DataDecoder::new(&bytes, 0);
let decoded = decoder.decode(offset).unwrap();
if let DataValue::Map(decoded_map) = decoded {
assert_eq!(
decoded_map.get("first_seen"),
Some(&DataValue::Timestamp(1727894671))
);
assert_eq!(
decoded_map.get("last_seen"),
Some(&DataValue::Timestamp(1727981071))
);
assert_eq!(
decoded_map.get("name"),
Some(&DataValue::String("test".to_string()))
);
} else {
panic!("Expected Map, got {decoded:?}");
}
}
}