<commit_before>use std::io::{self, BufRead, BufStream, Write};
use std::net::{TcpStream, ToSocketAddrs};

use message::{Message, OwnedMessage, ParseError};

pub struct Client {
    stream: BufStream<TcpStream>,
}

impl Client {

    pub fn new<A: ToSocketAddrs>(addr: A) -> io::Result<Client> {
        let stream = try!(TcpStream::connect(addr));
        Ok(Client{
            stream: BufStream::new(stream),
        })
    }

    fn send(&mut self, line: String) -> io::Result<()> {
        try!(write!(self.stream, "{}\r\n", line));
        self.stream.flush()
    }

    pub fn user(&mut self, username: &str, realname: &str) -> io::Result<()> {
        self.send(format!("USER {} 8 * :{}", username, realname))
    }

    pub fn nick(&mut self, nickname: &str) -> io::Result<()> {
        self.send(format!("NICK {}", nickname))
    }

    pub fn join(&mut self, channel: &str) -> io::Result<()> {
        self.send(format!("JOIN {}", channel))
    }

    pub fn read(&mut self) -> Result<OwnedMessage, ReadError> {
        let mut line = String::new();
        try!(self.stream.read_line(&mut line));
        let msg = try!(Message::parse(&line[..]));
        Ok(msg.to_owned())
    }

}

#[derive(Debug)]
pub enum ReadError {
    IoError(io::Error),
    ParseError(ParseError),
}

impl From<io::Error> for ReadError {
    fn from(err: io::Error) -> ReadError {
        ReadError::IoError(err)
    }
}

impl From<ParseError> for ReadError {
    fn from(err: ParseError) -> ReadError {
        ReadError::ParseError(err)
    }
}
<commit_msg>Add try_clone to client<commit_after>use std::io::{self, BufRead, BufStream, Write};
use std::net::{TcpStream, ToSocketAddrs};

use message::{Message, OwnedMessage, ParseError};

pub struct Client {
    stream: TcpStream,
    buf: BufStream<TcpStream>,
}

impl Client {

    pub fn new<A: ToSocketAddrs>(addr: A) -> io::Result<Client> {
        let stream = try!(TcpStream::connect(addr));
        let buf = BufStream::new(try!(stream.try_clone()));
        Ok(Client{
            stream: stream,
            buf: buf,
        })
    }

    pub fn send_raw<A: AsRef<str>>(&mut self, line: A) -> io::Result<()> {
        try!(write!(self.buf, "{}\r\n", line.as_ref()));
        self.buf.flush()
    }

    pub fn user(&mut self, username: &str, realname: &str) -> io::Result<()> {
        self.send_raw(format!("USER {} 8 * :{}", username, realname))
    }

    pub fn nick(&mut self, nickname: &str) -> io::Result<()> {
        self.send_raw(format!("NICK {}", nickname))
    }

    pub fn join(&mut self, channel: &str) -> io::Result<()> {
        self.send_raw(format!("JOIN {}", channel))
    }

    pub fn read(&mut self) -> Result<OwnedMessage, ReadError> {
        let mut line = String::new();
        try!(self.buf.read_line(&mut line));
        let msg = try!(Message::parse(&line[..]));
        Ok(msg.to_owned())
    }

    pub fn try_clone(&self) -> io::Result<Client> {
        let stream = try!(self.stream.try_clone());
        let buf = BufStream::new(try!(self.stream.try_clone()));
        Ok(Client{
            stream: stream,
            buf: buf,
        })
    }

}

#[derive(Debug)]
pub enum ReadError {
    IoError(io::Error),
    ParseError(ParseError),
}

impl From<io::Error> for ReadError {
    fn from(err: io::Error) -> ReadError {
        ReadError::IoError(err)
    }
}

impl From<ParseError> for ReadError {
    fn from(err: ParseError) -> ReadError {
        ReadError::ParseError(err)
    }
}
<|endoftext|>


<commit_before>//! This module contains the main client code, including the [`RocketClient`] type.
use crate::interpolation::*;
use crate::track::*;

use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use std::{
    convert::TryFrom,
    io::{Cursor, Read, Write},
    net::{TcpStream, ToSocketAddrs},
};
use thiserror::Error;

