<commit_before>/*
 * 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/.
 *
 * Copyright 2017 - Dario Ostuni <dario.ostuni@gmail.com>
 *
 */

use common::*;
use serde::ser::Serialize;
use serde_json;
use std::collections::HashMap;
use std::io::ErrorKind;
use std::net::{TcpListener, TcpStream};
use std::sync::mpsc::{Receiver, Sender, channel};
use std::thread;
use std::thread::JoinHandle;
use std::time::Duration;
use tungstenite::{Error, Message, WebSocket, accept};
use tungstenite::protocol::Role;

#[derive(Serialize, Deserialize, Debug)]
pub struct ServerConfig
{
    turn_time_ms: u64,
    server_port: u16,
}

enum ConnectionType
{
    Viewer,
    InGame(u64),
    InQueue,
    Unknown,
}

struct Connection
{
    role: ConnectionType,
    handle: JoinHandle<()>,
    socket: WebSocket<TcpStream>,
}

enum MessageResponse
{
    Mail(u64, Message),
    Disconnected(u64),
}

pub struct Server
{
    connections: HashMap<u64, Connection>,
    incoming_connections: Receiver<WebSocket<TcpStream>>,
    incoming_messages: Receiver<MessageResponse>,
    message_sender: Sender<MessageResponse>,
    seed: u64,
    turn_time: u64,
    game_id: u64,
}

impl Connection
{
    #[cfg_attr(rustfmt, rustfmt_skip)]
    pub fn new(ws: WebSocket<TcpStream>, sender: Sender<MessageResponse>, seed: u64) -> Connection
    {
        let mut ws = ws;
        Connection {
            socket: WebSocket::from_raw_socket(ws.get_ref().try_clone().unwrap(), Role::Server),
            role: ConnectionType::Unknown,
            handle: thread::spawn(move ||
                loop
                {
                    match ws.read_message()
                    {
                        Ok(x) => sender.send(MessageResponse::Mail(seed, x)).unwrap(),
                        Err(Error::ConnectionClosed) => {
                            sender.send(MessageResponse::Disconnected(seed)).unwrap();
                            break;
                        },
                        Err(_) => continue,
                    }
                }),
        }
    }
}

impl Server
{
    pub fn new(config: ServerConfig) -> Server
    {
        let (sen, rec) = channel();
        let server_port = config.server_port;
        thread::spawn(move || {
            let listener = TcpListener::bind(("0.0.0.0", server_port)).unwrap();
            for stream in listener.incoming()
            {
                match stream
                {
                    Ok(s) =>
                    {
                        match accept(s)
                        {
                            Ok(w) => sen.send(w).unwrap(),
                            Err(_) => error!("Error while accepting WebSocket connection"),
                        }
                    }
                    Err(why) => error!("Error while accepting connection: {}", why),
                }
            }
        });
        let (sen_msg, rec_msg) = channel();
        Server {
            connections: Default::default(),
            incoming_connections: rec,
            incoming_messages: rec_msg,
            message_sender: sen_msg,
            seed: 0,
            turn_time: config.turn_time_ms,
            game_id: 0,
        }
    }

    pub fn main(&mut self)
    {
        while let Ok(x) = self.incoming_connections.try_recv()
        {
            self.connections.insert(self.seed, Connection::new(x, self.message_sender.clone(), self.seed));
            self.seed += 1;
        }
        while let Ok(x) = self.incoming_messages.try_recv()
        {
            if let MessageResponse::Disconnected(id) = x
            {
                self.connections.remove(&id);
            }
            else if let MessageResponse::Mail(id, msg) = x
            {
                let msg = match msg.into_text()
                {
                    Ok(s) => s,
                    Err(why) =>
                    {
                        error!("Received garbage: {}", why);
                        continue;
                    }
                };
                let msg: ClientMessage = match serde_json::from_str(msg.as_str())
                {
                    Ok(s) => s,
                    Err(why) =>
                    {
                        error!("Received garbage: {}", why);
                        continue;
                    }
                };
                if let ClientMessage::HandShake(role) = msg
                {
                    self.handle_accept(id, role);
                }
                else if let ClientMessage::Command(command) = msg
                {
                    self.handle_command(id, command);
                }
                else
                {
                    unreachable!();
                }
            }
            else
            {
                unreachable!();
            }
        }
        thread::sleep(Duration::from_millis(self.turn_time));
    }

    fn handle_accept(&mut self, id: u64, role: ClientRole)
    {
        match role
        {
            ClientRole::Viewer =>
            {
                {
                    let v = self.connections.get_mut(&id);
                    if v.is_none()
                    {
                        return;
                    }
                    let v = v.unwrap();
                    v.role = ConnectionType::Viewer;
                }
                self.send_data(id, &ServerResponse::Ok);
            }
            ClientRole::Player(info) => unimplemented!(),
        };
    }

    fn handle_command(&mut self, id: u64, command: ClientCommand)
    {
        unimplemented!();
    }

    fn send_data<T: ?Sized + Serialize>(&mut self, id: u64, value: &T)
    {
        match serde_json::to_string(value)
        {
            Ok(s) =>
            {
                match self.connections.get_mut(&id)
                {
                    Some(conn) =>
                    {
                        match conn.socket.write_message(Message::text(s))
                        {
                            Ok(_) =>
                            {}
                            Err(Error::Io(x)) =>
                            {
                                match x.kind()
                                {
                                    ErrorKind::WouldBlock =>
                                    {
                                        match conn.socket.write_pending()
                                        {
                                            Ok(_) =>
                                            {}
                                            Err(why) => error!("Unable to flush WebSocket {}", why),
                                        }
                                    }
                                    _ => error!("I/O error"),
                                }
                            }
                            Err(Error::ConnectionClosed) => self.message_sender.send(MessageResponse::Disconnected(id)).unwrap(),
                            Err(_) => error!("Error while sending data"),
                        }
                    }
                    None => error!("Missing id"),
                }
            }
            Err(_) => unreachable!(),
        };
    }
}

impl ServerConfig
{
    pub fn new() -> ServerConfig
    {
        ServerConfig {
            turn_time_ms: 1000,
            server_port: 42000,
        }
    }
}
<commit_msg>Implemented wait queue for players<commit_after>/*
 * 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/.
 *
 * Copyright 2017 - Dario Ostuni <dario.ostuni@gmail.com>
 *
 */

use chasher::player_hash;
use common::*;
use serde::ser::Serialize;
use serde_json;
use std::collections::HashMap;
use std::io::ErrorKind;
use std::net::{TcpListener, TcpStream};
use std::sync::mpsc::{Receiver, Sender, channel};
use std::thread;
use std::thread::JoinHandle;
use std::time::Duration;
use tungstenite::{Error, Message, WebSocket, accept};
use tungstenite::protocol::Role;

#[derive(Serialize, Deserialize, Debug)]
pub struct ServerConfig
{
    tick_time_ms: u64,
    server_port: u16,
    game_start_ticks: u64,
}

enum ConnectionType
{
    Viewer,
    InGame(u64),
    InQueue(u64),
    Unknown,
}

struct Connection
{
    role: ConnectionType,
    handle: JoinHandle<()>,
    socket: WebSocket<TcpStream>,
}

enum MessageResponse
{
    Mail(u64, Message),
    Disconnected(u64),
}

pub struct Server
{
    connections: HashMap<u64, Connection>,
    incoming_connections: Receiver<WebSocket<TcpStream>>,
    incoming_messages: Receiver<MessageResponse>,
    message_sender: Sender<MessageResponse>,
    players: HashMap<u64, Player>,
    queue: HashMap<u64, Player>,
    seed: u64,
    tick_time: u64,
    game_id: u64,
    in_game: bool,
    game_start_ticks: u64,
}

impl Connection
{
    #[cfg_attr(rustfmt, rustfmt_skip)]
    pub fn new(ws: WebSocket<TcpStream>, sender: Sender<MessageResponse>, seed: u64) -> Connection
    {
        let mut ws = ws;
        Connection {
            socket: WebSocket::from_raw_socket(ws.get_ref().try_clone().unwrap(), Role::Server),
            role: ConnectionType::Unknown,
            handle: thread::spawn(move ||
                loop
                {
                    match ws.read_message()
                    {
                        Ok(x) => sender.send(MessageResponse::Mail(seed, x)).unwrap(),
                        Err(Error::ConnectionClosed) => {
                            sender.send(MessageResponse::Disconnected(seed)).unwrap();
                            break;
                        },
                        Err(_) => continue,
                    }
                }),
        }
    }
}

impl Server
{
    pub fn new(config: ServerConfig) -> Server
    {
        let (sen, rec) = channel();
        let server_port = config.server_port;
        thread::spawn(move || {
            let listener = TcpListener::bind(("0.0.0.0", server_port)).unwrap();
            for stream in listener.incoming()
            {
                match stream
                {
                    Ok(s) =>
                    {
                        match accept(s)
                        {
                            Ok(w) => sen.send(w).unwrap(),
                            Err(_) => error!("Error while accepting WebSocket connection"),
                        }
                    }
                    Err(why) => error!("Error while accepting connection: {}", why),
                }
            }
        });
        let (sen_msg, rec_msg) = channel();
        Server {
            connections: Default::default(),
            incoming_connections: rec,
            incoming_messages: rec_msg,
            message_sender: sen_msg,
            seed: 0,
            tick_time: config.tick_time_ms,
            game_id: 0,
            in_game: false,
            players: Default::default(),
            queue: Default::default(),
            game_start_ticks: config.game_start_ticks,
        }
    }

    pub fn main(&mut self)
    {
        while let Ok(x) = self.incoming_connections.try_recv()
        {
            self.connections.insert(self.seed, Connection::new(x, self.message_sender.clone(), self.seed));
            self.seed += 1;
        }
        while let Ok(x) = self.incoming_messages.try_recv()
        {
            if let MessageResponse::Disconnected(id) = x
            {
                self.connections.remove(&id);
            }
            else if let MessageResponse::Mail(id, msg) = x
            {
                let msg = match msg.into_text()
                {
                    Ok(s) => s,
                    Err(why) =>
                    {
                        error!("Received garbage: {}", why);
                        continue;
                    }
                };
                let msg: ClientMessage = match serde_json::from_str(msg.as_str())
                {
                    Ok(s) => s,
                    Err(why) =>
                    {
                        error!("Received garbage: {}", why);
                        continue;
                    }
                };
                if let ClientMessage::HandShake(role) = msg
                {
                    self.handle_accept(id, role);
                }
                else if let ClientMessage::Command(command) = msg
                {
                    self.handle_command(id, command);
                }
                else
                {
                    unreachable!();
                }
            }
            else
            {
                unreachable!();
            }
        }
        thread::sleep(Duration::from_millis(self.tick_time));
    }

