<commit_before>use std::fmt;
use std::io::{Cursor, Read, Write, ErrorKind};
use std::default::Default;

use rand;
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};

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

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
    }
}

/// 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
    }

    /// Test whether this is a control frame.
    #[inline]
    pub fn is_control(&self) -> bool {
        self.opcode.is_control()
    }

    /// 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.
    #[inline]
    pub fn set_rsv1(&mut self, has_rsv1: bool) -> &mut Frame {
        self.rsv1 = has_rsv1;
        self
    }

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

    /// Set the third reserved bit.
    #[inline]
    pub fn set_rsv3(&mut self, has_rsv3: bool) -> &mut Frame {
        self.rsv3 = 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(rand::random());
        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 payload = if let CloseCode::Empty = code {
            Vec::new()
        } else {
            let u: u16 = code.into();
            let raw = [(u >> 8) as u8, u as u8];
            [&raw, reason.as_bytes()].concat()
        };

        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() as u64 - cursor.position();
        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 let Some(length_nbytes) = match length {
            126 => Some(2),
            127 => Some(8),
            _ => None,
        } {
            match cursor.read_uint::<BigEndian>(length_nbytes) {
                Err(ref err) if err.kind() == ErrorKind::UnexpectedEof => {
                    cursor.set_position(initial);
                    return Ok(None);
                }
                Err(err) => {
                    return Err(Error::from(err));
                }
                Ok(read) => {
                    length = read;
                }
            };
            header_length += length_nbytes as u64;
        }
        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 + 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;
        }

        match self.payload.len() {
            len if len < 126 => {
                two |= len as u8;
            }
            len if len <= 65535 => {
                two |= 126;
            }
            _ => {
                two |= 127;
            }
        }
        try!(w.write(&[one, two]));

        if let Some(length_bytes) = match self.payload.len() {
            len if len < 126 => None,
            len if len <= 65535 => Some(2),
            _ => Some(8),
        } {
            try!(w.write_uint::<BigEndian>(self.payload.len() as u64, length_bytes));
        }

        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: {}
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 display_frame() {
        let f = Frame::message("hi there".into(), OpCode::Text, true);
        let view = format!("{}", f);
        view.contains("payload:");
    }
}
<commit_msg>Avoid integer overflow when parsing frame bytes.<commit_after>use std::fmt;
use std::io::{Cursor, Read, Write, ErrorKind};
use std::default::Default;

use rand;
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};

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

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
    }
}

/// 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
    }

    /// Test whether this is a control frame.
    #[inline]
    pub fn is_control(&self) -> bool {
        self.opcode.is_control()
    }

    /// 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.
    #[inline]
    pub fn set_rsv1(&mut self, has_rsv1: bool) -> &mut Frame {
        self.rsv1 = has_rsv1;
        self
    }

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

    /// Set the third reserved bit.
    #[inline]
    pub fn set_rsv3(&mut self, has_rsv3: bool) -> &mut Frame {
        self.rsv3 = 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(rand::random());
        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 payload = if let CloseCode::Empty = code {
            Vec::new()
        } else {
            let u: u16 = code.into();
            let raw = [(u >> 8) as u8, u as u8];
            [&raw, reason.as_bytes()].concat()
        };

        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() as u64 - cursor.position();
        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 let Some(length_nbytes) = match length {
            126 => Some(2),
            127 => Some(8),
            _ => None,
        } {
            match cursor.read_uint::<BigEndian>(length_nbytes) {
                Err(ref err) if err.kind() == ErrorKind::UnexpectedEof => {
                    cursor.set_position(initial);
                    return Ok(None);
                }
                Err(err) => {
                    return Err(Error::from(err));
                }
                Ok(read) => {
                    length = read;
                }
            };
            header_length += length_nbytes as u64;
        }
        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
        };

        match length.checked_add(header_length) {
            Some(l) if size < l => {
                cursor.set_position(initial);
                return Ok(None);
            }
            Some(_) => (),
            None => 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;
        }

        match self.payload.len() {
            len if len < 126 => {
                two |= len as u8;
            }
            len if len <= 65535 => {
                two |= 126;
            }
            _ => {
                two |= 127;
            }
        }
        try!(w.write(&[one, two]));

        if let Some(length_bytes) = match self.payload.len() {
            len if len < 126 => None,
            len if len <= 65535 => Some(2),
            _ => Some(8),
        } {
            try!(w.write_uint::<BigEndian>(self.payload.len() as u64, length_bytes));
        }

        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: {}
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 display_frame() {
        let f = Frame::message("hi there".into(), OpCode::Text, true);
        let view = format!("{}", f);
        view.contains("payload:");
    }
}
<|endoftext|>


<commit_before>/// Compute the real-valued digamma function.
///
/// The digamma function is defined by
///
/// ```math
///        d ln(Γ(x))
/// ψ(x) = ----------
///            dx
/// ```
///
/// where Γ is the Gamma function. The computation is based on an approximation
/// as described in the reference below.
///
/// ## Examples
///
/// ```
/// use special::digamma;
///
/// const EULER_MASCHERONI: f64 = 0.57721566490153286060651209008240243104215933593992;
/// assert!((digamma(1.0) + EULER_MASCHERONI).abs() < 1e-15);
/// ```
///
/// ## References
///
/// 1. M. J. Beal, Variational algorithms for approximate Bayesian inference.
///    University of London, 2003, pp. 265–266.
pub fn digamma(x: f64)-> f64 {
    macro_rules! evaluate_polynomial(
        ($x:expr, $coefficients:expr) => (
            $coefficients.iter().rev().fold(0.0, |sum, &c| $x * sum + c)
        );
    );

    if x <= 8.0 {
        return digamma(x + 1.0) - x.recip();
    }

    let inv_x = x.recip();
    let inv_x_2 = inv_x * inv_x;
    x.ln() - 0.5 * inv_x - inv_x_2 * evaluate_polynomial!(inv_x_2, [
        1.0 / 12.0, -1.0 / 120.0, 1.0 / 252.0, -1.0 / 240.0,
        5.0 / 660.0, -691.0 / 32760.0, 1.0 / 12.0, -3617.0 / 8160.0,
    ])
}

/// Compute the lower incomplete gamma function.
///
/// The code is based on a [C implementation][1] by John Burkardt. The original
/// algorithm was published in Applied Statistics and is known as
/// [Algorithm AS 239][2].
///
/// [1]: http://people.sc.fsu.edu/~jburkardt/c_src/asa239/asa239.html
/// [2]: http://www.jstor.org/stable/2347328
pub fn inc_gamma(x: f64, p: f64) -> f64 {
    debug_assert!(x >= 0.0 && p > 0.0);

    const ELIMIT: f64 = -88.0;
    const OFLO: f64 = 1.0e+37;
    const TOL: f64 = 1.0e-14;
    const XBIG: f64 = 1.0e+08;

    if x == 0.0 {
        return 0.0;
    }

    // For `p ≥ 1000`, the original algorithm uses an approximation shown below.
    // However, it introduces a substantial accuracy loss.
    //
    // ```
    // use m::erf;
    // use std::f64::consts::FRAC_1_SQRT_2;
    //
    // const PLIMIT: f64 = 1000.0;
    //
    // if PLIMIT < p {
    //     let pn1 = 3.0 * p.sqrt() * ((x / p).powf(1.0 / 3.0) + 1.0 / (9.0 * p) - 1.0);
    //     return 0.5 * (1.0 + unsafe { erf(FRAC_1_SQRT_2 * pn1) });
    // }
    // ```

    if XBIG < x {
        return 1.0;
    }

    if x <= 1.0 || x < p {
        let mut arg = p * x.ln() - x - ln_gamma(p + 1.0).0;
        let mut c = 1.0;
        let mut value = 1.0;
        let mut a = p;

        loop {
            a += 1.0;
            c *= x / a;
            value += c;

            if c <= TOL {
                break;
            }
        }

        arg += value.ln();

        if ELIMIT <= arg {
            return arg.exp();
        } else {
            return 0.0;
        };
    } else {
        let mut arg = p * x.ln() - x - ln_gamma(p).0;
        let mut a = 1.0 - p;
        let mut b = a + x + 1.0;
        let mut c = 0.0;
        let mut pn1 = 1.0;
        let mut pn2 = x;
        let mut pn3 = x + 1.0;
        let mut pn4 = x * b;
        let mut value = pn3 / pn4;

        loop {
            a += 1.0;
            b += 2.0;
            c += 1.0;
            let an = a * c;
            let pn5 = b * pn3 - an * pn1;
            let pn6 = b * pn4 - an * pn2;

            if pn6 != 0.0 {
                let rn = pn5 / pn6;
                if (value - rn).abs() <= TOL.min(TOL * rn) {
                    break;
                }
                value = rn;
            }

            pn1 = pn3;
            pn2 = pn4;
            pn3 = pn5;
            pn4 = pn6;

            if OFLO <= pn5.abs() {
                pn1 /= OFLO;
                pn2 /= OFLO;
                pn3 /= OFLO;
                pn4 /= OFLO;
            }
        }

        arg += value.ln();

        if ELIMIT <= arg {
            return 1.0 - arg.exp();
        } else {
            return 1.0;
        }
    }
}