#[derive(Debug, Error)]
/// The `Error` Type. This is the main error type.
pub enum Error {
    #[error("Failed to establish a TCP connection with the Rocket server")]
    /// Failure to connect to a rocket tracker. This can happen if the tracker is not running, the
    /// address isn't correct or other network-related reasons.
    Connect(#[source] std::io::Error),
    #[error("Handshake with the Rocket server failed")]
    /// Failure to transmit or receive greetings with the tracker
    Handshake(#[source] std::io::Error),
    #[error("The Rocket server greeting {0:?} wasn't correct")]
    /// Handshake was performed but the the received greeting wasn't correct
    HandshakeGreetingMismatch([u8; 12]),
    #[error("Cannot set Rocket's TCP connection to nonblocking mode")]
    /// Error from [`TcpStream::set_nonblocking`]
    SetNonblocking(#[source] std::io::Error),
    #[error("Rocket server disconnected")]
    /// Network IO error during operation
    IOError(#[source] std::io::Error),
}

#[derive(Debug)]
enum ClientState {
    New,
    Incomplete(usize),
    Complete,
}

#[derive(Debug, Copy, Clone)]
/// The `Event` Type. These are the various events from the tracker.
pub enum Event {
    /// The tracker changes row.
    SetRow(u32),
    /// The tracker pauses or unpauses.
    Pause(bool),
    /// The tracker asks us to save our track data.
    /// You may want to call [`RocketClient::save_tracks`] after receiving this event.
    SaveTracks,
}

enum ReceiveResult {
    Some(Event),
    None,
    Incomplete,
}

#[derive(Debug)]
/// The `RocketClient` type. This contains the connected socket and other fields.
pub struct RocketClient {
    stream: TcpStream,
    state: ClientState,
    cmd: Vec<u8>,
    tracks: Vec<Track>,
}

impl RocketClient {
    /// Construct a new RocketClient.
    ///
    /// This constructs a new Rocket client and connects to localhost on port 1338.
    ///
    /// # Errors
    ///
    /// [`Error::Connect`] if connection cannot be established, or [`Error::Handshake`]
    /// if the handshake fails.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # use rust_rocket::RocketClient;
    /// let mut rocket = RocketClient::new().unwrap();
    /// ```
    pub fn new() -> Result<Self, Error> {
        Self::connect(("localhost", 1338))
    }

    /// Construct a new RocketClient.
    ///
    /// This constructs a new Rocket client and connects to a specified host and port.
    ///
    /// # Errors
    ///
    /// [`Error::Connect`] if connection cannot be established, or [`Error::Handshake`]
    /// if the handshake fails.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # use rust_rocket::RocketClient;
    /// let mut rocket = RocketClient::connect(("localhost", 1338)).unwrap();
    /// ```
    pub fn connect(addr: impl ToSocketAddrs) -> Result<Self, Error> {
        let stream = TcpStream::connect(addr).map_err(Error::Connect)?;

        let mut rocket = Self {
            stream,
            state: ClientState::New,
            cmd: Vec::new(),
            tracks: Vec::new(),
        };

        rocket.handshake()?;

        rocket
            .stream
            .set_nonblocking(true)
            .map_err(Error::SetNonblocking)?;

        Ok(rocket)
    }

    /// Get track by name.
    ///
    /// If the track does not yet exist it will be created.
    ///
    /// # Errors
    ///
    /// This method can return an [`Error::IOError`] if Rocket tracker disconnects.
    ///
    /// # Panics
    ///
    /// Will panic if `name`'s length exceeds [`u32::MAX`].
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # use rust_rocket::RocketClient;
    /// # let mut rocket = RocketClient::new().unwrap();
    /// let track = rocket.get_track_mut("namespace:track").unwrap();
    /// track.get_value(3.5);
    /// ```
    pub fn get_track_mut(&mut self, name: &str) -> Result<&mut Track, Error> {
        if let Some((i, _)) = self
            .tracks
            .iter()
            .enumerate()
            .find(|(_, t)| t.get_name() == name)
        {
            Ok(&mut self.tracks[i])
        } else {
            // Send GET_TRACK message
            let mut buf = vec![2];
            buf.write_u32::<BigEndian>(u32::try_from(name.len()).expect("Track name too long"))
                .unwrap_or_else(|_|
                // Can writes to a vec fail? Consider changing to unreachable_unchecked in 1.0
                unreachable!());
            buf.extend_from_slice(&name.as_bytes());
            self.stream.write_all(&buf).map_err(Error::IOError)?;

            self.tracks.push(Track::new(name));
            Ok(self.tracks.last_mut().unwrap_or_else(||
                // tracks cannot be empty right after pushing into it, consider changing to
                // unreachable_unchecked in 1.0
                unreachable!()))
        }
    }

    /// Get track by name.
    ///
    /// You should use [`get_track_mut`](RocketClient::get_track_mut) to create a track.
    pub fn get_track(&self, name: &str) -> Option<&Track> {
        self.tracks.iter().find(|t| t.get_name() == name)
    }

    /// Create a clone of the tracks in the session which can then be serialized to a file in any
    /// format with a serde implementation.
    /// Tracks can be turned into a [`RocketPlayer`](crate::RocketPlayer::new) for playback.
    pub fn save_tracks(&self) -> Vec<Track> {
        self.tracks.clone()
    }

    /// Send a SetRow message.
    ///
    /// This changes the current row on the tracker side.
    ///
    /// # Errors
    ///
    /// This method can return an [`Error::IOError`] if Rocket tracker disconnects.
    pub fn set_row(&mut self, row: u32) -> Result<(), Error> {
        // Send SET_ROW message
        let mut buf = vec![3];
        buf.write_u32::<BigEndian>(row).unwrap_or_else(|_|
                // Can writes to a vec fail? Consider changing to unreachable_unchecked in 1.0
                unreachable!());
        self.stream.write_all(&buf).map_err(Error::IOError)
    }

    /// Poll for new events from the tracker.
    ///
    /// This polls from events from the tracker.
    /// You should call this fairly often your main loop.
    /// It is recommended to keep calling this as long as your receive `Some(Event)`.
    ///
    /// # Errors
    ///
    /// This method can return an [`Error::IOError`] if Rocket tracker disconnects.
    ///
    /// # Panics
    ///
    /// This method can panic on esoteric systems where `usize` is smaller than `u32`, and when
    /// there are more than `usize` tracks in use.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # use rust_rocket::RocketClient;
    /// # let mut rocket = RocketClient::new().unwrap();
    /// while let Some(event) = rocket.poll_events().unwrap() {
    ///     match event {
    ///         // Do something with the various events.
    ///         _ => (),
    ///     }
    /// }
    /// ```
    pub fn poll_events(&mut self) -> Result<Option<Event>, Error> {
        loop {
            let result = self.poll_event()?;
            match result {
                ReceiveResult::None => return Ok(None),
                ReceiveResult::Incomplete => (),
                ReceiveResult::Some(event) => return Ok(Some(event)),
            }
        }
    }

    fn poll_event(&mut self) -> Result<ReceiveResult, Error> {
        match self.state {
            ClientState::New => {
                let mut buf = [0; 1];
                match self.stream.read_exact(&mut buf) {
                    Ok(()) => {
                        self.cmd.extend_from_slice(&buf);
                        match self.cmd[0] {
                            0 => self.state = ClientState::Incomplete(4 + 4 + 4 + 1), //SET_KEY
                            1 => self.state = ClientState::Incomplete(4 + 4),         //DELETE_KEY
                            3 => self.state = ClientState::Incomplete(4),             //SET_ROW
                            4 => self.state = ClientState::Incomplete(1),             //PAUSE
                            5 => self.state = ClientState::Complete,                  //SAVE_TRACKS
                            _ => self.state = ClientState::Complete, // Error / Unknown
                        }
                        Ok(ReceiveResult::Incomplete)
                    }
                    Err(e) => match e.kind() {
                        std::io::ErrorKind::WouldBlock => Ok(ReceiveResult::None),
                        _ => Err(Error::IOError(e)),
                    },
                }
            }
            ClientState::Incomplete(bytes) => {
                let mut buf = vec![0; bytes];
                match self.stream.read(&mut buf) {
                    Ok(bytes_read) => {
                        self.cmd.extend_from_slice(&buf);
                        if bytes - bytes_read > 0 {
                            self.state = ClientState::Incomplete(bytes - bytes_read);
                        } else {
                            self.state = ClientState::Complete;
                        }
                        Ok(ReceiveResult::Incomplete)
                    }
                    Err(e) => match e.kind() {
                        std::io::ErrorKind::WouldBlock => Ok(ReceiveResult::None),
                        _ => Err(Error::IOError(e)),
                    },
                }
            }
            ClientState::Complete => {
                let mut result = ReceiveResult::None;
                {
                    // Following reads from cmd should never fail if above match arms are correct
                    let mut cursor = Cursor::new(&self.cmd);
                    let cmd = cursor.read_u8().unwrap();
                    match cmd {
                        0 => {
                            let track = &mut self.tracks
                                [usize::try_from(cursor.read_u32::<BigEndian>().unwrap()).unwrap()];
                            let row = cursor.read_u32::<BigEndian>().unwrap();
                            let value = cursor.read_f32::<BigEndian>().unwrap();
                            let interpolation = Interpolation::from(cursor.read_u8().unwrap());
                            let key = Key::new(row, value, interpolation);

                            track.set_key(key);
                        }
                        1 => {
                            let track = &mut self.tracks
                                [usize::try_from(cursor.read_u32::<BigEndian>().unwrap()).unwrap()];
                            let row = cursor.read_u32::<BigEndian>().unwrap();

                            track.delete_key(row);
                        }
                        3 => {
                            let row = cursor.read_u32::<BigEndian>().unwrap();
                            result = ReceiveResult::Some(Event::SetRow(row));
                        }
                        4 => {
                            let flag = cursor.read_u8().unwrap() == 1;
                            result = ReceiveResult::Some(Event::Pause(flag));
                        }
                        5 => {
                            result = ReceiveResult::Some(Event::SaveTracks);
                        }
                        _ => println!("Unknown {:?}", cmd),
                    }
                }

                self.cmd.clear();
                self.state = ClientState::New;

                Ok(result)
            }
        }
    }

    fn handshake(&mut self) -> Result<(), Error> {
        let client_greeting = b"hello, synctracker!";
        let server_greeting = b"hello, demo!";

        self.stream
            .write_all(client_greeting)
            .map_err(Error::Handshake)?;

        let mut buf = [0; 12];
        self.stream.read_exact(&mut buf).map_err(Error::Handshake)?;

        if &buf == server_greeting {
            Ok(())
        } else {
            Err(Error::HandshakeGreetingMismatch(buf))
        }
    }
}
<commit_msg>Fix comment about panics in poll_event<commit_after>//! This module contains the main client code, including the [`RocketClient`] type.
use crate::interpolation::*;
use crate::track::*;

use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use std::{
    convert::TryFrom,
    io::{Cursor, Read, Write},
    net::{TcpStream, ToSocketAddrs},
};
use thiserror::Error;

#[derive(Debug, Error)]
/// The `Error` Type. This is the main error type.
pub enum Error {
    #[error("Failed to establish a TCP connection with the Rocket server")]
    /// Failure to connect to a rocket tracker. This can happen if the tracker is not running, the
    /// address isn't correct or other network-related reasons.
    Connect(#[source] std::io::Error),
    #[error("Handshake with the Rocket server failed")]
    /// Failure to transmit or receive greetings with the tracker
    Handshake(#[source] std::io::Error),
    #[error("The Rocket server greeting {0:?} wasn't correct")]
    /// Handshake was performed but the the received greeting wasn't correct
    HandshakeGreetingMismatch([u8; 12]),
    #[error("Cannot set Rocket's TCP connection to nonblocking mode")]
    /// Error from [`TcpStream::set_nonblocking`]
    SetNonblocking(#[source] std::io::Error),
    #[error("Rocket server disconnected")]
    /// Network IO error during operation
    IOError(#[source] std::io::Error),
}

#[derive(Debug)]
enum ClientState {
    New,
    Incomplete(usize),
    Complete,
}

#[derive(Debug, Copy, Clone)]
/// The `Event` Type. These are the various events from the tracker.
pub enum Event {
    /// The tracker changes row.
    SetRow(u32),
    /// The tracker pauses or unpauses.
    Pause(bool),
    /// The tracker asks us to save our track data.
    /// You may want to call [`RocketClient::save_tracks`] after receiving this event.
    SaveTracks,
}

enum ReceiveResult {
    Some(Event),
    None,
    Incomplete,
}

#[derive(Debug)]
/// The `RocketClient` type. This contains the connected socket and other fields.
pub struct RocketClient {
    stream: TcpStream,
    state: ClientState,
    cmd: Vec<u8>,
    tracks: Vec<Track>,
}

impl RocketClient {
    /// Construct a new RocketClient.
    ///
    /// This constructs a new Rocket client and connects to localhost on port 1338.
    ///
    /// # Errors
    ///
    /// [`Error::Connect`] if connection cannot be established, or [`Error::Handshake`]
    /// if the handshake fails.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # use rust_rocket::RocketClient;
    /// let mut rocket = RocketClient::new().unwrap();
    /// ```
    pub fn new() -> Result<Self, Error> {
        Self::connect(("localhost", 1338))
    }

    /// Construct a new RocketClient.
    ///
    /// This constructs a new Rocket client and connects to a specified host and port.
    ///
    /// # Errors
    ///
    /// [`Error::Connect`] if connection cannot be established, or [`Error::Handshake`]
    /// if the handshake fails.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # use rust_rocket::RocketClient;
    /// let mut rocket = RocketClient::connect(("localhost", 1338)).unwrap();
    /// ```
    pub fn connect(addr: impl ToSocketAddrs) -> Result<Self, Error> {
        let stream = TcpStream::connect(addr).map_err(Error::Connect)?;

        let mut rocket = Self {
            stream,
            state: ClientState::New,
            cmd: Vec::new(),
            tracks: Vec::new(),
        };

        rocket.handshake()?;

        rocket
            .stream
            .set_nonblocking(true)
            .map_err(Error::SetNonblocking)?;

        Ok(rocket)
    }

    /// Get track by name.
    ///
    /// If the track does not yet exist it will be created.
    ///
    /// # Errors
    ///
    /// This method can return an [`Error::IOError`] if Rocket tracker disconnects.
    ///
    /// # Panics
    ///
    /// Will panic if `name`'s length exceeds [`u32::MAX`].
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # use rust_rocket::RocketClient;
    /// # let mut rocket = RocketClient::new().unwrap();
    /// let track = rocket.get_track_mut("namespace:track").unwrap();
    /// track.get_value(3.5);
    /// ```
    pub fn get_track_mut(&mut self, name: &str) -> Result<&mut Track, Error> {
        if let Some((i, _)) = self
            .tracks
            .iter()
            .enumerate()
            .find(|(_, t)| t.get_name() == name)
        {
            Ok(&mut self.tracks[i])
        } else {
            // Send GET_TRACK message
            let mut buf = vec![2];
            buf.write_u32::<BigEndian>(u32::try_from(name.len()).expect("Track name too long"))
                .unwrap_or_else(|_|
                // Can writes to a vec fail? Consider changing to unreachable_unchecked in 1.0
                unreachable!());
            buf.extend_from_slice(&name.as_bytes());
            self.stream.write_all(&buf).map_err(Error::IOError)?;

            self.tracks.push(Track::new(name));
            Ok(self.tracks.last_mut().unwrap_or_else(||
                // tracks cannot be empty right after pushing into it, consider changing to
                // unreachable_unchecked in 1.0
                unreachable!()))
        }
    }

    /// Get track by name.
    ///
    /// You should use [`get_track_mut`](RocketClient::get_track_mut) to create a track.
    pub fn get_track(&self, name: &str) -> Option<&Track> {
        self.tracks.iter().find(|t| t.get_name() == name)
    }

    /// Create a clone of the tracks in the session which can then be serialized to a file in any
    /// format with a serde implementation.
    /// Tracks can be turned into a [`RocketPlayer`](crate::RocketPlayer::new) for playback.
    pub fn save_tracks(&self) -> Vec<Track> {
        self.tracks.clone()
    }

    /// Send a SetRow message.
    ///
    /// This changes the current row on the tracker side.
    ///
    /// # Errors
    ///
    /// This method can return an [`Error::IOError`] if Rocket tracker disconnects.
    pub fn set_row(&mut self, row: u32) -> Result<(), Error> {
        // Send SET_ROW message
        let mut buf = vec![3];
        buf.write_u32::<BigEndian>(row).unwrap_or_else(|_|
                // Can writes to a vec fail? Consider changing to unreachable_unchecked in 1.0
                unreachable!());
        self.stream.write_all(&buf).map_err(Error::IOError)
    }

    /// Poll for new events from the tracker.
    ///
    /// This polls from events from the tracker.
    /// You should call this fairly often your main loop.
    /// It is recommended to keep calling this as long as your receive `Some(Event)`.
    ///
    /// # Errors
    ///
    /// This method can return an [`Error::IOError`] if Rocket tracker disconnects.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # use rust_rocket::RocketClient;
    /// # let mut rocket = RocketClient::new().unwrap();
    /// while let Some(event) = rocket.poll_events().unwrap() {
    ///     match event {
    ///         // Do something with the various events.
    ///         _ => (),
    ///     }
    /// }
    /// ```
    pub fn poll_events(&mut self) -> Result<Option<Event>, Error> {
        loop {
            let result = self.poll_event()?;
            match result {
                ReceiveResult::None => return Ok(None),
                ReceiveResult::Incomplete => (),
                ReceiveResult::Some(event) => return Ok(Some(event)),
            }
        }
    }

    fn poll_event(&mut self) -> Result<ReceiveResult, Error> {
        match self.state {
            ClientState::New => {
                let mut buf = [0; 1];
                match self.stream.read_exact(&mut buf) {
                    Ok(()) => {
                        self.cmd.extend_from_slice(&buf);
                        match self.cmd[0] {
                            0 => self.state = ClientState::Incomplete(4 + 4 + 4 + 1), //SET_KEY
                            1 => self.state = ClientState::Incomplete(4 + 4),         //DELETE_KEY
                            3 => self.state = ClientState::Incomplete(4),             //SET_ROW
                            4 => self.state = ClientState::Incomplete(1),             //PAUSE
                            5 => self.state = ClientState::Complete,                  //SAVE_TRACKS
                            _ => self.state = ClientState::Complete, // Error / Unknown
                        }
                        Ok(ReceiveResult::Incomplete)
                    }
                    Err(e) => match e.kind() {
                        std::io::ErrorKind::WouldBlock => Ok(ReceiveResult::None),
                        _ => Err(Error::IOError(e)),
                    },
                }
            }
            ClientState::Incomplete(bytes) => {
                let mut buf = vec![0; bytes];
                match self.stream.read(&mut buf) {
                    Ok(bytes_read) => {
                        self.cmd.extend_from_slice(&buf);
                        if bytes - bytes_read > 0 {
                            self.state = ClientState::Incomplete(bytes - bytes_read);
                        } else {
                            self.state = ClientState::Complete;
                        }
                        Ok(ReceiveResult::Incomplete)
                    }
                    Err(e) => match e.kind() {
                        std::io::ErrorKind::WouldBlock => Ok(ReceiveResult::None),
                        _ => Err(Error::IOError(e)),
                    },
                }
            }
            ClientState::Complete => {
                let mut result = ReceiveResult::None;
                {
                    // Following reads from cmd should never fail if above match arms are correct
                    let mut cursor = Cursor::new(&self.cmd);
                    let cmd = cursor.read_u8().unwrap();
                    match cmd {
                        0 => {
                            // usize::try_from(u32) will only panic if usize is smaller, and there
                            // are more than usize::MAX tracks in use, which isn't possible because
                            // I'd imagine Vec::push and everything else will panic first.
                            // If you're running this on a microcontroller, I'd love to see it!
                            let track = &mut self.tracks
                                [usize::try_from(cursor.read_u32::<BigEndian>().unwrap()).unwrap()];
                            let row = cursor.read_u32::<BigEndian>().unwrap();
                            let value = cursor.read_f32::<BigEndian>().unwrap();
                            let interpolation = Interpolation::from(cursor.read_u8().unwrap());
                            let key = Key::new(row, value, interpolation);

                            track.set_key(key);
                        }
                        1 => {
                            let track = &mut self.tracks
                                [usize::try_from(cursor.read_u32::<BigEndian>().unwrap()).unwrap()];
                            let row = cursor.read_u32::<BigEndian>().unwrap();

                            track.delete_key(row);
                        }
                        3 => {
                            let row = cursor.read_u32::<BigEndian>().unwrap();
                            result = ReceiveResult::Some(Event::SetRow(row));
                        }
                        4 => {
                            let flag = cursor.read_u8().unwrap() == 1;
                            result = ReceiveResult::Some(Event::Pause(flag));
                        }
                        5 => {
                            result = ReceiveResult::Some(Event::SaveTracks);
                        }
                        _ => println!("Unknown {:?}", cmd),
                    }
                }

                self.cmd.clear();
                self.state = ClientState::New;

                Ok(result)
            }
        }
    }

    fn handshake(&mut self) -> Result<(), Error> {
        let client_greeting = b"hello, synctracker!";
        let server_greeting = b"hello, demo!";

        self.stream
            .write_all(client_greeting)
            .map_err(Error::Handshake)?;

        let mut buf = [0; 12];
        self.stream.read_exact(&mut buf).map_err(Error::Handshake)?;

        if &buf == server_greeting {
            Ok(())
        } else {
            Err(Error::HandshakeGreetingMismatch(buf))
        }
    }
}
<|endoftext|>


<commit_before>
extern crate std;
extern crate core;
extern crate num;
extern crate uuid;

use std::str::SendStr;

use super::def::*;
use super::serialize::*;
use super::reader::*;
use std::collections::TreeMap;

pub static CQL_VERSION_STRINGS:  [&'static str, .. 3] = ["3.0.0", "3.0.0", "3.0.0"];
pub static CQL_MAX_SUPPORTED_VERSION:u8 = 0x03;

macro_rules! serialize_and_check_io_error(
    ($writer: expr, $obj: ident, $version: expr, $msg: expr) => {
        match $obj.serialize($writer, $version) {
            Err(e) => return Err(RCError { kind: SerializeError, desc: format!("{} -> {}", $msg, e.desc).into_maybe_owned()}),
            _ => ()
        }
    }
)

type PrepsStore = TreeMap<String, Box<CqlPreparedStat>>;

pub struct CqlClient {
    socket: std::io::net::tcp::TcpStream,
    pub version: u8,
    prepared: PrepsStore
}

impl CqlClient {

    fn new(socket: std::io::net::tcp::TcpStream, version: u8) -> CqlClient {
        CqlClient {socket: socket, version: version, prepared: TreeMap::new()}
    }

    fn build_auth<'a>(&self, creds: &'a Vec<SendStr>, stream: i8) -> CqlRequest<'a> {
        return CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: stream,
            opcode: OpcodeOptions,
            body: RequestCred(creds),
        };
    }

    fn build_options(&self) -> CqlRequest {
        return CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: 0x01,
            opcode: OpcodeOptions,
            body: RequestOptions,
        };
    }

    pub fn exec_query(&mut self, query_str: &str, con: Consistency::Consistency) -> RCResult<CqlResponse> {
        let q = CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: 0x01,
            opcode: OpcodeQuery,
            body: RequestQuery(query_str, con, 0)};

        let ref mut socket = self.socket;
        serialize_and_check_io_error!(socket, q, self.version, "Error serializing query");
        let res = read_and_check_io_error!(socket, read_cql_response, self.version, "Error reading query");
        Ok(res)
    }

    pub fn exec_prepared(&mut self, ps_id: &str, params: &[CqlValue], con: Consistency::Consistency) -> RCResult<CqlResponse> {
        let preps = match self.prepared.find(&String::from_str(ps_id)) {
            Some(ps) => &**ps,
            None => return Err(RCError::new(format!("Unknown prepared statement <{}>", ps_id), GenericError))
        };

        let q = CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: 0x01,
            opcode: OpcodeExecute,
            body: RequestExec(preps, params, con, 0x01),
        };

        serialize_and_check_io_error!(&mut self.socket, q, self.version, "Error serializing prepared statement execution");

        /* Code to debug prepared statements. Write to file the serialization of the request

        let path = Path::new("prepared_data.bin");
        let display = path.display();
        let mut file = match std::io::File::create(&path) {
            Err(why) => fail!("couldn't create {}: {}", display, why.desc),
            Ok(file) => file,
        };

        serialize_and_check_io_error!(&mut file, q, self.version, "Error serializing query");

        */

        Ok(read_and_check_io_error!(self.socket, read_cql_response, self.version, "Error reading prepared statement execution result"))
    }

    pub fn prepared_statement(&mut self, query_str: &str, query_id: &str) -> RCResult<()> {
        let q = CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: 0x01,
            opcode: OpcodePrepare,
            body: RequestPrepare(query_str),
        };

        serialize_and_check_io_error!(&mut self.socket, q, self.version, "Error serializing prepared statement");

        let res = read_and_check_io_error!(&mut self.socket, read_cql_response, self.version, "Error reading query");
        match res.body {
            ResultPrepared(preps) => {
                self.prepared.insert(String::from_str(query_id), preps);
                Ok(())
            },
            _ => Err(RCError::new("Response does not contain prepared statement", ReadError))
        }
    }
}

fn send_startup(socket: &mut std::io::TcpStream, version: u8, creds: Option<&Vec<SendStr>>) -> RCResult<()> {
    let body = CqlStringMap {
        pairs:vec![CqlPair{key: "CQL_VERSION", value: CQL_VERSION_STRINGS[(version-1) as uint]}],
    };
    let msg_startup = CqlRequest {
        version: version,
        flags: 0x00,
        stream: 0x01,
        opcode: OpcodeStartup,
        body: RequestStartup(body),
    };
    serialize_and_check_error!(socket, msg_startup, version, "Error serializing startup message");

    let response = read_and_check_io_error!(socket, read_cql_response, version, "Error reding response");
    match response.body {
        ResponseReady =>  Ok(()),
        ResponseAuth(_) => {
            match creds {
                Some(cred) => {
                    let msg_auth = CqlRequest {
                        version: version,
                        flags: 0x00,
                        stream: 0x01,
                        opcode: OpcodeOptions,
                        body: RequestCred(cred),
                    };
                    serialize_and_check_error!(socket, msg_auth, version, "Error serializing request (auth)");

                    let response = read_and_check_io_error!(socket, read_cql_response, version, "Error reding authenticaton response");
                    match response.body {
                        ResponseReady => Ok(()),
                        ResponseError(_, ref msg) => Err(RCError::new(format!("Error in authentication: {}", msg), ReadError)),
                        _ => Err(RCError::new("Server returned unknown message", ReadError))
                    }
                },
                None => Err(RCError::new("Credential should be provided for authentication", ReadError))
            }
        },
        ResponseError(_, ref msg) => Err(RCError::new(format!("Error connecting: {}", msg), ReadError)),
        _ => Err(RCError::new("Wrong response to startup", ReadError))
    }
}

pub fn connect(ip: &'static str, port: u16, creds:Option<&Vec<SendStr>>) -> RCResult<CqlClient> {

    let mut version = CQL_MAX_SUPPORTED_VERSION;

    while version >= 0x01 {
        let res = std::io::TcpStream::connect(ip.as_slice(), port);
        if res.is_err() {
            return Err(RCError::new(format!("Failed to connect to server at {}:{}", ip.as_slice(), port), ConnectionError));
        }

        let mut socket = res.unwrap();

        match send_startup(& mut socket, version, creds) {
            Ok(_) => return Ok(CqlClient::new(socket, version)),
            Err(e) => println!("Error connecting with protocol version v{}: {}", version, e.desc)
        }
        version -= 1;
    }
    Err(RCError::new("Unable to find suitable protocol version (v1, v2, v3)", ReadError))
}
<commit_msg>Refactored CqlClient -> Client<commit_after>
extern crate std;
extern crate core;
extern crate num;
extern crate uuid;

use std::str::SendStr;

use super::def::*;
use super::serialize::*;
use super::reader::*;
use std::collections::TreeMap;

pub static CQL_VERSION_STRINGS:  [&'static str, .. 3] = ["3.0.0", "3.0.0", "3.0.0"];
pub static CQL_MAX_SUPPORTED_VERSION:u8 = 0x03;

macro_rules! serialize_and_check_io_error(
    ($writer: expr, $obj: ident, $version: expr, $msg: expr) => {
        match $obj.serialize($writer, $version) {
            Err(e) => return Err(RCError { kind: SerializeError, desc: format!("{} -> {}", $msg, e.desc).into_maybe_owned()}),
            _ => ()
        }
    }
)

type PrepsStore = TreeMap<String, Box<CqlPreparedStat>>;

pub struct Client {
    socket: std::io::net::tcp::TcpStream,
    pub version: u8,
    prepared: PrepsStore
}

impl Client {

    fn new(socket: std::io::net::tcp::TcpStream, version: u8) -> Client {
        Client {socket: socket, version: version, prepared: TreeMap::new()}
    }

    fn build_auth<'a>(&self, creds: &'a Vec<SendStr>, stream: i8) -> CqlRequest<'a> {
        return CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: stream,
            opcode: OpcodeOptions,
            body: RequestCred(creds),
        };
    }

