websocket_sans_io/
frame_encoding.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
use nonmax::NonMaxU8;
use tinyvec::ArrayVec;

use crate::{FrameInfo, MAX_HEADER_LENGTH};

/// A low-level WebSocket frames decoder.
/// 
/// It lets to prepare frame headers and transform (mask) frame payloads when needed.
/// 
/// It does not validate that you supplied correct amount of payload bytes after headers or that headers make sense.
/// 
/// Example usage:
/// 
/// ```
#[doc=include_str!("../examples/encode_frame.rs")]
/// ```
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
pub struct WebsocketFrameEncoder {
    mask: [u8; 4],
    phase: Option<NonMaxU8>,
}

impl WebsocketFrameEncoder {
    /// Create new instance of WebsocketFrameEncoder
    pub const fn new() -> WebsocketFrameEncoder {
        WebsocketFrameEncoder {
            mask: [0; 4],
            phase: None,
        }
    }

    /// Serialize given frame header as bytes. You should write all those bytes to the socket 
    /// before starting to write payload contant (if any).
    /// 
    /// You can repeat the calls to `start_frame` to re-serialize the header of the same frame
    /// if you haven't yet called `transform_frame_payload`
    /// for that frame.
    /// 
    /// It does not validate the frame in any way and can allow you 
    /// to do nonsensial things such as starting conversation with `Continuation` frame
    /// or getting next frame header when current frame's payload is not completely written.
    /// 
    /// Use masked frames when you are client and unmasked frames when you are server.
    /// 
    /// Writing frame header with nonzero `frame_info.payload_length` means you are obligated to
    /// write this number of bytes before writing any new frame.
    /// 
    /// If you have large or unknown size of a WebSocket message, use frames with `fin=false` and
    /// [`crate::Opcode::Continuation`] frames with smaller payload lengths each.
    /// This also allows to interrupt the data transmissing to send a [`crate::Opcode::Ping`]
    /// or reply with a [`crate::Opcode::Pong`].
    #[inline]
    pub fn start_frame(&mut self, frame_info: &FrameInfo) -> ArrayVec<[u8; MAX_HEADER_LENGTH]> {
        if let Some(m) = frame_info.mask {
            self.mask = m;
            self.phase = Some(NonMaxU8::default());
        } else {
            self.phase = None;
        }
        encode_frame_header(frame_info)
    }

    /// Prepare this memory chunk to be transfitted to the socket as a part of WebSocket frame payload.
    /// 
    /// Call this after `start_frame`.
    /// 
    /// Chunks transformed by this method should be written to the socket in the same order as they
    /// are supplied to `transform_frame_payload`.
    /// 
    /// If you do not want to transmit the tranformed chunk or some of its trailing part, you can
    /// rollback the encoder state with [`WebsocketFrameEncoder::rollback_payload_transform`].
    #[inline]
    pub fn transform_frame_payload(&mut self, data: &mut [u8]) {
        if let Some(ref mut phase) = self.phase {
            let ph = phase.get();

            crate::masking::apply_mask(self.mask, data, ph);

            *phase = NonMaxU8::new( (ph + ((data.len() % 4) as u8)) % 4  ).unwrap();
        }
    }

    /// Undo transformation of this number of bytes.
    /// 
    /// Example:
    /// 
    /// ```
    #[doc=include_str!("../examples/encode_frame_with_rollback.rs")]
    /// ```
    #[inline]
    pub fn rollback_payload_transform(&mut self, n_bytes: usize) {
        if let Some(ref mut phase) = self.phase {
            let modulo = (n_bytes % 4) as u8;
            let newvalue = (phase.get() + 4 - modulo) % 4;
            *phase = NonMaxU8::new(newvalue).unwrap();
        }
    }

    /// Check if you can skip `transform_frame_payload` and just transfer payload as is.
    #[inline]
    pub const fn transform_needed(&self) -> bool {
        self.phase.is_some()
    }
}

/// Just encode the header to bytes without using any encoder instance.
/// 
/// May be useful when you do not need masking.
#[inline]
pub fn encode_frame_header(frame_info: &FrameInfo) -> ArrayVec<[u8; MAX_HEADER_LENGTH]> {
    debug_assert!(frame_info.reserved & 0x7 == frame_info.reserved);

    let mut ret: ArrayVec<_> = ArrayVec::new();

    ret.push(
        if frame_info.fin { 0x80 } else { 0x00 }
            | (frame_info.reserved << 4)
            | (frame_info.opcode as u8),
    );
    let mut second_byte = if frame_info.mask.is_some() {
        0x80
    } else {
        0x00
    };
    match frame_info.payload_length {
        x if x <= 0x7D => {
            second_byte |= x as u8;
            ret.push(second_byte);
        }
        #[allow(unused_comparisons)]
        x if x <= 0xFFFF => {
            second_byte |= 0x7E;
            ret.push(second_byte);
            ret.extend((x as u16).to_be_bytes());
        }
        #[cfg(feature = "large_frames")]
        x => {
            second_byte |= 0x7F;
            ret.push(second_byte);
            ret.extend((x as u64).to_be_bytes());
        }
        #[cfg(not(feature = "large_frames"))]
        _ => unreachable!(),
    };

    if let Some(mask) = frame_info.mask {
        ret.extend(mask);
    }

    ret
}