use alloc::{vec, vec::Vec};
use bytes::BufMut;
use core::time::Duration;
use crate::{
format::Format,
helpers::{take_buffer, take_buffer_iter, take_byte, take_byte_iter, take_num, take_num_iter},
Error, Packable, Unpackable,
};
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Extension {
Ext(i8, Vec<u8>),
Timestamp(Duration),
}
impl Extension {
pub const TIMESTAMP: i8 = -1;
}
impl Packable for Extension {
fn pack<T>(&self, buf: &mut T) -> usize
where
T: BufMut,
{
match self {
Extension::Ext(t, b) if b.len() == 1 => {
buf.put_slice(&[Format::FIXEXT1, *t as u8, b[0]]);
3
}
Extension::Ext(t, b) if b.len() == 2 => {
let mut s = [0u8; 4];
s[0] = Format::FIXEXT2;
s[1] = *t as u8;
s[2..].copy_from_slice(b);
buf.put_slice(&s);
4
}
Extension::Ext(t, b) if b.len() == 4 => {
let mut s = [0u8; 6];
s[0] = Format::FIXEXT4;
s[1] = *t as u8;
s[2..].copy_from_slice(b);
buf.put_slice(&s);
6
}
Extension::Ext(t, b) if b.len() == 8 => {
let mut s = [0u8; 10];
s[0] = Format::FIXEXT8;
s[1] = *t as u8;
s[2..].copy_from_slice(b);
buf.put_slice(&s);
10
}
Extension::Ext(t, b) if b.len() == 16 => {
let mut s = [0u8; 18];
s[0] = Format::FIXEXT16;
s[1] = *t as u8;
s[2..].copy_from_slice(b);
buf.put_slice(&s);
18
}
Extension::Ext(t, b) if b.len() <= u8::MAX as usize => {
buf.put_slice(&[Format::EXT8, b.len() as u8, *t as u8]);
buf.put_slice(b);
3 + b.len()
}
Extension::Ext(t, b) if b.len() <= u16::MAX as usize => {
let mut s = [0u8; 4];
s[0] = Format::EXT16;
s[1..3].copy_from_slice(&(b.len() as u16).to_be_bytes());
s[3] = *t as u8;
buf.put_slice(&s);
buf.put_slice(b);
4 + b.len()
}
Extension::Ext(t, b) if b.len() <= u32::MAX as usize => {
let mut s = [0u8; 6];
s[0] = Format::EXT32;
s[1..5].copy_from_slice(&(b.len() as u32).to_be_bytes());
s[5] = *t as u8;
buf.put_slice(&s);
buf.put_slice(b);
6 + b.len()
}
#[allow(unreachable_code)]
Extension::Ext(_, _) => {
#[cfg(feature = "strict")]
panic!("strict serialization enabled; the buffer is too large");
0
}
Extension::Timestamp(d) if d.as_secs() <= u32::MAX as u64 && d.subsec_nanos() == 0 => {
let mut s = [0u8; 6];
s[0] = Format::FIXEXT4;
s[1] = Self::TIMESTAMP as u8;
s[2..].copy_from_slice(&(d.as_secs() as u32).to_be_bytes());
buf.put_slice(&s);
6
}
Extension::Timestamp(d)
if d.as_secs() < 1u64 << 34 && d.subsec_nanos() < 1u32 << 30 =>
{
let secs = d.as_secs();
let secs_nanos = ((secs >> 32) & 0b11) as u32;
let secs = secs as u32;
let nanos = d.subsec_nanos() << 2;
let nanos = nanos | secs_nanos;
let mut s = [0u8; 10];
s[0] = Format::FIXEXT8;
s[1] = Self::TIMESTAMP as u8;
s[2..6].copy_from_slice(&nanos.to_be_bytes());
s[6..].copy_from_slice(&secs.to_be_bytes());
buf.put_slice(&s);
10
}
Extension::Timestamp(d) => {
let mut s = [0u8; 15];
s[0] = Format::EXT8;
s[1] = 12;
s[2] = Self::TIMESTAMP as u8;
s[3..7].copy_from_slice(&d.subsec_nanos().to_be_bytes());
s[7..].copy_from_slice(&d.as_secs().to_be_bytes());
buf.put_slice(&s);
15
}
}
}
}
impl Unpackable for Extension {
type Error = Error;
fn unpack_with_ofs(mut buf: &[u8]) -> Result<(usize, Self), Self::Error> {
let format = take_byte(&mut buf)?;
match format {
Format::FIXEXT1 => {
let t = take_byte(&mut buf)? as i8;
let x = take_byte(&mut buf)?;
Ok((3, Extension::Ext(t, vec![x])))
}
Format::FIXEXT2 => {
let t = take_byte(&mut buf)? as i8;
let b = take_buffer(&mut buf, 2)?;
Ok((4, Extension::Ext(t, b.to_vec())))
}
Format::FIXEXT4 => {
let t = take_byte(&mut buf)? as i8;
if t == Self::TIMESTAMP {
let secs = take_num(&mut buf, u32::from_be_bytes)?;
Ok((6, Extension::Timestamp(Duration::from_secs(secs as u64))))
} else {
let b = take_buffer(&mut buf, 4)?;
Ok((6, Extension::Ext(t, b.to_vec())))
}
}
Format::FIXEXT8 => {
let t = take_byte(&mut buf)? as i8;
if t == Self::TIMESTAMP {
let data = take_num(&mut buf, u64::from_be_bytes)?;
let nanos = (data >> 34) as u32;
let secs = data & ((1u64 << 34) - 1);
Ok((10, Extension::Timestamp(Duration::new(secs, nanos))))
} else {
let b = take_buffer(&mut buf, 8)?;
Ok((10, Extension::Ext(t, b.to_vec())))
}
}
Format::FIXEXT16 => {
let t = take_byte(&mut buf)? as i8;
let b = take_buffer(&mut buf, 16)?;
Ok((18, Extension::Ext(t, b.to_vec())))
}
Format::EXT8 => {
let len = take_byte(&mut buf)? as usize;
let t = take_byte(&mut buf)? as i8;
if len == 12 && t == Self::TIMESTAMP {
let nanos = take_num(&mut buf, u32::from_be_bytes)?;
let secs = take_num(&mut buf, u64::from_be_bytes)?;
Ok((15, Extension::Timestamp(Duration::new(secs, nanos))))
} else {
let b = take_buffer(&mut buf, len)?;
Ok((3 + len, Extension::Ext(t, b.to_vec())))
}
}
Format::EXT16 => {
let len = take_num(&mut buf, u16::from_be_bytes)? as usize;
let t = take_byte(&mut buf)? as i8;
let b = take_buffer(&mut buf, len)?;
Ok((4 + len, Extension::Ext(t, b.to_vec())))
}
Format::EXT32 => {
let len = take_num(&mut buf, u32::from_be_bytes)? as usize;
let t = take_byte(&mut buf)? as i8;
let b = take_buffer(&mut buf, len)?;
Ok((6 + len, Extension::Ext(t, b.to_vec())))
}
_ => Err(Error::InvalidExtension),
}
}
fn unpack_iter<I>(bytes: I) -> Result<(usize, Self), Self::Error>
where
I: IntoIterator<Item = u8>,
{
let mut bytes = bytes.into_iter();
let format = take_byte_iter(bytes.by_ref())?;
match format {
Format::FIXEXT1 => {
let t = take_byte_iter(bytes.by_ref())? as i8;
let x = take_byte_iter(bytes.by_ref())?;
Ok((3, Extension::Ext(t, vec![x])))
}
Format::FIXEXT2 => {
let t = take_byte_iter(bytes.by_ref())? as i8;
let b = take_buffer_iter(bytes.by_ref(), 2)?;
Ok((4, Extension::Ext(t, b)))
}
Format::FIXEXT4 => {
let t = take_byte_iter(bytes.by_ref())? as i8;
if t == Self::TIMESTAMP {
let secs = take_num_iter(bytes.by_ref(), u32::from_be_bytes)?;
Ok((6, Extension::Timestamp(Duration::from_secs(secs as u64))))
} else {
let b = take_buffer_iter(bytes.by_ref(), 4)?;
Ok((6, Extension::Ext(t, b)))
}
}
Format::FIXEXT8 => {
let t = take_byte_iter(bytes.by_ref())? as i8;
if t == Self::TIMESTAMP {
let data = take_num_iter(bytes.by_ref(), u64::from_be_bytes)?;
let nanos = (data >> 34) as u32;
let secs = data & ((1u64 << 34) - 1);
Ok((10, Extension::Timestamp(Duration::new(secs, nanos))))
} else {
let b = take_buffer_iter(bytes.by_ref(), 8)?;
Ok((10, Extension::Ext(t, b)))
}
}
Format::FIXEXT16 => {
let t = take_byte_iter(bytes.by_ref())? as i8;
let b = take_buffer_iter(bytes.by_ref(), 16)?;
Ok((18, Extension::Ext(t, b)))
}
Format::EXT8 => {
let len = take_byte_iter(bytes.by_ref())? as usize;
let t = take_byte_iter(bytes.by_ref())? as i8;
if len == 12 && t == Self::TIMESTAMP {
let nanos = take_num_iter(bytes.by_ref(), u32::from_be_bytes)?;
let secs = take_num_iter(bytes.by_ref(), u64::from_be_bytes)?;
Ok((15, Extension::Timestamp(Duration::new(secs, nanos))))
} else {
let b = take_buffer_iter(bytes.by_ref(), len)?;
Ok((3 + len, Extension::Ext(t, b)))
}
}
Format::EXT16 => {
let len = take_num_iter(bytes.by_ref(), u16::from_be_bytes)? as usize;
let t = take_byte_iter(bytes.by_ref())? as i8;
let b = take_buffer_iter(bytes.by_ref(), len)?;
Ok((4 + len, Extension::Ext(t, b)))
}
Format::EXT32 => {
let len = take_num_iter(bytes.by_ref(), u32::from_be_bytes)? as usize;
let t = take_byte_iter(bytes.by_ref())? as i8;
let b = take_buffer_iter(bytes.by_ref(), len)?;
Ok((6 + len, Extension::Ext(t, b)))
}
_ => Err(Error::InvalidExtension),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use proptest::prelude::*;
proptest! {
#[test]
fn extension_bytes(mut t: i8, b: Vec<u8>) {
if t == Extension::TIMESTAMP {
t -= 1;
}
let x = Extension::Ext(t, b);
let mut bytes = vec![];
x.pack(&mut bytes);
let (_, y) = Extension::unpack_with_ofs(&bytes).unwrap();
assert_eq!(x, y);
}
#[test]
fn extension_duration(d: Duration) {
let x = Extension::Timestamp(d);
let mut bytes = vec![];
x.pack(&mut bytes);
let (_, y) = Extension::unpack_with_ofs(&bytes).unwrap();
assert_eq!(x, y);
}
}
}