/// Compute the natural logarithm of the gamma function.
#[inline]
pub fn ln_gamma(x: f64) -> (f64, i32) {
    use m::lgamma_r;
    let mut sign: i32 = 0;
    let value = unsafe { lgamma_r(x, &mut sign) };
    (value, sign)
}

#[cfg(test)]
mod tests {
    use assert;

    #[test]
    fn digamma() {
        use std::f64::consts::{FRAC_PI_2, LN_2};
        const EULER_MASCHERONI: f64 = 0.57721566490153286060651209008240243104215933593992;
        assert_eq!(-FRAC_PI_2 - 3.0 * LN_2 - EULER_MASCHERONI, super::digamma(0.25));
    }

    #[test]
    fn inc_gamma_small_p() {
        let p = 4.2;
        let x = vec![
            0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5,
            5.0, 5.5, 6.0, 6.5, 7.0, 7.5, 8.0, 8.5, 9.0, 9.5,
        ];
        let y = vec![
            0.000000000000000e+00, 1.118818097571853e-03, 1.386936454406691e-02,
            5.173931341260211e-02, 1.186345016135447e-01, 2.092219100370761e-01,
            3.137074857845146e-01, 4.219736975484903e-01, 5.258832878449179e-01,
            6.200417684429613e-01, 7.016346974530381e-01, 7.698559854309833e-01,
            8.252531208548146e-01, 8.691540627528154e-01, 9.032348314835232e-01,
            9.292289044193606e-01, 9.487539837941871e-01, 9.632249236076943e-01,
            9.738240752336722e-01, 9.815062931491536e-01,
        ];

        let z = x.iter().map(|&x| super::inc_gamma(x, p)).collect::<Vec<_>>();
        assert::close(&z, &y, 1e-14);
    }

    #[test]
    fn inc_gamma_large_p() {
        let p = 1500.0;
        let x = vec![
            1400.0, 1410.0, 1420.0, 1430.0, 1440.0, 1450.0, 1460.0, 1470.0, 1480.0, 1490.0,
            1500.0, 1510.0, 1520.0, 1530.0, 1540.0, 1550.0, 1560.0, 1570.0, 1580.0, 1590.0,
        ];
        let y = vec![
            4.231080348517120e-03, 9.056183893278287e-03, 1.809200095269094e-02,
            3.380047876608895e-02, 5.917825204288418e-02, 9.731671903009434e-02,
            1.506860564044825e-01, 2.202940362119810e-01, 3.049926190158564e-01,
            4.012305141433336e-01, 5.034335611603484e-01, 6.049687028759408e-01,
            6.994146092246781e-01, 7.817404371226420e-01, 8.490443067241588e-01,
            9.006922020203059e-01, 9.379249194843459e-01, 9.631597573132408e-01,
            9.792520392189441e-01, 9.889149370075309e-01,
        ];

        let z = x.iter().map(|&x| super::inc_gamma(x, p)).collect::<Vec<_>>();
        assert::close(&z, &y, 1e-12);
    }
}
<commit_msg>Refine the description of inc_gamma<commit_after>/// Compute the real-valued digamma function.
///
/// The formula is as follows:
///
/// ```math
///        d ln(Γ(x))
/// ψ(x) = ----------
///            dx
/// ```
///
/// where Γ is the gamma function. The computation is based on an approximation
/// as described in the reference below.
///
/// ## Examples
///
/// ```
/// use special::digamma;
///
/// const EULER_MASCHERONI: f64 = 0.57721566490153286060651209008240243104215933593992;
/// assert!((digamma(1.0) + EULER_MASCHERONI).abs() < 1e-15);
/// ```
///
/// ## References
///
/// 1. M. J. Beal, Variational algorithms for approximate Bayesian inference.
///    University of London, 2003, pp. 265–266.
pub fn digamma(x: f64)-> f64 {
    macro_rules! evaluate_polynomial(
        ($x:expr, $coefficients:expr) => (
            $coefficients.iter().rev().fold(0.0, |sum, &c| $x * sum + c)
        );
    );

    if x <= 8.0 {
        return digamma(x + 1.0) - x.recip();
    }

    let inv_x = x.recip();
    let inv_x_2 = inv_x * inv_x;
    x.ln() - 0.5 * inv_x - inv_x_2 * evaluate_polynomial!(inv_x_2, [
        1.0 / 12.0, -1.0 / 120.0, 1.0 / 252.0, -1.0 / 240.0,
        5.0 / 660.0, -691.0 / 32760.0, 1.0 / 12.0, -3617.0 / 8160.0,
    ])
}

/// Compute the regularized lower incomplete gamma function.
///
/// The formula is as follows:
///
/// ```math
///                          ∞
///           γ(x, p)    1   /
/// P(x, p) = ------- = ---- | t^(p-1) e^(-t) dt
///            Γ(p)     Γ(p) /
///                          0
/// ```
///
/// where γ is the incomplete lower gamma function, and Γ is the complete gamma
/// function.
///
/// The code is based on a [C implementation][1] by John Burkardt. The original
/// algorithm was published in Applied Statistics and is known as
/// [Algorithm AS 239][2].
///
/// [1]: http://people.sc.fsu.edu/~jburkardt/c_src/asa239/asa239.html
/// [2]: http://www.jstor.org/stable/2347328
pub fn inc_gamma(x: f64, p: f64) -> f64 {
    debug_assert!(x >= 0.0 && p > 0.0);

    const ELIMIT: f64 = -88.0;
    const OFLO: f64 = 1.0e+37;
    const TOL: f64 = 1.0e-14;
    const XBIG: f64 = 1.0e+08;

    if x == 0.0 {
        return 0.0;
    }

    // For `p ≥ 1000`, the original algorithm uses an approximation shown below.
    // However, it introduces a substantial accuracy loss.
    //
    // ```
    // use m::erf;
    // use std::f64::consts::FRAC_1_SQRT_2;
    //
    // const PLIMIT: f64 = 1000.0;
    //
    // if PLIMIT < p {
    //     let pn1 = 3.0 * p.sqrt() * ((x / p).powf(1.0 / 3.0) + 1.0 / (9.0 * p) - 1.0);
    //     return 0.5 * (1.0 + unsafe { erf(FRAC_1_SQRT_2 * pn1) });
    // }
    // ```

    if XBIG < x {
        return 1.0;
    }

    if x <= 1.0 || x < p {
        let mut arg = p * x.ln() - x - ln_gamma(p + 1.0).0;
        let mut c = 1.0;
        let mut value = 1.0;
        let mut a = p;

        loop {
            a += 1.0;
            c *= x / a;
            value += c;

            if c <= TOL {
                break;
            }
        }

        arg += value.ln();

        if ELIMIT <= arg {
            return arg.exp();
        } else {
            return 0.0;
        };
    } else {
        let mut arg = p * x.ln() - x - ln_gamma(p).0;
        let mut a = 1.0 - p;
        let mut b = a + x + 1.0;
        let mut c = 0.0;
        let mut pn1 = 1.0;
        let mut pn2 = x;
        let mut pn3 = x + 1.0;
        let mut pn4 = x * b;
        let mut value = pn3 / pn4;

        loop {
            a += 1.0;
            b += 2.0;
            c += 1.0;
            let an = a * c;
            let pn5 = b * pn3 - an * pn1;
            let pn6 = b * pn4 - an * pn2;

            if pn6 != 0.0 {
                let rn = pn5 / pn6;
                if (value - rn).abs() <= TOL.min(TOL * rn) {
                    break;
                }
                value = rn;
            }

            pn1 = pn3;
            pn2 = pn4;
            pn3 = pn5;
            pn4 = pn6;

            if OFLO <= pn5.abs() {
                pn1 /= OFLO;
                pn2 /= OFLO;
                pn3 /= OFLO;
                pn4 /= OFLO;
            }
        }

        arg += value.ln();

        if ELIMIT <= arg {
            return 1.0 - arg.exp();
        } else {
            return 1.0;
        }
    }
}

/// Compute the natural logarithm of the gamma function.
#[inline]
pub fn ln_gamma(x: f64) -> (f64, i32) {
    use m::lgamma_r;
    let mut sign: i32 = 0;
    let value = unsafe { lgamma_r(x, &mut sign) };
    (value, sign)
}

#[cfg(test)]
mod tests {
    use assert;

    #[test]
    fn digamma() {
        use std::f64::consts::{FRAC_PI_2, LN_2};
        const EULER_MASCHERONI: f64 = 0.57721566490153286060651209008240243104215933593992;
        assert_eq!(-FRAC_PI_2 - 3.0 * LN_2 - EULER_MASCHERONI, super::digamma(0.25));
    }