    fn build_options(&self) -> CqlRequest {
        return CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: 0x01,
            opcode: OpcodeOptions,
            body: RequestOptions,
        };
    }

    pub fn exec_query(&mut self, query_str: &str, con: Consistency::Consistency) -> RCResult<CqlResponse> {
        let q = CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: 0x01,
            opcode: OpcodeQuery,
            body: RequestQuery(query_str, con, 0)};

        let ref mut socket = self.socket;
        serialize_and_check_io_error!(socket, q, self.version, "Error serializing query");
        let res = read_and_check_io_error!(socket, read_cql_response, self.version, "Error reading query");
        Ok(res)
    }

    pub fn exec_prepared(&mut self, ps_id: &str, params: &[CqlValue], con: Consistency::Consistency) -> RCResult<CqlResponse> {
        let preps = match self.prepared.find(&String::from_str(ps_id)) {
            Some(ps) => &**ps,
            None => return Err(RCError::new(format!("Unknown prepared statement <{}>", ps_id), GenericError))
        };

        let q = CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: 0x01,
            opcode: OpcodeExecute,
            body: RequestExec(preps, params, con, 0x01),
        };

        serialize_and_check_io_error!(&mut self.socket, q, self.version, "Error serializing prepared statement execution");

        /* Code to debug prepared statements. Write to file the serialization of the request

        let path = Path::new("prepared_data.bin");
        let display = path.display();
        let mut file = match std::io::File::create(&path) {
            Err(why) => fail!("couldn't create {}: {}", display, why.desc),
            Ok(file) => file,
        };

        serialize_and_check_io_error!(&mut file, q, self.version, "Error serializing query");

        */

        Ok(read_and_check_io_error!(self.socket, read_cql_response, self.version, "Error reading prepared statement execution result"))
    }

    pub fn prepared_statement(&mut self, query_str: &str, query_id: &str) -> RCResult<()> {
        let q = CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: 0x01,
            opcode: OpcodePrepare,
            body: RequestPrepare(query_str),
        };

        serialize_and_check_io_error!(&mut self.socket, q, self.version, "Error serializing prepared statement");

        let res = read_and_check_io_error!(&mut self.socket, read_cql_response, self.version, "Error reading query");
        match res.body {
            ResultPrepared(preps) => {
                self.prepared.insert(String::from_str(query_id), preps);
                Ok(())
            },
            _ => Err(RCError::new("Response does not contain prepared statement", ReadError))
        }
    }
}