    #[cfg_attr(feature = "cargo-clippy", allow(map_entry))]
    fn handle_accept(&mut self, id: u64, role: ClientRole)
    {
        match role
        {
            ClientRole::Viewer =>
            {
                {
                    let v = self.connections.get_mut(&id);
                    if v.is_none()
                    {
                        return;
                    }
                    let v = v.unwrap();
                    v.role = ConnectionType::Viewer;
                }
                self.send_data(id, &ServerResponse::Ok);
            }
            ClientRole::Player(info) =>
            {
                let user_game_id = player_hash(&info);
                if self.queue.contains_key(&user_game_id)
                {
                    self.send_data(id, &ServerResponse::Error(String::from("Username already taken")));
                }
                else
                {
                    {
                        let conn = self.connections.get_mut(&id);
                        if conn.is_none()
                        {
                            error!("Trying to register a non-existent WebSocket as player!");
                            return;
                        }
                        let conn = conn.unwrap();
                        conn.role = ConnectionType::InQueue(user_game_id);
                        self.queue.insert(user_game_id,
                                          Player {
                                              name: info.name,
                                              id: user_game_id,
                                              points: 0,
                                              position: Point { x: 0, y: 0 },
                                          });
                    }
                    self.send_data(id, &ServerResponse::Ok);
                }
            }
        };
    }

    fn handle_command(&mut self, id: u64, command: ClientCommand)
    {
        unimplemented!();
    }

    fn send_data<T: ?Sized + Serialize>(&mut self, id: u64, value: &T)
    {
        match serde_json::to_string(value)
        {
            Ok(s) =>
            {
                match self.connections.get_mut(&id)
                {
                    Some(conn) =>
                    {
                        match conn.socket.write_message(Message::text(s))
                        {
                            Ok(_) =>
                            {}
                            Err(Error::Io(x)) =>
                            {
                                match x.kind()
                                {
                                    ErrorKind::WouldBlock =>
                                    {
                                        match conn.socket.write_pending()
                                        {
                                            Ok(_) =>
                                            {}
                                            Err(why) => error!("Unable to flush WebSocket {}", why),
                                        }
                                    }
                                    _ => error!("I/O error"),
                                }
                            }
                            Err(Error::ConnectionClosed) => self.message_sender.send(MessageResponse::Disconnected(id)).unwrap(),
                            Err(_) => error!("Error while sending data"),
                        }
                    }
                    None => error!("Missing id"),
                }
            }
            Err(_) => unreachable!(),
        };
    }
}

impl ServerConfig
{
    pub fn new() -> ServerConfig
    {
        ServerConfig {
            tick_time_ms: 1000,
            server_port: 42000,
            game_start_ticks: 30,
        }
    }
}
<|endoftext|>


<commit_before>//! Signals are values that discretely change over time.

use crate::stream::Stream;
use crate::types::{MaybeOwned, SharedImpl, SharedSignal, Storage};
use parking_lot::Mutex;
use std::fmt;
use std::sync::{mpsc, Arc};

use self::SigValue::*;

/// Represents a discrete value that changes over time.
///
/// Signals are usually constructed by stream operations and can be read using the `sample` or
/// `take` methods. They update lazily when someone reads them.
#[derive(Clone, Debug)]
pub struct Signal<T>(SigValue<T>);

/// The content source of a signal.
#[derive(Clone)]
enum SigValue<T> {
    /// A signal with constant value.
    Constant(T),
    /// A signal that generates it's values from a function.
    ///
    /// This is produced by `Signal::from_fn`
    Dynamic(Arc<dyn Fn() -> T + Send + Sync>),
    /// A signal that contains shared data.
    ///
    /// This is mainly produced by stream methods or mapping other signals.
    Shared(Arc<dyn SharedSignal<T> + Send + Sync>),
    /// A signal that contains a signal, and allows sampling the inner signal directly.
    ///
    /// This is produced by `Signal::switch`
    Nested(Arc<dyn Fn() -> Signal<T> + Send + Sync>),
}

impl<T> Signal<T> {
    /// Creates a signal with constant value.
    ///
    /// The value is assumed to be constant, so changing it while it's stored on the
    /// signal is a logic error and will cause unexpected results.
    pub fn constant(val: T) -> Self {
        Signal(Constant(val))
    }

    /// Creates a signal that samples it's values from an external source.
    ///
    /// The closure is meant to sample a continuous value from the real world,
    /// so the signal value is assumed to be always changing.
    pub fn from_fn<F>(f: F) -> Self
    where
        F: Fn() -> T + Send + Sync + 'static,
    {
        Signal(Dynamic(Arc::new(f)))
    }

    /// Creates a new shared signal.
    pub(crate) fn shared<S>(storage: Arc<S>) -> Self
    where
        S: SharedSignal<T> + Send + Sync + 'static,
    {
        Signal(Shared(storage))
    }

    /// Moves the value out of the signal.
    ///
    /// To avoid leaving the internal storage empty (and causing panics), it must be filled with
    /// a default value. This marks the signal chain as changed, so most of time no one will
    /// notice anything.
    pub fn take(self) -> T
    where
        T: Default,
    {
        match self.0 {
            Constant(val) => val,
            Dynamic(f) => f(),
            Shared(s) => s.sample().replace(T::default()),
            Nested(f) => f().take(),
        }
    }

    /// Samples the value of the signal.
    ///
    /// This will clone the value stored in the signal.
    pub fn sample(&self) -> T
    where
        T: Clone,
    {
        match self.0 {
            Constant(ref val) => T::clone(val),
            Dynamic(ref f) => f(),
            Shared(ref s) => s.sample().get(),
            Nested(ref f) => f().sample(),
        }
    }

    /// Maps a signal with the provided function.
    ///
    /// The closure is called only when the parent signal changes.
    pub fn map<F, R>(&self, f: F) -> Signal<R>
    where
        F: Fn(T) -> R + Send + Sync + 'static,
        T: Clone + 'static,
        R: Send + 'static,
    {
        match self.0 {
            // constant signal: apply f once to produce another constant signal
            Constant(ref val) => Signal::constant(f(val.clone())),
            // dynamic signal: sample and apply f
            Dynamic(ref sf) => {
                let sf = sf.clone();
                Signal::from_fn(move || f(sf()))
            }
            // shared signal: create shared storage with source
            Shared(ref sig) => Signal::shared(
                SharedImpl {
                    storage: Storage::default(),
                    source: sig.clone(),
                    f,
                }
                .wrap(),
            ),
            // nested signal: extract signal, sample and apply f
            Nested(ref sf) => {
                let sf = sf.clone();
                Signal::from_fn(move || f(sf().sample()))
            }
        }
    }
}

impl<T: Send + 'static> Signal<T> {
    /// Samples the value of this signal every time the trigger stream fires.
    pub fn snapshot<S, F, R>(&self, trigger: &Stream<S>, f: F) -> Stream<R>
    where
        F: Fn(T, MaybeOwned<'_, S>) -> R + Send + Sync + 'static,
        T: Clone,
        S: 'static,
        R: 'static,
    {
        match self.0 {
            Constant(ref val) => {
                let val = Mutex::new(val.clone()); // need Sync for T
                trigger.map(move |t| f(val.lock().clone(), t))
            }
            Dynamic(ref sf) => {
                let sf = sf.clone();
                trigger.map(move |t| f(sf(), t))
            }
            Shared(ref sig) => {
                let sig = sig.clone();
                trigger.map(move |t| f(sig.sample().get(), t))
            }
            Nested(ref sf) => {
                let sf = sf.clone();
                trigger.map(move |t| f(sf().sample(), t))
            }
        }
    }

    /// Stores the last value sent to a channel.
    ///
    /// When sampled, the resulting signal consumes all the current values on the channel
    /// (using non-blocking operations) and returns the last value seen.
    #[inline]
    pub fn from_channel(initial: T, rx: mpsc::Receiver<T>) -> Self {
        Self::fold_channel(initial, rx, |_, v| v)
    }

    /// Creates a signal that folds the values from a channel.
    ///
    /// When sampled, the resulting signal consumes all the current values on the channel
    /// (using non-blocking operations) and folds them using the current signal value as the
    /// initial accumulator state.
    pub fn fold_channel<V, F>(initial: T, rx: mpsc::Receiver<V>, f: F) -> Self
    where
        F: Fn(T, V) -> T + Send + Sync + 'static,
        V: Send + 'static,
    {
        Signal::shared(
            SharedImpl {
                storage: Storage::new(initial),
                source: Mutex::new(rx),
                f,
            }
            .wrap(),
        )
    }
}

impl<T: Clone + 'static> Signal<Signal<T>> {
    /// Creates a new signal that samples the inner value of a nested signal.
    pub fn switch(&self) -> Signal<T> {
        match self.0 {
            // constant signal: just extract the inner signal
            Constant(ref sig) => Signal::clone(sig),
            // dynamic signal: re-label as nested
            Dynamic(ref f) => Signal(Nested(f.clone())),
            // shared signal: sample to extract the inner signal
            Shared(ref sig) => {
                let sig = sig.clone();
                Signal(Nested(Arc::new(move || sig.sample().get())))
            }
            // nested signal: remove one layer
            Nested(ref f) => {
                let f = f.clone();
                Signal(Nested(Arc::new(move || f().sample())))
            }
        }
    }
}

impl<T: Default> Default for Signal<T> {
    /// Creates a constant signal with T's default value.
    #[inline]
    fn default() -> Self {
        Signal::constant(T::default())
    }
}

impl<T> From<T> for Signal<T> {
    /// Creates a constant signal from T.
    #[inline]
    fn from(val: T) -> Self {
        Signal::constant(val)
    }
}

impl<T: fmt::Debug> fmt::Debug for SigValue<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            Constant(ref val) => write!(f, "Constant({:?})", val),
            Dynamic(ref rf) => write!(f, "Dynamic(Fn@{:p})", rf),
            Shared(ref rs) => write!(f, "Shared(SharedSignal@{:p})", rs),
            Nested(ref rf) => write!(f, "Nested(Fn@{:p})", rf),
        }
    }
}

impl<T: fmt::Display + Clone> fmt::Display for Signal<T> {
    /// Samples the signal and formats the value.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&self.sample(), f)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::RwLock;
    use std::time::Instant;

    #[test]
    fn signal_basic() {
        let signal = Signal::constant(42);
        let double = signal.map(|a| a * 2);
        let plusone = double.map(|a| a + 1);
        assert_eq!(signal.sample(), 42);
        assert_eq!(double.sample(), 84);
        assert_eq!(plusone.sample(), 85);
    }

    #[test]
    fn signal_shared() {
        let st = SharedImpl::new(1, ()).wrap();
        let signal = Signal::shared(st.clone());
        let double = signal.map(|a| a * 2);

        assert_eq!(signal.sample(), 1);
        assert_eq!(double.sample(), 2);
        st.set(42);
        assert_eq!(signal.sample(), 42);
        assert_eq!(double.sample(), 84);
    }

    #[test]
    fn signal_dynamic() {
        let t = Instant::now();
        let signal = Signal::from_fn(move || t);
        assert_eq!(signal.sample(), t);

        let n = Arc::new(RwLock::new(1));
        let cloned = n.clone();
        let signal = Signal::from_fn(move || *cloned.read().unwrap());
        let double = signal.map(|a| a * 2);
        let plusone = double.map(|a| a + 1);
        assert_eq!(signal.sample(), 1);
        assert_eq!(double.sample(), 2);
        assert_eq!(plusone.sample(), 3);
        *n.write().unwrap() = 13;
        assert_eq!(signal.sample(), 13);
        assert_eq!(double.sample(), 26);
        assert_eq!(plusone.sample(), 27);
    }