    #[test]
    fn inc_gamma_small_p() {
        let p = 4.2;
        let x = vec![
            0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5,
            5.0, 5.5, 6.0, 6.5, 7.0, 7.5, 8.0, 8.5, 9.0, 9.5,
        ];
        let y = vec![
            0.000000000000000e+00, 1.118818097571853e-03, 1.386936454406691e-02,
            5.173931341260211e-02, 1.186345016135447e-01, 2.092219100370761e-01,
            3.137074857845146e-01, 4.219736975484903e-01, 5.258832878449179e-01,
            6.200417684429613e-01, 7.016346974530381e-01, 7.698559854309833e-01,
            8.252531208548146e-01, 8.691540627528154e-01, 9.032348314835232e-01,
            9.292289044193606e-01, 9.487539837941871e-01, 9.632249236076943e-01,
            9.738240752336722e-01, 9.815062931491536e-01,
        ];

        let z = x.iter().map(|&x| super::inc_gamma(x, p)).collect::<Vec<_>>();
        assert::close(&z, &y, 1e-14);
    }

    #[test]
    fn inc_gamma_large_p() {
        let p = 1500.0;
        let x = vec![
            1400.0, 1410.0, 1420.0, 1430.0, 1440.0, 1450.0, 1460.0, 1470.0, 1480.0, 1490.0,
            1500.0, 1510.0, 1520.0, 1530.0, 1540.0, 1550.0, 1560.0, 1570.0, 1580.0, 1590.0,
        ];
        let y = vec![
            4.231080348517120e-03, 9.056183893278287e-03, 1.809200095269094e-02,
            3.380047876608895e-02, 5.917825204288418e-02, 9.731671903009434e-02,
            1.506860564044825e-01, 2.202940362119810e-01, 3.049926190158564e-01,
            4.012305141433336e-01, 5.034335611603484e-01, 6.049687028759408e-01,
            6.994146092246781e-01, 7.817404371226420e-01, 8.490443067241588e-01,
            9.006922020203059e-01, 9.379249194843459e-01, 9.631597573132408e-01,
            9.792520392189441e-01, 9.889149370075309e-01,
        ];

        let z = x.iter().map(|&x| super::inc_gamma(x, p)).collect::<Vec<_>>();
        assert::close(&z, &y, 1e-12);
    }
}
<|endoftext|>


<commit_before>extern crate rand;
use ::std::str;
use self::rand::Rng;

#[derive(Debug, Ord, PartialOrd, Eq, PartialEq, Copy, Clone)]
pub enum Suit {
    Clubs,
    Diamonds,
    Hearts,
    Spades,
    Stars,
}

#[derive(Debug, Eq, PartialEq, PartialOrd)]
pub struct Card {
    suit: Suit,
    value: &'static str,
}

impl Card {
    pub fn new(suit: Suit, value: &'static str) -> Card {
        Card { suit: suit, value: value }
    }
}

impl Ord for Card {
    fn cmp(&self, other: &Self) -> ::std::cmp::Ordering {
        self.value.cmp(&other.value)
    }
}

pub fn make_standard_deck() -> Vec<Card> {
    let mut result = Vec::new();
    let suits  = &[Suit::Clubs, Suit::Diamonds, Suit::Hearts, Suit::Spades];
    let values = &["2", "3", "4", "5", "6", "7", "8", "9", "10",
                   "Jack", "Queen", "King", "Ace"];

    for suit in suits {
        for value in values {
            result.push(Card::new(*suit, value));
        }
    }
    return result;
}

pub fn make_fifth_dimension_deck() -> Vec<Card> {
    let mut result = Vec::new();
    let suits  = &[Suit::Clubs, Suit::Diamonds, Suit::Hearts, Suit::Spades, Suit::Stars];
    let values = &["1", "2", "3", "4", "5", "6", "7", "8", "9", "10",
                   "Jack", "Princess", "Queen", "King", "Joker"];

    for suit in suits {
        for value in values {
            result.push(Card::new(*suit, value));
        }
    }
    return result;
}

pub fn shuffle(deck: &mut Vec<Card>) {
    rand::thread_rng().shuffle(deck);
}
<commit_msg>change formatting<commit_after>extern crate rand;
use ::std::str;
use self::rand::Rng;

#[derive(Debug, Ord, PartialOrd, Eq, PartialEq, Copy, Clone)]
pub enum Suit { Clubs, Diamonds, Hearts, Spades, Stars }

#[derive(Debug, Eq, PartialEq, PartialOrd)]
pub struct Card {
    suit: Suit,
    value: &'static str,
}

impl Card {
    pub fn new(suit: Suit, value: &'static str) -> Card {
        Card { suit: suit, value: value }
    }
}

impl Ord for Card {
    fn cmp(&self, other: &Self) -> ::std::cmp::Ordering {
        self.value.cmp(&other.value)
    }
}

pub fn make_standard_deck() -> Vec<Card> {
    let mut result = Vec::new();
    let suits  = &[Suit::Clubs, Suit::Diamonds, Suit::Hearts, Suit::Spades];
    let values = &["2", "3", "4", "5", "6", "7", "8", "9", "10",
                   "Jack", "Queen", "King", "Ace"];

    for suit in suits {
        for value in values {
            result.push(Card::new(*suit, value));
        }
    }
    return result;
}

pub fn make_fifth_dimension_deck() -> Vec<Card> {
    let mut result = Vec::new();
    let suits  = &[Suit::Clubs, Suit::Diamonds, Suit::Hearts, Suit::Spades, Suit::Stars];
    let values = &["1", "2", "3", "4", "5", "6", "7", "8", "9", "10",
                   "Jack", "Princess", "Queen", "King", "Joker"];

    for suit in suits {
        for value in values {
            result.push(Card::new(*suit, value));
        }
    }
    return result;
}

pub fn shuffle(deck: &mut Vec<Card>) {
    rand::thread_rng().shuffle(deck);
}
<|endoftext|>


<commit_before>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the
// terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not
// distributed with this file, You can
// obtain one at
// http://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible
// With Secondary Licenses", as defined by
// the Mozilla Public License, v. 2.0.

use std::{str};
use std::ffi::{CStr, CString};
use std::thread;
use std::net::TcpStream;
use std::sync::mpsc::{channel, Sender, Receiver};

use super::libc::{c_int, c_char};
use super::simple_stream::bstream::Bstream;

/// This client's writer channel
static mut writer_tx: *mut Sender<Vec<u8>> = 0 as *mut Sender<Vec<u8>>;

/// This client's kill channel
static mut kill_tx: *mut Sender<()> = 0 as *mut Sender<()>;

#[no_mangle]
pub extern "C" fn start(address: *const c_char,
    data_handler: extern fn(*const c_char),
    on_connect_handler: extern fn(),
    on_disconnect_handler: extern fn()) -> c_int {

    // TODO - adjust this to accept a log level adjustable by whoever is running
    // the application
    // super::init();
    //
    // trace!("Rust - start()");
    //
    // let mut r_address;
    // unsafe {
    //     r_address = CStr::from_ptr(address);
    // }
    // let s_address = r_address.to_bytes();
    // let host_address = match str::from_utf8(s_address) {
    //     Ok(safe_str) => safe_str,
    //     Err(_) => {
    //         error!("Invalid host address");
    //         return -1 as c_int;
    //     }
    // };
    //
    // // Create and register a way to kill this client
    // let (k_tx, kill_rx): (Sender<()>, Receiver<()>) = channel();
    // unsafe { *kill_tx = k_tx.clone(); }
    //
    // // Writer thread's channel
    // let (w_tx, writer_rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();
    // unsafe { *writer_tx = w_tx.clone(); }
    //
    // debug!("Attempting connect to: {}", host_address);
    //
    // let result = TcpStream::connect(host_address);
    // if result.is_err() {
    //     error!("Error connecting to {} - {}", host_address, result.unwrap_err());
    //     return -1 as c_int;
    // }
    // debug!("Connected");
    // on_connect_handler();
    //
    // let stream = result.unwrap();
    // let client = Bstream::new(stream);
    // let r_client = client.clone();
    // let w_client = client.clone();
    //
    // // Start the reader thread
    // thread::Builder::new()
    //     .name("ReaderThread".to_string())
    //     .spawn(move||{
    //         reader_thread(r_client, data_handler)
    //     }).unwrap();
    //
    // // Start the writer thread
    // thread::Builder::new()
    //     .name("WriterThread".to_string())
    //     .spawn(move||{
    //         writer_thread(writer_rx, w_client)
    //     }).unwrap();
    //
    // // Wait for the kill signal
    // match kill_rx.recv() {
    //     Ok(_) => { }
    //     Err(e) => {
    //         error!("Error on kill channel: {}", e);
    //         on_disconnect_handler();
    //         return -1 as c_int;
    //     }
    // };
    // on_disconnect_handler();

    // Exit out in standard C fashion
    0 as c_int
}

