use super::types::ObjectStatus;
use crate::error::CodecError;
use crate::types::read_bytes;
use crate::varint::VarInt;
use bytes::{Buf, BufMut};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u64)]
pub enum StreamType {
Fetch = 0x05,
SubgroupZero = 0x10,
SubgroupZeroExt = 0x11,
SubgroupFirstObj = 0x12,
SubgroupFirstObjExt = 0x13,
SubgroupExplicit = 0x14,
SubgroupExplicitExt = 0x15,
SubgroupZeroEog = 0x18,
SubgroupZeroEogExt = 0x19,
SubgroupFirstObjEog = 0x1A,
SubgroupFirstObjEogExt = 0x1B,
SubgroupExplicitEog = 0x1C,
SubgroupExplicitEogExt = 0x1D,
}
fn check_extensions_against_status(
status: ObjectStatus,
extensions: &[u8],
) -> Result<(), CodecError> {
if status == ObjectStatus::ObjectDoesNotExist && !extensions.is_empty() {
return Err(CodecError::ExtensionsOnNonExistentObject(extensions.len()));
}
Ok(())
}
impl StreamType {
pub fn from_id(id: u64) -> Option<Self> {
match id {
0x05 => Some(StreamType::Fetch),
0x10 => Some(StreamType::SubgroupZero),
0x11 => Some(StreamType::SubgroupZeroExt),
0x12 => Some(StreamType::SubgroupFirstObj),
0x13 => Some(StreamType::SubgroupFirstObjExt),
0x14 => Some(StreamType::SubgroupExplicit),
0x15 => Some(StreamType::SubgroupExplicitExt),
0x18 => Some(StreamType::SubgroupZeroEog),
0x19 => Some(StreamType::SubgroupZeroEogExt),
0x1A => Some(StreamType::SubgroupFirstObjEog),
0x1B => Some(StreamType::SubgroupFirstObjEogExt),
0x1C => Some(StreamType::SubgroupExplicitEog),
0x1D => Some(StreamType::SubgroupExplicitEogExt),
_ => None,
}
}
pub fn is_subgroup(&self) -> bool {
matches!(
self,
StreamType::SubgroupZero
| StreamType::SubgroupZeroExt
| StreamType::SubgroupFirstObj
| StreamType::SubgroupFirstObjExt
| StreamType::SubgroupExplicit
| StreamType::SubgroupExplicitExt
| StreamType::SubgroupZeroEog
| StreamType::SubgroupZeroEogExt
| StreamType::SubgroupFirstObjEog
| StreamType::SubgroupFirstObjEogExt
| StreamType::SubgroupExplicitEog
| StreamType::SubgroupExplicitEogExt
)
}
pub fn has_extensions(&self) -> bool {
matches!(
self,
StreamType::SubgroupZeroExt
| StreamType::SubgroupFirstObjExt
| StreamType::SubgroupExplicitExt
| StreamType::SubgroupZeroEogExt
| StreamType::SubgroupFirstObjEogExt
| StreamType::SubgroupExplicitEogExt
)
}
pub fn contains_end_of_group(&self) -> bool {
matches!(
self,
StreamType::SubgroupZeroEog
| StreamType::SubgroupZeroEogExt
| StreamType::SubgroupFirstObjEog
| StreamType::SubgroupFirstObjEogExt
| StreamType::SubgroupExplicitEog
| StreamType::SubgroupExplicitEogExt
)
}
pub fn writes_subgroup_id(&self) -> bool {
matches!(
self,
StreamType::SubgroupExplicit
| StreamType::SubgroupExplicitExt
| StreamType::SubgroupExplicitEog
| StreamType::SubgroupExplicitEogExt
)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u64)]
pub enum DatagramType {
Datagram = 0x00,
DatagramExt = 0x01,
DatagramEog = 0x02,
DatagramEogExt = 0x03,
DatagramStatus = 0x04,
DatagramStatusExt = 0x05,
}
impl DatagramType {
pub fn from_id(id: u64) -> Option<Self> {
match id {
0x00 => Some(DatagramType::Datagram),
0x01 => Some(DatagramType::DatagramExt),
0x02 => Some(DatagramType::DatagramEog),
0x03 => Some(DatagramType::DatagramEogExt),
0x04 => Some(DatagramType::DatagramStatus),
0x05 => Some(DatagramType::DatagramStatusExt),
_ => None,
}
}
pub fn has_extensions(&self) -> bool {
matches!(
self,
DatagramType::DatagramExt
| DatagramType::DatagramEogExt
| DatagramType::DatagramStatusExt
)
}
pub fn is_status(&self) -> bool {
matches!(self, DatagramType::DatagramStatus | DatagramType::DatagramStatusExt)
}
pub fn is_end_of_group(&self) -> bool {
matches!(self, DatagramType::DatagramEog | DatagramType::DatagramEogExt)
}
}
fn stream_type_error(raw: u64) -> CodecError {
if StreamType::from_id(raw).is_some() {
CodecError::InvalidField
} else {
CodecError::UnknownStreamType(raw)
}
}
fn datagram_type_error(raw: u64) -> CodecError {
if DatagramType::from_id(raw).is_some() {
CodecError::InvalidField
} else {
CodecError::UnknownDatagramType(raw)
}
}
fn read_extension_bytes(buf: &mut impl Buf, byte_len: u64) -> Result<Vec<u8>, CodecError> {
read_bytes(buf, byte_len as usize)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SubgroupHeader {
pub stream_type: StreamType,
pub track_alias: VarInt,
pub group_id: VarInt,
pub subgroup_id: VarInt,
pub publisher_priority: u8,
}
impl SubgroupHeader {
pub fn encode_stream(&self, buf: &mut impl BufMut) {
VarInt::from_usize(self.stream_type as usize).encode(buf);
self.encode(buf);
}
pub fn encode(&self, buf: &mut impl BufMut) {
self.track_alias.encode(buf);
self.group_id.encode(buf);
if self.stream_type.writes_subgroup_id() {
self.subgroup_id.encode(buf);
}
buf.put_u8(self.publisher_priority);
}
pub fn encode_checked(&self, buf: &mut impl BufMut) -> Result<(), CodecError> {
if !self.stream_type.writes_subgroup_id() && self.subgroup_id.into_inner() != 0 {
return Err(CodecError::InvalidField);
}
self.encode(buf);
Ok(())
}
pub fn decode(buf: &mut impl Buf) -> Result<Self, CodecError> {
Self::decode_with_type(StreamType::SubgroupExplicit, buf)
}
pub fn decode_with_type(
stream_type: StreamType,
buf: &mut impl Buf,
) -> Result<Self, CodecError> {
let track_alias = VarInt::decode(buf)?;
let group_id = VarInt::decode(buf)?;
let subgroup_id = match stream_type {
StreamType::SubgroupZero
| StreamType::SubgroupZeroExt
| StreamType::SubgroupZeroEog
| StreamType::SubgroupZeroEogExt => VarInt::from_usize(0),
StreamType::SubgroupExplicit
| StreamType::SubgroupExplicitExt
| StreamType::SubgroupExplicitEog
| StreamType::SubgroupExplicitEogExt => VarInt::decode(buf)?,
StreamType::SubgroupFirstObj
| StreamType::SubgroupFirstObjExt
| StreamType::SubgroupFirstObjEog
| StreamType::SubgroupFirstObjEogExt => VarInt::from_usize(0),
_ => return Err(CodecError::InvalidField),
};
if buf.remaining() < 1 {
return Err(CodecError::UnexpectedEnd);
}
let publisher_priority = buf.get_u8();
Ok(Self { stream_type, track_alias, group_id, subgroup_id, publisher_priority })
}
pub fn decode_stream(buf: &mut impl Buf) -> Result<Self, CodecError> {
let raw = VarInt::decode(buf)?.into_inner();
let stream_type = StreamType::from_id(raw).ok_or_else(|| stream_type_error(raw))?;
if !stream_type.is_subgroup() {
return Err(stream_type_error(raw));
}
Self::decode_with_type(stream_type, buf)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ObjectHeader {
pub object_id: VarInt,
pub extension_headers_length: VarInt,
pub extensions: Vec<u8>,
pub payload_length: VarInt,
pub object_status: ObjectStatus,
}
impl ObjectHeader {
pub fn encode(&self, buf: &mut impl BufMut) {
self.encode_with_extensions(false, buf);
}
pub fn encode_checked(&self, buf: &mut impl BufMut) -> Result<(), CodecError> {
self.encode_checked_with_extensions(false, buf)
}
pub fn encode_checked_with_extensions(
&self,
has_extensions: bool,
buf: &mut impl BufMut,
) -> Result<(), CodecError> {
if self.payload_length.into_inner() != 0 && self.object_status as usize != 0 {
return Err(CodecError::InvalidField);
}
if !has_extensions && !self.extensions.is_empty() {
return Err(CodecError::InvalidField);
}
self.encode_with_extensions(has_extensions, buf);
Ok(())
}
pub fn encode_with_extensions(&self, has_extensions: bool, buf: &mut impl BufMut) {
self.object_id.encode(buf);
if has_extensions {
VarInt::from_usize(self.extensions.len()).encode(buf);
buf.put_slice(&self.extensions);
}
self.payload_length.encode(buf);
if self.payload_length.into_inner() == 0 {
VarInt::from_usize(self.object_status as usize).encode(buf);
}
}
pub fn decode(buf: &mut impl Buf) -> Result<Self, CodecError> {
Self::decode_with_extensions(false, buf)
}
pub fn decode_with_extensions(
has_extensions: bool,
buf: &mut impl Buf,
) -> Result<Self, CodecError> {
let object_id = VarInt::decode(buf)?;
let (extension_headers_length, extensions) = if has_extensions {
let ehl = VarInt::decode(buf)?;
let ext = read_extension_bytes(buf, ehl.into_inner())?;
(ehl, ext)
} else {
(VarInt::from_usize(0), Vec::new())
};
let payload_length = VarInt::decode(buf)?;
let object_status = if payload_length.into_inner() == 0 {
let sv = VarInt::decode(buf)?.into_inner();
ObjectStatus::from_u64(sv).ok_or(CodecError::InvalidField)?
} else {
ObjectStatus::Normal
};
check_extensions_against_status(object_status, &extensions)?;
Ok(Self { object_id, extension_headers_length, extensions, payload_length, object_status })
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DatagramHeader {
pub track_alias: VarInt,
pub group_id: VarInt,
pub object_id: VarInt,
pub publisher_priority: u8,
pub extension_headers_length: VarInt,
pub extensions: Vec<u8>,
pub end_of_group: bool,
}
impl DatagramHeader {
pub fn encode(&self, buf: &mut impl BufMut) {
self.encode_with_extensions(false, buf);
}
pub fn encode_with_extensions(&self, has_extensions: bool, buf: &mut impl BufMut) {
self.track_alias.encode(buf);
self.group_id.encode(buf);
self.object_id.encode(buf);
buf.put_u8(self.publisher_priority);
if has_extensions {
VarInt::from_usize(self.extensions.len()).encode(buf);
buf.put_slice(&self.extensions);
}
}
pub fn encode_checked(&self, buf: &mut impl BufMut) -> Result<(), CodecError> {
if !self.extensions.is_empty() {
return Err(CodecError::InvalidField);
}
self.encode(buf);
Ok(())
}
pub fn decode(buf: &mut impl Buf) -> Result<Self, CodecError> {
Self::decode_with_extensions(false, buf)
}
pub fn decode_with_extensions(
has_extensions: bool,
buf: &mut impl Buf,
) -> Result<Self, CodecError> {
let track_alias = VarInt::decode(buf)?;
let group_id = VarInt::decode(buf)?;
let object_id = VarInt::decode(buf)?;
if buf.remaining() < 1 {
return Err(CodecError::UnexpectedEnd);
}
let publisher_priority = buf.get_u8();
let (extension_headers_length, extensions) = if has_extensions {
let ehl = VarInt::decode(buf)?;
if ehl.into_inner() == 0 {
return Err(CodecError::InvalidField);
}
let ext = read_extension_bytes(buf, ehl.into_inner())?;
(ehl, ext)
} else {
(VarInt::from_usize(0), Vec::new())
};
Ok(Self {
track_alias,
group_id,
object_id,
publisher_priority,
extension_headers_length,
extensions,
end_of_group: false,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DatagramStatusHeader {
pub track_alias: VarInt,
pub group_id: VarInt,
pub object_id: VarInt,
pub publisher_priority: u8,
pub extension_headers_length: VarInt,
pub extensions: Vec<u8>,
pub object_status: ObjectStatus,
}
impl DatagramStatusHeader {
pub fn encode(&self, buf: &mut impl BufMut) {
self.encode_with_extensions(false, buf);
}
pub fn encode_checked(&self, buf: &mut impl BufMut) -> Result<(), CodecError> {
if !self.extensions.is_empty() {
return Err(CodecError::InvalidField);
}
self.encode(buf);
Ok(())
}
pub fn encode_with_extensions(&self, has_extensions: bool, buf: &mut impl BufMut) {
self.track_alias.encode(buf);
self.group_id.encode(buf);
self.object_id.encode(buf);
buf.put_u8(self.publisher_priority);
if has_extensions {
VarInt::from_usize(self.extensions.len()).encode(buf);
buf.put_slice(&self.extensions);
}
VarInt::from_usize(self.object_status as usize).encode(buf);
}
pub fn decode(buf: &mut impl Buf) -> Result<Self, CodecError> {
Self::decode_with_extensions(false, buf)
}
pub fn decode_with_extensions(
has_extensions: bool,
buf: &mut impl Buf,
) -> Result<Self, CodecError> {
let track_alias = VarInt::decode(buf)?;
let group_id = VarInt::decode(buf)?;
let object_id = VarInt::decode(buf)?;
if buf.remaining() < 1 {
return Err(CodecError::UnexpectedEnd);
}
let publisher_priority = buf.get_u8();
let (extension_headers_length, extensions) = if has_extensions {
let ehl = VarInt::decode(buf)?;
if ehl.into_inner() == 0 {
return Err(CodecError::InvalidField);
}
let ext = read_extension_bytes(buf, ehl.into_inner())?;
(ehl, ext)
} else {
(VarInt::from_usize(0), Vec::new())
};
let sv = VarInt::decode(buf)?.into_inner();
let object_status = ObjectStatus::from_u64(sv).ok_or(CodecError::InvalidField)?;
check_extensions_against_status(object_status, &extensions)?;
Ok(Self {
track_alias,
group_id,
object_id,
publisher_priority,
extension_headers_length,
extensions,
object_status,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Datagram {
Payload(DatagramHeader),
Status(DatagramStatusHeader),
}
impl Datagram {
pub fn is_status(&self) -> bool {
matches!(self, Self::Status(_))
}
pub fn datagram_type(&self) -> DatagramType {
match self {
Self::Payload(header) => match (header.end_of_group, header.extensions.is_empty()) {
(false, true) => DatagramType::Datagram,
(false, false) => DatagramType::DatagramExt,
(true, true) => DatagramType::DatagramEog,
(true, false) => DatagramType::DatagramEogExt,
},
Self::Status(header) => {
if header.extensions.is_empty() {
DatagramType::DatagramStatus
} else {
DatagramType::DatagramStatusExt
}
}
}
}
pub fn decode(buf: &mut impl Buf) -> Result<Self, CodecError> {
let raw = VarInt::decode(buf)?.into_inner();
let datagram_type = DatagramType::from_id(raw).ok_or_else(|| datagram_type_error(raw))?;
let has_extensions = datagram_type.has_extensions();
if datagram_type.is_status() {
return Ok(Self::Status(DatagramStatusHeader::decode_with_extensions(
has_extensions,
buf,
)?));
}
let mut header = DatagramHeader::decode_with_extensions(has_extensions, buf)?;
header.end_of_group = datagram_type.is_end_of_group();
Ok(Self::Payload(header))
}
pub fn encode(&self, buf: &mut impl BufMut) {
let datagram_type = self.datagram_type();
let has_extensions = datagram_type.has_extensions();
VarInt::from_usize(datagram_type as usize).encode(buf);
match self {
Self::Payload(header) => header.encode_with_extensions(has_extensions, buf),
Self::Status(header) => header.encode_with_extensions(has_extensions, buf),
}
}
pub fn encode_checked(&self, buf: &mut impl BufMut) -> Result<(), CodecError> {
let mut body = Vec::with_capacity(64);
let datagram_type = self.datagram_type();
let has_extensions = datagram_type.has_extensions();
match self {
Self::Payload(header) => header.encode_with_extensions(has_extensions, &mut body),
Self::Status(header) => {
check_extensions_against_status(header.object_status, &header.extensions)?;
header.encode_with_extensions(has_extensions, &mut body);
}
}
VarInt::from_usize(datagram_type as usize).encode(buf);
buf.put_slice(&body);
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FetchHeader {
pub request_id: VarInt,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FetchObjectHeader {
pub group_id: VarInt,
pub subgroup_id: VarInt,
pub object_id: VarInt,
pub publisher_priority: u8,
pub extension_headers_length: VarInt,
pub extensions: Vec<u8>,
pub payload_length: VarInt,
pub object_status: ObjectStatus,
}
impl FetchHeader {
pub fn encode_stream(&self, buf: &mut impl BufMut) {
VarInt::from_usize(StreamType::Fetch as usize).encode(buf);
self.encode(buf);
}
pub fn encode(&self, buf: &mut impl BufMut) {
self.request_id.encode(buf);
}
pub fn decode(buf: &mut impl Buf) -> Result<Self, CodecError> {
let request_id = VarInt::decode(buf)?;
Ok(Self { request_id })
}
pub fn decode_stream(buf: &mut impl Buf) -> Result<Self, CodecError> {
let stream_type = VarInt::decode(buf)?.into_inner();
if stream_type != StreamType::Fetch as u64 {
return Err(stream_type_error(stream_type));
}
Self::decode(buf)
}
}
impl FetchObjectHeader {
pub fn encode(&self, buf: &mut impl BufMut) {
self.group_id.encode(buf);
self.subgroup_id.encode(buf);
self.object_id.encode(buf);
buf.put_u8(self.publisher_priority);
VarInt::from_usize(self.extensions.len()).encode(buf);
buf.put_slice(&self.extensions);
self.payload_length.encode(buf);
if self.payload_length.into_inner() == 0 {
VarInt::from_usize(self.object_status as usize).encode(buf);
}
}
pub fn encode_checked(&self, buf: &mut impl BufMut) -> Result<(), CodecError> {
if self.payload_length.into_inner() != 0 && self.object_status as usize != 0 {
return Err(CodecError::InvalidField);
}
self.encode(buf);
Ok(())
}
pub fn decode(buf: &mut impl Buf) -> Result<Self, CodecError> {
let group_id = VarInt::decode(buf)?;
let subgroup_id = VarInt::decode(buf)?;
let object_id = VarInt::decode(buf)?;
if buf.remaining() < 1 {
return Err(CodecError::UnexpectedEnd);
}
let publisher_priority = buf.get_u8();
let extension_headers_length = VarInt::decode(buf)?;
let extensions = read_extension_bytes(buf, extension_headers_length.into_inner())?;
let payload_length = VarInt::decode(buf)?;
let object_status = if payload_length.into_inner() == 0 {
let sv = VarInt::decode(buf)?.into_inner();
ObjectStatus::from_u64(sv).ok_or(CodecError::InvalidField)?
} else {
ObjectStatus::Normal
};
check_extensions_against_status(object_status, &extensions)?;
Ok(Self {
group_id,
subgroup_id,
object_id,
publisher_priority,
extension_headers_length,
extensions,
payload_length,
object_status,
})
}
}