websocket_rs/

codec.rs

//! FOR FUTURE ME (me talking to me in the future)
//! =============
//! Well oussama if you come next year and you did take a look at this code,
//! and start blaming yourself on why you wasted time doing this, like you
//! always do, then read the following:
//!
//! MOTIVATION
//! ==========
//! I tried to integrate a couple of websockets crates but boy it was just
//! plain painfull, way too complicated and time consuming when all i need is
//! get the data that the client sent and move on with my life.
//!
//! So i find it to be easier and straight forward (thanks to tokio) to simply
//! implement a decoder and encoder for websocket frames.
//!
//! This is in no way a full implementation of RFC 6455, but rather an
//! implementation in a controlled environment that only adds features required
//! by the application, and since the client will always be a legit browser
//! then ill safely make some assumptions and skip some checks (like validating
//! the frame structure).
//!
//! My minimal websocket implementation will only be used by xtermjs frontend
//! using AttachAddon.
//!
//! FEATURES
//! ========
//! ```ignore
//! [+] Receive opcode TEXT
//! [+] Pong when Ping is received
//! [+] Handle all payload lengths (**le 125**, **=126**, **=127**)
//! [+] Mask key
//! [+] Close websocket connection
//! [+] Read fragmented payload
//! [+] Extract websocket to its own crate
//! [ ] Add binary support
//! [ ] Send fragmented frame when the size reaches a treshold
//! [ ] Schedule ping/pong
//! [ ] Pong with application data included in ping
//! [ ] Keep track of connected clients
//! ```
//!
//! REFERENCE
//! =========
//! [RFC 6455 5.2](https://tools.ietf.org/html/rfc6455#section-5.2)
//!
//!  ```ignore
//!  0                   1                   2                   3
//!  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
//! +-+-+-+-+-------+-+-------------+-------------------------------+
//! |F|R|R|R| opcode|M| Payload len |    Extended payload length    |
//! |I|S|S|S|  (4)  |A|     (7)     |             (16/64)           |
//! |N|V|V|V|       |S|             |   (if payload len==126/127)   |
//! | |1|2|3|       |K|             |                               |
//! +-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
//! |     Extended payload length continued, if payload len == 127  |
//! + - - - - - - - - - - - - - - - +-------------------------------+
//! |                               |Masking-key, if MASK set to 1  |
//! +-------------------------------+-------------------------------+
//! | Masking-key (continued)       |          Payload Data         |
//! +-------------------------------- - - - - - - - - - - - - - - - +
//! :                     Payload Data continued ...                :
//! + - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - +
//! |                     Payload Data continued ...                |
//! +---------------------------------------------------------------+
//! ```
//!
//! ```ignore
//! FIN (1 bit): if 1 -> final fragment
//! RSV1 RSV2 RSV3 (3 bit): Must be 0 unless an extension defines them (SKIP)
//! Opcode (4 bit): Always assume its good (data sent from browser)
//!   - 0   -> continuation frame
//!   - 1   -> text frame
//!   - 2   -> binary frame
//!   - 3-7 -> reserved for further non-control frames
//!   - 8   -> connection closed
//!   - 9   -> ping
//!   - A   -> pong
//!   - B-F -> reserved for further control frames
//! Mask (1 bit): if 1 -> masking key is present. #section 5.3
//! Payload length (7 bit, 7+16 bit, 7+64 bit): length of the payload data
//!   - 0-125 -> thats the payload length
//!   - 126 -> next 2 bytes (UNSIGNED) are the payload length
//!   - 127 -> next 8 bytes (UNSIGNED) are the payload length
//! Masking key (0 or 4 bytes): present if mask bit is set to 1. #section 5.3
//! Payload data (x+y bytes): extension data + application data
//! Extension data (x bytes): is 0 unless an extension is negotiated (not in our case)
//! Application data (y bytes): is payload length - length of extension data
//! ```

use std::convert::TryInto;

use bytes::{ BytesMut, Buf, BufMut };
use tokio_util::codec::{ Decoder, Encoder };

use crate::{
    Message,
    Opcode,
    WebsocketError,
    WebsocketResult,
};

