zedmq/codec/
frame.rs

1use std::convert::TryInto;
2
3use super::{Command, Message};
4
5// -- FrameBuf
6
7/// An owned and growable frame.
8#[derive(Debug, PartialEq)]
9pub struct FrameBuf {
10    pub(crate) bytes: Vec<u8>,
11}
12
13impl AsRef<[u8]> for FrameBuf {
14    fn as_ref(&self) -> &[u8] {
15        self.bytes.as_slice()
16    }
17}
18
19impl From<Vec<u8>> for FrameBuf {
20    fn from(bytes: Vec<u8>) -> Self {
21        Self { bytes }
22    }
23}
24
25impl From<FrameBuf> for Vec<u8> {
26    fn from(FrameBuf { bytes }: FrameBuf) -> Self {
27        bytes
28    }
29}
30
31impl FrameBuf {
32    pub fn new(bytes: Vec<u8>) -> Self {
33        Self { bytes }
34    }
35
36    pub fn short_command<'a, I>(name: &str, properties: Option<I>) -> Self
37    where
38        I: IntoIterator<Item = (&'a str, &'a str)>,
39    {
40        assert_eq!(name.len(), name.as_bytes().len());
41
42        let mut bytes = vec![
43            // SHORT COMMAND
44            0x4,
45            // LENGTH OF ENTIRE FRAME
46            0x00,
47            // name as cstr length
48            name.len() as u8,
49        ];
50
51        let body: Vec<u8> = match properties {
52            None => vec![],
53            Some(it) => {
54                let mut payload = vec![];
55
56                // For every property:
57                //
58                //  1) push the name length as u8
59                //  2) extend with the name bytes
60                //  3) push the field length as u64 in network byte order
61                //  4) extend with the field bytes
62                for (st, field) in it.into_iter() {
63                    match st.len().try_into() {
64                        Ok(length) => payload.push(length),
65                        Err(_) => panic!("property names can not be longer than 255 bytes."),
66                    }
67
68                    payload.extend_from_slice(st.as_bytes());
69                    payload.extend_from_slice(&u32::to_be_bytes(field.len() as u32) as &[_]);
70                    payload.extend_from_slice(field.as_bytes());
71                }
72
73                payload
74            }
75        };
76
77        bytes.extend_from_slice(name.as_bytes());
78        bytes.extend_from_slice(body.as_slice());
79
80        // Size is size_of(frame) - 2
81        // (don't include the flags byte or the size byte itself.)
82        bytes[1] = (bytes.len() - 2).try_into().unwrap();
83
84        Self { bytes }
85    }
86
87    pub fn as_frame<'a>(&'a self) -> Frame<'a> {
88        Frame::new(self.bytes.as_slice())
89    }
90}
91
92// -- Frame<'a>
93
94/// A slice of frame (akin to `str` or `Path`)
95#[derive(Debug, PartialEq)]
96pub struct Frame<'a> {
97    pub(crate) bytes: &'a [u8],
98}
99
100impl<'a> Frame<'a> {
101    pub fn new(bytes: &'a [u8]) -> Self {
102        Self { bytes }
103    }
104
105    pub fn try_into_command(self) -> Option<Command<'a>> {
106        match self.kind()? {
107            FrameKind::Command => Some(Command { frame: self }),
108            _ => None,
109        }
110    }
111
112    pub fn try_into_message(self) -> Option<Message<'a>> {
113        match self.kind()? {
114            FrameKind::MessagePart => Some(Message {
115                frame: self,
116                is_last: false,
117            }),
118            FrameKind::MessageTail => Some(Message {
119                frame: self,
120                is_last: true,
121            }),
122            _ => None,
123        }
124    }
125
126    /// Get the size of the frame.
127    pub fn size(&self) -> Option<usize> {
128        match self.bytes.get(0)? {
129            0x0 | 0x1 | 0x4 => Some(*self.bytes.get(1)? as usize),
130
131            0x2 | 0x3 | 0x6 => {
132                let slice = self.bytes.get(1..)?.try_into().ok()?;
133                let size = u64::from_be_bytes(slice);
134                Some(size as usize)
135            }
136
137            _ => None,
138        }
139    }
140
141    /// The type of frame.
142    pub fn kind(&self) -> Option<FrameKind> {
143        let kind = match self.bytes.get(0)? {
144            // short/long message tail.
145            0x0 | 0x2 => FrameKind::MessageTail,
146
147            // short/long message part.
148            0x1 | 0x3 => FrameKind::MessagePart,
149
150            // short/long command.part.
151            // [frame_tag, ...(frame_size{1u8} | frame_size{8u8}), size, ...command_string{size}]
152            0x4 | 0x6 => FrameKind::Command,
153
154            // All the rest...
155            _ => return None,
156        };
157
158        Some(kind)
159    }
160}
161
162/// The various types a frame can be.
163#[derive(Debug, PartialEq, Copy, Clone)]
164pub enum FrameKind {
165    /// A command frame.
166    Command,
167
168    /// A single part (frame) of a message (multipart)
169    MessagePart,
170
171    /// The last frame of a multipart message or a singular frame of a non-multipart.
172    MessageTail,
173}
174
175impl<'a> From<&'a [u8]> for Frame<'a> {
176    fn from(bytes: &'a [u8]) -> Self {
177        Self { bytes }
178    }
179}