websocket_sans_io/
frame_encoding.rs

1use nonmax::NonMaxU8;
2use tinyvec::ArrayVec;
3
4use crate::{FrameInfo, MAX_HEADER_LENGTH};
5
6/// A low-level WebSocket frames decoder.
7/// 
8/// It lets to prepare frame headers and transform (mask) frame payloads when needed.
9/// 
10/// It does not validate that you supplied correct amount of payload bytes after headers or that headers make sense.
11/// 
12/// Example usage:
13/// 
14/// ```
15#[doc=include_str!("../examples/encode_frame.rs")]
16/// ```
17#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
18pub struct WebsocketFrameEncoder {
19    mask: [u8; 4],
20    phase: Option<NonMaxU8>,
21}
22
23impl WebsocketFrameEncoder {
24    /// Create new instance of WebsocketFrameEncoder
25    pub const fn new() -> WebsocketFrameEncoder {
26        WebsocketFrameEncoder {
27            mask: [0; 4],
28            phase: None,
29        }
30    }
31
32    /// Serialize given frame header as bytes. You should write all those bytes to the socket 
33    /// before starting to write payload contant (if any).
34    /// 
35    /// You can repeat the calls to `start_frame` to re-serialize the header of the same frame
36    /// if you haven't yet called `transform_frame_payload`
37    /// for that frame.
38    /// 
39    /// It does not validate the frame in any way and can allow you 
40    /// to do nonsensial things such as starting conversation with `Continuation` frame
41    /// or getting next frame header when current frame's payload is not completely written.
42    /// 
43    /// Use masked frames when you are client and unmasked frames when you are server.
44    /// 
45    /// Writing frame header with nonzero `frame_info.payload_length` means you are obligated to
46    /// write this number of bytes before writing any new frame.
47    /// 
48    /// If you have large or unknown size of a WebSocket message, use frames with `fin=false` and
49    /// [`crate::Opcode::Continuation`] frames with smaller payload lengths each.
50    /// This also allows to interrupt the data transmissing to send a [`crate::Opcode::Ping`]
51    /// or reply with a [`crate::Opcode::Pong`].
52    #[inline]
53    pub fn start_frame(&mut self, frame_info: &FrameInfo) -> ArrayVec<[u8; MAX_HEADER_LENGTH]> {
54        if let Some(m) = frame_info.mask {
55            self.mask = m;
56            self.phase = Some(NonMaxU8::default());
57        } else {
58            self.phase = None;
59        }
60        encode_frame_header(frame_info)
61    }
62
63    /// Prepare this memory chunk to be transfitted to the socket as a part of WebSocket frame payload.
64    /// 
65    /// Call this after `start_frame`.
66    /// 
67    /// Chunks transformed by this method should be written to the socket in the same order as they
68    /// are supplied to `transform_frame_payload`.
69    /// 
70    /// If you do not want to transmit the tranformed chunk or some of its trailing part, you can
71    /// rollback the encoder state with [`WebsocketFrameEncoder::rollback_payload_transform`].
72    #[inline]
73    pub fn transform_frame_payload(&mut self, data: &mut [u8]) {
74        if let Some(ref mut phase) = self.phase {
75            let ph = phase.get();
76
77            crate::masking::apply_mask(self.mask, data, ph);
78
79            *phase = NonMaxU8::new( (ph + ((data.len() % 4) as u8)) % 4  ).unwrap();
80        }
81    }
82
83    /// Undo transformation of this number of bytes.
84    /// 
85    /// Example:
86    /// 
87    /// ```
88    #[doc=include_str!("../examples/encode_frame_with_rollback.rs")]
89    /// ```
90    #[inline]
91    pub fn rollback_payload_transform(&mut self, n_bytes: usize) {
92        if let Some(ref mut phase) = self.phase {
93            let modulo = (n_bytes % 4) as u8;
94            let newvalue = (phase.get() + 4 - modulo) % 4;
95            *phase = NonMaxU8::new(newvalue).unwrap();
96        }
97    }
98
99    /// Check if you can skip `transform_frame_payload` and just transfer payload as is.
100    #[inline]
101    pub const fn transform_needed(&self) -> bool {
102        self.phase.is_some()
103    }
104}
105
106/// Just encode the header to bytes without using any encoder instance.
107/// 
108/// May be useful when you do not need masking.
109#[inline]
110pub fn encode_frame_header(frame_info: &FrameInfo) -> ArrayVec<[u8; MAX_HEADER_LENGTH]> {
111    debug_assert!(frame_info.reserved & 0x7 == frame_info.reserved);
112
113    let mut ret: ArrayVec<_> = ArrayVec::new();
114
115    ret.push(
116        if frame_info.fin { 0x80 } else { 0x00 }
117            | (frame_info.reserved << 4)
118            | (frame_info.opcode as u8),
119    );
120    let mut second_byte = if frame_info.mask.is_some() {
121        0x80
122    } else {
123        0x00
124    };
125    match frame_info.payload_length {
126        x if x <= 0x7D => {
127            second_byte |= x as u8;
128            ret.push(second_byte);
129        }
130        #[allow(unused_comparisons)]
131        x if x <= 0xFFFF => {
132            second_byte |= 0x7E;
133            ret.push(second_byte);
134            ret.extend((x as u16).to_be_bytes());
135        }
136        #[cfg(feature = "large_frames")]
137        x => {
138            second_byte |= 0x7F;
139            ret.push(second_byte);
140            ret.extend((x as u64).to_be_bytes());
141        }
142        #[cfg(not(feature = "large_frames"))]
143        _ => unreachable!(),
144    };
145
146    if let Some(mask) = frame_info.mask {
147        ret.extend(mask);
148    }
149
150    ret
151}
websocket_sans_io/frame_encoding.rs

websocket_sans_io/
frame_encoding.rs
