atem-connection-rs/atem-connection-rs/src/atem_lib/atem_packet.rs

use core::{fmt::Display, str};
use enumflags2::{BitFlags, bitflags};

use super::atem_field::{Field, FieldParsingError, RawField};

/// The "hello" packet to start communication with the ATEM
pub const COMMAND_CONNECT_HELLO: [u8; 20] = [
    0x10, 0x14, 0x53, 0xab, 0x00, 0x00, 0x00, 0x00, 0x00, 0x3a, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00,
];

#[derive(Debug)]
pub struct AtemPacket<T: AsRef<[u8]>> {
    buf: T,
}

impl<'a> TryFrom<&'a [u8]> for AtemPacket<&'a [u8]> {
    type Error = AtemPacketErr;
    fn try_from(buf: &'a [u8]) -> Result<Self, Self::Error> {
        AtemPacket::new_checked(buf)
    }
}

#[derive(Debug)]
pub enum AtemPacketErr {
    /// The packet was too short
    TooShort {
        got: usize,
    },
    /// The packet's stated and actual lengths were different
    LengthDiffers {
        expected: u16,
        got: usize,
    },
    InvalidFlags,
}

#[bitflags]
#[repr(u8)]
#[derive(PartialEq, Copy, Clone, Debug)]
pub enum PacketFlag {
    AckRequest = 0x1,
    NewSessionId = 0x2,
    IsRetransmit = 0x4,
    RetransmitRequest = 0x8,
    AckReply = 0x10,
}

impl<T: AsRef<[u8]>> AtemPacket<T> {
    pub fn new_checked(buf: T) -> Result<Self, AtemPacketErr> {
        let len = buf.as_ref().len();
        if len < 12 {
            return Err(AtemPacketErr::TooShort {
                got: buf.as_ref().len(),
            });
        }
        let p = Self { buf };
        if p.length() as usize != len {
            return Err(AtemPacketErr::LengthDiffers {
                expected: p.length(),
                got: len,
            });
        }
        // Check flags are valid
        let _: BitFlags<PacketFlag> = p
            .flags_raw()
            .try_into()
            .map_err(|_| AtemPacketErr::InvalidFlags)?;
        Ok(p)
    }
    pub fn length(&self) -> u16 {
        u16::from_be_bytes(self.buf.as_ref()[0..=1].try_into().unwrap()) & 0x07ff
    }

    fn flags_raw(&self) -> u8 {
        self.buf.as_ref()[0] >> 3
    }

    pub fn flags(&self) -> BitFlags<PacketFlag> {
        // We `unwrap` here, but given we check the flags in the constructor this should never panic.
        self.flags_raw().try_into().unwrap()
    }

    pub fn session_id(&self) -> u16 {
        u16::from_be_bytes(self.buf.as_ref()[2..=3].try_into().unwrap())
    }

    pub fn ack_number(&self) -> u16 {
        u16::from_be_bytes(self.buf.as_ref()[4..=5].try_into().unwrap())
    }

    pub fn remote_sequence_number(&self) -> u16 {
        u16::from_be_bytes(self.buf.as_ref()[9..=10].try_into().unwrap())
    }

    pub fn local_sequence_number(&self) -> u16 {
        u16::from_be_bytes(self.buf.as_ref()[10..=11].try_into().unwrap())
    }

    pub fn retransmit_request(&self) -> Option<u16> {
        self.flags()
            .contains(PacketFlag::RetransmitRequest)
            .then_some(u16::from_be_bytes([
                self.buf.as_ref()[6],
                self.buf.as_ref()[7],
            ]))
    }

    pub fn ack_reply(&self) -> Option<u16> {
        self.flags()
            .contains(PacketFlag::AckReply)
            .then_some(self.ack_number())
    }

    /// Get an iterator over the `Field`s in this packet.
    ///
    /// Returns None if this is a packet without fields.
    pub fn raw_fields(&self) -> RawFields {
        // TODO: do we only ever get newsessionid during the handshake (i.e. not in a packet with fields)?
        let has_fields = !self.flags().contains(PacketFlag::NewSessionId);
        RawFields::new(if has_fields { self.body() } else { &[] })
    }

    pub fn fields<'a, F: Field<'a>>(
        &'a self,
    ) -> impl Iterator<Item = Result<F, FieldParsingError>> + use<'a, F, T> {
        self.raw_fields().filter_map(|f| match F::try_from_raw(f) {
            Err(FieldParsingError::MismatchedFieldType) => None,
            x => Some(x),
        })
    }

    pub fn body(&self) -> &[u8] {
        &self.buf.as_ref()[12..]
    }
}

impl<T: AsRef<[u8]>> Display for AtemPacket<T> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(
            f,
            "len: {}, sid: {}, flags: {}, ack: {}, localseq: {}, remoteseq: {}",
            self.length(),
            self.session_id(),
            self.flags(),
            self.ack_number(),
            self.remote_sequence_number(),
            self.local_sequence_number(),
        )
    }
}

/// Created by [`AtemPacket::raw_fields`]
pub struct RawFields<'a> {
    data: &'a [u8],
    // The offset of the next field in the packet
    offset: usize,
}

impl<'a> RawFields<'a> {
    pub(crate) fn new(data: &'a [u8]) -> Self {
        Self { data, offset: 0 }
    }
}

impl<'a> Iterator for RawFields<'a> {
    type Item = RawField<'a>;
    fn next(&mut self) -> Option<Self::Item> {
        let remain = self.data.len() - self.offset;
        if remain == 0 {
            return None;
        } else if remain < 8 {
            // TODO: is 8 indeed the minimum size for something here? (i.e. no field data)
            panic!("Oh no");
        }

        let length =
            u16::from_be_bytes(self.data[self.offset..=self.offset + 1].try_into().unwrap());
        // TODO: sanity check length
        let r#type = str::from_utf8(&self.data[self.offset + 4..=self.offset + 7]).unwrap();
        let data = &self.data[self.offset + 8..self.offset + (length as usize)];
        self.offset += length as usize;
        Some(RawField { r#type, data })
    }
}