/// Writes the complete contents of buffer to the server
/// Returns -1 on error
#[no_mangle]
pub extern "C" fn hydrogen_write(buffer: *const c_char) -> c_int {
    trace!("Rust.write");

    let mut buf_as_cstr;
    unsafe {
        buf_as_cstr = CStr::from_ptr(buffer);
    }
    let buf_as_slice = buf_as_cstr.to_bytes();

    let mut n_buffer = Vec::<u8>::with_capacity(buf_as_slice.len());
    for byte in buf_as_slice.iter() {
        n_buffer.push(*byte);
    }

    unsafe {
        match (*writer_tx).send(n_buffer) {
            Ok(_) => { }
            Err(e) => {
                warn!("Error sending buffer: {}", e);
                let _ = (*kill_tx).send(());
                return -1 as c_int;
            }
        };
    }

    0 as c_int
}

/// Forever listens to incoming data and when a complete message is received,
/// the passed callback is hit
fn reader_thread(client: Bstream, handler: extern fn(*const c_char)) {
    trace!("Rust.reader_thread started");

    let mut reader = client.clone();
    loop {
        match reader.read() {
            Ok(buffer) => {
                // Launch the handler in a new thread
                thread::Builder::new()
                    .name("Reader-Worker".to_string())
                    .spawn(move||{
                        let slice = &buffer[..];
                        let c_buffer = CString::new(slice).unwrap();
                        handler(c_buffer.as_ptr());
                    }).unwrap();
            }
            Err(e) => {
                error!("Error: {}", e);
                break;
            }
        };
    }
    debug!("Rust.reader_thread finished");
    unsafe { let _ = (*kill_tx).send(()); }
}

/// Forever listens to Receiver<Vec<u8>> waiting on messages to come in
/// Once available, blocks until the entire message has been written
fn writer_thread(rx: Receiver<Vec<u8>>, client: Bstream) {
    trace!("Rust.writer_thread started");

    let mut writer = client.clone();
    loop {
        match rx.recv() {
            Ok(ref mut buffer) => {
                match writer.write(buffer) {
                    Ok(_) => { }
                    Err(e) => {
                        error!("Error: {}", e);
                        break;
                    }
                };
            }
            Err(e) => {
                error!("Error: {}", e);
                break;
            }
        };
    }

    debug!("Rust.writer_thread finished");
    unsafe { let _ = (*kill_tx).send(()); }
}

/// Drops the current connection and kills all current threads
#[no_mangle]
pub extern "C" fn kill() {
    unsafe { let _ = (*kill_tx).send(()); }
}
<commit_msg>Debugging<commit_after>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the
// terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not
// distributed with this file, You can
// obtain one at
// http://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible
// With Secondary Licenses", as defined by
// the Mozilla Public License, v. 2.0.

use std::{str};
use std::ffi::{CStr, CString};
use std::thread;
use std::net::TcpStream;
use std::sync::mpsc::{channel, Sender, Receiver};

use super::libc::{c_int, c_char};
use super::simple_stream::bstream::Bstream;

/// This client's writer channel
static mut writer_tx: *mut Sender<Vec<u8>> = 0 as *mut Sender<Vec<u8>>;

/// This client's kill channel
static mut kill_tx: *mut Sender<()> = 0 as *mut Sender<()>;

#[no_mangle]
pub extern "C" fn start(address: *const c_char,
    data_handler: extern fn(*const c_char),
    on_connect_handler: extern fn(),
    on_disconnect_handler: extern fn()) -> c_int {

    // TODO - adjust this to accept a log level adjustable by whoever is running
    // the application
    super::init();
    //
    // trace!("Rust - start()");
    //
    // let mut r_address;
    // unsafe {
    //     r_address = CStr::from_ptr(address);
    // }
    // let s_address = r_address.to_bytes();
    // let host_address = match str::from_utf8(s_address) {
    //     Ok(safe_str) => safe_str,
    //     Err(_) => {
    //         error!("Invalid host address");
    //         return -1 as c_int;
    //     }
    // };
    //
    // // Create and register a way to kill this client
    // let (k_tx, kill_rx): (Sender<()>, Receiver<()>) = channel();
    // unsafe { *kill_tx = k_tx.clone(); }
    //
    // // Writer thread's channel
    // let (w_tx, writer_rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();
    // unsafe { *writer_tx = w_tx.clone(); }
    //
    // debug!("Attempting connect to: {}", host_address);
    //
    // let result = TcpStream::connect(host_address);
    // if result.is_err() {
    //     error!("Error connecting to {} - {}", host_address, result.unwrap_err());
    //     return -1 as c_int;
    // }
    // debug!("Connected");
    // on_connect_handler();
    //
    // let stream = result.unwrap();
    // let client = Bstream::new(stream);
    // let r_client = client.clone();
    // let w_client = client.clone();
    //
    // // Start the reader thread
    // thread::Builder::new()
    //     .name("ReaderThread".to_string())
    //     .spawn(move||{
    //         reader_thread(r_client, data_handler)
    //     }).unwrap();
    //
    // // Start the writer thread
    // thread::Builder::new()
    //     .name("WriterThread".to_string())
    //     .spawn(move||{
    //         writer_thread(writer_rx, w_client)
    //     }).unwrap();
    //
    // // Wait for the kill signal
    // match kill_rx.recv() {
    //     Ok(_) => { }
    //     Err(e) => {
    //         error!("Error on kill channel: {}", e);
    //         on_disconnect_handler();
    //         return -1 as c_int;
    //     }
    // };
    // on_disconnect_handler();

    // Exit out in standard C fashion
    0 as c_int
}

/// Writes the complete contents of buffer to the server
/// Returns -1 on error
#[no_mangle]
pub extern "C" fn hydrogen_write(buffer: *const c_char) -> c_int {
    trace!("Rust.write");

    let mut buf_as_cstr;
    unsafe {
        buf_as_cstr = CStr::from_ptr(buffer);
    }
    let buf_as_slice = buf_as_cstr.to_bytes();

    let mut n_buffer = Vec::<u8>::with_capacity(buf_as_slice.len());
    for byte in buf_as_slice.iter() {
        n_buffer.push(*byte);
    }

    unsafe {
        match (*writer_tx).send(n_buffer) {
            Ok(_) => { }
            Err(e) => {
                warn!("Error sending buffer: {}", e);
                let _ = (*kill_tx).send(());
                return -1 as c_int;
            }
        };
    }

    0 as c_int
}

/// Forever listens to incoming data and when a complete message is received,
/// the passed callback is hit
fn reader_thread(client: Bstream, handler: extern fn(*const c_char)) {
    trace!("Rust.reader_thread started");

    let mut reader = client.clone();
    loop {
        match reader.read() {
            Ok(buffer) => {
                // Launch the handler in a new thread
                thread::Builder::new()
                    .name("Reader-Worker".to_string())
                    .spawn(move||{
                        let slice = &buffer[..];
                        let c_buffer = CString::new(slice).unwrap();
                        handler(c_buffer.as_ptr());
                    }).unwrap();
            }
            Err(e) => {
                error!("Error: {}", e);
                break;
            }
        };
    }
    debug!("Rust.reader_thread finished");
    unsafe { let _ = (*kill_tx).send(()); }
}

/// Forever listens to Receiver<Vec<u8>> waiting on messages to come in
/// Once available, blocks until the entire message has been written
fn writer_thread(rx: Receiver<Vec<u8>>, client: Bstream) {
    trace!("Rust.writer_thread started");

    let mut writer = client.clone();
    loop {
        match rx.recv() {
            Ok(ref mut buffer) => {
                match writer.write(buffer) {
                    Ok(_) => { }
                    Err(e) => {
                        error!("Error: {}", e);
                        break;
                    }
                };
            }
            Err(e) => {
                error!("Error: {}", e);
                break;
            }
        };
    }

    debug!("Rust.writer_thread finished");
    unsafe { let _ = (*kill_tx).send(()); }
}

/// Drops the current connection and kills all current threads
#[no_mangle]
pub extern "C" fn kill() {
    unsafe { let _ = (*kill_tx).send(()); }
}
<|endoftext|>


<commit_before>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the
// terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not
// distributed with this file, You can
// obtain one at
// http://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible
// With Secondary Licenses", as defined by
// the Mozilla Public License, v. 2.0.

use std::{str, ptr};
use std::ffi::{CStr, CString};
use std::thread;
use std::net::TcpStream;
use std::sync::mpsc::{channel, Sender, Receiver};

use super::libc::{c_int, c_char};
use super::simple_stream::bstream::Bstream;

extern "C" {
    fn register_writer_tx(tx: *mut Sender<Vec<u8>>);
    fn register_stop_tx(tx: *mut Sender<()>);
}