    #[test]
    fn signal_take() {
        let signal = Signal::constant(42);
        assert_eq!(signal.take(), 42);

        let signal = Signal::from_fn(move || 33);
        assert_eq!(signal.clone().take(), 33);
        assert_eq!(signal.take(), 33);

        let st = SharedImpl::new(123, ()).wrap();
        let signal = Signal::shared(st.clone());
        assert_eq!(signal.clone().take(), 123);
        assert_eq!(signal.clone().take(), 0);
        st.set(13);
        assert_eq!(signal.take(), 13);
    }
}
<commit_msg>updated doc comments<commit_after>//! Signals are values that discretely change over time.

use crate::stream::Stream;
use crate::types::{MaybeOwned, SharedImpl, SharedSignal, Storage};
use parking_lot::Mutex;
use std::fmt;
use std::sync::{mpsc, Arc};

use self::SigValue::*;

/// Represents a discrete value that changes over time.
///
/// Signals are usually constructed by stream operations and can be read using the `sample` or
/// `take` methods. They update lazily when someone reads them.
#[derive(Clone, Debug)]
pub struct Signal<T>(SigValue<T>);

/// The content source of a signal.
#[derive(Clone)]
enum SigValue<T> {
    /// A signal with constant value.
    Constant(T),
    /// A signal that generates it's values from a function.
    ///
    /// This is produced by `Signal::from_fn`
    Dynamic(Arc<dyn Fn() -> T + Send + Sync>),
    /// A signal that contains shared data.
    ///
    /// This is mainly produced by stream methods or mapping other signals.
    Shared(Arc<dyn SharedSignal<T> + Send + Sync>),
    /// A signal that contains a signal, and allows sampling the inner signal directly.
    ///
    /// This is produced by `Signal::switch`
    Nested(Arc<dyn Fn() -> Signal<T> + Send + Sync>),
}

impl<T> Signal<T> {
    /// Creates a signal with constant value.
    ///
    /// The value is assumed to be constant, so changing it while it's stored on the
    /// signal is a logic error and will cause unexpected results.
    pub fn constant(val: T) -> Self {
        Signal(Constant(val))
    }

    /// Creates a signal that samples it's values from an external source.
    ///
    /// The closure is meant to sample a continuous value from the real world,
    /// so the signal value is assumed to be always changing.
    pub fn from_fn<F>(f: F) -> Self
    where
        F: Fn() -> T + Send + Sync + 'static,
    {
        Signal(Dynamic(Arc::new(f)))
    }

    /// Creates a new shared signal.
    pub(crate) fn shared<S>(storage: Arc<S>) -> Self
    where
        S: SharedSignal<T> + Send + Sync + 'static,
    {
        Signal(Shared(storage))
    }

    /// Moves the value out of the signal.
    ///
    /// This is meant to be a clone-free alternative to `sample`, but with some drawbacks.
    /// Taking the value out of a signal can potentially cause unintended results.
    /// To avoid leaving the internal storage empty (and causing panics), it must be filled with
    /// a default value. This default value will only be observed when taking the content of a
    /// signal created by the methods:
    ///
    /// - `Stream::fold`
    /// - `Stream::fold_clone`
    /// - `Stream::hold`
    /// - `Stream::hold_if`
    /// - `Stream::collect`
    /// - `Stream::fold_channel`
    /// - `Signal::from_channel`
    ///
    /// The fold and collect methods will reset their folding after taking their accumulator.
    /// To avoid causing bad effects in those cases, it's recommended to map the signal with the
    /// identity function just to create a new independant storage.
    pub fn take(self) -> T
    where
        T: Default,
    {
        match self.0 {
            Constant(val) => val,
            Dynamic(f) => f(),
            Shared(s) => s.sample().replace(T::default()),
            Nested(f) => f().take(),
        }
    }

    /// Samples the value of the signal.
    ///
    /// This will clone the value stored in the signal.
    pub fn sample(&self) -> T
    where
        T: Clone,
    {
        match self.0 {
            Constant(ref val) => T::clone(val),
            Dynamic(ref f) => f(),
            Shared(ref s) => s.sample().get(),
            Nested(ref f) => f().sample(),
        }
    }

    /// Maps a signal with the provided function.
    pub fn map<F, R>(&self, f: F) -> Signal<R>
    where
        F: Fn(T) -> R + Send + Sync + 'static,
        T: Clone + 'static,
        R: Send + 'static,
    {
        match self.0 {
            // constant signal: apply f once to produce another constant signal
            Constant(ref val) => Signal::constant(f(val.clone())),
            // dynamic signal: sample and apply f
            Dynamic(ref sf) => {
                let sf = sf.clone();
                Signal::from_fn(move || f(sf()))
            }
            // shared signal: create shared storage with source
            Shared(ref sig) => Signal::shared(
                SharedImpl {
                    storage: Storage::default(),
                    source: sig.clone(),
                    f,
                }
                .wrap(),
            ),
            // nested signal: extract signal, sample and apply f
            Nested(ref sf) => {
                let sf = sf.clone();
                Signal::from_fn(move || f(sf().sample()))
            }
        }
    }
}

impl<T: Send + 'static> Signal<T> {
    /// Samples the value of this signal every time the trigger stream fires.
    pub fn snapshot<S, F, R>(&self, trigger: &Stream<S>, f: F) -> Stream<R>
    where
        F: Fn(T, MaybeOwned<'_, S>) -> R + Send + Sync + 'static,
        T: Clone,
        S: 'static,
        R: 'static,
    {
        match self.0 {
            Constant(ref val) => {
                let val = Mutex::new(val.clone()); // need Sync for T
                trigger.map(move |t| f(val.lock().clone(), t))
            }
            Dynamic(ref sf) => {
                let sf = sf.clone();
                trigger.map(move |t| f(sf(), t))
            }
            Shared(ref sig) => {
                let sig = sig.clone();
                trigger.map(move |t| f(sig.sample().get(), t))
            }
            Nested(ref sf) => {
                let sf = sf.clone();
                trigger.map(move |t| f(sf().sample(), t))
            }
        }
    }

    /// Stores the last value sent to a channel.
    ///
    /// When sampled, the resulting signal consumes all the current values on the channel
    /// (using `try_recv`) and returns the last value seen.
    #[inline]
    pub fn from_channel(initial: T, rx: mpsc::Receiver<T>) -> Self {
        Self::fold_channel(initial, rx, |_, v| v)
    }

    /// Creates a signal that folds the values from a channel.
    ///
    /// When sampled, the resulting signal consumes all the current values on the channel
    /// (using `try_recv`) and folds them using the current signal value as the
    /// initial accumulator state.
    pub fn fold_channel<V, F>(initial: T, rx: mpsc::Receiver<V>, f: F) -> Self
    where
        F: Fn(T, V) -> T + Send + Sync + 'static,
        V: Send + 'static,
    {
        Signal::shared(
            SharedImpl {
                storage: Storage::new(initial),
                source: Mutex::new(rx),
                f,
            }
            .wrap(),
        )
    }
}

impl<T: Clone + 'static> Signal<Signal<T>> {
    /// Creates a new signal that samples the inner value of a nested signal.
    pub fn switch(&self) -> Signal<T> {
        match self.0 {
            // constant signal: just extract the inner signal
            Constant(ref sig) => Signal::clone(sig),
            // dynamic signal: re-label as nested
            Dynamic(ref f) => Signal(Nested(f.clone())),
            // shared signal: sample to extract the inner signal
            Shared(ref sig) => {
                let sig = sig.clone();
                Signal(Nested(Arc::new(move || sig.sample().get())))
            }
            // nested signal: remove one layer
            Nested(ref f) => {
                let f = f.clone();
                Signal(Nested(Arc::new(move || f().sample())))
            }
        }
    }
}

impl<T: Default> Default for Signal<T> {
    /// Creates a constant signal with T's default value.
    #[inline]
    fn default() -> Self {
        Signal::constant(T::default())
    }
}

impl<T> From<T> for Signal<T> {
    /// Creates a constant signal from T.
    #[inline]
    fn from(val: T) -> Self {
        Signal::constant(val)
    }
}

impl<T: fmt::Debug> fmt::Debug for SigValue<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            Constant(ref val) => write!(f, "Constant({:?})", val),
            Dynamic(ref rf) => write!(f, "Dynamic(Fn@{:p})", rf),
            Shared(ref rs) => write!(f, "Shared(SharedSignal@{:p})", rs),
            Nested(ref rf) => write!(f, "Nested(Fn@{:p})", rf),
        }
    }
}

impl<T: fmt::Display + Clone> fmt::Display for Signal<T> {
    /// Samples the signal and formats the value.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&self.sample(), f)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::RwLock;
    use std::time::Instant;

    #[test]
    fn signal_basic() {
        let signal = Signal::constant(42);
        let double = signal.map(|a| a * 2);
        let plusone = double.map(|a| a + 1);
        assert_eq!(signal.sample(), 42);
        assert_eq!(double.sample(), 84);
        assert_eq!(plusone.sample(), 85);
    }

    #[test]
    fn signal_shared() {
        let st = SharedImpl::new(1, ()).wrap();
        let signal = Signal::shared(st.clone());
        let double = signal.map(|a| a * 2);

        assert_eq!(signal.sample(), 1);
        assert_eq!(double.sample(), 2);
        st.set(42);
        assert_eq!(signal.sample(), 42);
        assert_eq!(double.sample(), 84);
    }

    #[test]
    fn signal_dynamic() {
        let t = Instant::now();
        let signal = Signal::from_fn(move || t);
        assert_eq!(signal.sample(), t);

        let n = Arc::new(RwLock::new(1));
        let cloned = n.clone();
        let signal = Signal::from_fn(move || *cloned.read().unwrap());
        let double = signal.map(|a| a * 2);
        let plusone = double.map(|a| a + 1);
        assert_eq!(signal.sample(), 1);
        assert_eq!(double.sample(), 2);
        assert_eq!(plusone.sample(), 3);
        *n.write().unwrap() = 13;
        assert_eq!(signal.sample(), 13);
        assert_eq!(double.sample(), 26);
        assert_eq!(plusone.sample(), 27);
    }

    #[test]
    fn signal_take() {
        let signal = Signal::constant(42);
        assert_eq!(signal.take(), 42);

        let signal = Signal::from_fn(move || 33);
        assert_eq!(signal.clone().take(), 33);
        assert_eq!(signal.take(), 33);

        let st = SharedImpl::new(123, ()).wrap();
        let signal = Signal::shared(st.clone());
        assert_eq!(signal.clone().take(), 123);
        assert_eq!(signal.clone().take(), 0);
        st.set(13);
        assert_eq!(signal.take(), 13);
    }
}
<|endoftext|>