#[derive(Debug, Default)]
pub struct WebsocketCodec {
    fin: bool,
    opcode: Opcode,
    payload: BytesMut,
}

impl WebsocketCodec {
    #[inline(always)]
    fn next_frame(&mut self, buf: &mut BytesMut) -> Option<()> {
        if buf.len() < 6 {
            return None
        }

        self.fin = (0x80 & buf[0]) >> 7 != 0;

        self._set_opcode(0x0F & buf[0]);

        if self.opcode == Opcode::Text {
            let (payload_len, payload_idx, mask_key) = self._payload_meta(buf);

            if buf[payload_idx..].len() < payload_len {
                return None
            }

            self._unmask(buf, payload_len, payload_idx, mask_key);

            self.payload.extend_from_slice(&buf[payload_idx..payload_idx + payload_len]);
            buf.advance(payload_idx + payload_len);

            return Some(())
        }

        if
            self.opcode == Opcode::Ping
            || self.opcode == Opcode::Pong
            || self.opcode == Opcode::Close
        {
            buf.advance(buf.len());
            return Some(())
        }

        unimplemented!()
    }

    #[inline(always)]
    fn _unmask(&mut self, buf: &mut BytesMut, len: usize, idx: usize, key: [u8; 4]) {
        if (0x80 & buf[1]) >> 7 != 0 {
            for i in idx..(idx + len) {
                buf[i] ^= key[(i - idx) % 4];
            }
        }
    }

    #[inline(always)]
    fn _payload_meta(&self, buf: &BytesMut) -> (usize, usize, [u8; 4]) {
        let len = (0x7F & buf[1]) as usize;

        if len == 127 {
            return (
                usize::from_be_bytes(buf[2..10].try_into().unwrap()),
                14,
                buf[10..14].try_into().unwrap()
            )
        }

        if len == 126 {
            return (
                u16::from_be_bytes(buf[2..4].try_into().unwrap()) as usize,
                8,
                buf[4..8].try_into().unwrap()
            )
        }

        (len, 6, buf[2..6].try_into().unwrap())
    }

    #[inline(always)]
    fn _set_opcode(&mut self, opcode: u8) {
        if opcode == 1 { self.opcode = Opcode::Text; }
        else if opcode == 0x9 { self.opcode = Opcode::Ping; }
        else if opcode == 0xA { self.opcode = Opcode::Pong; }
        else if opcode == 0x8 { self.opcode = Opcode::Close; }
        else { self.opcode = Opcode::Unknown; }
    }
}

impl Decoder for WebsocketCodec {
    type Item = Message;
    type Error = WebsocketError;

    fn decode(&mut self, buf: &mut BytesMut) -> WebsocketResult<Option<Self::Item>> {
        while self.next_frame(buf).is_some() {
            if self.fin && (
                self.opcode == Opcode::Text
                || self.opcode == Opcode::Ping
                || self.opcode == Opcode::Pong
                || self.opcode == Opcode::Close
            ) {
                return Ok(Some(Message {
                    opcode: self.opcode,
                    payload: self.payload.to_bytes(),
                }))
            }
        }

        Ok(None)
    }
}

impl Encoder<Message> for WebsocketCodec {
    type Error = WebsocketError;

