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 {
Datagram = 0x01,
DatagramStatus = 0x02,
Subgroup = 0x04,
Fetch = 0x05,
}
impl StreamType {
pub fn from_id(id: u64) -> Option<Self> {
match id {
0x01 => Some(StreamType::Datagram),
0x02 => Some(StreamType::DatagramStatus),
0x04 => Some(StreamType::Subgroup),
0x05 => Some(StreamType::Fetch),
_ => None,
}
}
}
fn stream_type_error(raw: u64) -> CodecError {
if raw == StreamType::Subgroup as u64 || raw == StreamType::Fetch as u64 {
CodecError::InvalidField
} else {
CodecError::UnknownStreamType(raw)
}
}
fn datagram_type_error(raw: u64) -> CodecError {
if raw == StreamType::Datagram as u64 || raw == StreamType::DatagramStatus as u64 {
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)
}
fn encode_extensions(extensions: &[u8], buf: &mut impl BufMut) {
buf.put_slice(extensions);
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SubgroupHeader {
pub track_alias: VarInt,
pub group_id: VarInt,
pub subgroup_id: VarInt,
pub publisher_priority: u8,
}
#[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 SubgroupHeader {
pub fn encode_stream(&self, buf: &mut impl BufMut) {
VarInt::from_usize(StreamType::Subgroup 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);
self.subgroup_id.encode(buf);
buf.put_u8(self.publisher_priority);
}
pub fn decode(buf: &mut impl Buf) -> Result<Self, CodecError> {
let track_alias = VarInt::decode(buf)?;
let group_id = VarInt::decode(buf)?;
let subgroup_id = VarInt::decode(buf)?;
if buf.remaining() < 1 {
return Err(CodecError::UnexpectedEnd);
}
let publisher_priority = buf.get_u8();
Ok(Self { track_alias, group_id, subgroup_id, publisher_priority })
}
pub fn decode_stream(buf: &mut impl Buf) -> Result<Self, CodecError> {
let stream_type = VarInt::decode(buf)?.into_inner();
if stream_type != StreamType::Subgroup as u64 {
return Err(stream_type_error(stream_type));
}
Self::decode(buf)
}
}
fn check_end_of_track(object_id: VarInt, status: ObjectStatus) -> Result<(), CodecError> {
if status == ObjectStatus::EndOfTrack && object_id.into_inner() != 0 {
return Err(CodecError::EndOfTrackObjectId(object_id.into_inner()));
}
Ok(())
}
impl ObjectHeader {
pub fn encode(&self, buf: &mut impl BufMut) {
self.object_id.encode(buf);
VarInt::from_usize(self.extensions.len()).encode(buf);
encode_extensions(&self.extensions, buf);
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> {
check_end_of_track(self.object_id, self.object_status)?;
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 object_id = VarInt::decode(buf)?;
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 status_val = VarInt::decode(buf)?.into_inner();
ObjectStatus::from_u64(status_val).ok_or(CodecError::InvalidField)?
} else {
ObjectStatus::Normal
};
check_end_of_track(object_id, object_status)?;
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>,
}
impl DatagramHeader {
pub fn encode(&self, 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);
VarInt::from_usize(self.extensions.len()).encode(buf);
encode_extensions(&self.extensions, buf);
}
pub fn encode_checked(&self, buf: &mut impl BufMut) -> Result<(), CodecError> {
self.encode(buf);
Ok(())
}
pub fn decode(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 = VarInt::decode(buf)?;
let extensions = read_extension_bytes(buf, extension_headers_length.into_inner())?;
Ok(Self {
track_alias,
group_id,
object_id,
publisher_priority,
extension_headers_length,
extensions,
})
}
}
#[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.track_alias.encode(buf);
self.group_id.encode(buf);
self.object_id.encode(buf);
buf.put_u8(self.publisher_priority);
VarInt::from_usize(self.extensions.len()).encode(buf);
encode_extensions(&self.extensions, buf);
VarInt::from_usize(self.object_status as usize).encode(buf);
}
pub fn encode_checked(&self, buf: &mut impl BufMut) -> Result<(), CodecError> {
check_end_of_track(self.object_id, self.object_status)?;
self.encode(buf);
Ok(())
}
pub fn decode(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 = VarInt::decode(buf)?;
let extensions = read_extension_bytes(buf, extension_headers_length.into_inner())?;
let status_val = VarInt::decode(buf)?.into_inner();
let object_status = ObjectStatus::from_u64(status_val).ok_or(CodecError::InvalidField)?;
check_end_of_track(object_id, object_status)?;
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) -> StreamType {
match self {
Self::Payload(_) => StreamType::Datagram,
Self::Status(_) => StreamType::DatagramStatus,
}
}
pub fn decode(buf: &mut impl Buf) -> Result<Self, CodecError> {
let raw = VarInt::decode(buf)?.into_inner();
match StreamType::from_id(raw) {
Some(StreamType::Datagram) => Ok(Self::Payload(DatagramHeader::decode(buf)?)),
Some(StreamType::DatagramStatus) => {
Ok(Self::Status(DatagramStatusHeader::decode(buf)?))
}
_ => Err(datagram_type_error(raw)),
}
}
pub fn encode(&self, buf: &mut impl BufMut) {
VarInt::from_usize(self.datagram_type() as usize).encode(buf);
match self {
Self::Payload(header) => header.encode(buf),
Self::Status(header) => header.encode(buf),
}
}
pub fn encode_checked(&self, buf: &mut impl BufMut) -> Result<(), CodecError> {
let mut body = Vec::with_capacity(64);
match self {
Self::Payload(header) => header.encode_checked(&mut body)?,
Self::Status(header) => header.encode_checked(&mut body)?,
}
VarInt::from_usize(self.datagram_type() as usize).encode(buf);
buf.put_slice(&body);
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FetchHeader {
pub subscribe_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 object_status: ObjectStatus,
pub payload_length: VarInt,
}
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.subscribe_id.encode(buf);
}
pub fn decode(buf: &mut impl Buf) -> Result<Self, CodecError> {
let subscribe_id = VarInt::decode(buf)?;
Ok(Self { subscribe_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);
encode_extensions(&self.extensions, buf);
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> {
check_end_of_track(self.object_id, self.object_status)?;
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 status_val = VarInt::decode(buf)?.into_inner();
ObjectStatus::from_u64(status_val).ok_or(CodecError::InvalidField)?
} else {
ObjectStatus::Normal
};
check_end_of_track(object_id, object_status)?;
Ok(Self {
group_id,
subgroup_id,
object_id,
publisher_priority,
extension_headers_length,
extensions,
object_status,
payload_length,
})
}
}