fn send_startup(socket: &mut std::io::TcpStream, version: u8, creds: Option<&Vec<SendStr>>) -> RCResult<()> {
    let body = CqlStringMap {
        pairs:vec![CqlPair{key: "CQL_VERSION", value: CQL_VERSION_STRINGS[(version-1) as uint]}],
    };
    let msg_startup = CqlRequest {
        version: version,
        flags: 0x00,
        stream: 0x01,
        opcode: OpcodeStartup,
        body: RequestStartup(body),
    };
    serialize_and_check_error!(socket, msg_startup, version, "Error serializing startup message");

    let response = read_and_check_io_error!(socket, read_cql_response, version, "Error reding response");
    match response.body {
        ResponseReady =>  Ok(()),
        ResponseAuth(_) => {
            match creds {
                Some(cred) => {
                    let msg_auth = CqlRequest {
                        version: version,
                        flags: 0x00,
                        stream: 0x01,
                        opcode: OpcodeOptions,
                        body: RequestCred(cred),
                    };
                    serialize_and_check_error!(socket, msg_auth, version, "Error serializing request (auth)");

                    let response = read_and_check_io_error!(socket, read_cql_response, version, "Error reding authenticaton response");
                    match response.body {
                        ResponseReady => Ok(()),
                        ResponseError(_, ref msg) => Err(RCError::new(format!("Error in authentication: {}", msg), ReadError)),
                        _ => Err(RCError::new("Server returned unknown message", ReadError))
                    }
                },
                None => Err(RCError::new("Credential should be provided for authentication", ReadError))
            }
        },
        ResponseError(_, ref msg) => Err(RCError::new(format!("Error connecting: {}", msg), ReadError)),
        _ => Err(RCError::new("Wrong response to startup", ReadError))
    }
}