<commit_before>use std::io::{self, Buffer, Reader, Writer, IoResult, BufferedStream, standard_error};
use std::mem;
use std::collections::BTreeMap;
use std::rand::{thread_rng, Rng};
use std::{u64, u16};
use std::num::{Int, FromPrimitive, ToPrimitive};
use url::Url;

#[cfg(test)]
use test::Bencher;

use nonce::Nonce;
use message::{WSMessage, WSHeader, WS_MASK, WS_LEN, WS_LEN16, WS_LEN64, WS_OPTERM};
use stream::NetworkStream;

pub struct WebSocket<S = NetworkStream> {
    stream: Option<BufferedStream<S>>,
    pub url: Url,
    hostname: String,
    use_ssl: bool,
}

impl WebSocket {
    pub fn new(url: Url) -> WebSocket {
        let use_ssl = url.scheme[] == "wss";

        let port = match url.port() {
            Some(p) => p,
            None if use_ssl => 443,
            _ => 80
        };

        WebSocket {
            stream: None,
            hostname: format!("{}:{}", url.serialize_host().unwrap(), port),
            url: url,
            use_ssl: use_ssl,
        }
    }

    fn try_connect(&mut self) -> IoResult<()> {
        self.stream = Some(BufferedStream::new(try!(NetworkStream::connect(self.hostname[], self.use_ssl))));
        Ok(())
    }

    fn write_request(&mut self, nonce: &str) -> IoResult<()> {
        let s = match self.stream { Some(ref mut s) => s, None => return Err(standard_error(io::IoErrorKind::NotConnected)) };

        try!(s.write(format!("GET {} HTTP/1.1\r\n", self.url.serialize_path().unwrap_or("/".to_string())).as_bytes()));
        try!(s.write(format!("Host: {}\r\n", self.url.host().unwrap()).as_bytes()));
        try!(s.write(format!("Origin: {}\r\n", self.url.serialize_no_fragment()).as_bytes()));
        try!(s.write(format!("Sec-WebSocket-Key: {}\r\n", nonce).as_bytes()));

        try!(s.write(b"Upgrade: websocket\r\n"));
        try!(s.write(b"Connection: Upgrade\r\n"));
        try!(s.write(b"Sec-WebSocket-Protocol: chat, superchat\r\n"));
        try!(s.write(b"Sec-WebSocket-Version: 13\r\n"));
        try!(s.write(b"\r\n"));

        s.flush()
    }

    fn read_response(&mut self, nonce: &str) -> IoResult<()> {
        let spaces: &[_] = &[' ', '\t', '\r', '\n'];
        let s = match self.stream { Some(ref mut s) => s, None => return Err(standard_error(io::IoErrorKind::NotConnected)) };
        let status = try!(s.read_line())[].splitn(2, ' ').nth(1).and_then(|s| s.parse::<uint>());

        match status {
            Some(101) => (),
            _ => return Err(standard_error(io::InvalidInput))
        }

        let headers = s.lines().map(|r| r.unwrap_or("\r\n".to_string())) .take_while(|l| l[] != "\r\n")
            .map(|s| s[].splitn(1, ':').map(|s| s.trim_matches(spaces).to_string()).collect::<Vec<String>>())
            .map(|p| (p[0].to_string(), p[1].to_string()))
            .collect::<BTreeMap<String, String>>();

        try!(s.flush());

        let response = headers.get("Sec-WebSocket-Accept");
        match response {
            Some(r) if nonce == r[] => (),
            _ => return Err(standard_error(io::InvalidInput))
        }

        Ok(())
    }

    pub fn connect(&mut self) -> IoResult<()> {
        let mut nonce = Nonce::new();

        try!(self.try_connect());
        try!(self.write_request(nonce[]));

        nonce = nonce.encode();
        try!(self.read_response(nonce[]));

        Ok(())
    }

    fn read_header(&mut self) -> IoResult<WSHeader> {
        // XXX: this is a bug, WSHeader should accept u16
        Ok(WSHeader::from_bits_truncate(try!(self.read_be_u16())))
    }

    fn read_length(&mut self, header: &WSHeader) -> IoResult<uint> {
        let wslen = *header & WS_LEN;
        if wslen == WS_LEN16 { self.read_be_u16().map(|v| v as uint) }
        else if wslen == WS_LEN64 { self.read_be_u64().map(|v| v as uint) }
        else { Ok(wslen.bits() as uint) }
    }

    pub fn read_message(&mut self) -> IoResult<WSMessage> {
        let header = try!(self.read_header());
        let mut len = try!(self.read_length(&header));

        let mask = if header.contains(WS_MASK) {
            Some(try!(self.read_be_u32()))
        } else {
            None
        };

        // If this is the terminating frame (close command),
        // first two bytes of data MUST BE u16 status code
        let mut status = if header.contains(WS_OPTERM) {
            // compensate length of status code
            len = len - 2;
            Some(try!(self.read_be_u16()))
        } else {
            None
        };

        let mut data = try!(self.read_exact(len));

        // If we have mask, decrypt data
        if let Some(mut m) = mask {
            // decrypt status if present
            if let Some(s) = status {
                status = Some(s ^ (m & 0xffff) as u16);
                // compensate the usage of two mask bytes
                m = m.rotate_right(16);
            }
            data = WebSocket::mask_data(data[], m);
        }

        Ok(WSMessage { header: header, data: data, status: status.and_then(FromPrimitive::from_u16) })
    }

    fn mask_data(data: &[u8], mask: u32) -> Vec<u8> {
        data.iter().enumerate().map(|(i, b)| *b ^ (mask >> ((i % 4) << 3) & 0xff) as u8).collect::<Vec<u8>>()
    }

    pub fn send_message(&mut self, msg: &WSMessage) -> IoResult<()> {
        let mut len = msg.data.len();
        let mut hdr = msg.header - WS_LEN;

        // If we have status set, the data length is increased by status size
        if msg.status.is_some() {
            len = len + 2;
        }

        // Encode and send length along with header
        if len < u16::MAX as uint {
            hdr = hdr | WSHeader::from_bits_truncate(len as u16);
            try!(self.write_be_u16(hdr.bits()));

        } else if len < u64::MAX as uint {
            hdr = hdr | WS_LEN16;
            try!(self.write_be_u16(hdr.bits()));
            try!(self.write_be_u16(len as u16));

        } else {
            hdr = hdr | WS_LEN64;
            try!(self.write_be_u16(hdr.bits()));
            try!(self.write_be_u64(len as u64));
        }

        // If user required masking, encrypt all data
        if hdr.contains(WS_MASK) {
            // Generate and send random mask
            let mut mask = thread_rng().gen::<u32>();
            try!(self.write_be_u32(mask));

            // Encrypt status code if present
            if let Some(status) = msg.status {
                try!(self.write_be_u16(status.to_u16().unwrap() ^ (mask & 0xffff) as u16));
                // compensate for mask already used for status encryption
                mask = mask.rotate_right(16);
            }

            try!(self.write(WebSocket::mask_data(msg.data[], mask)[]));
        } else {
            // Send status code if present
            if let Some(status) = msg.status {
                try!(self.write_be_u16(status.to_u16().unwrap()));
            }
            try!(self.write(msg.data[]));
        }

        self.flush()
    }

    pub fn iter(&mut self) -> WSMessages {
        WSMessages { sock: self }
    }
}

impl Reader for WebSocket {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
        match self.stream {
            Some(ref mut s) => s.read(buf),
            None => Err(standard_error(io::IoErrorKind::NotConnected))
        }
    }
}

impl Writer for WebSocket {
    fn write(&mut self, buf: &[u8]) -> IoResult<()> {
        match self.stream {
            Some(ref mut s) => s.write(buf),
            None => Err(standard_error(io::IoErrorKind::NotConnected))
        }
    }

    fn flush(&mut self) -> IoResult<()> {
        match self.stream {
            Some(ref mut s) => s.flush(),
            None => Err(standard_error(io::IoErrorKind::NotConnected))
        }
    }
}

impl Buffer for WebSocket {
    fn fill_buf<'a>(&'a mut self) -> IoResult<&'a [u8]> {
        match self.stream {
            Some(ref mut s) => s.fill_buf(),
            None => Err(standard_error(io::IoErrorKind::NotConnected))
        }
    }

    fn consume(&mut self, amt: uint) {
        match self.stream {
            Some(ref mut s) => s.consume(amt),
            None => ()
        }
    }
}

pub struct WSMessages<'a> {
    sock: &'a mut WebSocket
}

pub struct WSDefragMessages<'a> {
    underlying: &'a mut WSMessages<'a>,
    buffer: WSMessage
}

impl<'a> WSMessages<'a> {
    pub fn defrag(&'a mut self) -> WSDefragMessages<'a> {
        WSDefragMessages{ underlying: self, buffer: WSMessage{ header: WSHeader::empty(), data: Vec::new(), status: None } }
    }
}

impl<'a> Iterator for WSMessages<'a> {
    type Item = WSMessage;
    fn next(&mut self) -> Option<WSMessage> {
        self.sock.read_message().ok()
    }
}

impl<'a> WSDefragMessages<'a> {
    fn popbuf(&mut self) -> Option<WSMessage> {
        if self.buffer.data.is_empty() {
            None
        } else {
            let mut buf = WSMessage{ header: WSHeader::empty(), data: Vec::new(), status: None };
            mem::swap(&mut self.buffer, &mut buf);
            Some(buf)
        }
    }

    fn swapbuf(&mut self, msg: &mut WSMessage) {
        mem::swap(&mut self.buffer, msg);
    }
}

impl<'a> Iterator for WSDefragMessages<'a> {
    type Item = WSMessage;
    fn next(&mut self) -> Option<WSMessage> {
        loop {
            match self.underlying.next() {
                None => return self.popbuf(),
                Some(mut msg) => {
                    if msg.is_whole() {
                        return Some(msg);
                    } else if msg.is_first() {
                        self.swapbuf(&mut msg);
                    } else if msg.is_more() {
                        self.buffer.push(msg);
                    } else if msg.is_last() {
                        self.buffer.push(msg);
                        return self.popbuf().map(|v| v.last());
                    }
                }
            }
        }
    }
}

<commit_msg>support protocols, extensions and version options<commit_after>use std::io::{self, Buffer, Reader, Writer, IoResult, BufferedStream, standard_error};
use std::mem;
use std::collections::BTreeMap;
use std::rand::{thread_rng, Rng};
use std::{u64, u16};
use std::num::{Int, FromPrimitive, ToPrimitive};
use std::slice::SliceConcatExt;
use url::Url;

#[cfg(test)]
use test::Bencher;

use nonce::Nonce;
use message::{WSMessage, WSHeader, WS_MASK, WS_LEN, WS_LEN16, WS_LEN64, WS_OPTERM};
use stream::NetworkStream;

pub struct WebSocket<S = NetworkStream> {
    stream: Option<BufferedStream<S>>,
    pub url: Url,
    hostname: String,
    use_ssl: bool,
    version: uint,
    extensions: Option<Vec<String>>,
    protocols: Option<Vec<String>>
}

impl WebSocket {
    pub fn with_options(url: Url, version: uint, protocols: Option<&[&str]>, extensions: Option<&[&str]>) -> WebSocket {
        let use_ssl = url.scheme[] == "wss";

        let port = match url.port() {
            Some(p) => p,
            None if use_ssl => 443,
            _ => 80
        };

        WebSocket {
            stream: None,
            hostname: format!("{}:{}", url.serialize_host().unwrap(), port),
            url: url,
            use_ssl: use_ssl,
            version: version,
            extensions: extensions.map(|v| v.iter().map(|v| v.to_string()).collect()),
            protocols: protocols.map(|v| v.iter().map(|v| v.to_string()).collect())
        }
    }

