<commit_before>// Copyright (c) Carl-Erwin Griffith
//
// Permission is hereby granted, free of charge, to any
// person obtaining a copy of this software and associated
// documentation files (the "Software"), to deal in the
// Software without restriction, including without
// limitation the rights to use, copy, modify, merge,
// publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software
// is furnished to do so, subject to the following
// conditions:
//
// The above copyright notice and this permission notice
// shall be included in all copies or substantial portions
// of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
// ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
// SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
// IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER

use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};

use crate::core::codepointinfo::CodepointInfo;

pub type LineCellIndex = usize;

/// A LineCell encapsulates code point information (CodepoinInfo).

/// The displayed Lines are composed of LineCell
#[derive(Hash, Default, Debug, Clone, Eq, PartialEq)]
pub struct LineCell {
    pub cpi: CodepointInfo,
    pub is_used: bool,
}

impl LineCell {
    pub fn new() -> Self {
        LineCell {
            cpi: CodepointInfo::new(),
            is_used: false,
        }
    }
}

/// A Line is an array of LineCell and some metadata.

/// A Screen is composed of Line(s)
#[derive(Default, Debug, Clone, Eq, PartialEq)]
pub struct Line {
    pub cells: Vec<LineCell>,
    pub nb_cells: usize,
    pub max_width: usize,
    pub start_index: usize,
    width: usize,
    pub read_only: bool,
    pub metadata: bool,
    hash_cache: u64,
    hash_unclipped_cache: u64,
}

impl Line {
    pub fn new(width: LineCellIndex) -> Self {
        assert_eq!(width > 0, true);

        let mut cells = Vec::with_capacity(width);
        for _ in 0..width {
            cells.push(LineCell::new());
        }

        Line {
            cells,
            nb_cells: 0,
            max_width: width,
            start_index: 0,
            width,
            read_only: false,
            metadata: false,
            hash_cache: 0,
            hash_unclipped_cache: 0,
        }
    }

    pub fn hash(&mut self) -> u64 {
        if self.hash_cache != 0 {
            return self.hash_cache;
        }

        let mut s = DefaultHasher::new();

        for i in 0..self.width() {
            self.cells[self.start_index + i].hash(&mut s);
        }

        self.hash_cache = s.finish();
        self.hash_cache
    }

    pub fn hash_unclipped(&mut self) -> u64 {
        if self.hash_unclipped_cache != 0 {
            return self.hash_unclipped_cache;
        }

        let mut s = DefaultHasher::new();

        for i in 0..self.max_width {
            self.cells[i].hash(&mut s);
        }
        self.hash_unclipped_cache = s.finish();
        self.hash_unclipped_cache
    }

    pub fn width(&self) -> usize {
        self.width - self.start_index
    }

    pub fn max_width(&self) -> usize {
        self.max_width
    }

    pub fn resize(&mut self, width: LineCellIndex) {
        self.cells.resize(width, LineCell::new());
        self.nb_cells = 0;
        self.max_width = width;
        self.start_index = 0;
        self.width = width;
        self.read_only = false;
        self.metadata = false;
        self.hash_cache = 0;
    }

    /// returns (start_index, width) tupple
    /// [ 0 <= start_index < width <= screen.max_width() ]
    pub fn clipping(&mut self) -> (usize, usize) {
        (self.start_index, self.width)
    }

    /// [ 0 <= start_index < width <= screen.max_width() ]
    pub fn set_clipping(&mut self, start_index: usize, width: usize) {
        assert!(start_index < self.max_width);
        assert!(width <= self.max_width);
        assert!(start_index + width <= self.max_width);

        self.start_index = start_index;
        self.width = start_index + width;
        self.hash_cache = 0;
        self.nb_cells = 0;
        self.read_only = false;
    }

    pub fn clear_clip_width(&mut self) {
        self.set_clipping(0, self.max_width);
    }

    pub fn push(&mut self, cpi: CodepointInfo) -> (bool, LineCellIndex) {
        if self.nb_cells < self.width() && !self.read_only {
            self.cells[self.start_index + self.nb_cells].cpi = cpi;
            self.cells[self.start_index + self.nb_cells].is_used = true;

            self.nb_cells += 1;

            if self.nb_cells == self.width {
                self.read_only = true;
            }

            (true, self.nb_cells)
        } else {
            self.read_only = true;
            (false, self.nb_cells)
        }
    }

    pub fn clear(&mut self) {
        for i in 0..self.width() {
            self.cells[self.start_index + i] = LineCell::new();
        }

        self.hash_cache = 0;
        self.nb_cells = 0;
        self.read_only = false;
        self.metadata = false;
        self.hash_cache = 0;
    }

    pub fn get_first_cpi(&self) -> Option<&CodepointInfo> {
        if self.start_index < self.width {
            Some(&self.cells[self.start_index].cpi)
        } else {
            // internal error
            None
        }
    }

    pub fn get_last_cpi(&self) -> Option<&CodepointInfo> {
        if self.nb_cells > 0 {
            Some(&self.cells[self.start_index + self.nb_cells - 1].cpi)
        } else {
            None
        }
    }

    pub fn get_unclipped_cpi(&self, index: LineCellIndex) -> Option<&CodepointInfo> {
        if index < self.max_width() {
            Some(&self.cells[index].cpi)
        } else {
            None
        }
    }

    pub fn get_mut_unclipped_cpi(&mut self, index: LineCellIndex) -> Option<&mut CodepointInfo> {
        if index < self.max_width() {
            Some(&mut self.cells[index].cpi)
        } else {
            None
        }
    }

    pub fn get_cpi(&self, index: LineCellIndex) -> Option<&CodepointInfo> {
        if index < self.width() {
            Some(&self.cells[self.start_index + index].cpi)
        } else {
            None
        }
    }

    pub fn get_mut_cpi(&mut self, index: LineCellIndex) -> Option<&mut CodepointInfo> {
        if index < self.width() {
            Some(&mut self.cells[self.start_index + index].cpi)
        } else {
            None
        }
    }

    pub fn at(&mut self, index: LineCellIndex) -> Option<&mut CodepointInfo> {
        self.get_mut_cpi(index)
    }

    pub fn get_used_cpi(&self, index: LineCellIndex) -> Option<&CodepointInfo> {
        if index < self.nb_cells {
            self.get_cpi(index)
        } else {
            None
        }
    }

    pub fn get_mut_used_cpi(&mut self, index: LineCellIndex) -> Option<&mut CodepointInfo> {
        if index < self.nb_cells {
            self.get_mut_cpi(index)
        } else {
            None
        }
    }

    pub fn get_used_cpi_clipped(
        &self,
        index: LineCellIndex,
    ) -> (Option<&CodepointInfo>, LineCellIndex) {
        if self.nb_cells == 0 {
            return (None, 0);
        }

        let index = ::std::cmp::min(index, self.nb_cells - 1);
        (self.get_used_cpi(index), index)
    }
}
<commit_msg>api soundness : usize <-> LineCellIndex<commit_after>// Copyright (c) Carl-Erwin Griffith
//
// Permission is hereby granted, free of charge, to any
// person obtaining a copy of this software and associated
// documentation files (the "Software"), to deal in the
// Software without restriction, including without
// limitation the rights to use, copy, modify, merge,
// publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software
// is furnished to do so, subject to the following
// conditions:
//
// The above copyright notice and this permission notice
// shall be included in all copies or substantial portions
// of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
// ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
// SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
// IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER

use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};

use crate::core::codepointinfo::CodepointInfo;

pub type LineCellIndex = usize;

/// A LineCell encapsulates code point information (CodepoinInfo).

/// The displayed Lines are composed of LineCell
#[derive(Hash, Default, Debug, Clone, Eq, PartialEq)]
pub struct LineCell {
    pub cpi: CodepointInfo,
    pub is_used: bool,
}

impl LineCell {
    pub fn new() -> Self {
        LineCell {
            cpi: CodepointInfo::new(),
            is_used: false,
        }
    }
}

/// A Line is an array of LineCell and some metadata.

/// A Screen is composed of Line(s)
#[derive(Default, Debug, Clone, Eq, PartialEq)]
pub struct Line {
    pub cells: Vec<LineCell>,
    pub nb_cells: usize,
    pub max_width: usize,
    pub start_index: usize,
    width: usize,
    pub read_only: bool,
    pub metadata: bool,
    hash_cache: u64,
    hash_unclipped_cache: u64,
}

impl Line {
    pub fn new(width: usize) -> Self {
        assert_eq!(width > 0, true);

        let mut cells = Vec::with_capacity(width);
        for _ in 0..width {
            cells.push(LineCell::new());
        }

        Line {
            cells,
            nb_cells: 0,
            max_width: width,
            start_index: 0,
            width,
            read_only: false,
            metadata: false,
            hash_cache: 0,
            hash_unclipped_cache: 0,
        }
    }

    pub fn hash(&mut self) -> u64 {
        if self.hash_cache != 0 {
            return self.hash_cache;
        }

        let mut s = DefaultHasher::new();

        for i in 0..self.width() {
            self.cells[self.start_index + i].hash(&mut s);
        }

        self.hash_cache = s.finish();
        self.hash_cache
    }

    pub fn hash_unclipped(&mut self) -> u64 {
        if self.hash_unclipped_cache != 0 {
            return self.hash_unclipped_cache;
        }

        let mut s = DefaultHasher::new();

        for i in 0..self.max_width {
            self.cells[i].hash(&mut s);
        }
        self.hash_unclipped_cache = s.finish();
        self.hash_unclipped_cache
    }

    pub fn width(&self) -> usize {
        self.width - self.start_index
    }

    pub fn max_width(&self) -> usize {
        self.max_width
    }

    pub fn resize(&mut self, width: usize) {
        self.cells.resize(width, LineCell::new());
        self.nb_cells = 0;
        self.max_width = width;
        self.start_index = 0;
        self.width = width;
        self.read_only = false;
        self.metadata = false;
        self.hash_cache = 0;
    }

    /// returns (start_index, width) tupple
    /// [ 0 <= start_index < width <= screen.max_width() ]
    pub fn clipping(&mut self) -> (usize, usize) {
        (self.start_index, self.width)
    }

    /// [ 0 <= start_index < width <= screen.max_width() ]
    pub fn set_clipping(&mut self, start_index: LineCellIndex, width: usize) {
        assert!(start_index < self.max_width);
        assert!(width <= self.max_width);
        assert!(start_index + width <= self.max_width);

        self.start_index = start_index;
        self.width = start_index + width;
        self.hash_cache = 0;
        self.nb_cells = 0;
        self.read_only = false;
    }

    pub fn clear_clip_width(&mut self) {
        self.set_clipping(0, self.max_width);
    }