pub fn connect(ip: &'static str, port: u16, creds:Option<&Vec<SendStr>>) -> RCResult<Client> {

    let mut version = CQL_MAX_SUPPORTED_VERSION;

    while version >= 0x01 {
        let res = std::io::TcpStream::connect(ip.as_slice(), port);
        if res.is_err() {
            return Err(RCError::new(format!("Failed to connect to server at {}:{}", ip.as_slice(), port), ConnectionError));
        }

        let mut socket = res.unwrap();

        match send_startup(& mut socket, version, creds) {
            Ok(_) => return Ok(Client::new(socket, version)),
            Err(e) => println!("Error connecting with protocol version v{}: {}", version, e.desc)
        }
        version -= 1;
    }
    Err(RCError::new("Unable to find suitable protocol version (v1, v2, v3)", ReadError))
}
<|endoftext|>


<commit_before>use std::io;
use std::net::SocketAddr;

use futures::{Future, Sink, Stream};

use tokio_core::net::TcpStream;
use tokio_core::reactor::Core;
use tokio_io::AsyncRead;

use super::resp;

/// TODO: comeback and optimise this number
const DEFAULT_BUFFER_SIZE:usize = 100;

fn connect(addr: &SocketAddr, core: &Core) -> Box<Future<Item=ClientConnection, Error=io::Error>> {
    let handle = core.handle();
    let client_con = TcpStream::connect(addr, &handle).map(move |socket| {
        let framed = socket.framed(resp::RespCodec);
        let (write_f, read_f) = framed.split();
        let write_b = write_f.buffer(DEFAULT_BUFFER_SIZE);
        ClientConnection {
            sender: ClientSink(Box::new(write_b)),
            receiver: ClientStream(Box::new(read_f))
        }
    });
    Box::new(client_con)
}

/// TODO - is the boxing necessary?  It makes the type signature much simpler
struct ClientSink(Box<Sink<SinkItem=resp::RespValue, SinkError=io::Error>>);
struct ClientStream(Box<Stream<Item=resp::RespValue, Error=io::Error>>);

/// A low-level client connection representing a sender and a receiver.
///
/// The two halves operate independently from one another
struct ClientConnection {
    sender: ClientSink,
    receiver: ClientStream
}

#[cfg(test)]
mod test {
    use std::io;

    use futures::{Future, Sink, Stream};
    use futures::stream;

    use tokio_core::reactor::Core;

    use super::resp;

    #[test]
    fn can_connect() {
        let mut core = Core::new().unwrap();
        let addr = "127.0.0.1:6379".parse().unwrap();

        let connection = super::connect(&addr, &core).and_then(|connection| {
            let a = connection.sender.0.send(["PING", "TEST"].as_ref().into());
            let b = connection.receiver.0.take(1).collect();
            a.join(b)
        });
        let (_, values) = core.run(connection).unwrap();
        assert_eq!(values.len(), 1);
        assert_eq!(values[0], "TEST".into());
    }

    #[test]
    fn complex_test() {
        let mut core = Core::new().unwrap();
        let addr = "127.0.0.1:6379".parse().unwrap();
        let connection = super::connect(&addr, &core).and_then(|connection| {
            let mut ops = Vec::<resp::RespValue>::new();
            ops.push(["FLUSH"].as_ref().into());
            ops.extend((0..1000).map(|i| ["SADD", "test_set", &format!("VALUE: {}", i)].as_ref().into()));
            ops.push(["SMEMBERS", "test_set"].as_ref().into());
            let ops_r:Vec<Result<resp::RespValue, io::Error>> = ops.into_iter().map(Result::Ok).collect();
            let send = connection.sender.0.send_all(stream::iter(ops_r));
            let receive = connection.receiver.0.skip(1001).take(1).collect();
            send.join(receive)
        });
        let (_, values) = core.run(connection).unwrap();
        assert_eq!(values.len(), 1);
        let values = match &values[0] {
            &resp::RespValue::Array(ref values) => values.clone(),
            _ => panic!("Not an array")
        };
        assert_eq!(values.len(), 1000);
    }
}<commit_msg>Whitespace<commit_after>use std::io;
use std::net::SocketAddr;

use futures::{Future, Sink, Stream};

use tokio_core::net::TcpStream;
use tokio_core::reactor::Core;
use tokio_io::AsyncRead;

use super::resp;

/// TODO: comeback and optimise this number
const DEFAULT_BUFFER_SIZE:usize = 100;

fn connect(addr: &SocketAddr, core: &Core) -> Box<Future<Item=ClientConnection, Error=io::Error>> {
    let handle = core.handle();
    let client_con = TcpStream::connect(addr, &handle).map(move |socket| {
        let framed = socket.framed(resp::RespCodec);
        let (write_f, read_f) = framed.split();
        let write_b = write_f.buffer(DEFAULT_BUFFER_SIZE);
        ClientConnection {
            sender: ClientSink(Box::new(write_b)),
            receiver: ClientStream(Box::new(read_f))
        }
    });
    Box::new(client_con)
}