    #[inline] pub fn new(url: Url) -> WebSocket {
        WebSocket::with_options(url, 1, None, None)
    }

    fn try_connect(&mut self) -> IoResult<()> {
        self.stream = Some(BufferedStream::new(try!(NetworkStream::connect(self.hostname[], self.use_ssl))));
        Ok(())
    }

    fn write_request(&mut self, nonce: &str) -> IoResult<()> {
        let s = match self.stream { Some(ref mut s) => s, None => return Err(standard_error(io::IoErrorKind::NotConnected)) };

        try!(s.write(format!("GET {} HTTP/1.1\r\n", self.url.serialize_path().unwrap_or("/".to_string())).as_bytes()));
        try!(s.write(format!("Host: {}\r\n", self.url.host().unwrap()).as_bytes()));
        try!(s.write(format!("Origin: {}\r\n", self.url.serialize_no_fragment()).as_bytes()));
        try!(s.write(format!("Sec-WebSocket-Key: {}\r\n", nonce).as_bytes()));

        try!(s.write(b"Upgrade: websocket\r\n"));
        try!(s.write(b"Connection: Upgrade\r\n"));
        try!(s.write(format!("Sec-WebSocket-Version: {}\r\n", self.version).as_bytes()));
        if let Some(ref protos) = self.protocols {
            try!(s.write(format!("Sec-WebSocket-Protocol: {}\r\n", protos.connect(", ")).as_bytes()));
        }
        if let Some(ref exts) = self.extensions {
            try!(s.write(format!("Sec-WebSocket-Extensions: {}\r\n", exts.connect(", ")).as_bytes()));
        }
        try!(s.write(b"\r\n"));

        s.flush()
    }

    fn read_response(&mut self, nonce: &str) -> IoResult<()> {
        let spaces: &[_] = &[' ', '\t', '\r', '\n'];
        let s = match self.stream { Some(ref mut s) => s, None => return Err(standard_error(io::IoErrorKind::NotConnected)) };
        let status = try!(s.read_line())[].splitn(2, ' ').nth(1).and_then(|s| s.parse::<uint>());

        match status {
            Some(101) => (),
            _ => return Err(standard_error(io::InvalidInput))
        }

        let headers = s.lines().map(|r| r.unwrap_or("\r\n".to_string())) .take_while(|l| l[] != "\r\n")
            .map(|s| s[].splitn(1, ':').map(|s| s.trim_matches(spaces).to_string()).collect::<Vec<String>>())
            .map(|p| (p[0].to_string(), p[1].to_string()))
            .collect::<BTreeMap<String, String>>();

        try!(s.flush());

        let response = headers.get("Sec-WebSocket-Accept");
        match response {
            Some(r) if nonce == r[] => (),
            _ => return Err(standard_error(io::InvalidInput))
        }

        Ok(())
    }

    pub fn connect(&mut self) -> IoResult<()> {
        let mut nonce = Nonce::new();

        try!(self.try_connect());
        try!(self.write_request(nonce[]));

        nonce = nonce.encode();
        try!(self.read_response(nonce[]));

        Ok(())
    }

    fn read_header(&mut self) -> IoResult<WSHeader> {
        // XXX: this is a bug, WSHeader should accept u16
        Ok(WSHeader::from_bits_truncate(try!(self.read_be_u16())))
    }

    fn read_length(&mut self, header: &WSHeader) -> IoResult<uint> {
        let wslen = *header & WS_LEN;
        if wslen == WS_LEN16 { self.read_be_u16().map(|v| v as uint) }
        else if wslen == WS_LEN64 { self.read_be_u64().map(|v| v as uint) }
        else { Ok(wslen.bits() as uint) }
    }

    pub fn read_message(&mut self) -> IoResult<WSMessage> {
        let header = try!(self.read_header());
        let mut len = try!(self.read_length(&header));

        let mask = if header.contains(WS_MASK) {
            Some(try!(self.read_be_u32()))
        } else {
            None
        };

        // If this is the terminating frame (close command),
        // first two bytes of data MUST BE u16 status code
        let mut status = if header.contains(WS_OPTERM) {
            // compensate length of status code
            len = len - 2;
            Some(try!(self.read_be_u16()))
        } else {
            None
        };

        let mut data = try!(self.read_exact(len));

        // If we have mask, decrypt data
        if let Some(mut m) = mask {
            // decrypt status if present
            if let Some(s) = status {
                status = Some(s ^ (m & 0xffff) as u16);
                // compensate the usage of two mask bytes
                m = m.rotate_right(16);
            }
            data = WebSocket::mask_data(data[], m);
        }

        Ok(WSMessage { header: header, data: data, status: status.and_then(FromPrimitive::from_u16) })
    }

    fn mask_data(data: &[u8], mask: u32) -> Vec<u8> {
        data.iter().enumerate().map(|(i, b)| *b ^ (mask >> ((i % 4) << 3) & 0xff) as u8).collect::<Vec<u8>>()
    }

    pub fn send_message(&mut self, msg: &WSMessage) -> IoResult<()> {
        let mut len = msg.data.len();
        let mut hdr = msg.header - WS_LEN;

        // If we have status set, the data length is increased by status size
        if msg.status.is_some() {
            len = len + 2;
        }

        // Encode and send length along with header
        if len < u16::MAX as uint {
            hdr = hdr | WSHeader::from_bits_truncate(len as u16);
            try!(self.write_be_u16(hdr.bits()));

        } else if len < u64::MAX as uint {
            hdr = hdr | WS_LEN16;
            try!(self.write_be_u16(hdr.bits()));
            try!(self.write_be_u16(len as u16));

        } else {
            hdr = hdr | WS_LEN64;
            try!(self.write_be_u16(hdr.bits()));
            try!(self.write_be_u64(len as u64));
        }

        // If user required masking, encrypt all data
        if hdr.contains(WS_MASK) {
            // Generate and send random mask
            let mut mask = thread_rng().gen::<u32>();
            try!(self.write_be_u32(mask));

            // Encrypt status code if present
            if let Some(status) = msg.status {
                try!(self.write_be_u16(status.to_u16().unwrap() ^ (mask & 0xffff) as u16));
                // compensate for mask already used for status encryption
                mask = mask.rotate_right(16);
            }

            try!(self.write(WebSocket::mask_data(msg.data[], mask)[]));
        } else {
            // Send status code if present
            if let Some(status) = msg.status {
                try!(self.write_be_u16(status.to_u16().unwrap()));
            }
            try!(self.write(msg.data[]));
        }

        self.flush()
    }

    pub fn iter(&mut self) -> WSMessages {
        WSMessages { sock: self }
    }
}

impl Reader for WebSocket {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
        match self.stream {
            Some(ref mut s) => s.read(buf),
            None => Err(standard_error(io::IoErrorKind::NotConnected))
        }
    }
}

impl Writer for WebSocket {
    fn write(&mut self, buf: &[u8]) -> IoResult<()> {
        match self.stream {
            Some(ref mut s) => s.write(buf),
            None => Err(standard_error(io::IoErrorKind::NotConnected))
        }
    }

    fn flush(&mut self) -> IoResult<()> {
        match self.stream {
            Some(ref mut s) => s.flush(),
            None => Err(standard_error(io::IoErrorKind::NotConnected))
        }
    }
}

impl Buffer for WebSocket {
    fn fill_buf<'a>(&'a mut self) -> IoResult<&'a [u8]> {
        match self.stream {
            Some(ref mut s) => s.fill_buf(),
            None => Err(standard_error(io::IoErrorKind::NotConnected))
        }
    }

    fn consume(&mut self, amt: uint) {
        match self.stream {
            Some(ref mut s) => s.consume(amt),
            None => ()
        }
    }
}

pub struct WSMessages<'a> {
    sock: &'a mut WebSocket
}

pub struct WSDefragMessages<'a> {
    underlying: &'a mut WSMessages<'a>,
    buffer: WSMessage
}

impl<'a> WSMessages<'a> {
    pub fn defrag(&'a mut self) -> WSDefragMessages<'a> {
        WSDefragMessages{ underlying: self, buffer: WSMessage{ header: WSHeader::empty(), data: Vec::new(), status: None } }
    }
}

impl<'a> Iterator for WSMessages<'a> {
    type Item = WSMessage;
    fn next(&mut self) -> Option<WSMessage> {
        self.sock.read_message().ok()
    }
}

impl<'a> WSDefragMessages<'a> {
    fn popbuf(&mut self) -> Option<WSMessage> {
        if self.buffer.data.is_empty() {
            None
        } else {
            let mut buf = WSMessage{ header: WSHeader::empty(), data: Vec::new(), status: None };
            mem::swap(&mut self.buffer, &mut buf);
            Some(buf)
        }
    }

    fn swapbuf(&mut self, msg: &mut WSMessage) {
        mem::swap(&mut self.buffer, msg);
    }
}

impl<'a> Iterator for WSDefragMessages<'a> {
    type Item = WSMessage;
    fn next(&mut self) -> Option<WSMessage> {
        loop {
            match self.underlying.next() {
                None => return self.popbuf(),
                Some(mut msg) => {
                    if msg.is_whole() {
                        return Some(msg);
                    } else if msg.is_first() {
                        self.swapbuf(&mut msg);
                    } else if msg.is_more() {
                        self.buffer.push(msg);
                    } else if msg.is_last() {
                        self.buffer.push(msg);
                        return self.popbuf().map(|v| v.last());
                    }
                }
            }
        }
    }
}

<|endoftext|>


<commit_before>//! # Socket code around `libc`
//!
//! ## Notes
//!
//! This module provides a low level one-to-one mapping between `libc` system call wrappers
//! with defaults specific to netlink sockets as well as a higher level API for simplification
//! of netlink code.

use std::io;
use std::os::unix::io::{AsRawFd,IntoRawFd,RawFd};
use std::marker::PhantomData;
use std::mem::{zeroed,size_of};

use buffering::copy::{StreamReadBuffer,StreamWriteBuffer};
use libc::{self,c_int,c_void};
use mio::{self,Evented};
#[cfg(feature = "stream")]
use tokio::io::AsyncRead;
#[cfg(feature = "stream")]
use tokio::prelude::{Async,Stream};

use {Nl,MAX_NL_LENGTH};
use err::{NlError,Nlmsgerr};
use consts::{self,AddrFamily,NlFamily,GenlId,CtrlCmd,CtrlAttr,CtrlAttrMcastGrp,NlmF};
use genl::Genlmsghdr;
use nlattr::Nlattr;
use nl::Nlmsghdr;

/// Handle for the socket file descriptor
pub struct NlSocket<T, P> {
    fd: c_int,
    poll: mio::Poll,
    data_type: PhantomData<T>,
    data_payload: PhantomData<P>,
}

impl<T, P> NlSocket<T, P> where T: Nl, P: Nl {
    /// Wrapper around `socket()` syscall filling in the netlink-specific information
    pub fn new(proto: NlFamily) -> Result<Self, io::Error> {
        let fd = match unsafe {
            libc::socket(AddrFamily::Netlink.into(), libc::SOCK_RAW, proto.into())
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }?;
        let poll = mio::Poll::new()?;
        let socket = NlSocket { fd, poll, data_type: PhantomData, data_payload: PhantomData };
        socket.register(&socket.poll, mio::Token(0), mio::Ready::readable(), mio::PollOpt::edge())?;
        Ok(socket)
    }