    fn encode(&mut self, msg: Message, dst: &mut BytesMut) -> WebsocketResult<()> {
        if msg.opcode == Opcode::Pong {
            // TODO RFC 5.5.3 must have identical application data sent by the
            // client.
            dst.extend_from_slice(&[138, 0]);
        }

        if msg.opcode == Opcode::Text {
            dst.put_u8(0x81);

            if msg.payload.len() <= 125 {
                dst.put_u8(msg.payload.len() as u8);
            } else if msg.payload.len() == 126 {
                dst.put_u8(126);
                dst.put_u16(msg.payload.len() as u16);
            } else {
                dst.put_u8(127);
                dst.put_u64(msg.payload.len() as u64);
            }

            dst.put_slice(&msg.payload[..]);
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    extern crate test;

    use super::*;
    use bytes::BytesMut;
    use std::str;
    use test::Bencher;

    macro_rules! frame {
        ($buf:expr, $payload:expr, $opcode:expr) => {
            let mut buf: BytesMut = $buf.as_bytes()
                .chunks(2)
                .map(|s| u8::from_str_radix(unsafe { str::from_utf8_unchecked(s) }, 16).unwrap())
                .collect();

            let mut wsc = WebsocketCodec::default();

            if let Some(_) = wsc.next_frame(&mut buf) {
                let payload = unsafe { str::from_utf8_unchecked(&wsc.payload) };
                assert!(payload == $payload);
                assert!(wsc.opcode == $opcode);
            } else {
                assert!(false, "failed to extract frame from '{:x}'")
            }
        }
    }

    #[test]
    fn text_frame_with_no_payload() {
        frame!("8180e0350bbc", "", Opcode::Text);
    }

    #[test]
    fn text_frame_with_payload_len_eq_1() {
        frame!("8181f3dba99092", "a", Opcode::Text);
    }

    #[test]
    fn text_frame_with_hello_world_payload() {
        frame!("818bb013fc40d876902cdf338b2fc27f98", "hello world", Opcode::Text);
    }

    #[test]
    fn text_frame_payload_125() {
        frame!(
            "81fd16f9d90b5795b52b6291b87f369eb562628dbc7965d9b0783697b67f369eb66772d99f6a7f8bf96265d9bf646395f52b7797bd2b7096ac673690aa2b7098b0792cd99164609cab2b6291ab64639eb12b6291bc2b7096be2b7797bd2b7090b57f7e80f96a7f8bf72b36adb16e659cf97d7f96b56e788df96f7395b06c7e8daa2b7e",
            "All that glitters is not gold Fair is foul, and foul is fair: Hover through the fog and filthy air.  These violent delights h",
            Opcode::Text
        );
    }

    #[bench]
    fn bench_text_frame_payload_125(b: &mut Bencher) {
        // 1,570 ns/iter (+/- 50)
        b.iter(|| text_frame_payload_125())
    }

    #[test]
    fn text_frame_long_payload() {
        frame!(
            "81fe020d6742ba33262ed613132adb474725d65a1336df411462d340472cd5474725d55f0348fc520e309a5a1462dc5c122e9613062cde13012dcf5f472bc9130123d3415d62f25c1127c813132ac85c1225d213132adf13012ddd13062cde13012bd6470f3b9a520e309439332adf400262cc5a082edf5d1362de560b2bdd5b13319a5b0634df13112bd55f022cce13022cde40496c94392f27d65f472bc913022fca471e62db5d0362db5f0b62ce5b0262de56112bd6404723c856472adf41026cb0711e62ce5b0262ca410e21d15a09259a5c0162d74a4736d2460a20c91f4711d55e0236d25a09259a440e21d1560362ce5b0e319a44063b9a50082fdf404962f543022c96130b2dd958146e9a640f2ddf4502309a58092dd9581463b0670f279a5f0626c313032dce5b4732c85c1327c9474736d55c472fcf500f6e9a5e0236d25a0929c91d6d00c856112bce4a472bc913132adf13142dcf5f472ddc13102bce1d6d0bdc130a37c95a0462d8564736d2564724d55c0362d555472ed545026e9a430b23c313082c9439292dcd130e319a470f279a440e2cce561562d555472dcf414726d340042dd447022cce1d6d05d55c0362d45a002ace1f4725d55c0362d45a002ace124732db41132bd454472bc9131437d95b4731cd5602369a400830c85c106e9a670f23ce132e62c95b062ed6131423c313002dd557472cd3540f369a470e2ed6130e369a510262d75c1530d54449",
            "All that glitters is not gold
Fair is foul, and foul is fair: Hover through the fog and filthy air.
These violent delights have violent ends...
Hell is empty and all the devils are here.
By the pricking of my thumbs, Something wicked this way comes. Open, locks, Whoever knocks!
The lady doth protest too much, methinks.
Brevity is the soul of wit.
If music be the food of love, play on.
Now is the winter of our discontent.
Good night, good night! parting is such sweet sorrow, That I shall say good night till it be morrow.",
            Opcode::Text
        );
    }

    #[bench]
    fn bench_text_frame_long_payload(b: &mut Bencher) {
        // 6,795 ns/iter (+/- 172)
        b.iter(|| text_frame_long_payload())
    }
}