/// TODO - is the boxing necessary?  It makes the type signature much simpler
struct ClientSink(Box<Sink<SinkItem=resp::RespValue, SinkError=io::Error>>);
struct ClientStream(Box<Stream<Item=resp::RespValue, Error=io::Error>>);

/// A low-level client connection representing a sender and a receiver.
///
/// The two halves operate independently from one another
struct ClientConnection {
    sender: ClientSink,
    receiver: ClientStream
}

#[cfg(test)]
mod test {
    use std::io;

    use futures::{Future, Sink, Stream};
    use futures::stream;

    use tokio_core::reactor::Core;

    use super::resp;

    #[test]
    fn can_connect() {
        let mut core = Core::new().unwrap();
        let addr = "127.0.0.1:6379".parse().unwrap();

        let connection = super::connect(&addr, &core).and_then(|connection| {
            let a = connection.sender.0.send(["PING", "TEST"].as_ref().into());
            let b = connection.receiver.0.take(1).collect();
            a.join(b)
        });
        let (_, values) = core.run(connection).unwrap();
        assert_eq!(values.len(), 1);
        assert_eq!(values[0], "TEST".into());
    }

    #[test]
    fn complex_test() {
        let mut core = Core::new().unwrap();
        let addr = "127.0.0.1:6379".parse().unwrap();
        let connection = super::connect(&addr, &core).and_then(|connection| {
            let mut ops = Vec::<resp::RespValue>::new();
            ops.push(["FLUSH"].as_ref().into());
            ops.extend((0..1000).map(|i| ["SADD", "test_set", &format!("VALUE: {}", i)].as_ref().into()));
            ops.push(["SMEMBERS", "test_set"].as_ref().into());
            let ops_r:Vec<Result<resp::RespValue, io::Error>> = ops.into_iter().map(Result::Ok).collect();
            let send = connection.sender.0.send_all(stream::iter(ops_r));
            let receive = connection.receiver.0.skip(1001).take(1).collect();
            send.join(receive)
        });
        let (_, values) = core.run(connection).unwrap();
        assert_eq!(values.len(), 1);
        let values = match &values[0] {
            &resp::RespValue::Array(ref values) => values.clone(),
            _ => panic!("Not an array")
        };
        assert_eq!(values.len(), 1000);
    }
}<|endoftext|>


<commit_before>use std::io::{BufRead, BufReader, Write};
use std::net::TcpStream;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Receiver, Sender, TryRecvError};
use std::thread;

use chatmap::{ChatMap, ClientInfo};
use message::Message;

pub fn create_client(stream: TcpStream, sender_to_router: Sender<Message>, chat_map: &Arc<Mutex<ChatMap>>) {
    println!("New client");
    let mut stream = stream;
    stream.write(&"Welcome to Smazy\n".to_string().into_bytes()).unwrap();
    stream.write(&"Please enter a username:\n".to_string().into_bytes()).unwrap();
    stream.flush().unwrap();
    let mut reader = BufReader::new(stream.try_clone().unwrap());
    let mut username = "".to_string();
    match reader.read_line(&mut username) {
        Ok(n) => {
            let username = username.trim().to_string();
            let (sender_to_client, receiver_from_router) = channel();
            {
                let mut guard = chat_map.lock().unwrap();
                if n > 0 && !guard.contains_key(&username) {
                    let client_info = ClientInfo{partner: None, sender_to_client: sender_to_client};
                    guard.insert(username.to_string(), client_info);
                    println!("New user: {}", username.to_string());
                }
            }
            let partner = choose_chat_partner(stream.try_clone().unwrap(), username.to_string(), chat_map);
            chat(stream, username, partner, sender_to_router, receiver_from_router);
        },
        Err(e) => println!("{}", e)
    }
}

fn choose_chat_partner(stream: TcpStream, username: String, chat_map: &Arc<Mutex<ChatMap>>) -> String {
    let mut stream = stream;
    stream.write(&"Here are the users available to chat:\n".to_string().into_bytes()).unwrap();
    {
        let guard = chat_map.lock().unwrap();
        for (name, client_info) in guard.iter() {
            if name.as_str() != username.as_str() && client_info.partner == None {
                stream.write(&name.to_string().into_bytes()).unwrap();
                stream.write(&"\n".to_string().into_bytes()).unwrap();
            }
        }
    }
    stream.write(&"Select who you want to chat with:\n".to_string().into_bytes()).unwrap();
    let mut reader = BufReader::new(stream.try_clone().unwrap());
    let mut partner = "".to_string();
    match reader.read_line(&mut partner) {
        Ok(_) => {
            let partner = partner.trim().to_string();
            let success: bool;
            {
                let mut guard = chat_map.lock().unwrap();
                match guard.get_mut(&partner) {
                    Some(clientinfo) => {
                        if partner.as_str() != username.as_str() && clientinfo.partner == None {
                            success = true;
                            clientinfo.partner = Some(username.to_string());
                        } else {
                            success = false;
                        }
                    },
                    None => success = false
                }
            }

            if success {
                // TODO: Changing yourself to be chatting means that your partner can't choose
                //       to chat with you, since it sees you already have a partner
                //       Need to notify partner when someone wants to chat with them
                //let mut guard = chat_map.lock().unwrap();
                //guard.get_mut(&username).unwrap().partner = Some(partner.to_string());
            } else {
                return choose_chat_partner(stream, username, chat_map);
            }
        }
        Err(e) => println!("{}", e)
    }
    partner.trim().to_string()
}

// TODO: Receiving a message while in the middle of typing a message
//       inserts received message into middle of your message
fn chat(stream: TcpStream, username: String, partner: String, sender_to_router: Sender<Message>, receiver_from_router: Receiver<Message>) {
    // Send messages
    {
        let username = username.clone();
        let reader = BufReader::new(stream.try_clone().unwrap());
        thread::spawn(move|| {
            let mut lines = reader.lines();
            while let Some(Ok(line)) = lines.next() {
                println!("{}",line);
                let msg = Message::new("Date".to_string(), username.to_string(), partner.to_string(), line.to_string());
                sender_to_router.send(msg).unwrap();
            }
        });
    }

    // Receive Messages
    {
        let username = username.clone();
        let mut stream = stream;
        thread::spawn(move|| {
            loop {
                match receiver_from_router.try_recv() {
                    Ok(msg) => {
                        println!("User {} received message: {}", &username, msg.content());
                        stream.write(&msg.from().to_string().into_bytes()).unwrap();
                        stream.write(&": ".to_string().into_bytes()).unwrap();
                        stream.write(&msg.content().to_string().into_bytes()).unwrap();
                        stream.write(&"\n".to_string().into_bytes()).unwrap();
                    }
                    Err(TryRecvError::Empty) => continue,
                    Err(TryRecvError::Disconnected) => panic!("User {} disconnected from router", &username)
                }
            }
        });
    }
}<commit_msg>Able to chat between clients, still no notification when someone requests to chat with you<commit_after>use std::io::{BufRead, BufReader, Write};
use std::net::TcpStream;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Receiver, Sender, TryRecvError};
use std::thread;

use chatmap::{ChatMap, ClientInfo};
use message::Message;

pub fn create_client(stream: TcpStream, sender_to_router: Sender<Message>, chat_map: &Arc<Mutex<ChatMap>>) {
    println!("New client");
    let mut stream = stream;
    stream.write(&"Welcome to Smazy\n".to_string().into_bytes()).unwrap();
    stream.write(&"Please enter a username:\n".to_string().into_bytes()).unwrap();
    stream.flush().unwrap();
    let mut reader = BufReader::new(stream.try_clone().unwrap());
    let mut username = "".to_string();
    match reader.read_line(&mut username) {
        Ok(n) => {
            let username = username.trim().to_string();
            let (sender_to_client, receiver_from_router) = channel();
            {
                let mut guard = chat_map.lock().unwrap();
                if n > 0 && !guard.contains_key(&username) {
                    let client_info = ClientInfo{partner: None, sender_to_client: sender_to_client};
                    guard.insert(username.to_string(), client_info);
                    println!("New user: {}", username.to_string());
                }
            }
            let partner = choose_chat_partner(stream.try_clone().unwrap(), username.to_string(), chat_map);
            chat(stream, username, partner, sender_to_router, receiver_from_router);
        },
        Err(e) => println!("{}", e)
    }
}