/// Connects to the provided address, (eg "123.123.123.123:3000") and registers
/// the on data received handler
#[no_mangle]
pub extern "C" fn connect(address: *const c_char,
    handler: extern fn(*const c_char, c_int),
    on_connect_handler: extern fn(),
    on_disconnect_handler: extern fn()) -> c_int {

    let mut r_address;
    unsafe {
        r_address = CStr::from_ptr(address);
    }
    let s_address = r_address.to_bytes();
    let host_address = match str::from_utf8(s_address) {
        Ok(safe_str) => safe_str,
        Err(_) => {
            println!("Invalid host address");
            return -1 as c_int;
        }
    };

    // Create and register a way to kill this client
    let (k_tx, kill_rx): (Sender<()>, Receiver<()>) = channel();
    let kill_tx = k_tx.clone();
    let mut k_tx_ptr = Box::new(k_tx);
    unsafe {
        register_stop_tx(&mut *k_tx_ptr);
    }

    // Writer thread's channel
    let (w_tx, w_rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();
    let mut w_tx_ptr = Box::new(w_tx);
    unsafe {
        register_writer_tx(&mut *w_tx_ptr);
    }

    println!("host_address: {}", host_address);

    // let result = TcpStream::connect(host_address);
    // if result.is_err() {
    //     println!("Error connecting to {} - {}", host_address, result.unwrap_err());
    //     return -1 as c_int;
    // }
    // println!("Connected");
    // on_connect_handler();

    // let stream = result.unwrap();
    // let client = Bstream::new(stream);
    //
    // let r_client = client.clone();
    // let r_kill_tx = kill_tx.clone();
    //
    // let w_client = client.clone();
    // let w_kill_tx = kill_tx.clone();
    //
    // // Start the reader thread
    // thread::Builder::new()
    //     .name("ReaderThread".to_string())
    //     .spawn(move||{
    //         reader_thread(r_client, handler, r_kill_tx)
    //     }).unwrap();
    //
    // // Start the writer thread
    // thread::Builder::new()
    //     .name("WriterThread".to_string())
    //     .spawn(move||{
    //         writer_thread(w_rx, w_client, w_kill_tx)
    //     }).unwrap();
    //
    // // Wait for the kill signal
    // match kill_rx.recv() {
    //     Ok(_) => { }
    //     Err(e) => {
    //         println!("Error on kill channel: {}", e);
    //         return -1 as c_int;
    //     }
    // };
    // on_disconnect_handler();

    // Exit out in standard C fashion
    0 as c_int
}

/// Writes the complete contents of buffer to the server
/// Returns -1 on error
#[no_mangle]
pub extern "C" fn send_to_writer(w_tx: *mut Sender<Vec<u8>>,
                                 buffer: *const c_char,
                                 count: c_int,
                                 k_tx: *mut Sender<()>) -> c_int {
    if count < 1 {
        println!("Error - count must be greater than zero");
        return -1 as c_int;
    }

    let num_elts = count as usize;
    let mut n_buffer = Vec::<u8>::with_capacity(num_elts);
    for x in 0..num_elts as isize {
        unsafe {
            n_buffer.push(ptr::read(buffer.offset(x)) as u8);
        }
    }

    unsafe {
        match (*w_tx).send(n_buffer) {
            Ok(_) => { }
            Err(e) => {
                println!("Error sending buffer: {}", e);
                let _ = (*k_tx).send(());
                return -1 as c_int;
            }
        };
    }

    0 as c_int
}

/// Forever listens to incoming data and when a complete message is received,
/// the passed callback is hit
fn reader_thread(client: Bstream,
                 event_handler: extern fn(*const c_char, c_int),
                 kill_tx: Sender<()>) {

    let mut reader = client.clone();
    loop {
        match reader.read() {
            Ok(buffer) => {
                // Launch the handler in a new thread
                thread::Builder::new()
                    .name("Reader-Worker".to_string())
                    .spawn(move||{
                        let slice = &buffer[..];
                        let c_buffer = CString::new(slice).unwrap();
                        event_handler(c_buffer.as_ptr(), buffer.len() as c_int);
                    }).unwrap();
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Reader thread finished");
    let _ = kill_tx.send(());
}

/// Forever listens to Receiver<Vec<u8>> waiting on messages to come in
/// Once available, blocks until the entire message has been written
fn writer_thread(rx: Receiver<Vec<u8>>, client: Bstream, kill_tx: Sender<()>) {
    let mut writer = client.clone();
    loop {
        match rx.recv() {
            Ok(ref mut buffer) => {
                match writer.write(buffer) {
                    Ok(_) => { }
                    Err(e) => {
                        println!("Error: {}", e);
                        break;
                    }
                };
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Writer thread finished");
    let _ = kill_tx.send(());
}
<commit_msg>Debugging<commit_after>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the
// terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not
// distributed with this file, You can
// obtain one at
// http://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible
// With Secondary Licenses", as defined by
// the Mozilla Public License, v. 2.0.

use std::{str, ptr};
use std::ffi::{CStr, CString};
use std::thread;
use std::net::TcpStream;
use std::sync::mpsc::{channel, Sender, Receiver};

use super::libc::{c_int, c_char};
use super::simple_stream::bstream::Bstream;

extern "C" {
    fn register_writer_tx(tx: *mut Sender<Vec<u8>>);
    fn register_stop_tx(tx: *mut Sender<()>);
}

/// Connects to the provided address, (eg "123.123.123.123:3000") and registers
/// the on data received handler
#[no_mangle]
pub extern "C" fn connect(address: *const c_char,
    handler: extern fn(*const c_char, c_int),
    on_connect_handler: extern fn(),
    on_disconnect_handler: extern fn()) -> c_int {

    let mut r_address;
    unsafe {
        r_address = CStr::from_ptr(address);
    }
    let s_address = r_address.to_bytes();
    let host_address = match str::from_utf8(s_address) {
        Ok(safe_str) => safe_str,
        Err(_) => {
            println!("Invalid host address");
            return -1 as c_int;
        }
    };

    // Create and register a way to kill this client
    let (k_tx, kill_rx): (Sender<()>, Receiver<()>) = channel();
    let kill_tx = k_tx.clone();
    let mut k_tx_ptr = Box::new(k_tx);
    unsafe {
        register_stop_tx(&mut *k_tx_ptr);
    }

    // Writer thread's channel
    let (w_tx, w_rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();
    let mut w_tx_ptr = Box::new(w_tx);
    unsafe {
        register_writer_tx(&mut *w_tx_ptr);
    }

    // println!("host_address: {}", host_address);

    // let result = TcpStream::connect(host_address);
    // if result.is_err() {
    //     println!("Error connecting to {} - {}", host_address, result.unwrap_err());
    //     return -1 as c_int;
    // }
    // println!("Connected");
    // on_connect_handler();

    // let stream = result.unwrap();
    // let client = Bstream::new(stream);
    //
    // let r_client = client.clone();
    // let r_kill_tx = kill_tx.clone();
    //
    // let w_client = client.clone();
    // let w_kill_tx = kill_tx.clone();
    //
    // // Start the reader thread
    // thread::Builder::new()
    //     .name("ReaderThread".to_string())
    //     .spawn(move||{
    //         reader_thread(r_client, handler, r_kill_tx)
    //     }).unwrap();
    //
    // // Start the writer thread
    // thread::Builder::new()
    //     .name("WriterThread".to_string())
    //     .spawn(move||{
    //         writer_thread(w_rx, w_client, w_kill_tx)
    //     }).unwrap();
    //
    // // Wait for the kill signal
    // match kill_rx.recv() {
    //     Ok(_) => { }
    //     Err(e) => {
    //         println!("Error on kill channel: {}", e);
    //         return -1 as c_int;
    //     }
    // };
    // on_disconnect_handler();

    // Exit out in standard C fashion
    0 as c_int
}

/// Writes the complete contents of buffer to the server
/// Returns -1 on error
#[no_mangle]
pub extern "C" fn send_to_writer(w_tx: *mut Sender<Vec<u8>>,
                                 buffer: *const c_char,
                                 count: c_int,
                                 k_tx: *mut Sender<()>) -> c_int {
    if count < 1 {
        println!("Error - count must be greater than zero");
        return -1 as c_int;
    }

    let num_elts = count as usize;
    let mut n_buffer = Vec::<u8>::with_capacity(num_elts);
    for x in 0..num_elts as isize {
        unsafe {
            n_buffer.push(ptr::read(buffer.offset(x)) as u8);
        }
    }

    unsafe {
        match (*w_tx).send(n_buffer) {
            Ok(_) => { }
            Err(e) => {
                println!("Error sending buffer: {}", e);
                let _ = (*k_tx).send(());
                return -1 as c_int;
            }
        };
    }

    0 as c_int
}

/// Forever listens to incoming data and when a complete message is received,
/// the passed callback is hit
fn reader_thread(client: Bstream,
                 event_handler: extern fn(*const c_char, c_int),
                 kill_tx: Sender<()>) {

    let mut reader = client.clone();
    loop {
        match reader.read() {
            Ok(buffer) => {
                // Launch the handler in a new thread
                thread::Builder::new()
                    .name("Reader-Worker".to_string())
                    .spawn(move||{
                        let slice = &buffer[..];
                        let c_buffer = CString::new(slice).unwrap();
                        event_handler(c_buffer.as_ptr(), buffer.len() as c_int);
                    }).unwrap();
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Reader thread finished");
    let _ = kill_tx.send(());
}

/// Forever listens to Receiver<Vec<u8>> waiting on messages to come in
/// Once available, blocks until the entire message has been written
fn writer_thread(rx: Receiver<Vec<u8>>, client: Bstream, kill_tx: Sender<()>) {
    let mut writer = client.clone();
    loop {
        match rx.recv() {
            Ok(ref mut buffer) => {
                match writer.write(buffer) {
                    Ok(_) => { }
                    Err(e) => {
                        println!("Error: {}", e);
                        break;
                    }
                };
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Writer thread finished");
    let _ = kill_tx.send(());
}
<|endoftext|>


<commit_before>pub mod messages;

use address::Address;
use mio;
use network;
use self::messages::{CommMessage};
use std::collections::{HashMap, HashSet};
use std::io;
use std::sync::mpsc;
use std::thread;

#[derive(Debug)]
pub enum Task {
    HandleNetworkEvent(network::Event),
    ScheduleMessageDelivery(Address, CommMessage)
}

#[derive(Debug)]
pub enum ScheduledTask {
    DeliverMessage(Address, CommMessage),
}

pub struct Client {
    address: Address,
    network_commands: Option<network::TaskSender>,
    received: HashSet<Address>,
    pending_deliveries: HashMap<Address, mio::Timeout>,
    delivered: HashMap<Address, usize>,
    acknowledgements: HashMap<Address, CommMessage>
}

impl Client {
    pub fn new(address: Address) -> Client {
        Client {
            address: address,
            network_commands: None,
            received: HashSet::new(),
            pending_deliveries: HashMap::new(),
            delivered: HashMap::new(),
            acknowledgements: HashMap::new()
        }
    }

    pub fn run(mut self, mut network: network::Network, headless: bool) {
        let mut event_loop = mio::EventLoop::new().unwrap();
        let (event_sender, event_receiver) = mpsc::channel();
        network.register_event_listener(event_sender);
        let notify_channel = event_loop.channel();
        let sender = self.address;

        thread::spawn(move|| {
            for event in event_receiver.iter() {
                notify_channel
                    .send(Task::HandleNetworkEvent(event))
                    .unwrap_or_else(|err| info!("Couldn't handle network event: {:?}", err));
            }
        });

        self.network_commands = Some(network.run());

        let notify_channel = event_loop.channel();
        thread::spawn(move || event_loop.run(&mut self).unwrap());

        info!("Running client at {}", sender);

        if headless {
            loop { thread::park(); }
        } else {
            loop {
                let mut line = String::new();
                io::stdin().read_line(&mut line).unwrap();
                let parts: Vec<&str> = line.splitn(2, ' ').collect();
                let recipient = Address::from_str(parts[0]);
                let message_text = parts[1].trim().to_string();

                let message = messages::create_text_message(recipient, sender, message_text);
                notify_channel
                    .send(Task::ScheduleMessageDelivery(recipient, message))
                    .unwrap_or_else(|err| info!("Couldn't schedule message delivery: {:?}", err));
            }
        }
    }

    fn handle_networking_event(&mut self, event: network::Event, event_loop: &mut mio::EventLoop<Client>) {
        match event {
            network::Event::ReceivedPacket(data) => {
                let comm_message = messages::decode(data);
                let CommMessage { recipient, .. } = comm_message.clone();

                if let Some(ref text_message) = comm_message.text_message {
                    if !self.received.contains(&text_message.id) {
                        self.received.insert(text_message.id);
                        if recipient == self.address {
                            println!("{}: {}", text_message.sender, text_message.text);
                            let ack = messages::create_message_acknowledgement(text_message.sender, text_message.id);
                            self.schedule_message_delivery(text_message.sender, ack, event_loop);
                        } else {
                            if let Some(ack) = self.acknowledgements.remove(&text_message.id) {
                                self.schedule_message_delivery(recipient, ack.clone(), event_loop);
                                self.acknowledgements.insert(text_message.id, ack);
                            } else {
                                self.schedule_message_delivery(recipient, comm_message.clone(), event_loop);
                            }
                        }
                    }
                }

                if let Some(ref ack) = comm_message.message_acknowledgement {
                    if let None = self.acknowledgements.insert(ack.id, comm_message.clone()) {
                        self.pending_deliveries.remove(&ack.id).map(|p| event_loop.clear_timeout(p));
                        if recipient == self.address {
                            debug!("ack {}", ack.id);
                        } else {
                            self.schedule_message_delivery(recipient, comm_message.clone(), event_loop);
                        }
                    }
                }
            }
        }
    }

    fn schedule_message_delivery(&mut self, recipient: Address, message: CommMessage, event_loop: &mut mio::EventLoop<Client>) {
        if let Some(ref text_message) = message.text_message {
            if !self.pending_deliveries.contains_key(&text_message.id) {
                let delivered = self.delivered.entry(text_message.id).or_insert(0);
                let delay = (2u64.pow(*delivered as u32) - 1) * 1000;
                debug!("Delivery with delay {:?}", delay);
                let timeout = event_loop.timeout_ms(ScheduledTask::DeliverMessage(recipient, message.clone()), delay).unwrap();
                self.pending_deliveries.insert(text_message.id, timeout);
                *delivered += 1;
            }
        }

        if let Some(_) = message.message_acknowledgement {
            event_loop.timeout_ms(ScheduledTask::DeliverMessage(recipient, message.clone()), 0).unwrap();
        }
    }

    fn deliver_message(&mut self, recipient: Address, message: CommMessage) {
        if let Some(ref commands) = self.network_commands {
            if let Some(ref text_message) = message.text_message {
                self.pending_deliveries.remove(&text_message.id);
            }
            commands.send(network::OneshotTask::SendPacket(recipient, message.encode())).unwrap();
        }
    }
}

impl mio::Handler for Client {
    type Timeout = ScheduledTask;
    type Message = Task;

    fn notify(&mut self, event_loop: &mut mio::EventLoop<Client>, task: Task) {
        match task {
            Task::HandleNetworkEvent(event) => self.handle_networking_event(event, event_loop),
            Task::ScheduleMessageDelivery(recipient, message) => self.schedule_message_delivery(recipient, message, event_loop)
        }
    }

    fn timeout(&mut self, _event_loop: &mut mio::EventLoop<Client>, task: ScheduledTask) {
        match task {
            ScheduledTask::DeliverMessage(recipient, message) => self.deliver_message(recipient, message)
        }
    }
}
<commit_msg>Make message's keep geting delivered with exp backoff<commit_after>pub mod messages;

use address::Address;
use mio;
use network;
use self::messages::{CommMessage};
use std::collections::{HashMap, HashSet};
use std::io;
use std::sync::mpsc;
use std::thread;

#[derive(Debug)]
pub enum Task {
    HandleNetworkEvent(network::Event),
    ScheduleMessageDelivery(Address, CommMessage)
}

#[derive(Debug)]
pub enum ScheduledTask {
    DeliverMessage(Address, CommMessage),
}

pub struct Client {
    address: Address,
    network_commands: Option<network::TaskSender>,
    received: HashSet<Address>,
    pending_deliveries: HashMap<Address, mio::Timeout>,
    delivered: HashMap<Address, usize>,
    acknowledgements: HashMap<Address, CommMessage>
}

impl Client {
    pub fn new(address: Address) -> Client {
        Client {
            address: address,
            network_commands: None,
            received: HashSet::new(),
            pending_deliveries: HashMap::new(),
            delivered: HashMap::new(),
            acknowledgements: HashMap::new()
        }
    }

    pub fn run(mut self, mut network: network::Network, headless: bool) {
        let mut event_loop = mio::EventLoop::new().unwrap();
        let (event_sender, event_receiver) = mpsc::channel();
        network.register_event_listener(event_sender);
        let notify_channel = event_loop.channel();
        let sender = self.address;

        thread::spawn(move|| {
            for event in event_receiver.iter() {
                notify_channel
                    .send(Task::HandleNetworkEvent(event))
                    .unwrap_or_else(|err| info!("Couldn't handle network event: {:?}", err));
            }
        });

        self.network_commands = Some(network.run());

        let notify_channel = event_loop.channel();
        thread::spawn(move || event_loop.run(&mut self).unwrap());

        info!("Running client at {}", sender);

        if headless {
            loop { thread::park(); }
        } else {
            loop {
                let mut line = String::new();
                io::stdin().read_line(&mut line).unwrap();
                let parts: Vec<&str> = line.splitn(2, ' ').collect();
                let recipient = Address::from_str(parts[0]);
                let message_text = parts[1].trim().to_string();

                let message = messages::create_text_message(recipient, sender, message_text);
                notify_channel
                    .send(Task::ScheduleMessageDelivery(recipient, message))
                    .unwrap_or_else(|err| info!("Couldn't schedule message delivery: {:?}", err));
            }
        }
    }

    fn handle_networking_event(&mut self, event: network::Event, event_loop: &mut mio::EventLoop<Client>) {
        match event {
            network::Event::ReceivedPacket(data) => {
                let comm_message = messages::decode(data);
                let CommMessage { recipient, .. } = comm_message.clone();

                if let Some(ref text_message) = comm_message.text_message {
                    if !self.received.contains(&text_message.id) {
                        self.received.insert(text_message.id);
                        if recipient == self.address {
                            println!("{}: {}", text_message.sender, text_message.text);
                            let ack = messages::create_message_acknowledgement(text_message.sender, text_message.id);
                            self.schedule_message_delivery(text_message.sender, ack, event_loop);
                        } else {
                            if let Some(ack) = self.acknowledgements.remove(&text_message.id) {
                                self.schedule_message_delivery(recipient, ack.clone(), event_loop);
                                self.acknowledgements.insert(text_message.id, ack);
                            } else {
                                self.schedule_message_delivery(recipient, comm_message.clone(), event_loop);
                            }
                        }
                    }
                }

                if let Some(ref ack) = comm_message.message_acknowledgement {
                    if let None = self.acknowledgements.insert(ack.id, comm_message.clone()) {
                        self.pending_deliveries.remove(&ack.id).map(|p| event_loop.clear_timeout(p));
                        if recipient == self.address {
                            debug!("ack {}", ack.id);
                        } else {
                            self.schedule_message_delivery(recipient, comm_message.clone(), event_loop);
                        }
                    }
                }
            }
        }
    }

    fn schedule_message_delivery(&mut self, recipient: Address, message: CommMessage, event_loop: &mut mio::EventLoop<Client>) {
        if let Some(ref text_message) = message.text_message {
            if !self.pending_deliveries.contains_key(&text_message.id) {
                let delivered = self.delivered.entry(text_message.id).or_insert(0);
                let delay = (2u64.pow(*delivered as u32) - 1) * 1000;
                debug!("Delivery with delay {:?}", delay);
                let timeout = event_loop.timeout_ms(ScheduledTask::DeliverMessage(recipient, message.clone()), delay).unwrap();
                self.pending_deliveries.insert(text_message.id, timeout);
                *delivered += 1;
            }
        }

        if let Some(_) = message.message_acknowledgement {
            event_loop.timeout_ms(ScheduledTask::DeliverMessage(recipient, message.clone()), 0).unwrap();
        }
    }

    fn deliver_message(&mut self, recipient: Address, message: CommMessage, event_loop: &mut mio::EventLoop<Client>) {
        if let Some(ref text_message) = message.text_message {
            self.pending_deliveries.remove(&text_message.id);
            self.schedule_message_delivery(recipient, message.clone(), event_loop);
        }
        if let Some(ref commands) = self.network_commands {
            commands.send(network::OneshotTask::SendPacket(recipient, message.encode())).unwrap();
        }
    }
}

impl mio::Handler for Client {
    type Timeout = ScheduledTask;
    type Message = Task;

    fn notify(&mut self, event_loop: &mut mio::EventLoop<Client>, task: Task) {
        match task {
            Task::HandleNetworkEvent(event) => self.handle_networking_event(event, event_loop),
            Task::ScheduleMessageDelivery(recipient, message) => self.schedule_message_delivery(recipient, message, event_loop)
        }
    }

    fn timeout(&mut self, event_loop: &mut mio::EventLoop<Client>, task: ScheduledTask) {
        match task {
            ScheduledTask::DeliverMessage(recipient, message) => self.deliver_message(recipient, message, event_loop)
        }
    }
}
<|endoftext|>


<commit_before>//! # client
//!
//! Ssh2-config implementation with a ssh2 client

/**
 * MIT License
 *
 * ssh2-config - Copyright (c) 2021 Christian Visintin
 *
 * 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 DEALINGS IN THE
 * SOFTWARE.
 */
use dirs::home_dir;
use ssh2::{MethodType, Session};
use ssh2_config::{HostParams, SshConfig};
use std::env::args;
use std::fs::File;
use std::io::BufReader;
use std::net::{SocketAddr, TcpStream, ToSocketAddrs};
use std::path::{Path, PathBuf};
use std::process::exit;
use std::time::Duration;

pub fn client() {
    // get args
    let args: Vec<String> = args().collect();
    let address = match args.get(1) {
        Some(addr) => addr.to_string(),
        None => {
            usage();
            exit(255)
        }
    };
    // check path
    let config_path = match args.get(2) {
        Some(p) => PathBuf::from(p),
        None => {
            let mut p = home_dir().expect("Failed to get home_dir for guest OS");
            p.extend(Path::new(".ssh/config"));
            p
        }
    };
    // Open config file
    let config = read_config(config_path.as_path());
    println!("Config file: {:?}", config);

    let params = config.query(address.as_str());
    println!("Params: {:?}", params);

    connect(address.as_str(), &params);
}

fn usage() {
    eprintln!("Usage: cargo run --example client -- <address:port> [config-path]");
}

fn read_config(p: &Path) -> SshConfig {
    let mut reader = match File::open(p) {
        Ok(f) => BufReader::new(f),
        Err(err) => panic!("Could not open file '{}': {}", p.display(), err),
    };
    match SshConfig::default().parse(&mut reader) {
        Ok(config) => config,
        Err(err) => panic!("Failed to parse configuration: {}", err),
    }
}

fn connect(host: &str, params: &HostParams) {
    // Resolve host
    let host = params.host_name.as_deref().unwrap_or(host);
    let port = params.port.unwrap_or(22);

    let formatted_host : String;
    let host: &str = if host.contains(':') {
        host
    } else {
        formatted_host = format!("{}:{}", host, port);
        &*formatted_host
    };

    println!("Connecting to host {}...", host);
    let socket_addresses: Vec<_> = host.to_socket_addrs().expect("Could not parse host").collect();
    // Try addresses
    let stream =
        socket_addresses.iter().find_map(|socket_addr|
            TcpStream::connect_timeout(
                socket_addr,
                params.connect_timeout.unwrap_or(Duration::from_secs(30)),
            ).map(|s| {
                println!("Established connection with {}", socket_addr);
                s
            }).ok()).expect("No suitable socket address found; connection timeout");

    let mut session: Session = match Session::new() {
        Ok(s) => s,
        Err(err) => {
            panic!("Could not create session: {}", err);
        }
    };
    // Configure session
    configure_session(&mut session, params);
    // Connect
    session.set_tcp_stream(stream);
    if let Err(err) = session.handshake() {
        panic!("Handshake failed: {}", err);
    }
    // Get username
    let username = match params.user.as_ref() {
        Some(u) => {
            println!("Using username '{}'", u);
            u.clone()
        }
        None => match std::env::var("USER") {
            Ok(val) => val,
            Err(e) => {
                println!("No USER env variable and no user in .ssh/config (consider using %i in .config)");
                exit(1)
            }
        },
    };
    for i in (params.identity_file.as_ref()).unwrap().iter() {
        if let Ok(_) = session.userauth_pubkey_file(&username, None, i, None) {
            break;
        }
    }
    if let Some(banner) = session.banner() {
        println!("{}", banner);
    }
    println!("Connection OK!");
    if let Err(err) = session.disconnect(None, "mandi mandi!", None) {
        panic!("Disconnection failed: {}", err);
    }
}

fn configure_session(session: &mut Session, params: &HostParams) {
    println!("Configuring session...");
    if let Some(compress) = params.compression {
        println!("compression: {}", compress);
        session.set_compress(compress);
    }
    if params.tcp_keep_alive.unwrap_or(false) && params.server_alive_interval.is_some() {
        let interval = params.server_alive_interval.unwrap().as_secs() as u32;
        println!("keepalive interval: {} seconds", interval);
        session.set_keepalive(true, interval);
    }
    // algos
    if let Some(algos) = params.kex_algorithms.as_deref() {
        if let Err(err) = session.method_pref(MethodType::Kex, algos.join(",").as_str()) {
            panic!("Could not set KEX algorithms: {}", err);
        }
    }
    if let Some(algos) = params.host_key_algorithms.as_deref() {
        if let Err(err) = session.method_pref(MethodType::HostKey, algos.join(",").as_str()) {
            panic!("Could not set host key algorithms: {}", err);
        }
    }
    if let Some(algos) = params.ciphers.as_deref() {
        if let Err(err) = session.method_pref(MethodType::CryptCs, algos.join(",").as_str()) {
            panic!("Could not set crypt algorithms (client-server): {}", err);
        }
        if let Err(err) = session.method_pref(MethodType::CryptSc, algos.join(",").as_str()) {
            panic!("Could not set crypt algorithms (server-client): {}", err);
        }
    }
    if let Some(algos) = params.mac.as_deref() {
        if let Err(err) = session.method_pref(MethodType::MacCs, algos.join(",").as_str()) {
            panic!("Could not set MAC algorithms (client-server): {}", err);
        }
        if let Err(err) = session.method_pref(MethodType::MacSc, algos.join(",").as_str()) {
            panic!("Could not set MAC algorithms (server-client): {}", err);
        }
    }
}
<commit_msg>src/client: listen on port 23 as in the Go implementation<commit_after>//! # client
//!
//! Ssh2-config implementation with a ssh2 client

/**
 * MIT License
 *
 * ssh2-config - Copyright (c) 2021 Christian Visintin
 *
 * 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 DEALINGS IN THE
 * SOFTWARE.
 */
use dirs::home_dir;
use ssh2::{MethodType, Session};
use ssh2_config::{HostParams, SshConfig};
use std::env::args;
use std::fs::File;
use std::io::BufReader;
use std::net::{SocketAddr, TcpStream, ToSocketAddrs};
use std::path::{Path, PathBuf};
use std::process::exit;
use std::time::Duration;

pub fn client() {
    // get args
    let args: Vec<String> = args().collect();
    let address = match args.get(1) {
        Some(addr) => addr.to_string(),
        None => {
            usage();
            exit(255)
        }
    };
    // check path
    let config_path = match args.get(2) {
        Some(p) => PathBuf::from(p),
        None => {
            let mut p = home_dir().expect("Failed to get home_dir for guest OS");
            p.extend(Path::new(".ssh/config"));
            p
        }
    };
    // Open config file
    let config = read_config(config_path.as_path());
    println!("Config file: {:?}", config);

    let params = config.query(address.as_str());
    println!("Params: {:?}", params);

    connect(address.as_str(), &params);
}

fn usage() {
    eprintln!("Usage: cargo run --example client -- <address:port> [config-path]");
}

fn read_config(p: &Path) -> SshConfig {
    let mut reader = match File::open(p) {
        Ok(f) => BufReader::new(f),
        Err(err) => panic!("Could not open file '{}': {}", p.display(), err),
    };
    match SshConfig::default().parse(&mut reader) {
        Ok(config) => config,
        Err(err) => panic!("Failed to parse configuration: {}", err),
    }
}

fn connect(host: &str, params: &HostParams) {
    // Resolve host
    let host = params.host_name.as_deref().unwrap_or(host);
    let port = params.port.unwrap_or(23);

    let formatted_host : String;
    let host: &str = if host.contains(':') {
        host
    } else {
        formatted_host = format!("{}:{}", host, port);
        &*formatted_host
    };

    println!("Connecting to host {}...", host);
    let socket_addresses: Vec<_> = host.to_socket_addrs().expect("Could not parse host").collect();
    // Try addresses
    let stream =
        socket_addresses.iter().find_map(|socket_addr|
            TcpStream::connect_timeout(
                socket_addr,
                params.connect_timeout.unwrap_or(Duration::from_secs(30)),
            ).map(|s| {
                println!("Established connection with {}", socket_addr);
                s
            }).ok()).expect("No suitable socket address found; connection timeout");

    let mut session: Session = match Session::new() {
        Ok(s) => s,
        Err(err) => {
            panic!("Could not create session: {}", err);
        }
    };
    // Configure session
    configure_session(&mut session, params);
    // Connect
    session.set_tcp_stream(stream);
    if let Err(err) = session.handshake() {
        panic!("Handshake failed: {}", err);
    }
    // Get username
    let username = match params.user.as_ref() {
        Some(u) => {
            println!("Using username '{}'", u);
            u.clone()
        }
        None => match std::env::var("USER") {
            Ok(val) => val,
            Err(e) => {
                println!("No USER env variable and no user in .ssh/config (consider using %i in .config)");
                exit(1)
            }
        },
    };
    for i in (params.identity_file.as_ref()).unwrap().iter() {
        if let Ok(_) = session.userauth_pubkey_file(&username, None, i, None) {
            break;
        }
    }
    if let Some(banner) = session.banner() {
        println!("{}", banner);
    }
    println!("Connection OK!");
    if let Err(err) = session.disconnect(None, "mandi mandi!", None) {
        panic!("Disconnection failed: {}", err);
    }
}

fn configure_session(session: &mut Session, params: &HostParams) {
    println!("Configuring session...");
    if let Some(compress) = params.compression {
        println!("compression: {}", compress);
        session.set_compress(compress);
    }
    if params.tcp_keep_alive.unwrap_or(false) && params.server_alive_interval.is_some() {
        let interval = params.server_alive_interval.unwrap().as_secs() as u32;
        println!("keepalive interval: {} seconds", interval);
        session.set_keepalive(true, interval);
    }
    // algos
    if let Some(algos) = params.kex_algorithms.as_deref() {
        if let Err(err) = session.method_pref(MethodType::Kex, algos.join(",").as_str()) {
            panic!("Could not set KEX algorithms: {}", err);
        }
    }
    if let Some(algos) = params.host_key_algorithms.as_deref() {
        if let Err(err) = session.method_pref(MethodType::HostKey, algos.join(",").as_str()) {
            panic!("Could not set host key algorithms: {}", err);
        }
    }
    if let Some(algos) = params.ciphers.as_deref() {
        if let Err(err) = session.method_pref(MethodType::CryptCs, algos.join(",").as_str()) {
            panic!("Could not set crypt algorithms (client-server): {}", err);
        }
        if let Err(err) = session.method_pref(MethodType::CryptSc, algos.join(",").as_str()) {
            panic!("Could not set crypt algorithms (server-client): {}", err);
        }
    }
    if let Some(algos) = params.mac.as_deref() {
        if let Err(err) = session.method_pref(MethodType::MacCs, algos.join(",").as_str()) {
            panic!("Could not set MAC algorithms (client-server): {}", err);
        }
        if let Err(err) = session.method_pref(MethodType::MacSc, algos.join(",").as_str()) {
            panic!("Could not set MAC algorithms (server-client): {}", err);
        }
    }
}
<|endoftext|>


<commit_before>/// A color
#[derive(Copy, Clone)]
pub struct Color {
    pub data: u32,
}

impl Color {
    /// Create a new color from RGB
    pub fn rgb(r: u8, g: u8, b: u8) -> Self {
        Color { data: 0xFF000000 | ((r as u32) << 16) | ((g as u32) << 8) | (b as u32) }
    }

    /// Set the alpha
    pub fn rgba(r: u8, g: u8, b: u8, a: u8) -> Self {
        Color { data: ((a as u32) << 24) | ((r as u32) << 16) | ((g as u32) << 8) | (b as u32) }
    }
}
<commit_msg>Add more specific color methods for `Color`<commit_after>/// A color
#[derive(Copy, Clone)]
pub struct Color {
    pub data: u32,
}

impl Color {
    /// Create a new color from RGB
    pub fn rgb(r: u8, g: u8, b: u8) -> Self {
        Color { data: 0xFF000000 | ((r as u32) << 16) | ((g as u32) << 8) | (b as u32) }
    }

    /// Set the alpha
    pub fn rgba(r: u8, g: u8, b: u8, a: u8) -> Self {
        Color { data: ((a as u32) << 24) | ((r as u32) << 16) | ((g as u32) << 8) | (b as u32) }
    }

    /// Create a new black color
    pub fn black() -> Self {
        Color { data: 0xFF000000 }
    }

    /// Create a new white color
    pub fn white() -> Self {
        Color { data: 0xFFFFFFFF }
    }

    /// Create a new gray color
    pub fn gray() -> Self {
        Color { data: 0xFF808080 }
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::hash::Hash;

use serde::{Deserialize, Serialize};

#[derive(Debug, Hash, PartialEq, Eq, PartialOrd, Deserialize, Serialize)]
pub struct Id(uuid::Uuid);

#[derive(Debug, Deserialize, Serialize)]
pub struct Collection {
    pub id: Id,
    pub name: String,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct CollectionPages(HashMap<Id, crate::page::Id>);
<commit_msg>inline `pages` field on `Collection`<commit_after>use std::collections::HashMap;
use std::hash::Hash;

use serde::{Deserialize, Serialize};

#[derive(Debug, Hash, PartialEq, Eq, PartialOrd, Deserialize, Serialize)]
pub struct Id(uuid::Uuid);

#[derive(Debug, Deserialize, Serialize)]
pub struct Collection {
    pub id: Id,
    pub name: String,
    pub pages: HashMap<Id, crate::page::Id>,
}
<|endoftext|>