    /// Set underlying socket file descriptor to be blocking
    pub fn block(&mut self) -> Result<&mut Self, io::Error> {
        match unsafe { libc::fcntl(self.fd, libc::F_SETFL,
                                   libc::fcntl(self.fd, libc::F_GETFL, 0) & !libc::O_NONBLOCK) } {
            i if i < 0 => return Err(io::Error::last_os_error()),
            _ => Ok(self),
        }
    }

    /// Set underlying socket file descriptor to be non blocking
    pub fn nonblock(&mut self) -> Result<&mut Self, io::Error> {
        match unsafe { libc::fcntl(self.fd, libc::F_SETFL,
                                   libc::fcntl(self.fd, libc::F_GETFL, 0) | libc::O_NONBLOCK) } {
            i if i < 0 => return Err(io::Error::last_os_error()),
            _ => Ok(self),
        }
    }

    /// Determines if underlying file descriptor is blocking - `Stream` feature will throw an
    /// error if this function returns false
    pub fn is_blocking(&self) -> Result<bool, io::Error> {
        let is_blocking = match unsafe { libc::fcntl(self.fd, libc::F_GETFL, 0) } {
            i if i >= 0 => i & libc::O_NONBLOCK == 0,
            _ => return Err(io::Error::last_os_error()),
        };
        Ok(is_blocking)
    }