fn choose_chat_partner(stream: TcpStream, username: String, chat_map: &Arc<Mutex<ChatMap>>) -> String {
    let mut stream = stream;

    stream.write(&"Here are the users available to chat:\n".to_string().into_bytes()).unwrap();
    {
        let guard = chat_map.lock().unwrap();
        for (name, client_info) in guard.iter() {
            if name.as_str() != username.as_str() && client_info.partner == None {
                stream.write(&name.to_string().into_bytes()).unwrap();
                stream.write(&"\n".to_string().into_bytes()).unwrap();
            }
        }
    }

    stream.write(&"Select who you want to chat with:\n".to_string().into_bytes()).unwrap();

    let mut reader = BufReader::new(stream.try_clone().unwrap());
    let mut partner = "".to_string();

    match reader.read_line(&mut partner) {
        Ok(_) => {
            partner = partner.trim().to_string();

            if !try_select_partner(chat_map, username.to_string(), partner.to_string()) {
                return choose_chat_partner(stream, username, chat_map);
            }
        }
        Err(e) => println!("{}", e)
    }

    partner
}

fn try_select_partner(chat_map: &Arc<Mutex<ChatMap>>, username: String, partner: String) -> bool {
    let success: bool;
    {
        let mut guard = chat_map.lock().unwrap();

        match guard.get_mut(&partner) {
            Some(clientinfo) => {
                if partner.as_str() != username.as_str()
                    && (clientinfo.partner == None
                        || clientinfo.partner == Some(username.to_string())) {
                    clientinfo.partner = Some(username.to_string());
                    success = true;
                } else {
                    success = false;
                }
            },
            None => success = false
        }
    }

    if success {
        let mut guard = chat_map.lock().unwrap();
        println!("Assigning partner {} to {}", partner, username);
        guard.get_mut(&username).unwrap().partner = Some(partner.to_string());
    }
    success
}

// TODO: Receiving a message while in the middle of typing a message
//       inserts received message into middle of your message
fn chat(stream: TcpStream, username: String, partner: String, sender_to_router: Sender<Message>, receiver_from_router: Receiver<Message>) {
    // Send messages
    {
        let username = username.clone();
        let reader = BufReader::new(stream.try_clone().unwrap());
        thread::spawn(move|| {
            let mut lines = reader.lines();
            while let Some(Ok(line)) = lines.next() {
                println!("{}",line);
                let msg = Message::new("Date".to_string(), username.to_string(), partner.to_string(), line.to_string());
                sender_to_router.send(msg).unwrap();
            }
        });
    }

    // Receive Messages
    {
        let username = username.clone();
        let mut stream = stream;
        thread::spawn(move|| {
            loop {
                match receiver_from_router.try_recv() {
                    Ok(msg) => {
                        println!("User {} received message: {}", &username, msg.content());
                        stream.write(&msg.from().to_string().into_bytes()).unwrap();
                        stream.write(&": ".to_string().into_bytes()).unwrap();
                        stream.write(&msg.content().to_string().into_bytes()).unwrap();
                        stream.write(&"\n".to_string().into_bytes()).unwrap();
                    }
                    Err(TryRecvError::Empty) => continue,
                    Err(TryRecvError::Disconnected) => panic!("User {} disconnected from router", &username)
                }
            }
        });
    }
}<|endoftext|>


<commit_before>//! Producer and consumer configuration.
//!
//! ## C library configuration
//!
//! The Rust library will forward all the configuration to the C library. The most frequently
//! used parameters are listed here.
//!
//! ### Frequently used parameters
//!
//! For producer-specific and consumer-specific parameters check the producer and consumer modules
//! documentation. The full list of available parameters is available in the  [librdkafka
//! documentation](https://github.com/edenhill/librdkafka/blob/master/CONFIGURATION.md).
//!
//! - `client.id` (rdkafka): Client identifier.
//! - `bootstrap.servers`: Initial list of brokers as a CSV list of broker host or host:port.
//! - `message.max.bytes` (1000000): Maximum message size.
//! - `debug`: A comma-separated list of debug contexts to enable. Use 'all' to print all the debugging information.
//! - `statistics.interval.ms` (0 - disabled): how often the statistic callback specified in the `Context` will be called.
//!

use std::collections::HashMap;
use std::ffi::CString;
use std::mem;

use log::{log_enabled, trace, Level};

use rdkafka_sys as rdsys;
use rdkafka_sys::types::*;

use crate::client::ClientContext;
use crate::error::{IsError, KafkaError, KafkaResult};
use crate::util::ErrBuf;

/// The log levels supported by librdkafka.
#[derive(Copy, Clone, Debug)]
pub enum RDKafkaLogLevel {
    /// Higher priority then Level::Error from the log crate.
    Emerg = 0,
    /// Higher priority then Level::Error from the log crate.
    Alert = 1,
    /// Higher priority then Level::Error from the log crate.
    Critical = 2,
    /// Equivalent to Level::Error from the log crate.
    Error = 3,
    /// Equivalent to Level::Warning from the log crate.
    Warning = 4,
    /// Higher priority then Level::Info from the log crate.
    Notice = 5,
    /// Equivalent to Level::Info from the log crate.
    Info = 6,
    /// Equivalent to Level::Debug from the log crate.
    Debug = 7,
}

impl RDKafkaLogLevel {
    pub(crate) fn from_int(level: i32) -> RDKafkaLogLevel {
        match level {
            0 => RDKafkaLogLevel::Emerg,
            1 => RDKafkaLogLevel::Alert,
            2 => RDKafkaLogLevel::Critical,
            3 => RDKafkaLogLevel::Error,
            4 => RDKafkaLogLevel::Warning,
            5 => RDKafkaLogLevel::Notice,
            6 => RDKafkaLogLevel::Info,
            _ => RDKafkaLogLevel::Debug,
        }
    }
}

//
// ********** CLIENT CONFIG **********
//

/// A native rdkafka-sys client config.
pub struct NativeClientConfig {
    ptr: *mut RDKafkaConf,
}

impl NativeClientConfig {
    /// Wraps a pointer to an `RDKafkaConfig` object and returns a new `NativeClientConfig`.
    pub(crate) unsafe fn from_ptr(ptr: *mut RDKafkaConf) -> NativeClientConfig {
        NativeClientConfig { ptr }
    }

    /// Returns the pointer to the librdkafka RDKafkaConf structure.
    pub fn ptr(&self) -> *mut RDKafkaConf {
        self.ptr
    }

    /// Returns the pointer to the librdkafka RDKafkaConf structure. This method should be used when
    /// the native pointer is intended to be moved. The destructor won't be executed automatically;
    /// the caller should take care of deallocating the resource when no longer needed.
    pub fn ptr_move(self) -> *mut RDKafkaConf {
        let ptr = self.ptr;
        mem::forget(self);
        ptr
    }
}

impl Drop for NativeClientConfig {
    fn drop(&mut self) {
        trace!("Drop NativeClientConfig {:p}", self.ptr());
        unsafe { rdsys::rd_kafka_conf_destroy(self.ptr) }
    }
}

/// Client configuration.
#[derive(Clone)]
pub struct ClientConfig {
    conf_map: HashMap<String, String>,
    /// The librdkafka logging level. Refer to `RDKafkaLogLevel` for the list of available levels.
    pub log_level: RDKafkaLogLevel,
}

impl Default for ClientConfig {
    fn default() -> Self {
        Self::new()
    }
}

impl ClientConfig {
    /// Creates a new empty configuration.
    pub fn new() -> ClientConfig {
        ClientConfig {
            conf_map: HashMap::new(),
            log_level: log_level_from_global_config(),
        }
    }

    /// Sets a new parameter in the configuration.
    pub fn set<'a>(&'a mut self, key: &str, value: &str) -> &'a mut ClientConfig {
        self.conf_map.insert(key.to_string(), value.to_string());
        self
    }

    /// Sets the log level of the client. If not specified, the log level will be calculated based
    /// on the global log level of the log crate.
    pub fn set_log_level(&mut self, log_level: RDKafkaLogLevel) -> &mut ClientConfig {
        self.log_level = log_level;
        self
    }

    /// Returns the native rdkafka-sys configuration.
    pub fn create_native_config(&self) -> KafkaResult<NativeClientConfig> {
        let conf = unsafe { rdsys::rd_kafka_conf_new() };
        let mut err_buf = ErrBuf::new();
        for (key, value) in &self.conf_map {
            let key_c = CString::new(key.to_string())?;
            let value_c = CString::new(value.to_string())?;
            let ret = unsafe {
                rdsys::rd_kafka_conf_set(
                    conf,
                    key_c.as_ptr(),
                    value_c.as_ptr(),
                    err_buf.as_mut_ptr(),
                    err_buf.len(),
                )
            };
            if ret.is_error() {
                return Err(KafkaError::ClientConfig(
                    ret,
                    err_buf.to_string(),
                    key.to_string(),
                    value.to_string(),
                ));
            }
        }
        Ok(unsafe { NativeClientConfig::from_ptr(conf) })
    }

    /// Uses the current configuration to create a new Consumer or Producer.
    pub fn create<T: FromClientConfig>(&self) -> KafkaResult<T> {
        T::from_config(self)
    }

    /// Uses the current configuration and the provided context to create a new Consumer or Producer.
    pub fn create_with_context<C, T>(&self, context: C) -> KafkaResult<T>
    where
        C: ClientContext,
        T: FromClientConfigAndContext<C>,
    {
        T::from_config_and_context(self, context)
    }
}