    pub fn push(&mut self, cpi: CodepointInfo) -> (bool, LineCellIndex) {
        if self.nb_cells < self.width() && !self.read_only {
            self.cells[self.start_index + self.nb_cells].cpi = cpi;
            self.cells[self.start_index + self.nb_cells].is_used = true;

            self.nb_cells += 1;

            if self.nb_cells == self.width {
                self.read_only = true;
            }

            (true, self.nb_cells)
        } else {
            self.read_only = true;
            (false, self.nb_cells)
        }
    }

    pub fn clear(&mut self) {
        for i in 0..self.width() {
            self.cells[self.start_index + i] = LineCell::new();
        }

        self.hash_cache = 0;
        self.nb_cells = 0;
        self.read_only = false;
        self.metadata = false;
        self.hash_cache = 0;
    }

    pub fn get_first_cpi(&self) -> Option<&CodepointInfo> {
        if self.start_index < self.width {
            Some(&self.cells[self.start_index].cpi)
        } else {
            // internal error
            None
        }
    }

    pub fn get_last_cpi(&self) -> Option<&CodepointInfo> {
        if self.nb_cells > 0 {
            Some(&self.cells[self.start_index + self.nb_cells - 1].cpi)
        } else {
            None
        }
    }

    pub fn get_unclipped_cpi(&self, index: LineCellIndex) -> Option<&CodepointInfo> {
        if index < self.max_width() {
            Some(&self.cells[index].cpi)
        } else {
            None
        }
    }

    pub fn get_mut_unclipped_cpi(&mut self, index: LineCellIndex) -> Option<&mut CodepointInfo> {
        if index < self.max_width() {
            Some(&mut self.cells[index].cpi)
        } else {
            None
        }
    }

    pub fn get_cpi(&self, index: LineCellIndex) -> Option<&CodepointInfo> {
        if index < self.width() {
            Some(&self.cells[self.start_index + index].cpi)
        } else {
            None
        }
    }

    pub fn get_mut_cpi(&mut self, index: LineCellIndex) -> Option<&mut CodepointInfo> {
        if index < self.width() {
            Some(&mut self.cells[self.start_index + index].cpi)
        } else {
            None
        }
    }

    pub fn at(&mut self, index: LineCellIndex) -> Option<&mut CodepointInfo> {
        self.get_mut_cpi(index)
    }

    pub fn get_used_cpi(&self, index: LineCellIndex) -> Option<&CodepointInfo> {
        if index < self.nb_cells {
            self.get_cpi(index)
        } else {
            None
        }
    }

    pub fn get_mut_used_cpi(&mut self, index: LineCellIndex) -> Option<&mut CodepointInfo> {
        if index < self.nb_cells {
            self.get_mut_cpi(index)
        } else {
            None
        }
    }

