ws_frame/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2
3//! # ws-frame
4//! 
5//! A library for decoding WebSocket
6//! ([RFC6455](https://tools.ietf.org/html/rfc6455)) frames.
7//! 
8//! # Example
9//! ```
10//! use ws_frame::{Frame, Opcode};
11//! 
12//! let buf = [0b10100010, 0b00000001, 0b00000010];
13//! let mut f = Frame::empty();
14//! 
15//! if f.decode(&buf).is_complete() {
16//!     if Opcode::Ping == f.head.unwrap().op {
17//!         println!("Pong!")
18//!     }
19//! }
20//! ```
21
22#[cfg(feature = "std")]
23extern crate std as core;
24
25use byteorder::{BigEndian, ByteOrder};
26use core::convert::TryInto;
27
28mod iter;
29use iter::Bytes;
30
31macro_rules! unwrap {
32    ($e:expr) => {
33        match $e {
34            Some(t) => t,
35            None => return Status::Partial,
36        }
37    };
38}
39
40/// The result of a successful decode pass.
41///
42/// `Complete` is used when the buffer
43/// contained the complete value. `Partial` is used when decoding did not reach
44/// the end of the expected value, but no invalid data was found.
45#[derive(Debug, PartialEq)]
46pub enum Status {
47    /// The completed result.
48    ///
49    /// Contains the amount of bytes decoded.
50    Complete(usize),
51    /// A partial result.
52    Partial,
53}
54
55impl Status {
56    /// Convenience method to check if status is complete.
57    #[inline]
58    pub fn is_complete(&self) -> bool {
59        match *self {
60            Status::Complete(..) => true,
61            Status::Partial => false,
62        }
63    }
64
65    /// Convenience method to check if status is partial.
66    #[inline]
67    pub fn is_partial(&self) -> bool {
68        match *self {
69            Status::Complete(..) => false,
70            Status::Partial => true,
71        }
72    }
73
74    /// Convenience method to unwrap a Complete value. Panics if the status is
75    /// partial.
76    #[inline]
77    pub fn unwrap(self) -> usize {
78        match self {
79            Status::Complete(len) => len,
80            Status::Partial => panic!("Tried to unwrap Status::Partial"),
81        }
82    }
83}
84
85#[derive(Debug, PartialEq)]
86pub enum Opcode {
87    Continue,
88    Text,
89    Binary,
90    Close,
91    Ping,
92    Pong,
93    Reserved,
94}
95
96impl From<u8> for Opcode {
97    fn from(opcode: u8) -> Opcode {
98        match opcode {
99            0 => Opcode::Continue,
100            1 => Opcode::Text,
101            2 => Opcode::Binary,
102            8 => Opcode::Close,
103            9 => Opcode::Ping,
104            10 => Opcode::Pong,
105            _ => Opcode::Reserved,
106        }
107    }
108}
109
110#[derive(Debug, PartialEq)]
111pub struct Head {
112    pub op: Opcode,
113    pub finished: bool,
114    pub rsv: [bool; 3],
115}
116
117/// A decoded Frame.
118///
119/// The optional values will be `None` if a decode was not complete, and did
120/// not decode the associated property. This allows you to inspect the parts
121/// that could be decoded, before reading more.
122///
123/// # Example
124/// ```
125/// use ws_frame::Frame;
126/// 
127/// let buf = &[0b10000010, 0b00000001];
128/// let mut f = Frame::empty();
129/// 
130/// if f.decode(buf).is_partial() {
131///     match f.head {
132///         Some(head) => assert_eq!([false; 3], head.rsv),
133///         None => {
134///             // read more and decode again
135///         }
136///     }
137/// }
138/// ```
139#[derive(Debug, PartialEq)]
140pub struct Frame<'buf> {
141    /// The head section of a frame.
142    pub head: Option<Head>,
143    /// An optional mask key to apply over the payload.
144    pub mask: Option<[u8; 4]>,
145    /// The payload section of a frame.
146    ///
147    /// An empty payload is represented as `Some(&[])`.
148    pub payload: Option<&'buf [u8]>,
149}
150
151impl<'buf> Frame<'buf> {
152    /// Creates a new `Frame`.
153    pub const fn empty() -> Self {
154        Self {
155            head: None,
156            mask: None,
157            payload: None,
158        }
159    }
160    /// Try to decode a buffer of bytes into this `Frame`.
161    pub fn decode(&mut self, buf: &'buf [u8]) -> Status {
162        let mut bytes = Bytes::new(buf);
163
164        let first = unwrap!(bytes.next());
165        let rsv_bits = first >> 4 & 0x7u8;
166
167        let mut rsv = [false; 3];
168        for i in 0..3 {
169            rsv[2 - i] = rsv_bits >> i & 0x1u8 == 1u8;
170        }
171
172        self.head = Some(Head {
173            op: Opcode::from(first & 0xF),
174            finished: first_bit(first),
175            rsv,
176        });
177
178        let second = unwrap!(bytes.next());
179        let len = match second & 0x3F {
180            126 => unwrap!(bytes.slice_to(4).map(BigEndian::read_u64)),
181            // TODO validate most-sig bit == 0
182            127 => unwrap!(bytes.slice_to(8).map(BigEndian::read_u64)),
183            l => l as u64,
184        };
185
186        if first_bit(second) {
187            let mut mask = [0; 4];
188            mask.copy_from_slice(unwrap!(bytes.slice_to(4)));
189            self.mask = Some(mask);
190        }
191
192        self.payload = Some(unwrap!(bytes.slice_to(len.try_into().unwrap())));
193        Status::Complete(bytes.pos())
194    }
195}
196
197#[inline]
198fn first_bit(byte: u8) -> bool {
199    byte >> 7 == 1u8
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205
206    #[test]
207    fn it_works() {
208        const BYTES: &[u8] = &[0b10100010, 0b00000011, 0b00000001, 0b00000010, 0b00000011];
209        let mut f = Frame::empty();
210        assert_eq!(Status::Complete(BYTES.len()), f.decode(BYTES));
211
212        let head = f.head.unwrap();
213        assert!(head.finished);
214        assert_eq!([false, true, false], head.rsv);
215        assert_eq!(3, f.payload.unwrap().len());
216    }
217}