/// Return the log level
fn log_level_from_global_config() -> RDKafkaLogLevel {
    if log_enabled!(target: "librdkafka", Level::Debug) {
        RDKafkaLogLevel::Debug
    } else if log_enabled!(target: "librdkafka", Level::Info) {
        RDKafkaLogLevel::Info
    } else if log_enabled!(target: "librdkafka", Level::Warn) {
        RDKafkaLogLevel::Warning
    } else {
        RDKafkaLogLevel::Error
    }
}

/// Create a new client based on the provided configuration.
pub trait FromClientConfig: Sized {
    /// Create a client from client configuration. The default client context will be used.
    fn from_config(_: &ClientConfig) -> KafkaResult<Self>;
}

/// Create a new client based on the provided configuration and context.
pub trait FromClientConfigAndContext<C: ClientContext>: Sized {
    /// Create a client from client configuration and a client context.
    fn from_config_and_context(_: &ClientConfig, _: C) -> KafkaResult<Self>;
}
<commit_msg>derive Debug for ClientConfig<commit_after>//! Producer and consumer configuration.
//!
//! ## C library configuration
//!
//! The Rust library will forward all the configuration to the C library. The most frequently
//! used parameters are listed here.
//!
//! ### Frequently used parameters
//!
//! For producer-specific and consumer-specific parameters check the producer and consumer modules
//! documentation. The full list of available parameters is available in the  [librdkafka
//! documentation](https://github.com/edenhill/librdkafka/blob/master/CONFIGURATION.md).
//!
//! - `client.id` (rdkafka): Client identifier.
//! - `bootstrap.servers`: Initial list of brokers as a CSV list of broker host or host:port.
//! - `message.max.bytes` (1000000): Maximum message size.
//! - `debug`: A comma-separated list of debug contexts to enable. Use 'all' to print all the debugging information.
//! - `statistics.interval.ms` (0 - disabled): how often the statistic callback specified in the `Context` will be called.
//!

use std::collections::HashMap;
use std::ffi::CString;
use std::mem;

use log::{log_enabled, trace, Level};

use rdkafka_sys as rdsys;
use rdkafka_sys::types::*;

use crate::client::ClientContext;
use crate::error::{IsError, KafkaError, KafkaResult};
use crate::util::ErrBuf;

/// The log levels supported by librdkafka.
#[derive(Copy, Clone, Debug)]
pub enum RDKafkaLogLevel {
    /// Higher priority then Level::Error from the log crate.
    Emerg = 0,
    /// Higher priority then Level::Error from the log crate.
    Alert = 1,
    /// Higher priority then Level::Error from the log crate.
    Critical = 2,
    /// Equivalent to Level::Error from the log crate.
    Error = 3,
    /// Equivalent to Level::Warning from the log crate.
    Warning = 4,
    /// Higher priority then Level::Info from the log crate.
    Notice = 5,
    /// Equivalent to Level::Info from the log crate.
    Info = 6,
    /// Equivalent to Level::Debug from the log crate.
    Debug = 7,
}

impl RDKafkaLogLevel {
    pub(crate) fn from_int(level: i32) -> RDKafkaLogLevel {
        match level {
            0 => RDKafkaLogLevel::Emerg,
            1 => RDKafkaLogLevel::Alert,
            2 => RDKafkaLogLevel::Critical,
            3 => RDKafkaLogLevel::Error,
            4 => RDKafkaLogLevel::Warning,
            5 => RDKafkaLogLevel::Notice,
            6 => RDKafkaLogLevel::Info,
            _ => RDKafkaLogLevel::Debug,
        }
    }
}

//
// ********** CLIENT CONFIG **********
//

/// A native rdkafka-sys client config.
pub struct NativeClientConfig {
    ptr: *mut RDKafkaConf,
}

impl NativeClientConfig {
    /// Wraps a pointer to an `RDKafkaConfig` object and returns a new `NativeClientConfig`.
    pub(crate) unsafe fn from_ptr(ptr: *mut RDKafkaConf) -> NativeClientConfig {
        NativeClientConfig { ptr }
    }

    /// Returns the pointer to the librdkafka RDKafkaConf structure.
    pub fn ptr(&self) -> *mut RDKafkaConf {
        self.ptr
    }

    /// Returns the pointer to the librdkafka RDKafkaConf structure. This method should be used when
    /// the native pointer is intended to be moved. The destructor won't be executed automatically;
    /// the caller should take care of deallocating the resource when no longer needed.
    pub fn ptr_move(self) -> *mut RDKafkaConf {
        let ptr = self.ptr;
        mem::forget(self);
        ptr
    }
}

impl Drop for NativeClientConfig {
    fn drop(&mut self) {
        trace!("Drop NativeClientConfig {:p}", self.ptr());
        unsafe { rdsys::rd_kafka_conf_destroy(self.ptr) }
    }
}

/// Client configuration.
#[derive(Clone, Debug)]
pub struct ClientConfig {
    conf_map: HashMap<String, String>,
    /// The librdkafka logging level. Refer to `RDKafkaLogLevel` for the list of available levels.
    pub log_level: RDKafkaLogLevel,
}

impl Default for ClientConfig {
    fn default() -> Self {
        Self::new()
    }
}

impl ClientConfig {
    /// Creates a new empty configuration.
    pub fn new() -> ClientConfig {
        ClientConfig {
            conf_map: HashMap::new(),
            log_level: log_level_from_global_config(),
        }
    }

    /// Sets a new parameter in the configuration.
    pub fn set<'a>(&'a mut self, key: &str, value: &str) -> &'a mut ClientConfig {
        self.conf_map.insert(key.to_string(), value.to_string());
        self
    }

    /// Sets the log level of the client. If not specified, the log level will be calculated based
    /// on the global log level of the log crate.
    pub fn set_log_level(&mut self, log_level: RDKafkaLogLevel) -> &mut ClientConfig {
        self.log_level = log_level;
        self
    }

    /// Returns the native rdkafka-sys configuration.
    pub fn create_native_config(&self) -> KafkaResult<NativeClientConfig> {
        let conf = unsafe { rdsys::rd_kafka_conf_new() };
        let mut err_buf = ErrBuf::new();
        for (key, value) in &self.conf_map {
            let key_c = CString::new(key.to_string())?;
            let value_c = CString::new(value.to_string())?;
            let ret = unsafe {
                rdsys::rd_kafka_conf_set(
                    conf,
                    key_c.as_ptr(),
                    value_c.as_ptr(),
                    err_buf.as_mut_ptr(),
                    err_buf.len(),
                )
            };
            if ret.is_error() {
                return Err(KafkaError::ClientConfig(
                    ret,
                    err_buf.to_string(),
                    key.to_string(),
                    value.to_string(),
                ));
            }
        }
        Ok(unsafe { NativeClientConfig::from_ptr(conf) })
    }

    /// Uses the current configuration to create a new Consumer or Producer.
    pub fn create<T: FromClientConfig>(&self) -> KafkaResult<T> {
        T::from_config(self)
    }

    /// Uses the current configuration and the provided context to create a new Consumer or Producer.
    pub fn create_with_context<C, T>(&self, context: C) -> KafkaResult<T>
    where
        C: ClientContext,
        T: FromClientConfigAndContext<C>,
    {
        T::from_config_and_context(self, context)
    }
}

/// Return the log level
fn log_level_from_global_config() -> RDKafkaLogLevel {
    if log_enabled!(target: "librdkafka", Level::Debug) {
        RDKafkaLogLevel::Debug
    } else if log_enabled!(target: "librdkafka", Level::Info) {
        RDKafkaLogLevel::Info
    } else if log_enabled!(target: "librdkafka", Level::Warn) {
        RDKafkaLogLevel::Warning
    } else {
        RDKafkaLogLevel::Error
    }
}

/// Create a new client based on the provided configuration.
pub trait FromClientConfig: Sized {
    /// Create a client from client configuration. The default client context will be used.
    fn from_config(_: &ClientConfig) -> KafkaResult<Self>;
}

/// Create a new client based on the provided configuration and context.
pub trait FromClientConfigAndContext<C: ClientContext>: Sized {
    /// Create a client from client configuration and a client context.
    fn from_config_and_context(_: &ClientConfig, _: C) -> KafkaResult<Self>;
}
<|endoftext|>