    pub fn get_used_cpi_clipped(
        &self,
        index: LineCellIndex,
    ) -> (Option<&CodepointInfo>, LineCellIndex) {
        if self.nb_cells == 0 {
            return (None, 0);
        }

        let index = ::std::cmp::min(index, self.nb_cells - 1);
        (self.get_used_cpi(index), index)
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::mem::transmute;
use std::io::{Cursor, Read, Write};
use std::default::Default;
use std::iter::FromIterator;

use rand;
use mio::TryRead;

use result::{Result, Error, Kind};
use protocol::{OpCode, CloseCode};

fn apply_mask(buf: &mut [u8], mask: &[u8; 4]) {
    let iter = buf.iter_mut().zip(mask.iter().cycle());
    for (byte, &key) in iter {
        *byte ^= key
    }
}

#[inline]
fn generate_mask() -> [u8; 4] {
    unsafe { transmute(rand::random::<u32>()) }
}

/// A struct representing a WebSocket frame.
#[derive(Debug, Clone)]
pub struct Frame {
    finished: bool,
    rsv1: bool,
    rsv2: bool,
    rsv3: bool,
    opcode: OpCode,

    mask: Option<[u8; 4]>,

    payload: Vec<u8>,
}

impl Frame {

    /// Get the length of the frame.
    /// This is the length of the header + the length of the payload.
    #[inline]
    pub fn len(&self) -> usize {
        let mut header_length = 2;
        let payload_len = self.payload().len();
        if payload_len > 125 {
            if payload_len <= u16::max_value() as usize {
                header_length += 2;
            } else {
                header_length += 8;
            }
        }

        if self.is_masked() {
            header_length += 4;
        }

        header_length + payload_len
    }

    /// Test whether the frame is a final frame.
    #[inline]
    pub fn is_final(&self) -> bool {
        self.finished
    }

    /// Test whether the first reserved bit is set.
    #[inline]
    pub fn has_rsv1(&self) -> bool {
        self.rsv1
    }

    /// Test whether the second reserved bit is set.
    #[inline]
    pub fn has_rsv2(&self) -> bool {
        self.rsv2
    }

    /// Test whether the third reserved bit is set.
    #[inline]
    pub fn has_rsv3(&self) -> bool {
        self.rsv3
    }

    /// Get the OpCode of the frame.
    #[inline]
    pub fn opcode(&self) -> OpCode {
        self.opcode
    }

    /// Get a reference to the frame's payload.
    #[inline]
    pub fn payload(&self) -> &Vec<u8> {
        &self.payload
    }

    // Test whether the frame is masked.
    #[doc(hidden)]
    #[inline]
    pub fn is_masked(&self) -> bool {
        self.mask.is_some()
    }

    // Get an optional reference to the frame's mask.
    #[doc(hidden)]
    #[allow(dead_code)]
    #[inline]
    pub fn mask(&self) -> Option<&[u8; 4]> {
        self.mask.as_ref()
    }

    /// Make this frame a final frame.
    #[allow(dead_code)]
    #[inline]
    pub fn set_final(&mut self, is_final: bool) -> &mut Frame {
        self.finished = is_final;
        self
    }

    /// Set the first reserved bit.
    #[allow(dead_code)]
    #[inline]
    pub fn set_rsv1(&mut self, has_rsv1: bool) -> &mut Frame {
        self.rsv1 = has_rsv1;
        self
    }

    /// Set the second reserved bit.
    #[allow(dead_code)]
    #[inline]
    pub fn set_rsv2(&mut self, has_rsv2: bool) -> &mut Frame {
        self.rsv2 = has_rsv2;
        self
    }

    /// Set the third reserved bit.
    #[allow(dead_code)]
    #[inline]
    pub fn set_rsv3(&mut self, has_rsv3: bool) -> &mut Frame {
        self.rsv1 = has_rsv3;
        self
    }

    /// Set the OpCode.
    #[allow(dead_code)]
    #[inline]
    pub fn set_opcode(&mut self, opcode: OpCode) -> &mut Frame {
        self.opcode = opcode;
        self
    }

    /// Edit the frame's payload.
    #[allow(dead_code)]
    #[inline]
    pub fn payload_mut(&mut self) -> &mut Vec<u8> {
        &mut self.payload
    }

    // Generate a new mask for this frame.
    //
    // This method simply generates and stores the mask. It does not change the payload data.
    // Instead, the payload data will be masked with the generated mask when the frame is sent
    // to the other endpoint.
    #[doc(hidden)]
    #[inline]
    pub fn set_mask(&mut self) -> &mut Frame {
        self.mask = Some(generate_mask());
        self
    }

    // This method unmasks the payload and should only be called on frames that are actually
    // masked. In other words, those frames that have just been received from a client endpoint.
    #[doc(hidden)]
    #[inline]
    pub fn remove_mask(&mut self) -> &mut Frame {
        self.mask.and_then(|mask| {
            Some(apply_mask(&mut self.payload, &mask))
        });
        self.mask = None;
        self
    }

    /// Consume the frame into its payload.
    pub fn into_data(self) -> Vec<u8> {
        self.payload
    }

    /// Create a new data frame.
    #[inline]
    pub fn message(data: Vec<u8>, code: OpCode, finished: bool) -> Frame {
        debug_assert!(match code {
            OpCode::Text | OpCode::Binary | OpCode::Continue => true,
            _ => false,
        }, "Invalid opcode for data frame.");

        Frame {
            finished: finished,
            opcode: code,
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Pong control frame.
    #[inline]
    pub fn pong(data: Vec<u8>) -> Frame {
        Frame {
            opcode: OpCode::Pong,
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Ping control frame.
    #[inline]
    pub fn ping(data: Vec<u8>) -> Frame {
        Frame {
            opcode: OpCode::Ping,
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Close control frame.
    #[inline]
    pub fn close(code: CloseCode, reason: &str) -> Frame {
        let raw: [u8; 2] = unsafe {
            let u: u16 = code.into();
            transmute(u.to_be())
        };

        let payload = if let CloseCode::Empty = code {
            Vec::new()
        } else {
            Vec::from_iter(
                raw[..].iter()
                       .chain(reason.as_bytes().iter())
                       .map(|&b| b))
        };

        Frame {
            payload: payload,
            .. Frame::default()
        }
    }

    /// Parse the input stream into a frame.
    pub fn parse(cursor: &mut Cursor<Vec<u8>>) -> Result<Option<Frame>> {
        let size = cursor.get_ref().len() - cursor.position() as usize;
        let initial = cursor.position();
        trace!("Position in buffer {}", initial);

        let mut head = [0u8; 2];
        if try!(cursor.read(&mut head)) != 2 {
            cursor.set_position(initial);
            return Ok(None)
        }

        trace!("Parsed headers {:?}", head);

        let first = head[0];
        let second = head[1];
        trace!("First: {:b}", first);
        trace!("Second: {:b}", second);

        let finished = first & 0x80 != 0;

        let rsv1 = first & 0x40 != 0;
        let rsv2 = first & 0x20 != 0;
        let rsv3 = first & 0x10 != 0;

        let opcode = OpCode::from(first & 0x0F);
        trace!("Opcode: {:?}", opcode);

        let masked = second & 0x80 != 0;
        trace!("Masked: {:?}", masked);

        let mut header_length = 2;

        let mut length = (second & 0x7F) as u64;

        if length == 126 {
            let mut length_bytes = [0u8; 2];
            if try!(cursor.read(&mut length_bytes)) != 2 {
                cursor.set_position(initial);
                return Ok(None)
            }

            length = unsafe {
                let mut wide: u16 = transmute(length_bytes);
                wide = u16::from_be(wide);
                wide
            } as u64;
            header_length += 2;
        } else if length == 127 {
            let mut length_bytes = [0u8; 8];
            if try!(cursor.read(&mut length_bytes)) != 8 {
                cursor.set_position(initial);
                return Ok(None)
            }

            unsafe { length = transmute(length_bytes); }
            length = u64::from_be(length);
            header_length += 8;
        }
        trace!("Payload length: {}", length);

        let mask = if masked {
            let mut mask_bytes = [0u8; 4];
            if try!(cursor.read(&mut mask_bytes)) != 4 {
                cursor.set_position(initial);
                return Ok(None)
            } else {
                header_length += 4;
                Some(mask_bytes)
            }
        } else {
            None
        };

        if size < length as usize + header_length {
            cursor.set_position(initial);
            return Ok(None)
        }

        let mut data = Vec::with_capacity(length as usize);
        if length > 0 {
            if let Some(read) = try!(cursor.try_read_buf(&mut data)) {
                debug_assert!(read == length as usize, "Read incorrect payload length!");
            }
        }

        // Disallow bad opcode
        if let OpCode::Bad = opcode {
            return Err(Error::new(Kind::Protocol, format!("Encountered invalid opcode: {}", first & 0x0F)))
        }

        // control frames must have length <= 125
        match opcode {
            OpCode::Ping | OpCode::Pong if length > 125 => {
                return Err(Error::new(Kind::Protocol, format!("Rejected WebSocket handshake.Received control frame with length: {}.", length)))
            }
            OpCode::Close if length > 125 => {
                debug!("Received close frame with payload length exceeding 125. Morphing to protocol close frame.");
                return Ok(Some(Frame::close(CloseCode::Protocol, "Received close frame with payload length exceeding 125.")))
            }
            _ => ()
        }

        let frame = Frame {
            finished: finished,
            rsv1: rsv1,
            rsv2: rsv2,
            rsv3: rsv3,
            opcode: opcode,
            mask: mask,
            payload: data,
        };

        Ok(Some(frame))
    }

    /// Write a frame out to a buffer
    pub fn format<W>(&mut self, w: &mut W) -> Result<()>
        where W: Write
    {
        let mut one = 0u8;
        let code: u8 = self.opcode.into();
        if self.is_final() {
            one |= 0x80;
        }
        if self.has_rsv1() {
            one |= 0x40;
        }
        if self.has_rsv2() {
            one |= 0x20;
        }
        if self.has_rsv3() {
            one |= 0x10;
        }
        one |= code;

        let mut two = 0u8;

        if self.is_masked() {
            two |= 0x80;
        }

        if self.payload.len() < 126 {
            two |= self.payload.len() as u8;
            let headers = [one, two];
            try!(w.write(&headers));
        } else if self.payload.len() <= 65535 {
            two |= 126;
            let length_bytes: [u8; 2] = unsafe {
                let short = self.payload.len() as u16;
                transmute(short.to_be())
            };
            let headers = [one, two, length_bytes[0], length_bytes[1]];
            try!(w.write(&headers));
        } else {
            two |= 127;
            let length_bytes: [u8; 8] = unsafe {
                let long = self.payload.len() as u64;
                transmute(long.to_be())
            };
            let headers = [
                one,
                two,
                length_bytes[0],
                length_bytes[1],
                length_bytes[2],
                length_bytes[3],
                length_bytes[4],
                length_bytes[5],
                length_bytes[6],
                length_bytes[7],
            ];
            try!(w.write(&headers));
        }

        if self.is_masked() {
            let mask = self.mask.take().unwrap();
            apply_mask(&mut self.payload, &mask);
            try!(w.write(&mask));
        }

        try!(w.write(&self.payload));
        Ok(())
    }
}

impl Default for Frame {
    fn default() -> Frame {
        Frame {
            finished: true,
            rsv1: false,
            rsv2: false,
            rsv3: false,
            opcode: OpCode::Close,
            mask: None,
            payload: Vec::new(),
        }
    }
}

impl fmt::Display for Frame {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
            "
<FRAME>
final: {}
reserved: {} {} {}
opcode: {}
mask: {}
length: {}
payload length: {}
payload: {:?}
            ",
            self.finished,
            self.rsv1,
            self.rsv2,
            self.rsv3,
            self.opcode,
            self.mask.map(|mask| format!("{:?}", mask)).unwrap_or("NONE".into()),
            self.len(),
            self.payload.len(),
            self.payload)
    }
}

mod test {
    #![allow(unused_imports, unused_variables, dead_code)]
    use super::*;
    use protocol::OpCode;

    #[test]
    fn test_display_frame() {
        let f = Frame::message("hi there".into(), OpCode::Text, true);
        let view = format!("{}", f);
        view.contains("payload:");
    }
}
<commit_msg>refactor(frame): Format payload as hexadecimal<commit_after>use std::fmt;
use std::mem::transmute;
use std::io::{Cursor, Read, Write};
use std::default::Default;
use std::iter::FromIterator;

use rand;
use mio::TryRead;

use result::{Result, Error, Kind};
use protocol::{OpCode, CloseCode};

fn apply_mask(buf: &mut [u8], mask: &[u8; 4]) {
    let iter = buf.iter_mut().zip(mask.iter().cycle());
    for (byte, &key) in iter {
        *byte ^= key
    }
}

#[inline]
fn generate_mask() -> [u8; 4] {
    unsafe { transmute(rand::random::<u32>()) }
}

/// A struct representing a WebSocket frame.
#[derive(Debug, Clone)]
pub struct Frame {
    finished: bool,
    rsv1: bool,
    rsv2: bool,
    rsv3: bool,
    opcode: OpCode,

    mask: Option<[u8; 4]>,

    payload: Vec<u8>,
}

impl Frame {

    /// Get the length of the frame.
    /// This is the length of the header + the length of the payload.
    #[inline]
    pub fn len(&self) -> usize {
        let mut header_length = 2;
        let payload_len = self.payload().len();
        if payload_len > 125 {
            if payload_len <= u16::max_value() as usize {
                header_length += 2;
            } else {
                header_length += 8;
            }
        }

        if self.is_masked() {
            header_length += 4;
        }

        header_length + payload_len
    }

    /// Test whether the frame is a final frame.
    #[inline]
    pub fn is_final(&self) -> bool {
        self.finished
    }

    /// Test whether the first reserved bit is set.
    #[inline]
    pub fn has_rsv1(&self) -> bool {
        self.rsv1
    }

    /// Test whether the second reserved bit is set.
    #[inline]
    pub fn has_rsv2(&self) -> bool {
        self.rsv2
    }

    /// Test whether the third reserved bit is set.
    #[inline]
    pub fn has_rsv3(&self) -> bool {
        self.rsv3
    }

    /// Get the OpCode of the frame.
    #[inline]
    pub fn opcode(&self) -> OpCode {
        self.opcode
    }

    /// Get a reference to the frame's payload.
    #[inline]
    pub fn payload(&self) -> &Vec<u8> {
        &self.payload
    }

    // Test whether the frame is masked.
    #[doc(hidden)]
    #[inline]
    pub fn is_masked(&self) -> bool {
        self.mask.is_some()
    }

    // Get an optional reference to the frame's mask.
    #[doc(hidden)]
    #[allow(dead_code)]
    #[inline]
    pub fn mask(&self) -> Option<&[u8; 4]> {
        self.mask.as_ref()
    }

    /// Make this frame a final frame.
    #[allow(dead_code)]
    #[inline]
    pub fn set_final(&mut self, is_final: bool) -> &mut Frame {
        self.finished = is_final;
        self
    }

    /// Set the first reserved bit.
    #[allow(dead_code)]
    #[inline]
    pub fn set_rsv1(&mut self, has_rsv1: bool) -> &mut Frame {
        self.rsv1 = has_rsv1;
        self
    }

    /// Set the second reserved bit.
    #[allow(dead_code)]
    #[inline]
    pub fn set_rsv2(&mut self, has_rsv2: bool) -> &mut Frame {
        self.rsv2 = has_rsv2;
        self
    }

    /// Set the third reserved bit.
    #[allow(dead_code)]
    #[inline]
    pub fn set_rsv3(&mut self, has_rsv3: bool) -> &mut Frame {
        self.rsv1 = has_rsv3;
        self
    }

    /// Set the OpCode.
    #[allow(dead_code)]
    #[inline]
    pub fn set_opcode(&mut self, opcode: OpCode) -> &mut Frame {
        self.opcode = opcode;
        self
    }

    /// Edit the frame's payload.
    #[allow(dead_code)]
    #[inline]
    pub fn payload_mut(&mut self) -> &mut Vec<u8> {
        &mut self.payload
    }

    // Generate a new mask for this frame.
    //
    // This method simply generates and stores the mask. It does not change the payload data.
    // Instead, the payload data will be masked with the generated mask when the frame is sent
    // to the other endpoint.
    #[doc(hidden)]
    #[inline]
    pub fn set_mask(&mut self) -> &mut Frame {
        self.mask = Some(generate_mask());
        self
    }

    // This method unmasks the payload and should only be called on frames that are actually
    // masked. In other words, those frames that have just been received from a client endpoint.
    #[doc(hidden)]
    #[inline]
    pub fn remove_mask(&mut self) -> &mut Frame {
        self.mask.and_then(|mask| {
            Some(apply_mask(&mut self.payload, &mask))
        });
        self.mask = None;
        self
    }

    /// Consume the frame into its payload.
    pub fn into_data(self) -> Vec<u8> {
        self.payload
    }

    /// Create a new data frame.
    #[inline]
    pub fn message(data: Vec<u8>, code: OpCode, finished: bool) -> Frame {
        debug_assert!(match code {
            OpCode::Text | OpCode::Binary | OpCode::Continue => true,
            _ => false,
        }, "Invalid opcode for data frame.");

        Frame {
            finished: finished,
            opcode: code,
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Pong control frame.
    #[inline]
    pub fn pong(data: Vec<u8>) -> Frame {
        Frame {
            opcode: OpCode::Pong,
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Ping control frame.
    #[inline]
    pub fn ping(data: Vec<u8>) -> Frame {
        Frame {
            opcode: OpCode::Ping,
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Close control frame.
    #[inline]
    pub fn close(code: CloseCode, reason: &str) -> Frame {
        let raw: [u8; 2] = unsafe {
            let u: u16 = code.into();
            transmute(u.to_be())
        };

        let payload = if let CloseCode::Empty = code {
            Vec::new()
        } else {
            Vec::from_iter(
                raw[..].iter()
                       .chain(reason.as_bytes().iter())
                       .map(|&b| b))
        };

        Frame {
            payload: payload,
            .. Frame::default()
        }
    }

    /// Parse the input stream into a frame.
    pub fn parse(cursor: &mut Cursor<Vec<u8>>) -> Result<Option<Frame>> {
        let size = cursor.get_ref().len() - cursor.position() as usize;
        let initial = cursor.position();
        trace!("Position in buffer {}", initial);

        let mut head = [0u8; 2];
        if try!(cursor.read(&mut head)) != 2 {
            cursor.set_position(initial);
            return Ok(None)
        }

        trace!("Parsed headers {:?}", head);

        let first = head[0];
        let second = head[1];
        trace!("First: {:b}", first);
        trace!("Second: {:b}", second);

        let finished = first & 0x80 != 0;

        let rsv1 = first & 0x40 != 0;
        let rsv2 = first & 0x20 != 0;
        let rsv3 = first & 0x10 != 0;

        let opcode = OpCode::from(first & 0x0F);
        trace!("Opcode: {:?}", opcode);

        let masked = second & 0x80 != 0;
        trace!("Masked: {:?}", masked);

        let mut header_length = 2;

        let mut length = (second & 0x7F) as u64;

        if length == 126 {
            let mut length_bytes = [0u8; 2];
            if try!(cursor.read(&mut length_bytes)) != 2 {
                cursor.set_position(initial);
                return Ok(None)
            }

            length = unsafe {
                let mut wide: u16 = transmute(length_bytes);
                wide = u16::from_be(wide);
                wide
            } as u64;
            header_length += 2;
        } else if length == 127 {
            let mut length_bytes = [0u8; 8];
            if try!(cursor.read(&mut length_bytes)) != 8 {
                cursor.set_position(initial);
                return Ok(None)
            }

            unsafe { length = transmute(length_bytes); }
            length = u64::from_be(length);
            header_length += 8;
        }
        trace!("Payload length: {}", length);

        let mask = if masked {
            let mut mask_bytes = [0u8; 4];
            if try!(cursor.read(&mut mask_bytes)) != 4 {
                cursor.set_position(initial);
                return Ok(None)
            } else {
                header_length += 4;
                Some(mask_bytes)
            }
        } else {
            None
        };

        if size < length as usize + header_length {
            cursor.set_position(initial);
            return Ok(None)
        }

        let mut data = Vec::with_capacity(length as usize);
        if length > 0 {
            if let Some(read) = try!(cursor.try_read_buf(&mut data)) {
                debug_assert!(read == length as usize, "Read incorrect payload length!");
            }
        }

        // Disallow bad opcode
        if let OpCode::Bad = opcode {
            return Err(Error::new(Kind::Protocol, format!("Encountered invalid opcode: {}", first & 0x0F)))
        }

        // control frames must have length <= 125
        match opcode {
            OpCode::Ping | OpCode::Pong if length > 125 => {
                return Err(Error::new(Kind::Protocol, format!("Rejected WebSocket handshake.Received control frame with length: {}.", length)))
            }
            OpCode::Close if length > 125 => {
                debug!("Received close frame with payload length exceeding 125. Morphing to protocol close frame.");
                return Ok(Some(Frame::close(CloseCode::Protocol, "Received close frame with payload length exceeding 125.")))
            }
            _ => ()
        }

        let frame = Frame {
            finished: finished,
            rsv1: rsv1,
            rsv2: rsv2,
            rsv3: rsv3,
            opcode: opcode,
            mask: mask,
            payload: data,
        };

        Ok(Some(frame))
    }

    /// Write a frame out to a buffer
    pub fn format<W>(&mut self, w: &mut W) -> Result<()>
        where W: Write
    {
        let mut one = 0u8;
        let code: u8 = self.opcode.into();
        if self.is_final() {
            one |= 0x80;
        }
        if self.has_rsv1() {
            one |= 0x40;
        }
        if self.has_rsv2() {
            one |= 0x20;
        }
        if self.has_rsv3() {
            one |= 0x10;
        }
        one |= code;

        let mut two = 0u8;

        if self.is_masked() {
            two |= 0x80;
        }

        if self.payload.len() < 126 {
            two |= self.payload.len() as u8;
            let headers = [one, two];
            try!(w.write(&headers));
        } else if self.payload.len() <= 65535 {
            two |= 126;
            let length_bytes: [u8; 2] = unsafe {
                let short = self.payload.len() as u16;
                transmute(short.to_be())
            };
            let headers = [one, two, length_bytes[0], length_bytes[1]];
            try!(w.write(&headers));
        } else {
            two |= 127;
            let length_bytes: [u8; 8] = unsafe {
                let long = self.payload.len() as u64;
                transmute(long.to_be())
            };
            let headers = [
                one,
                two,
                length_bytes[0],
                length_bytes[1],
                length_bytes[2],
                length_bytes[3],
                length_bytes[4],
                length_bytes[5],
                length_bytes[6],
                length_bytes[7],
            ];
            try!(w.write(&headers));
        }

        if self.is_masked() {
            let mask = self.mask.take().unwrap();
            apply_mask(&mut self.payload, &mask);
            try!(w.write(&mask));
        }

        try!(w.write(&self.payload));
        Ok(())
    }
}

impl Default for Frame {
    fn default() -> Frame {
        Frame {
            finished: true,
            rsv1: false,
            rsv2: false,
            rsv3: false,
            opcode: OpCode::Close,
            mask: None,
            payload: Vec::new(),
        }
    }
}

impl fmt::Display for Frame {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
            "
<FRAME>
final: {}
reserved: {} {} {}
opcode: {}
mask: {}
length: {}
payload length: {}
payload: 0x{}
            ",
            self.finished,
            self.rsv1,
            self.rsv2,
            self.rsv3,
            self.opcode,
            self.mask.map(|mask| format!("{:?}", mask)).unwrap_or("NONE".into()),
            self.len(),
            self.payload.len(),
            self.payload.iter().map(|byte| format!("{:x}", byte)).collect::<String>())
    }
}

mod test {
    #![allow(unused_imports, unused_variables, dead_code)]
    use super::*;
    use protocol::OpCode;

    #[test]
    fn test_display_frame() {
        let f = Frame::message("hi there".into(), OpCode::Text, true);
        let view = format!("{}", f);
        view.contains("payload:");
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::mem::transmute;
use std::io::{Cursor, Read, Write};
use std::default::Default;
use std::iter::FromIterator;

use rand;
use mio::TryRead;

use result::{Result, Error, Kind};
use protocol::{OpCode, CloseCode};

fn apply_mask(buf: &mut [u8], mask: &[u8; 4]) {
    let iter = buf.iter_mut().zip(mask.iter().cycle());
    for (byte, &key) in iter {
        *byte ^= key
    }
}

#[inline]
fn generate_mask() -> [u8; 4] {
    unsafe { transmute(rand::random::<u32>()) }
}

/// A struct representing a WebSocket frame.
#[derive(Debug, Clone)]
pub struct Frame {
    finished: bool,
    rsv1: bool,
    rsv2: bool,
    rsv3: bool,
    opcode: OpCode,

    mask: Option<[u8; 4]>,

    payload: Vec<u8>,
}

impl Frame {

    /// Get the length of the frame.
    /// This is the length of the header + the length of the payload.
    #[inline]
    pub fn len(&self) -> usize {
        let mut header_length = 2;
        let payload_len = self.payload().len();
        if payload_len > 125 {
            if payload_len <= u16::max_value() as usize {
                header_length += 2;
            } else {
                header_length += 8;
            }
        }

        if self.is_masked() {
            header_length += 4;
        }

        header_length + payload_len
    }

    /// Test whether the frame is a final frame.
    #[inline]
    pub fn is_final(&self) -> bool {
        self.finished
    }

    /// Test whether the first reserved bit is set.
    #[inline]
    pub fn has_rsv1(&self) -> bool {
        self.rsv1
    }

    /// Test whether the second reserved bit is set.
    #[inline]
    pub fn has_rsv2(&self) -> bool {
        self.rsv2
    }

    /// Test whether the third reserved bit is set.
    #[inline]
    pub fn has_rsv3(&self) -> bool {
        self.rsv3
    }

    /// Get the OpCode of the frame.
    #[inline]
    pub fn opcode(&self) -> OpCode {
        self.opcode
    }

    /// Get a reference to the frame's payload.
    #[inline]
    pub fn payload(&self) -> &Vec<u8> {
        &self.payload
    }

    // Test whether the frame is masked.
    #[doc(hidden)]
    #[inline]
    pub fn is_masked(&self) -> bool {
        self.mask.is_some()
    }

    // Get an optional reference to the frame's mask.
    #[doc(hidden)]
    #[allow(dead_code)]
    #[inline]
    pub fn mask(&self) -> Option<&[u8; 4]> {
        self.mask.as_ref()
    }

    /// Make this frame a final frame.
    #[allow(dead_code)]
    #[inline]
    pub fn set_final(&mut self, is_final: bool) -> &mut Frame {
        self.finished = is_final;
        self
    }

    /// Set the first reserved bit.
    #[allow(dead_code)]
    #[inline]
    pub fn set_rsv1(&mut self, has_rsv1: bool) -> &mut Frame {
        self.rsv1 = has_rsv1;
        self
    }

    /// Set the second reserved bit.
    #[allow(dead_code)]
    #[inline]
    pub fn set_rsv2(&mut self, has_rsv2: bool) -> &mut Frame {
        self.rsv2 = has_rsv2;
        self
    }

    /// Set the third reserved bit.
    #[allow(dead_code)]
    #[inline]
    pub fn set_rsv3(&mut self, has_rsv3: bool) -> &mut Frame {
        self.rsv1 = has_rsv3;
        self
    }

    /// Set the OpCode.
    #[allow(dead_code)]
    #[inline]
    pub fn set_opcode(&mut self, opcode: OpCode) -> &mut Frame {
        self.opcode = opcode;
        self
    }

    /// Edit the frame's payload.
    #[allow(dead_code)]
    #[inline]
    pub fn payload_mut(&mut self) -> &mut Vec<u8> {
        &mut self.payload
    }

    // Generate a new mask for this frame.
    //
    // This method simply generates and stores the mask. It does not change the payload data.
    // Instead, the payload data will be masked with the generated mask when the frame is sent
    // to the other endpoint.
    #[doc(hidden)]
    #[inline]
    pub fn set_mask(&mut self) -> &mut Frame {
        self.mask = Some(generate_mask());
        self
    }

    // This method unmasks the payload and should only be called on frames that are actually
    // masked. In other words, those frames that have just been received from a client endpoint.
    #[doc(hidden)]
    #[inline]
    pub fn remove_mask(&mut self) -> &mut Frame {
        self.mask.and_then(|mask| {
            Some(apply_mask(&mut self.payload, &mask))
        });
        self.mask = None;
        self
    }

    /// Consume the frame into its payload.
    pub fn into_data(self) -> Vec<u8> {
        self.payload
    }

    /// Create a new data frame.
    #[inline]
    pub fn message(data: Vec<u8>, code: OpCode, finished: bool) -> Frame {
        debug_assert!(match code {
            OpCode::Text | OpCode::Binary | OpCode::Continue => true,
            _ => false,
        }, "Invalid opcode for data frame.");

        Frame {
            finished: finished,
            opcode: code,
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Pong control frame.
    #[inline]
    pub fn pong(data: Vec<u8>) -> Frame {
        Frame {
            opcode: OpCode::Pong,
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Ping control frame.
    #[inline]
    pub fn ping(data: Vec<u8>) -> Frame {
        Frame {
            opcode: OpCode::Ping,
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Close control frame.
    #[inline]
    pub fn close(code: CloseCode, reason: &str) -> Frame {
        let raw: [u8; 2] = unsafe {
            let u: u16 = code.into();
            transmute(u.to_be())
        };

        let payload = if let CloseCode::Empty = code {
            Vec::new()
        } else {
            Vec::from_iter(
                raw[..].iter()
                       .chain(reason.as_bytes().iter())
                       .map(|&b| b))
        };

        Frame {
            payload: payload,
            .. Frame::default()
        }
    }

    /// Parse the input stream into a frame.
    pub fn parse(cursor: &mut Cursor<Vec<u8>>) -> Result<Option<Frame>> {
        let size = cursor.get_ref().len() - cursor.position() as usize;
        let initial = cursor.position();
        trace!("Position in buffer {}", initial);

        let mut head = [0u8; 2];
        if try!(cursor.read(&mut head)) != 2 {
            cursor.set_position(initial);
            return Ok(None)
        }

        trace!("Parsed headers {:?}", head);

        let first = head[0];
        let second = head[1];
        trace!("First: {:b}", first);
        trace!("Second: {:b}", second);

        let finished = first & 0x80 != 0;

        let rsv1 = first & 0x40 != 0;
        let rsv2 = first & 0x20 != 0;
        let rsv3 = first & 0x10 != 0;

        let opcode = OpCode::from(first & 0x0F);
        trace!("Opcode: {:?}", opcode);
        if let OpCode::Bad = opcode {
            return Err(Error::new(Kind::Protocol, format!("Encountered invalid opcode: {}", first & 0x0F)))
        }
        let masked = second & 0x80 != 0;
        trace!("Masked: {:?}", masked);

        let mut header_length = 2;

        let mut length = (second & 0x7F) as u64;

        if length == 126 {
            let mut length_bytes = [0u8; 2];
            if try!(cursor.read(&mut length_bytes)) != 2 {
                cursor.set_position(initial);
                return Ok(None)
            }

            length = unsafe {
                let mut wide: u16 = transmute(length_bytes);
                wide = u16::from_be(wide);
                wide
            } as u64;
            header_length += 2;
        } else if length == 127 {
            let mut length_bytes = [0u8; 8];
            if try!(cursor.read(&mut length_bytes)) != 8 {
                cursor.set_position(initial);
                return Ok(None)
            }

            unsafe { length = transmute(length_bytes); }
            length = u64::from_be(length);
            header_length += 8;
        }
        trace!("Payload length: {}", length);

        // control frames must have length <= 125
        match opcode {
            OpCode::Close | OpCode::Ping | OpCode::Pong if length > 125 => {
                return Err(Error::new(Kind::Protocol, format!("Rejected WebSocket handshake.Received control frame with length: {}.", length)))
            }
            _ => ()
        }

        let mask = if masked {
            let mut mask_bytes = [0u8; 4];
            if try!(cursor.read(&mut mask_bytes)) != 4 {
                cursor.set_position(initial);
                return Ok(None)
            } else {
                header_length += 4;
                Some(mask_bytes)
            }
        } else {
            None
        };

        if size < length as usize + header_length {
            cursor.set_position(initial);
            return Ok(None)
        }

        let mut data = Vec::with_capacity(length as usize);
        if length > 0 {
            // It is ok to unwrap here because this cursor won't block
            let read = try!(cursor.try_read_buf(&mut data)).unwrap();
            debug_assert!(read == length as usize, "Read incorrect payload length!");
        }

        let frame = Frame {
            finished: finished,
            rsv1: rsv1,
            rsv2: rsv2,
            rsv3: rsv3,
            opcode: opcode,
            mask: mask,
            payload: data,
        };
        Ok(Some(frame))
    }

    /// Write a frame out to a buffer
    pub fn format<W>(&mut self, w: &mut W) -> Result<()>
        where W: Write
    {
        let mut one = 0u8;
        let code: u8 = self.opcode.into();
        if self.is_final() {
            one |= 0x80;
        }
        if self.has_rsv1() {
            one |= 0x40;
        }
        if self.has_rsv2() {
            one |= 0x20;
        }
        if self.has_rsv3() {
            one |= 0x10;
        }
        one |= code;

        let mut two = 0u8;

        if self.is_masked() {
            two |= 0x80;
        }

        if self.payload.len() < 126 {
            two |= self.payload.len() as u8;
            let headers = [one, two];
            try!(w.write(&headers));
        } else if self.payload.len() <= 65535 {
            two |= 126;
            let length_bytes: [u8; 2] = unsafe {
                let short = self.payload.len() as u16;
                transmute(short.to_be())
            };
            let headers = [one, two, length_bytes[0], length_bytes[1]];
            try!(w.write(&headers));
        } else {
            two |= 127;
            let length_bytes: [u8; 8] = unsafe {
                transmute(self.payload.len().to_be())
            };
            let headers = [
                one,
                two,
                length_bytes[0],
                length_bytes[1],
                length_bytes[2],
                length_bytes[3],
                length_bytes[4],
                length_bytes[5],
                length_bytes[6],
                length_bytes[7],
            ];
            try!(w.write(&headers));
        }

        if self.is_masked() {
            let mask = self.mask.take().unwrap();
            apply_mask(&mut self.payload, &mask);
            try!(w.write(&mask));
        }

        try!(w.write(&self.payload));
        Ok(())
    }
}

impl Default for Frame {
    fn default() -> Frame {
        Frame {
            finished: true,
            rsv1: false,
            rsv2: false,
            rsv3: false,
            opcode: OpCode::Close,
            mask: None,
            payload: Vec::new(),
        }
    }
}

impl fmt::Display for Frame {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
            "
<FRAME>
final: {}
reserved: {} {} {}
opcode: {}
mask: {}
payload: {:?}
            ",
            self.finished,
            self.rsv1,
            self.rsv2,
            self.rsv3,
            self.opcode,
            self.mask.map(|mask| format!("{:?}", mask)).unwrap_or("NONE".into()),
            self.payload)
    }
}

mod test {
    #![allow(unused_imports, unused_variables, dead_code)]
    use super::*;
    use protocol::OpCode;

    #[test]
    fn test_display_frame() {
        let f = Frame::message("hi there".into(), OpCode::Text, true);
        let view = format!("{}", f);
        view.contains("payload:");
    }
}
<commit_msg>fix(frame): fail to compile on 32bit, close #20<commit_after>use std::fmt;
use std::mem::transmute;
use std::io::{Cursor, Read, Write};
use std::default::Default;
use std::iter::FromIterator;

use rand;
use mio::TryRead;

use result::{Result, Error, Kind};
use protocol::{OpCode, CloseCode};

fn apply_mask(buf: &mut [u8], mask: &[u8; 4]) {
    let iter = buf.iter_mut().zip(mask.iter().cycle());
    for (byte, &key) in iter {
        *byte ^= key
    }
}

#[inline]
fn generate_mask() -> [u8; 4] {
    unsafe { transmute(rand::random::<u32>()) }
}

/// A struct representing a WebSocket frame.
#[derive(Debug, Clone)]
pub struct Frame {
    finished: bool,
    rsv1: bool,
    rsv2: bool,
    rsv3: bool,
    opcode: OpCode,

    mask: Option<[u8; 4]>,

    payload: Vec<u8>,
}

impl Frame {

    /// Get the length of the frame.
    /// This is the length of the header + the length of the payload.
    #[inline]
    pub fn len(&self) -> usize {
        let mut header_length = 2;
        let payload_len = self.payload().len();
        if payload_len > 125 {
            if payload_len <= u16::max_value() as usize {
                header_length += 2;
            } else {
                header_length += 8;
            }
        }

        if self.is_masked() {
            header_length += 4;
        }

        header_length + payload_len
    }

    /// Test whether the frame is a final frame.
    #[inline]
    pub fn is_final(&self) -> bool {
        self.finished
    }

    /// Test whether the first reserved bit is set.
    #[inline]
    pub fn has_rsv1(&self) -> bool {
        self.rsv1
    }

    /// Test whether the second reserved bit is set.
    #[inline]
    pub fn has_rsv2(&self) -> bool {
        self.rsv2
    }

    /// Test whether the third reserved bit is set.
    #[inline]
    pub fn has_rsv3(&self) -> bool {
        self.rsv3
    }

    /// Get the OpCode of the frame.
    #[inline]
    pub fn opcode(&self) -> OpCode {
        self.opcode
    }

    /// Get a reference to the frame's payload.
    #[inline]
    pub fn payload(&self) -> &Vec<u8> {
        &self.payload
    }

    // Test whether the frame is masked.
    #[doc(hidden)]
    #[inline]
    pub fn is_masked(&self) -> bool {
        self.mask.is_some()
    }

    // Get an optional reference to the frame's mask.
    #[doc(hidden)]
    #[allow(dead_code)]
    #[inline]
    pub fn mask(&self) -> Option<&[u8; 4]> {
        self.mask.as_ref()
    }

    /// Make this frame a final frame.
    #[allow(dead_code)]
    #[inline]
    pub fn set_final(&mut self, is_final: bool) -> &mut Frame {
        self.finished = is_final;
        self
    }

    /// Set the first reserved bit.
    #[allow(dead_code)]
    #[inline]
    pub fn set_rsv1(&mut self, has_rsv1: bool) -> &mut Frame {
        self.rsv1 = has_rsv1;
        self
    }

    /// Set the second reserved bit.
    #[allow(dead_code)]
    #[inline]
    pub fn set_rsv2(&mut self, has_rsv2: bool) -> &mut Frame {
        self.rsv2 = has_rsv2;
        self
    }

    /// Set the third reserved bit.
    #[allow(dead_code)]
    #[inline]
    pub fn set_rsv3(&mut self, has_rsv3: bool) -> &mut Frame {
        self.rsv1 = has_rsv3;
        self
    }

    /// Set the OpCode.
    #[allow(dead_code)]
    #[inline]
    pub fn set_opcode(&mut self, opcode: OpCode) -> &mut Frame {
        self.opcode = opcode;
        self
    }

    /// Edit the frame's payload.
    #[allow(dead_code)]
    #[inline]
    pub fn payload_mut(&mut self) -> &mut Vec<u8> {
        &mut self.payload
    }

    // Generate a new mask for this frame.
    //
    // This method simply generates and stores the mask. It does not change the payload data.
    // Instead, the payload data will be masked with the generated mask when the frame is sent
    // to the other endpoint.
    #[doc(hidden)]
    #[inline]
    pub fn set_mask(&mut self) -> &mut Frame {
        self.mask = Some(generate_mask());
        self
    }

    // This method unmasks the payload and should only be called on frames that are actually
    // masked. In other words, those frames that have just been received from a client endpoint.
    #[doc(hidden)]
    #[inline]
    pub fn remove_mask(&mut self) -> &mut Frame {
        self.mask.and_then(|mask| {
            Some(apply_mask(&mut self.payload, &mask))
        });
        self.mask = None;
        self
    }

    /// Consume the frame into its payload.
    pub fn into_data(self) -> Vec<u8> {
        self.payload
    }

    /// Create a new data frame.
    #[inline]
    pub fn message(data: Vec<u8>, code: OpCode, finished: bool) -> Frame {
        debug_assert!(match code {
            OpCode::Text | OpCode::Binary | OpCode::Continue => true,
            _ => false,
        }, "Invalid opcode for data frame.");

        Frame {
            finished: finished,
            opcode: code,
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Pong control frame.
    #[inline]
    pub fn pong(data: Vec<u8>) -> Frame {
        Frame {
            opcode: OpCode::Pong,
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Ping control frame.
    #[inline]
    pub fn ping(data: Vec<u8>) -> Frame {
        Frame {
            opcode: OpCode::Ping,
            payload: data,
            .. Frame::default()
        }
    }

    /// Create a new Close control frame.
    #[inline]
    pub fn close(code: CloseCode, reason: &str) -> Frame {
        let raw: [u8; 2] = unsafe {
            let u: u16 = code.into();
            transmute(u.to_be())
        };

        let payload = if let CloseCode::Empty = code {
            Vec::new()
        } else {
            Vec::from_iter(
                raw[..].iter()
                       .chain(reason.as_bytes().iter())
                       .map(|&b| b))
        };

        Frame {
            payload: payload,
            .. Frame::default()
        }
    }

    /// Parse the input stream into a frame.
    pub fn parse(cursor: &mut Cursor<Vec<u8>>) -> Result<Option<Frame>> {
        let size = cursor.get_ref().len() - cursor.position() as usize;
        let initial = cursor.position();
        trace!("Position in buffer {}", initial);

        let mut head = [0u8; 2];
        if try!(cursor.read(&mut head)) != 2 {
            cursor.set_position(initial);
            return Ok(None)
        }

        trace!("Parsed headers {:?}", head);

        let first = head[0];
        let second = head[1];
        trace!("First: {:b}", first);
        trace!("Second: {:b}", second);

        let finished = first & 0x80 != 0;

        let rsv1 = first & 0x40 != 0;
        let rsv2 = first & 0x20 != 0;
        let rsv3 = first & 0x10 != 0;

        let opcode = OpCode::from(first & 0x0F);
        trace!("Opcode: {:?}", opcode);
        if let OpCode::Bad = opcode {
            return Err(Error::new(Kind::Protocol, format!("Encountered invalid opcode: {}", first & 0x0F)))
        }
        let masked = second & 0x80 != 0;
        trace!("Masked: {:?}", masked);

        let mut header_length = 2;

        let mut length = (second & 0x7F) as u64;

        if length == 126 {
            let mut length_bytes = [0u8; 2];
            if try!(cursor.read(&mut length_bytes)) != 2 {
                cursor.set_position(initial);
                return Ok(None)
            }

            length = unsafe {
                let mut wide: u16 = transmute(length_bytes);
                wide = u16::from_be(wide);
                wide
            } as u64;
            header_length += 2;
        } else if length == 127 {
            let mut length_bytes = [0u8; 8];
            if try!(cursor.read(&mut length_bytes)) != 8 {
                cursor.set_position(initial);
                return Ok(None)
            }

            unsafe { length = transmute(length_bytes); }
            length = u64::from_be(length);
            header_length += 8;
        }
        trace!("Payload length: {}", length);

        // control frames must have length <= 125
        match opcode {
            OpCode::Close | OpCode::Ping | OpCode::Pong if length > 125 => {
                return Err(Error::new(Kind::Protocol, format!("Rejected WebSocket handshake.Received control frame with length: {}.", length)))
            }
            _ => ()
        }

        let mask = if masked {
            let mut mask_bytes = [0u8; 4];
            if try!(cursor.read(&mut mask_bytes)) != 4 {
                cursor.set_position(initial);
                return Ok(None)
            } else {
                header_length += 4;
                Some(mask_bytes)
            }
        } else {
            None
        };

        if size < length as usize + header_length {
            cursor.set_position(initial);
            return Ok(None)
        }

        let mut data = Vec::with_capacity(length as usize);
        if length > 0 {
            // It is ok to unwrap here because this cursor won't block
            let read = try!(cursor.try_read_buf(&mut data)).unwrap();
            debug_assert!(read == length as usize, "Read incorrect payload length!");
        }

        let frame = Frame {
            finished: finished,
            rsv1: rsv1,
            rsv2: rsv2,
            rsv3: rsv3,
            opcode: opcode,
            mask: mask,
            payload: data,
        };
        Ok(Some(frame))
    }

    /// Write a frame out to a buffer
    pub fn format<W>(&mut self, w: &mut W) -> Result<()>
        where W: Write
    {
        let mut one = 0u8;
        let code: u8 = self.opcode.into();
        if self.is_final() {
            one |= 0x80;
        }
        if self.has_rsv1() {
            one |= 0x40;
        }
        if self.has_rsv2() {
            one |= 0x20;
        }
        if self.has_rsv3() {
            one |= 0x10;
        }
        one |= code;

        let mut two = 0u8;

        if self.is_masked() {
            two |= 0x80;
        }

        if self.payload.len() < 126 {
            two |= self.payload.len() as u8;
            let headers = [one, two];
            try!(w.write(&headers));
        } else if self.payload.len() <= 65535 {
            two |= 126;
            let length_bytes: [u8; 2] = unsafe {
                let short = self.payload.len() as u16;
                transmute(short.to_be())
            };
            let headers = [one, two, length_bytes[0], length_bytes[1]];
            try!(w.write(&headers));
        } else {
            two |= 127;
            let length_bytes: [u8; 8] = unsafe {
                let long = self.payload.len() as u64;
                transmute(long.to_be())
            };
            let headers = [
                one,
                two,
                length_bytes[0],
                length_bytes[1],
                length_bytes[2],
                length_bytes[3],
                length_bytes[4],
                length_bytes[5],
                length_bytes[6],
                length_bytes[7],
            ];
            try!(w.write(&headers));
        }

        if self.is_masked() {
            let mask = self.mask.take().unwrap();
            apply_mask(&mut self.payload, &mask);
            try!(w.write(&mask));
        }

        try!(w.write(&self.payload));
        Ok(())
    }
}

impl Default for Frame {
    fn default() -> Frame {
        Frame {
            finished: true,
            rsv1: false,
            rsv2: false,
            rsv3: false,
            opcode: OpCode::Close,
            mask: None,
            payload: Vec::new(),
        }
    }
}

impl fmt::Display for Frame {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
            "
<FRAME>
final: {}
reserved: {} {} {}
opcode: {}
mask: {}
payload: {:?}
            ",
            self.finished,
            self.rsv1,
            self.rsv2,
            self.rsv3,
            self.opcode,
            self.mask.map(|mask| format!("{:?}", mask)).unwrap_or("NONE".into()),
            self.payload)
    }
}

mod test {
    #![allow(unused_imports, unused_variables, dead_code)]
    use super::*;
    use protocol::OpCode;

    #[test]
    fn test_display_frame() {
        let f = Frame::message("hi there".into(), OpCode::Text, true);
        let view = format!("{}", f);
        view.contains("payload:");
    }
}
<|endoftext|>


<commit_before>use cpu::common_defs::OpcodeExecInfo;
use cpu::CpuState;
use cpu::DecodeRegister;
use cpu::common_defs::address_mode::AddressMode;
use cpu::common_defs::opcode_class::OpcodeClass;
use memory::Memory;

macro_rules! set_zs {
    ($cpu_state:expr,$val:expr) => (set_z!($cpu_state,$val);set_s!($cpu_state,$val););
}
macro_rules! set_z {
    ($cpu_state:expr,$val:expr) => ($cpu_state.Z = $val == 0;);
}
macro_rules! set_s {
    ($cpu_state:expr,$val:expr) => ($cpu_state.S = $val >= 0x80;);
}
macro_rules! stack_push8 {
    ($cpu_state:expr,$mem:expr,$val:expr) => (try!($mem.write8(0x0100 | $cpu_state.sp,$val)); $cpu_state.sp-=1;);
}
macro_rules! stack_push16 {
    ($cpu_state:expr,$mem:expr,$val:expr) => (try!($mem.write16(0x0100 as u16 | $cpu_state.sp as u16 -1,$val)); $cpu_state.sp-=2;);
}
macro_rules! stack_pull8 {
    ($cpu_state:expr,$mem:expr) => {
        {
            $cpu_state.sp+=1;
            $mem.read8(0x0100 | $cpu_state.sp;)
        }
    }
}
macro_rules! stack_pull16 {
    ($cpu_state:expr,$mem:expr) => {
        {
            $cpu_state.sp+=2;
            $mem.read16(0x0100 as u16 | $cpu_state.sp as u16 -1)
        }
    }
}

#[derive(Debug)]
pub enum ExecutionError {
  MemoryError(::memory::MemoryError),
  UnexpectedOpcode(String),
  UnexpectedAddressMode(String)
}
impl From<::memory::MemoryError> for ExecutionError {
    fn from(err: ::memory::MemoryError) -> ExecutionError {
        ExecutionError::MemoryError(err)
    }
}

pub struct CpuExecutor {
    op_table: Vec<OpcodeExecInfo>,
}

impl CpuExecutor {
	/// Construct a new CpuExecutor.
    pub fn new(opcodes: Vec<OpcodeExecInfo> ) -> CpuExecutor {
        CpuExecutor {
            op_table:  opcodes,
        }
    }

	/// Reset the given cpu_state.
    pub fn power_on(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) {
        self.reset(cpu_state,mem);
    }
	/// Reset the given cpu_state.
    pub fn reset(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) {
        cpu_state.pc = mem.read16(0xFFFC).unwrap();
        cpu_state.sp = 0xFD;
        cpu_state.pack_flags(0x24);
    }

    /// Execute a single instruction in the given memory and cpu_state context.
    pub fn step(self: &CpuExecutor, cpu_state: &mut CpuState,mem:&mut Memory) -> Result<(),ExecutionError> {
        try!(self.fetch_and_decode(cpu_state,mem));
        try!(self.execute(cpu_state,mem));
        Ok(())
    }

	/// Fetch the next instruction and perform address resolution.
    pub fn fetch_and_decode(self: &CpuExecutor, cpu_state: &mut CpuState,mem:&mut Memory) -> Result<(),ExecutionError> {
        cpu_state.instruction_register = mem.read8(cpu_state.pc).unwrap();
        cpu_state.decode_register = try!(self.decode(cpu_state,mem));
        cpu_state.pc = cpu_state.pc+1;
        Ok(())
    }

    /// Perform address resolution, returning the info in a DecodeRegister.
    fn decode(self: &CpuExecutor, cpu_state: &CpuState, mem: &Memory) -> Result<DecodeRegister,ExecutionError> {
        let mut dr = DecodeRegister {
            info : self.op_table[cpu_state.instruction_register as usize].clone(),
            ..Default::default()
        };
        match dr.info.address_mode {
            // no explicit addresses for the following modes
            AddressMode::Accumulator  => { },
            AddressMode::Implied      => { },

            // explicit addresses from here on out
            AddressMode::Absolute => {
                dr.addr_final  = Some(mem.read16(cpu_state.pc+1).unwrap());
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::AbsoluteX       => {
                dr.addr_final  = Some( ((mem.read16(cpu_state.pc+1).unwrap() as u32 + cpu_state.x as u32) % 65536) as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::AbsoluteY       => {
                dr.addr_final  = Some( ((mem.read16(cpu_state.pc+1).unwrap() as u32 + cpu_state.y as u32) % 65536) as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Immediate       => {
                dr.addr_final  = Some(cpu_state.pc+1);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Indirect        => {
                dr.addr_intermediate = Some(mem.read16(cpu_state.pc+1).unwrap());
                dr.addr_final        = Some(mem.read16(dr.addr_intermediate.unwrap()).unwrap());
            },
            AddressMode::IndexedIndirect => {
                dr.addr_intermediate = Some( (mem.read8(cpu_state.pc+1).unwrap() + cpu_state.x) as u16);
                dr.addr_final        = Some(mem.read16(dr.addr_intermediate.unwrap()).unwrap());
                dr.value_final       = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::IndirectIndexed => {
                dr.addr_init         = Some(mem.read8(cpu_state.pc+1).unwrap() as u16);
                dr.addr_intermediate = Some(mem.read16(dr.addr_init.unwrap()).unwrap());
                dr.addr_final        = Some( ((dr.addr_intermediate.unwrap() as u32 + cpu_state.y as u32) % 65536) as u16);
                dr.value_final       = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Relative        => {
                dr.value_final = Some(mem.read8(cpu_state.pc+1).unwrap());

                // NOTE: relative is from the end of the current instruction and relative
                // instructions are 2 bytes long, so we add 2 before adding in the specified offset
                //
                if dr.value_final.unwrap() < 0x80 { dr.addr_final  = Some(cpu_state.pc + 2 + dr.value_final.unwrap() as u16); }
                else                              { dr.addr_final  = Some(cpu_state.pc + 2 + dr.value_final.unwrap() as u16 - 0x100);}
            },
            AddressMode::ZeroPage        => {
                dr.addr_final  = Some(mem.read8(cpu_state.pc+1).unwrap() as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPageX       => {
                dr.addr_final  = Some( (mem.read8(cpu_state.pc+1).unwrap() as u16 + cpu_state.x as u16) % 256);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPageY       => {
                dr.addr_final  = Some( (mem.read8(cpu_state.pc+1).unwrap() as u16 + cpu_state.y as u16) % 256);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            _ => { return Err(ExecutionError::UnexpectedAddressMode(format!("unrecognized addressing mode '{:?}' while decoding instruction_register!",dr.info.address_mode))); }
        }

        return Ok(dr);
    }

    /// Perform the current instruction, returning the CpuState after execution.
    pub fn execute(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) -> Result<(),ExecutionError> {
    	// Figure out which opcode is being executed.
    	match cpu_state.decode_register.info.opcode_class {
    		OpcodeClass::ASL => {
				cpu_state.a = cpu_state.a << 1;
				cpu_state.C = (128 & cpu_state.a) > 0;
    			if cpu_state.a == 0 { cpu_state.Z = true; }

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::BCC => {
                if !cpu_state.C { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else            { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BCS => {
                if cpu_state.C { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else           { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BEQ => {
                if cpu_state.Z { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else           { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BIT => {
                cpu_state.V = (cpu_state.decode_register.value_final.unwrap() >> 6) & 1 > 0;
                set_z!(cpu_state,cpu_state.decode_register.value_final.unwrap() & cpu_state.a);
                set_s!(cpu_state,cpu_state.decode_register.value_final.unwrap());

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::BNE => {
                if !cpu_state.Z { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else            { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BPL => {
                if !cpu_state.S { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else            { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BVC => {
                if !cpu_state.V { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else            { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BVS => {
                if cpu_state.V { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else           { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::CLC => {
                cpu_state.C = false;
    		},
    		OpcodeClass::LDA => {
    			cpu_state.a = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.a);

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::LDX => {
    			cpu_state.x = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.x);

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::LDY => {
    			cpu_state.y = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.y);

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::LSR => {
    			cpu_state.a = cpu_state.a >> 1;
    			cpu_state.C = 1 & cpu_state.a > 0;
    			if cpu_state.a == 0 { cpu_state.Z = true; }

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::JMP => {
    			cpu_state.pc = cpu_state.decode_register.addr_final.unwrap();
    		},
    		OpcodeClass::JSR => {
                stack_push16!(cpu_state,mem,cpu_state.pc-1);
    			cpu_state.pc = cpu_state.decode_register.addr_final.unwrap();
    		},
    		OpcodeClass::NOP => {
            },
    		OpcodeClass::SEC => {
                cpu_state.C = true;
    		},
    		OpcodeClass::STA => {
                try!(mem.write8(cpu_state.decode_register.addr_final.unwrap(),cpu_state.a));
                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::STX => {
                try!(mem.write8(cpu_state.decode_register.addr_final.unwrap(),cpu_state.x));
                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},

			_ => { return Err(ExecutionError::UnexpectedOpcode(format!("Unrecognised opcode class: {:?}", cpu_state.decode_register.info.opcode_class)));}

    	}
        Ok(())
    }
}

<commit_msg>fixed JSR implementation<commit_after>use cpu::common_defs::OpcodeExecInfo;
use cpu::CpuState;
use cpu::DecodeRegister;
use cpu::common_defs::address_mode::AddressMode;
use cpu::common_defs::opcode_class::OpcodeClass;
use memory::Memory;

macro_rules! set_zs {
    ($cpu_state:expr,$val:expr) => (set_z!($cpu_state,$val);set_s!($cpu_state,$val););
}
macro_rules! set_z {
    ($cpu_state:expr,$val:expr) => ($cpu_state.Z = $val == 0;);
}
macro_rules! set_s {
    ($cpu_state:expr,$val:expr) => ($cpu_state.S = $val >= 0x80;);
}
macro_rules! stack_push8 {
    ($cpu_state:expr,$mem:expr,$val:expr) => (try!($mem.write8(0x0100 | $cpu_state.sp,$val)); $cpu_state.sp-=1;);
}
macro_rules! stack_push16 {
    ($cpu_state:expr,$mem:expr,$val:expr) => (try!($mem.write16(0x0100 as u16 | $cpu_state.sp as u16 -1,$val)); $cpu_state.sp-=2;);
}
macro_rules! stack_pull8 {
    ($cpu_state:expr,$mem:expr) => {
        {
            $cpu_state.sp+=1;
            $mem.read8(0x0100 | $cpu_state.sp;)
        }
    }
}
macro_rules! stack_pull16 {
    ($cpu_state:expr,$mem:expr) => {
        {
            $cpu_state.sp+=2;
            $mem.read16(0x0100 as u16 | $cpu_state.sp as u16 -1)
        }
    }
}

#[derive(Debug)]
pub enum ExecutionError {
  MemoryError(::memory::MemoryError),
  UnexpectedOpcode(String),
  UnexpectedAddressMode(String)
}
impl From<::memory::MemoryError> for ExecutionError {
    fn from(err: ::memory::MemoryError) -> ExecutionError {
        ExecutionError::MemoryError(err)
    }
}

pub struct CpuExecutor {
    op_table: Vec<OpcodeExecInfo>,
}

impl CpuExecutor {
	/// Construct a new CpuExecutor.
    pub fn new(opcodes: Vec<OpcodeExecInfo> ) -> CpuExecutor {
        CpuExecutor {
            op_table:  opcodes,
        }
    }

	/// Reset the given cpu_state.
    pub fn power_on(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) {
        self.reset(cpu_state,mem);
    }
	/// Reset the given cpu_state.
    pub fn reset(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) {
        cpu_state.pc = mem.read16(0xFFFC).unwrap();
        cpu_state.sp = 0xFD;
        cpu_state.pack_flags(0x24);
    }

    /// Execute a single instruction in the given memory and cpu_state context.
    pub fn step(self: &CpuExecutor, cpu_state: &mut CpuState,mem:&mut Memory) -> Result<(),ExecutionError> {
        try!(self.fetch_and_decode(cpu_state,mem));
        try!(self.execute(cpu_state,mem));
        Ok(())
    }

	/// Fetch the next instruction and perform address resolution.
    pub fn fetch_and_decode(self: &CpuExecutor, cpu_state: &mut CpuState,mem:&mut Memory) -> Result<(),ExecutionError> {
        cpu_state.instruction_register = mem.read8(cpu_state.pc).unwrap();
        cpu_state.decode_register = try!(self.decode(cpu_state,mem));
        cpu_state.pc = cpu_state.pc+1;
        Ok(())
    }

    /// Perform address resolution, returning the info in a DecodeRegister.
    fn decode(self: &CpuExecutor, cpu_state: &CpuState, mem: &Memory) -> Result<DecodeRegister,ExecutionError> {
        let mut dr = DecodeRegister {
            info : self.op_table[cpu_state.instruction_register as usize].clone(),
            ..Default::default()
        };
        match dr.info.address_mode {
            // no explicit addresses for the following modes
            AddressMode::Accumulator  => { },
            AddressMode::Implied      => { },

            // explicit addresses from here on out
            AddressMode::Absolute => {
                dr.addr_final  = Some(mem.read16(cpu_state.pc+1).unwrap());
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::AbsoluteX       => {
                dr.addr_final  = Some( ((mem.read16(cpu_state.pc+1).unwrap() as u32 + cpu_state.x as u32) % 65536) as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::AbsoluteY       => {
                dr.addr_final  = Some( ((mem.read16(cpu_state.pc+1).unwrap() as u32 + cpu_state.y as u32) % 65536) as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Immediate       => {
                dr.addr_final  = Some(cpu_state.pc+1);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Indirect        => {
                dr.addr_intermediate = Some(mem.read16(cpu_state.pc+1).unwrap());
                dr.addr_final        = Some(mem.read16(dr.addr_intermediate.unwrap()).unwrap());
            },
            AddressMode::IndexedIndirect => {
                dr.addr_intermediate = Some( (mem.read8(cpu_state.pc+1).unwrap() + cpu_state.x) as u16);
                dr.addr_final        = Some(mem.read16(dr.addr_intermediate.unwrap()).unwrap());
                dr.value_final       = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::IndirectIndexed => {
                dr.addr_init         = Some(mem.read8(cpu_state.pc+1).unwrap() as u16);
                dr.addr_intermediate = Some(mem.read16(dr.addr_init.unwrap()).unwrap());
                dr.addr_final        = Some( ((dr.addr_intermediate.unwrap() as u32 + cpu_state.y as u32) % 65536) as u16);
                dr.value_final       = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Relative        => {
                dr.value_final = Some(mem.read8(cpu_state.pc+1).unwrap());

                // NOTE: relative is from the end of the current instruction and relative
                // instructions are 2 bytes long, so we add 2 before adding in the specified offset
                //
                if dr.value_final.unwrap() < 0x80 { dr.addr_final  = Some(cpu_state.pc + 2 + dr.value_final.unwrap() as u16); }
                else                              { dr.addr_final  = Some(cpu_state.pc + 2 + dr.value_final.unwrap() as u16 - 0x100);}
            },
            AddressMode::ZeroPage        => {
                dr.addr_final  = Some(mem.read8(cpu_state.pc+1).unwrap() as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPageX       => {
                dr.addr_final  = Some( (mem.read8(cpu_state.pc+1).unwrap() as u16 + cpu_state.x as u16) % 256);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPageY       => {
                dr.addr_final  = Some( (mem.read8(cpu_state.pc+1).unwrap() as u16 + cpu_state.y as u16) % 256);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            _ => { return Err(ExecutionError::UnexpectedAddressMode(format!("unrecognized addressing mode '{:?}' while decoding instruction_register!",dr.info.address_mode))); }
        }

        return Ok(dr);
    }

    /// Perform the current instruction, returning the CpuState after execution.
    pub fn execute(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) -> Result<(),ExecutionError> {
    	// Figure out which opcode is being executed.
    	match cpu_state.decode_register.info.opcode_class {
    		OpcodeClass::ASL => {
				cpu_state.a = cpu_state.a << 1;
				cpu_state.C = (128 & cpu_state.a) > 0;
    			if cpu_state.a == 0 { cpu_state.Z = true; }

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::BCC => {
                if !cpu_state.C { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else            { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BCS => {
                if cpu_state.C { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else           { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BEQ => {
                if cpu_state.Z { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else           { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BIT => {
                cpu_state.V = (cpu_state.decode_register.value_final.unwrap() >> 6) & 1 > 0;
                set_z!(cpu_state,cpu_state.decode_register.value_final.unwrap() & cpu_state.a);
                set_s!(cpu_state,cpu_state.decode_register.value_final.unwrap());

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::BNE => {
                if !cpu_state.Z { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else            { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BPL => {
                if !cpu_state.S { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else            { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BVC => {
                if !cpu_state.V { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else            { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BVS => {
                if cpu_state.V { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else           { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::CLC => {
                cpu_state.C = false;
    		},
    		OpcodeClass::LDA => {
    			cpu_state.a = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.a);

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::LDX => {
    			cpu_state.x = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.x);

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::LDY => {
    			cpu_state.y = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.y);

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::LSR => {
    			cpu_state.a = cpu_state.a >> 1;
    			cpu_state.C = 1 & cpu_state.a > 0;
    			if cpu_state.a == 0 { cpu_state.Z = true; }

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::JMP => {
    			cpu_state.pc = cpu_state.decode_register.addr_final.unwrap();
    		},
    		OpcodeClass::JSR => {
                stack_push16!(cpu_state,mem,cpu_state.pc+1);
    			cpu_state.pc = cpu_state.decode_register.addr_final.unwrap();
    		},
    		OpcodeClass::NOP => {
            },
    		OpcodeClass::SEC => {
                cpu_state.C = true;
    		},
    		OpcodeClass::STA => {
                try!(mem.write8(cpu_state.decode_register.addr_final.unwrap(),cpu_state.a));
                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::STX => {
                try!(mem.write8(cpu_state.decode_register.addr_final.unwrap(),cpu_state.x));
                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},

			_ => { return Err(ExecutionError::UnexpectedOpcode(format!("Unrecognised opcode class: {:?}", cpu_state.decode_register.info.opcode_class)));}

    	}
        Ok(())
    }
}

<|endoftext|>


<commit_before>use fragment::matching;
use ExclusionPattern;
use walkdir::{DirEntry, Error, WalkDir};
use std::path::PathBuf;
use IndexedPath;

pub struct Index {
    path: PathBuf,
    entries: Vec<IndexedPath>,
}

impl Index {
    pub fn new(path: PathBuf) -> Index {
        Index {
            path: path,
            entries: Vec::new()
        }
    }

    /// Finds all files inside and beneath the index path
    /// and adds them to the index entries vector.
    pub fn populate(&mut self, exclusions: Option<Vec<ExclusionPattern>>, case_sensitive: bool) {
        // The entries listed by read_dir include the root index path; we want
        // relative paths, so we get this length so that we can strip it.
        let prefix_length = match self.path.to_str() {
            Some(path) => path.len() + 1,
            None => return,
        };

        // Start indexing at the specified path.
        let filtered_entries = WalkDir::new(&self.path).into_iter().filter_entry(|entry| {
            if let Some(ref exclusions) = exclusions {
                !exclusions.iter().any(|exclusion| {
                    exclusion.matches(entry.path().to_string_lossy().as_ref())
                })
            } else {
                true
            }
        });

        for entry in filtered_entries {
            relative_entry_path(entry, prefix_length).map(|entry_path| {
                self.entries.push(
                    IndexedPath::new(entry_path, case_sensitive)
                );
            });
        }
    }

    pub fn find(&self, term: &str, limit: usize) -> Vec<&PathBuf> {
        matching::find(term, &self.entries, limit)
            .into_iter()
            .map(|result| &result.path)
            .collect()
    }
}

/// Transforms a DirEntry object into an optional relative path string,
/// returning None if any errors occur or if the entry is not a file.
fn relative_entry_path(entry: Result<DirEntry, Error>, prefix_length: usize) -> Option<String> {
    entry.ok().and_then(|e| {
        // Limit path look-ups to files.
        e.path().metadata().ok().and_then(|metadata| {
            if metadata.is_file() {
                // Map the absolute path to a relative one.
                e.path().to_str().map(|path| path[prefix_length..].to_string())
            } else {
                None
            }
        })
    })
}

#[cfg(test)]
mod tests {
    use super::{Index, IndexedPath, ExclusionPattern};
    use std::path::PathBuf;

    #[test]
    fn populate_respects_exclusions() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath::new("root_file", true)];
        index.populate(Some(vec![ExclusionPattern::new("**/directory").unwrap()]), true);

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn populate_lowercases_entries_when_case_sensitive_is_false() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath::new("directory/Capitalized_file", false),
                                    IndexedPath::new("directory/nested_file", false),
                                    IndexedPath::new("root_file", false)];
        index.populate(None, false);
        index.entries.sort();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn populate_lowercases_entries_when_case_sensitive_is_true() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath::new("directory/Capitalized_file", true),
                                    IndexedPath::new("directory/nested_file", true),
                                    IndexedPath::new("root_file", true)];
        index.populate(None, true);
        index.entries.sort();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn find_defers_to_matching_module() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        index.populate(None, true);
        let term = "root";
        let limit = 5;

        // Get a string version of the results (PathBuf doesn't implement the display trait).
        let results: Vec<String> = index.find(term, limit)
                                        .iter()
                                        .map(|r| r.to_string_lossy().into_owned())
                                        .collect();

        assert_eq!(results, vec!["root_file".to_string()]);
    }
}
<commit_msg>Update index find method to return Path references<commit_after>use fragment::matching;
use ExclusionPattern;
use walkdir::{DirEntry, Error, WalkDir};
use std::path::{Path, PathBuf};
use IndexedPath;

pub struct Index {
    path: PathBuf,
    entries: Vec<IndexedPath>,
}

impl Index {
    pub fn new(path: PathBuf) -> Index {
        Index {
            path: path,
            entries: Vec::new()
        }
    }

    /// Finds all files inside and beneath the index path
    /// and adds them to the index entries vector.
    pub fn populate(&mut self, exclusions: Option<Vec<ExclusionPattern>>, case_sensitive: bool) {
        // The entries listed by read_dir include the root index path; we want
        // relative paths, so we get this length so that we can strip it.
        let prefix_length = match self.path.to_str() {
            Some(path) => path.len() + 1,
            None => return,
        };

        // Start indexing at the specified path.
        let filtered_entries = WalkDir::new(&self.path).into_iter().filter_entry(|entry| {
            if let Some(ref exclusions) = exclusions {
                !exclusions.iter().any(|exclusion| {
                    exclusion.matches(entry.path().to_string_lossy().as_ref())
                })
            } else {
                true
            }
        });

        for entry in filtered_entries {
            relative_entry_path(entry, prefix_length).map(|entry_path| {
                self.entries.push(
                    IndexedPath::new(entry_path, case_sensitive)
                );
            });
        }
    }

    pub fn find(&self, term: &str, limit: usize) -> Vec<&Path> {
        matching::find(term, &self.entries, limit)
            .into_iter()
            .map(|result| result.as_path())
            .collect()
    }
}

/// Transforms a DirEntry object into an optional relative path string,
/// returning None if any errors occur or if the entry is not a file.
fn relative_entry_path(entry: Result<DirEntry, Error>, prefix_length: usize) -> Option<String> {
    entry.ok().and_then(|e| {
        // Limit path look-ups to files.
        e.path().metadata().ok().and_then(|metadata| {
            if metadata.is_file() {
                // Map the absolute path to a relative one.
                e.path().to_str().map(|path| path[prefix_length..].to_string())
            } else {
                None
            }
        })
    })
}

#[cfg(test)]
mod tests {
    use super::{Index, IndexedPath, ExclusionPattern};
    use std::path::{Path, PathBuf};

    #[test]
    fn populate_respects_exclusions() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath::new("root_file", true)];
        index.populate(Some(vec![ExclusionPattern::new("**/directory").unwrap()]), true);

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn populate_lowercases_entries_when_case_sensitive_is_false() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath::new("directory/Capitalized_file", false),
                                    IndexedPath::new("directory/nested_file", false),
                                    IndexedPath::new("root_file", false)];
        index.populate(None, false);
        index.entries.sort();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn populate_lowercases_entries_when_case_sensitive_is_true() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath::new("directory/Capitalized_file", true),
                                    IndexedPath::new("directory/nested_file", true),
                                    IndexedPath::new("root_file", true)];
        index.populate(None, true);
        index.entries.sort();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn find_defers_to_matching_module() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        index.populate(None, true);
        let term = "root";
        let limit = 5;
        let results = index.find(term, limit);

        assert_eq!(results, vec![Path::new("root_file")]);
    }
}
<|endoftext|>