    /// Use this function to bind to a netlink ID and subscribe to groups. See netlink(7)
    /// man pages for more information on netlink IDs and groups.
    pub fn bind(&mut self, pid: Option<u32>, groups: Vec<u32>) -> Result<(), io::Error> {
        let mut nladdr = unsafe { zeroed::<libc::sockaddr_nl>() };
        nladdr.nl_family = libc::c_int::from(AddrFamily::Netlink) as u16;
        nladdr.nl_pid = pid.unwrap_or(0);
        nladdr.nl_groups = groups.into_iter().fold(0, |acc, next| {
            acc | (1 << (next - 1))
        });
        match unsafe {
            libc::bind(self.fd, &nladdr as *const _ as *const libc::sockaddr,
                       size_of::<libc::sockaddr_nl>() as u32)
        } {
            i if i >= 0 => Ok(()),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Send message encoded as byte slice to the netlink ID specified in the netlink header
    /// (`neli::nl::Nlmsghdr`)
    pub fn send<B>(&mut self, buf: B, flags: i32) -> Result<libc::ssize_t, io::Error> where B: AsRef<[u8]> {
        match unsafe {
            libc::send(self.fd, buf.as_ref() as *const _ as *const c_void, buf.as_ref().len(), flags)
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Convenience function to send an `Nlmsghdr` struct
    pub fn send_nl(&mut self, msg: Nlmsghdr<T, P>) -> Result<(), NlError> {
        let mut mem = StreamWriteBuffer::new_growable(Some(msg.asize()));
        msg.serialize(&mut mem)?;
        self.send(mem, 0)?;
        Ok(())
    }

    /// Receive message encoded as byte slice from the netlink socket
    pub fn recv<'a, B>(&mut self, mut buf: B, flags: i32) -> Result<libc::ssize_t, io::Error> where B: AsMut<[u8]> {
        match unsafe {
            libc::recv(self.fd, buf.as_mut() as *mut _ as *mut c_void, buf.as_mut().len(), flags)
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Convenience function to receive an `Nlmsghdr` struct
    pub fn recv_nl(&mut self, buf_sz: Option<usize>) -> Result<Nlmsghdr<T, P>, NlError> {
        let mut mem = vec![0; buf_sz.unwrap_or(MAX_NL_LENGTH)];
        let mem_read = self.recv(&mut mem, 0)?;
        mem.truncate(mem_read as usize);
        Ok(Nlmsghdr::<T, P>::deserialize(&mut StreamReadBuffer::new(mem))?)
    }

    /// Convenience function to receive an `Nlmsghdr` struct with function type parameters
    /// that determine deserialization type
    pub fn recv_nl_typed<TT, PP>(&mut self, buf_sz: Option<usize>)
            -> Result<Nlmsghdr<TT, PP>, NlError> where TT: Nl + Into<u16> + From<u16>, PP: Nl {
        let mut mem = vec![0; buf_sz.unwrap_or(MAX_NL_LENGTH)];
        let mem_read = self.recv(&mut mem, 0)?;
        mem.truncate(mem_read as usize);
        Ok(Nlmsghdr::<TT, PP>::deserialize(&mut StreamReadBuffer::new(mem))?)
    }

    /// Consume an ACK and return an error if an ACK is not found
    pub fn recv_ack(&mut self, buf_sz: Option<usize>) -> Result<(), NlError> {
        let ack = self.recv_nl_typed::<consts::Nlmsg, Nlmsgerr<consts::Nlmsg>>(buf_sz)?;
        if ack.nl_type == consts::Nlmsg::Error && ack.nl_payload.error == 0 {
            Ok(())
        } else {
            Err(NlError::NoAck)
        }
    }

    /// Equivalent of `socket` and `bind` calls.
    pub fn connect(proto: NlFamily, pid: Option<u32>, groups: Vec<u32>)
                   -> Result<Self, io::Error> {
        let mut s = try!(NlSocket::new(proto));
        try!(s.bind(pid, groups));
        Ok(s)
    }
}

impl NlSocket<GenlId, Genlmsghdr<CtrlCmd>> {
    /// Create generic netlink resolution socket
    pub fn new_genl() -> Result<NlSocket<GenlId, Genlmsghdr<CtrlCmd>>, io::Error> {
        Self::connect(NlFamily::Generic, None, Vec::new())
    }

    fn get_genl_family(&mut self, family_name: &str)
            -> Result<Nlmsghdr<GenlId, Genlmsghdr<CtrlCmd>>, NlError> {
        let attrs = vec![Nlattr::new_str_payload(None, CtrlAttr::FamilyName, family_name)?];
        let genlhdr = Genlmsghdr::new(CtrlCmd::Getfamily, 2, attrs)?;
        let nlhdr = Nlmsghdr::new(None, GenlId::Ctrl,
                               vec![NlmF::Request], None, None, genlhdr);
        self.send_nl(nlhdr)?;

        Ok(self.recv_nl(Some(4096))?)
    }

    /// Convenience function for resolving a `&str` containing the multicast group name to a
    /// numeric netlink ID
    pub fn resolve_genl_family(&mut self, family_name: &str) -> Result<u16, NlError> {
        let nlhdr = self.get_genl_family(family_name)?;
        let mut handle = nlhdr.nl_payload.get_attr_handle::<CtrlAttr>();
        Ok(handle.get_payload_with::<u16>(CtrlAttr::FamilyId, None)?)
    }

    /// Convenience function for resolving a `&str` containing the multicast group name to a
    /// numeric netlink ID
    pub fn resolve_nl_mcast_group(&mut self, family_name: &str, mcast_name: &str)
            -> Result<u32, NlError> {
        let nlhdr = self.get_genl_family(family_name)?;
        let mut handle = nlhdr.nl_payload.get_attr_handle::<CtrlAttr>();
        let mut mcast_groups = handle.get_nested_attributes::<u16>(CtrlAttr::McastGroups)?;
        mcast_groups.parse_nested_attributes()?;
        let mut id = None;
        if let Some(iter) = mcast_groups.iter() {
            for attribute in iter {
                let attribute_len = attribute.nla_len;
                let mut handle = attribute.get_attr_handle();
                let string = handle.get_payload_with::<String>(CtrlAttrMcastGrp::Name,
                    Some(attribute_len as usize))?;
                if string.as_str() == mcast_name {
                    id = handle.get_payload_with::<u32>(CtrlAttrMcastGrp::Id, None).ok();
                }
            }
        }
        id.ok_or(NlError::new("Failed to resolve multicast group ID"))
    }
}

impl<T, P> AsRawFd for NlSocket<T, P> {
    fn as_raw_fd(&self) -> RawFd {
        self.fd
    }
}

impl<T, P> IntoRawFd for NlSocket<T, P> {
    fn into_raw_fd(self) -> RawFd {
        self.fd
    }
}

impl<T, P> io::Read for NlSocket<T, P> where T: Nl, P: Nl {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match unsafe {
            libc::recv(self.fd, buf as *mut _ as *mut c_void, buf.len(), 0)
        } {
            i if i >= 0 => Ok(i as usize),
            _ => Err(io::Error::last_os_error()),
        }
    }
}

#[cfg(feature = "stream")]
impl<T, P> AsyncRead for NlSocket<T, P> where T: Nl, P: Nl { }

#[cfg(feature = "stream")]
impl<T, P> Stream for NlSocket<T, P> where T: Nl, P: Nl {
    type Item = Nlmsghdr<T, P>;
    type Error = io::Error;

    fn poll(&mut self) -> Result<Async<Option<Self::Item>>, Self::Error> {
        if self.is_blocking()? {
            self.nonblock()?;
        }
        let mut mem = vec![0; MAX_NL_LENGTH];
        let bytes_written = match self.poll_read(mem.as_mut_slice()) {
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Ok(Async::Ready(0)) => return Ok(Async::Ready(None)),
            Ok(Async::Ready(i)) => i,
            Err(e) => return Err(e),
        };
        mem.truncate(bytes_written);
        let mut mem_read = StreamReadBuffer::new(mem);
        let hdr = match Nlmsghdr::<T, P>::deserialize(&mut mem_read) {
            Ok(h) => h,
            Err(_) => return Err(io::Error::from(io::ErrorKind::Other)),
        };
        Ok(Async::Ready(Some(hdr)))
    }
}

impl<T, P> Evented for NlSocket<T, P> {
    fn register(&self, poll: &mio::Poll, token: mio::Token, interest: mio::Ready,
                opts: mio::PollOpt) -> io::Result<()> {
        poll.register(&mio::unix::EventedFd(&self.as_raw_fd()), token, interest, opts)
    }

    fn reregister(&self, poll: &mio::Poll, token: mio::Token, interest: mio::Ready,
                  opts: mio::PollOpt) -> io::Result<()> {
        poll.reregister(&mio::unix::EventedFd(&self.as_raw_fd()), token, interest, opts)
    }

    fn deregister(&self, poll: &mio::Poll) -> io::Result<()> {
        poll.deregister(&mio::unix::EventedFd(&self.as_raw_fd()))
    }
}

impl<T, P> Drop for NlSocket<T, P> {
    /// Closes underlying file descriptor to avoid file descriptor leaks.
    fn drop(&mut self) {
        unsafe { libc::close(self.fd); }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use consts::{CtrlCmd,Nlmsg};
    use genl::Genlmsghdr;

    #[test]
    fn test_socket_creation() {
       NlSocket::<Nlmsg, Genlmsghdr<CtrlCmd>>::connect(NlFamily::Generic, None, Vec::new()).unwrap();
    }

    #[ignore]
    #[test]
    fn test_genl_family_resolve() {
        let mut sock = NlSocket::new_genl().unwrap();
        assert_eq!(19, sock.resolve_genl_family("nl80211").unwrap());
    }

    #[ignore]
    #[test]
    fn test_nl_mcast_group_resolve() {
        let mut sock = NlSocket::new_genl().unwrap();
        assert_eq!(5, sock.resolve_nl_mcast_group("nl80211", "mlme").unwrap());
    }
}
<commit_msg>Rework tokio support to mirror TcpStream and follow docs - testing pending<commit_after>//! # Socket code around `libc`
//!
//! ## Notes
//!
//! This module provides a low level one-to-one mapping between `libc` system call wrappers
//! with defaults specific to netlink sockets as well as a higher level API for simplification
//! of netlink code.

use std::io;
use std::os::unix::io::{AsRawFd,IntoRawFd,RawFd};
use std::marker::PhantomData;
use std::mem::{zeroed,size_of};

use buffering::copy::{StreamReadBuffer,StreamWriteBuffer};
use libc::{self,c_int,c_void};

use {Nl,MAX_NL_LENGTH};
use err::{NlError,Nlmsgerr};
use consts::{self,AddrFamily,NlFamily,GenlId,CtrlCmd,CtrlAttr,CtrlAttrMcastGrp,NlmF};
use genl::Genlmsghdr;
use nlattr::Nlattr;
use nl::Nlmsghdr;

/// Handle for the socket file descriptor
pub struct NlSocket<T, P> {
    fd: c_int,
    data_type: PhantomData<T>,
    data_payload: PhantomData<P>,
}

impl<T, P> NlSocket<T, P> {
    /// Wrapper around `socket()` syscall filling in the netlink-specific information
    pub fn new(proto: NlFamily) -> Result<Self, io::Error> {
        let fd = match unsafe {
            libc::socket(AddrFamily::Netlink.into(), libc::SOCK_RAW, proto.into())
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }?;
        Ok(NlSocket { fd, data_type: PhantomData, data_payload: PhantomData })
    }

    /// Set underlying socket file descriptor to be blocking
    pub fn block(&mut self) -> Result<&mut Self, io::Error> {
        match unsafe { libc::fcntl(self.fd, libc::F_SETFL,
                                   libc::fcntl(self.fd, libc::F_GETFL, 0) & !libc::O_NONBLOCK) } {
            i if i < 0 => return Err(io::Error::last_os_error()),
            _ => Ok(self),
        }
    }

    /// Set underlying socket file descriptor to be non blocking
    pub fn nonblock(&mut self) -> Result<&mut Self, io::Error> {
        match unsafe { libc::fcntl(self.fd, libc::F_SETFL,
                                   libc::fcntl(self.fd, libc::F_GETFL, 0) | libc::O_NONBLOCK) } {
            i if i < 0 => return Err(io::Error::last_os_error()),
            _ => Ok(self),
        }
    }

    /// Determines if underlying file descriptor is blocking - `Stream` feature will throw an
    /// error if this function returns false
    pub fn is_blocking(&self) -> Result<bool, io::Error> {
        let is_blocking = match unsafe { libc::fcntl(self.fd, libc::F_GETFL, 0) } {
            i if i >= 0 => i & libc::O_NONBLOCK == 0,
            _ => return Err(io::Error::last_os_error()),
        };
        Ok(is_blocking)
    }

    /// Use this function to bind to a netlink ID and subscribe to groups. See netlink(7)
    /// man pages for more information on netlink IDs and groups.
    pub fn bind(&mut self, pid: Option<u32>, groups: Vec<u32>) -> Result<(), io::Error> {
        let mut nladdr = unsafe { zeroed::<libc::sockaddr_nl>() };
        nladdr.nl_family = libc::c_int::from(AddrFamily::Netlink) as u16;
        nladdr.nl_pid = pid.unwrap_or(0);
        nladdr.nl_groups = groups.into_iter().fold(0, |acc, next| {
            acc | (1 << (next - 1))
        });
        match unsafe {
            libc::bind(self.fd, &nladdr as *const _ as *const libc::sockaddr,
                       size_of::<libc::sockaddr_nl>() as u32)
        } {
            i if i >= 0 => Ok(()),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Send message encoded as byte slice to the netlink ID specified in the netlink header
    /// (`neli::nl::Nlmsghdr`)
    pub fn send<B>(&mut self, buf: B, flags: i32) -> Result<libc::ssize_t, io::Error> where B: AsRef<[u8]> {
        match unsafe {
            libc::send(self.fd, buf.as_ref() as *const _ as *const c_void, buf.as_ref().len(), flags)
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Receive message encoded as byte slice from the netlink socket
    pub fn recv<'a, B>(&mut self, mut buf: B, flags: i32) -> Result<libc::ssize_t, io::Error> where B: AsMut<[u8]> {
        match unsafe {
            libc::recv(self.fd, buf.as_mut() as *mut _ as *mut c_void, buf.as_mut().len(), flags)
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Equivalent of `socket` and `bind` calls.
    pub fn connect(proto: NlFamily, pid: Option<u32>, groups: Vec<u32>)
                   -> Result<Self, io::Error> {
        let mut s = try!(NlSocket::new(proto));
        try!(s.bind(pid, groups));
        Ok(s)
    }
}

impl<T, P> NlSocket<T, P> where T: Nl, P: Nl {
    /// Convenience function to send an `Nlmsghdr` struct
    pub fn send_nl(&mut self, msg: Nlmsghdr<T, P>) -> Result<(), NlError> {
        let mut mem = StreamWriteBuffer::new_growable(Some(msg.asize()));
        msg.serialize(&mut mem)?;
        self.send(mem, 0)?;
        Ok(())
    }

    /// Convenience function to receive an `Nlmsghdr` struct
    pub fn recv_nl(&mut self, buf_sz: Option<usize>) -> Result<Nlmsghdr<T, P>, NlError> {
        let mut mem = vec![0; buf_sz.unwrap_or(MAX_NL_LENGTH)];
        let mem_read = self.recv(&mut mem, 0)?;
        mem.truncate(mem_read as usize);
        Ok(Nlmsghdr::<T, P>::deserialize(&mut StreamReadBuffer::new(mem))?)
    }

    /// Convenience function to receive an `Nlmsghdr` struct with function type parameters
    /// that determine deserialization type
    pub fn recv_nl_typed<TT, PP>(&mut self, buf_sz: Option<usize>)
            -> Result<Nlmsghdr<TT, PP>, NlError> where TT: Nl + Into<u16> + From<u16>, PP: Nl {
        let mut mem = vec![0; buf_sz.unwrap_or(MAX_NL_LENGTH)];
        let mem_read = self.recv(&mut mem, 0)?;
        mem.truncate(mem_read as usize);
        Ok(Nlmsghdr::<TT, PP>::deserialize(&mut StreamReadBuffer::new(mem))?)
    }

    /// Consume an ACK and return an error if an ACK is not found
    pub fn recv_ack(&mut self, buf_sz: Option<usize>) -> Result<(), NlError> {
        let ack = self.recv_nl_typed::<consts::Nlmsg, Nlmsgerr<consts::Nlmsg>>(buf_sz)?;
        if ack.nl_type == consts::Nlmsg::Error && ack.nl_payload.error == 0 {
            Ok(())
        } else {
            Err(NlError::NoAck)
        }
    }
}

impl NlSocket<GenlId, Genlmsghdr<CtrlCmd>> {
    /// Create generic netlink resolution socket
    pub fn new_genl() -> Result<NlSocket<GenlId, Genlmsghdr<CtrlCmd>>, io::Error> {
        Self::connect(NlFamily::Generic, None, Vec::new())
    }

    fn get_genl_family(&mut self, family_name: &str)
            -> Result<Nlmsghdr<GenlId, Genlmsghdr<CtrlCmd>>, NlError> {
        let attrs = vec![Nlattr::new_str_payload(None, CtrlAttr::FamilyName, family_name)?];
        let genlhdr = Genlmsghdr::new(CtrlCmd::Getfamily, 2, attrs)?;
        let nlhdr = Nlmsghdr::new(None, GenlId::Ctrl,
                               vec![NlmF::Request], None, None, genlhdr);
        self.send_nl(nlhdr)?;

        Ok(self.recv_nl(Some(4096))?)
    }

    /// Convenience function for resolving a `&str` containing the multicast group name to a
    /// numeric netlink ID
    pub fn resolve_genl_family(&mut self, family_name: &str) -> Result<u16, NlError> {
        let nlhdr = self.get_genl_family(family_name)?;
        let mut handle = nlhdr.nl_payload.get_attr_handle::<CtrlAttr>();
        Ok(handle.get_payload_with::<u16>(CtrlAttr::FamilyId, None)?)
    }

    /// Convenience function for resolving a `&str` containing the multicast group name to a
    /// numeric netlink ID
    pub fn resolve_nl_mcast_group(&mut self, family_name: &str, mcast_name: &str)
            -> Result<u32, NlError> {
        let nlhdr = self.get_genl_family(family_name)?;
        let mut handle = nlhdr.nl_payload.get_attr_handle::<CtrlAttr>();
        let mut mcast_groups = handle.get_nested_attributes::<u16>(CtrlAttr::McastGroups)?;
        mcast_groups.parse_nested_attributes()?;
        let mut id = None;
        if let Some(iter) = mcast_groups.iter() {
            for attribute in iter {
                let attribute_len = attribute.nla_len;
                let mut handle = attribute.get_attr_handle();
                let string = handle.get_payload_with::<String>(CtrlAttrMcastGrp::Name,
                    Some(attribute_len as usize))?;
                if string.as_str() == mcast_name {
                    id = handle.get_payload_with::<u32>(CtrlAttrMcastGrp::Id, None).ok();
                }
            }
        }
        id.ok_or(NlError::new("Failed to resolve multicast group ID"))
    }
}

impl<T, P> AsRawFd for NlSocket<T, P> {
    fn as_raw_fd(&self) -> RawFd {
        self.fd
    }
}

impl<T, P> IntoRawFd for NlSocket<T, P> {
    fn into_raw_fd(self) -> RawFd {
        self.fd
    }
}

impl<T, P> io::Read for NlSocket<T, P> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match unsafe {
            libc::recv(self.fd, buf as *mut _ as *mut c_void, buf.len(), 0)
        } {
            i if i >= 0 => Ok(i as usize),
            _ => Err(io::Error::last_os_error()),
        }
    }
}

/// Tokio-specific features for neli
#[cfg(feature = "stream")]
pub mod tokio {
    use super::*;

    use std::io::{ErrorKind,Read};

    use mio::{self,Evented,Ready};
    use tokio::prelude::{Async,Stream};
    use tokio::reactor::PollEvented2;

    /// Tokio-enabled Netlink socket struct
    pub struct NlSocket<T, P>(PollEvented2<super::NlSocket<T, P>>);

    impl<T, P> NlSocket<T, P> {
        /// Setup NlSocket for use with tokio - set to nonblocking state and wrap in polling mechanism
        pub fn new(mut sock: super::NlSocket<T, P>) -> io::Result<Self> {
            if sock.is_blocking()? {
                sock.nonblock()?;
            }
            Ok(NlSocket(PollEvented2::new(sock)))
        }
    }

    impl<T, P> Stream for NlSocket<T, P> where T: Nl, P: Nl {
        type Item = Nlmsghdr<T, P>;
        type Error = io::Error;

        fn poll(&mut self) -> Result<Async<Option<Self::Item>>, Self::Error> {
            self.0.poll_read_ready(Ready::readable())?;

            let mut mem = vec![0; MAX_NL_LENGTH];
            let bytes_written = match self.0.read(mem.as_mut_slice()) {
                Ok(i) if i > 0 => i,
                Ok(i) if i == 0 => return Ok(Async::Ready(None)),
                Ok(_) => {
                    match io::Error::last_os_error().kind() {
                        ErrorKind::WouldBlock => {
                            self.0.clear_read_ready(Ready::readable())?;
                            return Ok(Async::NotReady);
                        },
                        e => return Err(io::Error::from(e)),
                    }
                },
                Err(e) => return Err(e),
            };
            mem.truncate(bytes_written);
            let mut mem_read = StreamReadBuffer::new(mem);
            let hdr = match Nlmsghdr::<T, P>::deserialize(&mut mem_read) {
                Ok(h) => h,
                Err(_) => return Err(io::Error::from(io::ErrorKind::Other)),
            };
            Ok(Async::Ready(Some(hdr)))
        }
    }

    impl<T, P> Evented for super::NlSocket<T, P> {
        fn register(&self, poll: &mio::Poll, token: mio::Token, interest: mio::Ready,
                    opts: mio::PollOpt) -> io::Result<()> {
            poll.register(&mio::unix::EventedFd(&self.as_raw_fd()), token, interest, opts)
        }

        fn reregister(&self, poll: &mio::Poll, token: mio::Token, interest: mio::Ready,
                      opts: mio::PollOpt) -> io::Result<()> {
            poll.reregister(&mio::unix::EventedFd(&self.as_raw_fd()), token, interest, opts)
        }

        fn deregister(&self, poll: &mio::Poll) -> io::Result<()> {
            poll.deregister(&mio::unix::EventedFd(&self.as_raw_fd()))
        }
    }
}

impl<T, P> Drop for NlSocket<T, P> {
    /// Closes underlying file descriptor to avoid file descriptor leaks.
    fn drop(&mut self) {
        unsafe { libc::close(self.fd); }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use consts::{CtrlCmd,Nlmsg};
    use genl::Genlmsghdr;

    #[test]
    fn test_socket_creation() {
       NlSocket::<Nlmsg, Genlmsghdr<CtrlCmd>>::connect(NlFamily::Generic, None, Vec::new()).unwrap();
    }

    #[ignore]
    #[test]
    fn test_genl_family_resolve() {
        let mut sock = NlSocket::new_genl().unwrap();
        assert_eq!(19, sock.resolve_genl_family("nl80211").unwrap());
    }

    #[ignore]
    #[test]
    fn test_nl_mcast_group_resolve() {
        let mut sock = NlSocket::new_genl().unwrap();
        assert_eq!(5, sock.resolve_nl_mcast_group("nl80211", "mlme").unwrap());
    }
}
<|endoftext|>


<commit_before>
use std::iter;
use std::ascii::AsciiExt;

static STANDARD_CHARS: &'static[u8] = b"abcdefghijklmnopqrstuvwxyz234567";

pub fn encode(b: &[u8]) -> String {
    let mut out = Vec::new();
    let mut n: u64 = 0;

    for i in 0 .. b.len() {
        // Calculate the shift (32, 24, 16, 8, 0)
        let shift = 8 * (4 - (i % 5));
        n |= (b[i] as u64) << shift;

        // If the shift was zero, or this is the last character of input, extract
        // each 5-bit group from n.
        if shift == 0 || (i == b.len()-1) {
            // Determine number of output groups, based on the shift
            let groups = match shift {
                32 => 2,
                24 => 4,
                16 => 5,
                8  => 7,
                _  => 8
            };

            // Pull out each 5-bit index
            for j in iter::range_step_inclusive(35, (35 - (groups * 5)), -5) {
                let ndx = (n >> j as u32) & 31;
                out.push(STANDARD_CHARS[ndx as usize]);
            }

            // Reset n
            n = 0;
        }
    }

    let result = match String::from_utf8(out) {
        Ok(s) => s,
        Err(e) => panic!(e)
    };
    result
}

pub fn decode(b: &[u8]) -> Vec<u8> {
    let mut out: Vec<u8> = Vec::new();
    let mut n: u64 = 0;

    for i in iter::range_inclusive(0, b.len()) {
        // Use the character to generate an index
        let c = b[i];
        let idx = match c as char {
            'A' ... 'Z' => c - 65, // A..Z => 65..90
            'a' ... 'z' => c - 97, // a..z => 97..122
            _ => 26 + (c - 50)     // 2..7 => 50..55
        };

        // Calculate the shift for this 5-bit group (35, 25, ..., 0)
        let shift = 35 - (5 * (i % 8));
        n |= (idx as u64 & 31) << shift;

        // If the shift is zero or this is the last character of input, extract
        // each 8-bit group from n.
        if shift == 0 || (i == b.len()-1) {
            // Determine number of output groups, based on shift
            let groups = match shift {
                35 => 0,
                30 => 1,
                25 => 2,
                20 => 2,
                15 => 3,
                10 => 4,
                5  => 4,
                _  => 5
            };

            // Pull out each 8-bit group
            for j in iter::range_step_inclusive(32, (32 - (groups * 8)), -8) {
                let v = ((n >> j as u8) & 255) as u8;
                out.push(v);
            }

            // Reset n
            n = 0;
        }
    }

    out
}

macro_rules! check_encode(
    ($input:expr, $expected:expr) => ( {
        let x = encode($input.as_bytes());
        assert!(x.eq_ignore_ascii_case($expected));
    });
    );

macro_rules! check_decode(
    ($expected:expr, $input:expr) => (
        assert_eq!($expected, decode($input.as_bytes()));
        );
);

#[test]
fn encode_tests() {
    check_encode!("", "");
    check_encode!("f", "MY");
    check_encode!("foo", "MZXw6");
    check_encode!("foob", "MZXW6YQ");
    check_encode!("fooba", "MZXW6YTB");
    check_encode!("foobar", "MZXW6YTBOI");
}

#[test]
fn decode_tests() {
    check_decode!("", "");
    check_decode!("f", "MY");
    check_decode!("foo", "MZXW6");
    check_decode!("foob", "MZxw6YQ");
    check_decode!("fooba", "MZXW6YTB");
    check_decode!("foobar", "MZXW6YTBOI");
}

<commit_msg>Fix test conversion<commit_after>
use std::iter;
use std::ascii::AsciiExt;

static STANDARD_CHARS: &'static[u8] = b"abcdefghijklmnopqrstuvwxyz234567";

pub fn encode(b: &[u8]) -> String {
    let mut out = Vec::new();
    let mut n: u64 = 0;

    for i in 0 .. b.len() {
        // Calculate the shift (32, 24, 16, 8, 0)
        let shift = 8 * (4 - (i % 5));
        n |= (b[i] as u64) << shift;

        // If the shift was zero, or this is the last character of input, extract
        // each 5-bit group from n.
        if shift == 0 || (i == b.len()-1) {
            // Determine number of output groups, based on the shift
            let groups = match shift {
                32 => 2,
                24 => 4,
                16 => 5,
                8  => 7,
                _  => 8
            };

            // Pull out each 5-bit index
            for j in iter::range_step_inclusive(35, (35 - (groups * 5)), -5) {
                let ndx = (n >> j as u32) & 31;
                out.push(STANDARD_CHARS[ndx as usize]);
            }

            // Reset n
            n = 0;
        }
    }

    let result = match String::from_utf8(out) {
        Ok(s) => s,
        Err(e) => panic!(e)
    };
    result
}

pub fn decode(b: &[u8]) -> Vec<u8> {
    let mut out: Vec<u8> = Vec::new();
    let mut n: u64 = 0;

    for i in iter::range_inclusive(0, b.len()) {
        // Use the character to generate an index
        let c = b[i];
        let idx = match c as char {
            'A' ... 'Z' => c - 65, // A..Z => 65..90
            'a' ... 'z' => c - 97, // a..z => 97..122
            _ => 26 + (c - 50)     // 2..7 => 50..55
        };

        // Calculate the shift for this 5-bit group (35, 25, ..., 0)
        let shift = 35 - (5 * (i % 8));
        n |= (idx as u64 & 31) << shift;

        // If the shift is zero or this is the last character of input, extract
        // each 8-bit group from n.
        if shift == 0 || (i == b.len()-1) {
            // Determine number of output groups, based on shift
            let groups = match shift {
                35 => 0,
                30 => 1,
                25 => 2,
                20 => 2,
                15 => 3,
                10 => 4,
                5  => 4,
                _  => 5
            };

            // Pull out each 8-bit group
            for j in iter::range_step_inclusive(32, (32 - (groups * 8)), -8) {
                let v = ((n >> j as u8) & 255) as u8;
                out.push(v);
            }

            // Reset n
            n = 0;
        }
    }

    out
}

macro_rules! check_encode(
    ($input:expr, $expected:expr) => ( {
        let x = encode($input.as_bytes());
        assert!(x.eq_ignore_ascii_case($expected));
    });
    );

macro_rules! check_decode(
    ($expected:expr, $input:expr) => ( {
        let d = String::from_utf8(decode($input.as_bytes()));
        assert_eq!($expected, d.ok().unwrap());
    });
    );

#[test]
fn encode_tests() {
    check_encode!("", "");
    check_encode!("f", "MY");
    check_encode!("foo", "MZXw6");
    check_encode!("foob", "MZXW6YQ");
    check_encode!("fooba", "MZXW6YTB");
    check_encode!("foobar", "MZXW6YTBOI");
}

#[test]
fn decode_tests() {
    check_decode!("", "");
    check_decode!("f", "MY");
    check_decode!("foo", "MZXW6");
    check_decode!("foob", "MZxw6YQ");
    check_decode!("fooba", "MZXW6YTB");
    check_decode!("foobar", "MZXW6YTBOI");
}

<|endoftext|>